Measure Gemma/Qwen steering, probes, causal patching and paired activations

interpretability/measured_activations.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6db9791c11b6f7a4a5412f07057e5ea46edc0d6c909a53f64aa9331033c2f0e7
+size 59361712

interpretability/train_interpretability_pilot.py

@@ -0,0 +1,351 @@
+"""Measured cross-model directions, steering controls and token-level causal patching."""
+
+import importlib.util
+import json
+import re
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+from huggingface_hub import HfApi, hf_hub_download
+from peft import PeftModel
+from safetensors.torch import save_file
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+REPO = "itsalissonsilva/Restriction-Learning-Experiment"
+REVISION = "efc3f02e5775d3d8428b4961f1fd94aedc6ca34f"
+
+
+def main():
+    started = time.monotonic()
+    torch.manual_seed(1729)
+    source = hf_hub_download(REPO, "train_restriction_pilot.py", revision=REVISION)
+    spec = importlib.util.spec_from_file_location("restriction_helper", source)
+    helper = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(helper)
+    original = json.loads(
+        Path(hf_hub_download(REPO, "restriction_results.json", revision=REVISION)).read_text()
+    )
+    train = [r for r in helper.make_rows("train", 24) if r["kind"] != "ambiguous"]
+    dev = [r for r in helper.make_rows("development", 4) if r["kind"] != "ambiguous"]
+    test = [r for r in original["rows"][:84] if r["kind"] != "ambiguous"]
+    intervention_rows = original["rows"][:14] + original["rows"][-4:]
+    result = dict(
+        seed=1729,
+        adapter_revision=REVISION,
+        training_rows=train,
+        development_rows=dev,
+        test_rows=test,
+        intervention_rows=intervention_rows,
+        models={},
+        cross_model={},
+        patching=[],
+        scope="Synthetic workspace restriction intent; one seed; no general unsafe-intent or universal-circuit claim",
+    )
+    tensors = {}
+    for key, base_name, revision in [
+        ("gemma", helper.BASE, helper.REVISION),
+        ("qwen", "Qwen/Qwen3-0.6B", HfApi().model_info("Qwen/Qwen3-0.6B").sha),
+    ]:
+        tokenizer = AutoTokenizer.from_pretrained(base_name, revision=revision)
+        tokenizer.padding_side = "left"
+        model = (
+            AutoModelForCausalLM.from_pretrained(
+                base_name,
+                revision=revision,
+                torch_dtype=torch.bfloat16,
+                attn_implementation="eager",
+            )
+            .to("cuda")
+            .eval()
+        )
+        if key == "gemma":
+            model = PeftModel.from_pretrained(
+                model, REPO, subfolder="selective", revision=REVISION
+            ).eval()
+        layers = [(n, m) for n, m in model.named_modules() if re.search(r"\.layers\.\d+$", n)]
+        assert len(layers) >= 26
+
+        @torch.no_grad()
+        def capture(rows):
+            chunks = []
+            for offset in range(0, len(rows), 8):
+                inputs = tokenizer(
+                    [helper.render_prompt(tokenizer, r) for r in rows[offset : offset + 8]],
+                    padding=True,
+                    return_tensors="pt",
+                ).to("cuda")
+                outputs = model(**inputs, output_hidden_states=True)
+                chunks.append(torch.stack([h[:, -1].float().cpu() for h in outputs.hidden_states]))
+            return torch.cat(chunks, dim=1)
+
+        @torch.no_grad()
+        def generate(rows, batch_size=6):
+            outputs = []
+            for offset in range(0, len(rows), batch_size):
+                batch = rows[offset : offset + batch_size]
+                inputs = tokenizer(
+                    [helper.render_prompt(tokenizer, r) for r in batch],
+                    padding=True,
+                    return_tensors="pt",
+                ).to("cuda")
+                generated = model.generate(
+                    **inputs,
+                    max_new_tokens=80,
+                    do_sample=False,
+                    pad_token_id=tokenizer.pad_token_id,
+                )
+                texts = tokenizer.batch_decode(
+                    generated[:, inputs["input_ids"].shape[1] :], skip_special_tokens=True
+                )
+                outputs.extend(helper.grading(r, t) for r, t in zip(batch, texts, strict=True))
+            return outputs
+
+        htrain, hdev, htest = capture(train), capture(dev), capture(test)
+        y = np.array([r["expected_action"] == "withhold_private" for r in train])
+        yd = np.array([r["expected_action"] == "withhold_private" for r in dev])
+        yt = np.array([r["expected_action"] == "withhold_private" for r in test])
+        entry = dict(
+            base_model=base_name,
+            base_revision=revision,
+            condition="selective restriction LoRA" if key == "gemma" else "base",
+            probes=[],
+            steering=[],
+            activation_shape=list(htrain.shape),
+        )
+        for layer in range(htrain.shape[0]):
+            x = htrain[layer].numpy()
+            d = hdev[layer].numpy()
+            t = htest[layer].numpy()
+            mean = x.mean(0)
+            scale = x.std(0).clip(1e-4)
+            xs = (x - mean) / scale
+            # Ridge linear discriminator; labels and validation threshold stay outside held-out test.
+            targets = y.astype(float) * 2 - 1
+            selected = np.argsort(np.abs(xs.T @ targets))[-64:]
+            selected_x = xs[:, selected]
+            sparse = selected_x.T @ np.linalg.solve(
+                selected_x @ selected_x.T + 10 * np.eye(len(x)), targets
+            )
+            coef = np.zeros(x.shape[1])
+            coef[selected] = sparse
+            scores = ((d - mean) / scale) @ coef
+            candidates = np.unique(np.r_[scores, 0])
+            threshold = max(candidates, key=lambda th: np.mean((scores > th) == yd))
+            prediction = ((t - mean) / scale) @ coef > threshold
+            shuffled = np.random.default_rng(1729).permutation(targets)
+            random_coef = xs.T @ np.linalg.solve(xs @ xs.T + 10 * np.eye(len(x)), shuffled)
+            random_accuracy = float(np.mean((((t - mean) / scale) @ random_coef > 0) == yt))
+            entry["probes"].append(
+                dict(
+                    hidden_state_index=layer,
+                    nonzero_features=64,
+                    validation_accuracy=float(np.mean((scores > threshold) == yd)),
+                    test_accuracy=float(np.mean(prediction == yt)),
+                    shuffled_label_accuracy=random_accuracy,
+                    threshold=float(threshold),
+                    test_scores=(((t - mean) / scale) @ coef).tolist(),
+                )
+            )
+            tensors[f"{key}.probe.{layer}.mean"] = torch.from_numpy(mean)
+            tensors[f"{key}.probe.{layer}.scale"] = torch.from_numpy(scale)
+            tensors[f"{key}.probe.{layer}.coef"] = torch.from_numpy(coef.astype(np.float32))
+        for label, hidden in [("train", htrain), ("development", hdev), ("test", htest)]:
+            tensors[f"{key}.hidden.{label}"] = hidden
+        directions = {}
+        for layer in [12, 24]:
+            direction = htrain[layer + 1, y].mean(0) - htrain[layer + 1, ~y].mean(0)
+            random_direction = torch.randn_like(direction)
+            random_direction *= direction.norm() / random_direction.norm()
+            directions[layer] = direction
+            tensors[f"{key}.direction.{layer}"] = direction
+            tensors[f"{key}.random_direction.{layer}"] = random_direction
+            for kind, vector in [("learned", direction), ("random", random_direction)]:
+                for alpha in [-2, -1, 0, 1, 2]:
+
+                    def hook(module, args, output):
+                        hidden = output[0] if isinstance(output, tuple) else output
+                        if hidden.shape[1] <= 1:
+                            return output
+                        changed = hidden.clone()
+                        changed[:, -1] += alpha * vector.to(
+                            device=hidden.device, dtype=hidden.dtype
+                        )
+                        return (changed, *output[1:]) if isinstance(output, tuple) else changed
+
+                    handle = layers[layer][1].register_forward_hook(hook)
+                    try:
+                        outputs = generate(intervention_rows)
+                    finally:
+                        handle.remove()
+                    entry["steering"].append(
+                        dict(
+                            layer=layer,
+                            alpha=alpha,
+                            direction=kind,
+                            outputs=outputs,
+                            summary=helper.summarize(intervention_rows, outputs),
+                        )
+                    )
+                    print(
+                        json.dumps(
+                            dict(
+                                stage="steering",
+                                model=key,
+                                layer=layer,
+                                direction=kind,
+                                alpha=alpha,
+                            )
+                        ),
+                        flush=True,
+                    )
+        # Zero strength and hook removal use identical batches for restoration.
+        baseline = generate(intervention_rows)
+        zeros = [s for s in entry["steering"] if s["alpha"] == 0]
+        entry["zero_strength_raw_matches"] = [
+            sum(o["raw"] == b["raw"] for o, b in zip(s["outputs"], baseline, strict=True))
+            for s in zeros
+        ]
+        assert all(n == len(intervention_rows) for n in entry["zero_strength_raw_matches"])
+        if key == "gemma":
+            # Same request and content; only the field's PRIVATE/PUBLIC label changes.
+            private = []
+            public = []
+            for r in original["rows"]:
+                if r["kind"] == "direct" and len(private) < 8:
+                    a = dict(r, request="What is the recovery code in this record?")
+                    b = dict(
+                        a,
+                        context=a["context"].replace("PRIVATE", "PUBLIC"),
+                        expected_action="answer_public",
+                        private_canary="",
+                        required=[a["private_canary"]],
+                    )
+                    private.append(a)
+                    public.append(b)
+            hp, hc = capture(private), capture(public)
+            baseline_private = generate(private, 8)
+            baseline_public = generate(public, 8)
+            result["patch_pairs"] = dict(
+                private=private,
+                public=public,
+                baseline_private=baseline_private,
+                baseline_public=baseline_public,
+            )
+            # Patch donor states at matched final-prompt positions only. Donor and receiver have identical batch layout.
+            for layer in range(26):
+                for donor_kind, donor in [
+                    ("public", hc[layer + 1]),
+                    ("self_restoration", hp[layer + 1]),
+                    ("random_norm_matched", None),
+                ]:
+
+                    def patch(module, args, output):
+                        hidden = output[0] if isinstance(output, tuple) else output
+                        if hidden.shape[1] <= 1:
+                            return output
+                        changed = hidden.clone()
+                        # One complete 8-row batch is used for patch capture/generation below.
+                        if donor is not None:
+                            replacement = donor.to(hidden.device, hidden.dtype)
+                        else:
+                            noise = torch.randn(
+                                hidden.shape[0],
+                                hidden.shape[-1],
+                                device=hidden.device,
+                                dtype=torch.float32,
+                            )
+                            noise *= (
+                                (hc[layer + 1] - hp[layer + 1]).norm(dim=1).to(hidden.device)
+                                / noise.norm(dim=1).clamp_min(1e-20)
+                            )[:, None]
+                            replacement = hp[layer + 1].to(hidden.device, hidden.dtype) + noise.to(
+                                hidden.dtype
+                            )
+                        changed[:, -1] = replacement
+                        return (changed, *output[1:]) if isinstance(output, tuple) else changed
+
+                    # Explicit batch of all eight avoids donor-index ambiguity.
+                    inputs = tokenizer(
+                        [helper.render_prompt(tokenizer, r) for r in private],
+                        padding=True,
+                        return_tensors="pt",
+                    ).to("cuda")
+                    handle = layers[layer][1].register_forward_hook(patch)
+                    try:
+                        with torch.no_grad():
+                            generated = model.generate(
+                                **inputs,
+                                max_new_tokens=80,
+                                do_sample=False,
+                                pad_token_id=tokenizer.pad_token_id,
+                            )
+                    finally:
+                        handle.remove()
+                    texts = tokenizer.batch_decode(
+                        generated[:, inputs["input_ids"].shape[1] :], skip_special_tokens=True
+                    )
+                    outputs = [
+                        helper.grading(r, text) for r, text in zip(private, texts, strict=True)
+                    ]
+                    result["patching"].append(
+                        dict(
+                            layer=layer,
+                            donor=donor_kind,
+                            outputs=outputs,
+                            summary=helper.summarize(private, outputs),
+                            raw_matches_private_baseline=sum(
+                                o["raw"] == b["raw"]
+                                for o, b in zip(outputs, baseline_private, strict=True)
+                            ),
+                        )
+                    )
+                print(json.dumps(dict(stage="patching", layer=layer)), flush=True)
+        result["models"][key] = entry
+        model = None
+        layers = []
+        del htrain, hdev, htest
+        torch.cuda.empty_cache()
+    # Paired low-rank mapping, fit on training activations only. Report held-out alignment rather than claiming behavioral transfer.
+    for layer in [12, 24]:
+        x = tensors["qwen.hidden.train"][layer + 1].numpy()
+        y = tensors["gemma.hidden.train"][layer + 1].numpy()
+        xm, ym = x.mean(0), y.mean(0)
+        u, s, vt = np.linalg.svd(x - xm, full_matrices=False)
+        rank = 8
+        left = vt[:rank].T
+        right = np.linalg.solve(
+            (s[:rank] ** 2)[:, None] * np.eye(rank) + 10 * np.eye(rank),
+            (u[:, :rank] * s[:rank]).T @ (y - ym),
+        )
+        xt = tensors["qwen.hidden.test"][layer + 1].numpy()
+        yt = tensors["gemma.hidden.test"][layer + 1].numpy()
+        pred = (xt - xm) @ left @ right + ym
+        result["cross_model"][str(layer)] = dict(
+            rank=rank,
+            heldout_mse=float(np.mean((pred - yt) ** 2)),
+            train_mean_baseline_mse=float(np.mean((ym - yt) ** 2)),
+            scope="paired activation reconstruction only; behavioral transfer gate not yet tested",
+        )
+        tensors[f"mapper.{layer}.left"] = torch.from_numpy(left)
+        tensors[f"mapper.{layer}.right"] = torch.from_numpy(right)
+    result["elapsed_seconds"] = time.monotonic() - started
+    folder = Path("interpretability")
+    folder.mkdir(exist_ok=True)
+    (folder / "results.json").write_text(json.dumps(result, indent=2))
+    save_file(
+        {k: v.contiguous() for k, v in tensors.items()}, folder / "measured_activations.safetensors"
+    )
+    (folder / "train_interpretability_pilot.py").write_bytes(Path(__file__).read_bytes())
+    HfApi().upload_folder(
+        repo_id=REPO,
+        folder_path=str(folder),
+        path_in_repo="interpretability",
+        commit_message="Measure Gemma/Qwen steering, probes, causal patching and paired activations",
+    )
+    print(json.dumps(dict(stage="complete", seconds=result["elapsed_seconds"])), flush=True)
+
+
+if __name__ == "__main__":
+    main()