"""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()