"""Restriction-only answer likelihood, hidden-state drift and adapter weight spectra.""" import importlib.util import json import time from pathlib import Path import torch from huggingface_hub import HfApi, hf_hub_download from peft import PeftModel 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) report = json.loads(Path(hf_hub_download(REPO, "restriction_results.json", revision=REVISION)).read_text()) rows = report["rows"][:report["test_rows"]] tokenizer = AutoTokenizer.from_pretrained(helper.BASE, revision=helper.REVISION) tokenizer.padding_side = "left" base = AutoModelForCausalLM.from_pretrained(helper.BASE, revision=helper.REVISION, torch_dtype=torch.bfloat16, attn_implementation="eager").to("cuda") model = PeftModel.from_pretrained(base, REPO, subfolder="selective", revision=REVISION).eval() modules = {name: m for name, m in model.named_modules() if hasattr(m, "lora_A") and "default" in m.lora_A} scales = {name: m.scaling["default"] for name, m in modules.items()} def scale(alpha): for name, module in modules.items(): module.scaling["default"] = scales[name] * alpha @torch.no_grad() def likelihood(): values = [] for offset in range(0, len(rows), 4): sequences, labels = [], [] for row in rows[offset:offset+4]: prefix = tokenizer(helper.render_prompt(tokenizer, row), add_special_tokens=False)["input_ids"] target = tokenizer(json.dumps(dict(action=row["expected_action"], response=row["expected_response"])) + tokenizer.eos_token, add_special_tokens=False)["input_ids"] sequences.append(prefix + target) labels.append([-100]*len(prefix) + target) length = max(map(len, sequences)) ids = torch.tensor([[tokenizer.pad_token_id]*(length-len(s))+s for s in sequences], device="cuda") mask = torch.tensor([[0]*(length-len(s))+[1]*len(s) for s in sequences], device="cuda") targets = torch.tensor([[-100]*(length-len(s))+s for s in labels], device="cuda")[:, 1:] logits = model(input_ids=ids, attention_mask=mask).logits[:, :-1].float() losses = torch.nn.functional.cross_entropy(logits.transpose(1, 2), targets, reduction="none", ignore_index=-100) values.extend((losses.sum(1)/(targets != -100).sum(1)).cpu().tolist()) return values @torch.no_grad() def hidden(alpha): scale(alpha) result = [] 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) result.append(torch.stack([h[:, -1].float().cpu() for h in outputs.hidden_states])) return torch.cat(result, dim=1) result = dict(source_revision=REVISION, base_model=helper.BASE, base_revision=helper.REVISION, ids=[r["id"] for r in rows], group_ids=[r["group_id"] for r in rows], scope="84 previously opened held-out policy cases; 12 workspace groups; descriptive follow-up, one seed; no model tuning", strength_sweep=[], layer_ablations=[], activations=[], spectra=[]) for alpha in [0, .5, 1, 1.5, 2]: scale(alpha) values = likelihood() result["strength_sweep"].append(dict(alpha=alpha, nll=values)) print(json.dumps(dict(stage="strength_nll", alpha=alpha, mean=sum(values)/len(values))), flush=True) scale(1) for layer in range(26): for name, module in modules.items(): if f".layers.{layer}." in name: module.scaling["default"] = 0 result["layer_ablations"].append(dict(layer=layer, nll=likelihood())) scale(1) print(json.dumps(dict(stage="layer_nll", layer=layer)), flush=True) h0, h1 = hidden(0), hidden(1) for i in range(h0.shape[0]): relative = torch.linalg.vector_norm(h1[i]-h0[i], dim=1)/torch.linalg.vector_norm(h0[i], dim=1).clamp_min(1e-8) cosine = torch.nn.functional.cosine_similarity(h0[i], h1[i], dim=1) result["activations"].append(dict(hidden_state_index=i, relative_l2=relative.tolist(), cosine=cosine.tolist(), mean_relative_l2=float(relative.mean()), mean_cosine=float(cosine.mean()))) scale(1) for name, module in modules.items(): a, b = module.lora_A["default"].weight.float(), module.lora_B["default"].weight.float() def spectrum(aa, bb): qa, ra = torch.linalg.qr(aa.T, mode="reduced") qb, rb = torch.linalg.qr(bb, mode="reduced") return torch.linalg.svdvals(rb @ ra.T)*abs(scales[name]) singular = spectrum(a, b) ar, br = torch.randn_like(a), torch.randn_like(b) random_singular = spectrum(ar, br) random_singular *= torch.linalg.vector_norm(singular)/torch.linalg.vector_norm(random_singular).clamp_min(1e-20) result["spectra"].append(dict(module=name, rank=a.shape[0], singular_values=singular.tolist(), random_control_singular_values=random_singular.tolist(), delta_frobenius=float(torch.linalg.vector_norm(singular)))) restored = likelihood() full = next(x["nll"] for x in result["strength_sweep"] if x["alpha"] == 1) result["restoration_max_nll_difference"] = max(abs(a-b) for a,b in zip(full, restored, strict=True)) assert result["restoration_max_nll_difference"] == 0 result["elapsed_seconds"] = time.monotonic()-started folder = Path("restriction_analysis") folder.mkdir(exist_ok=True) (folder/"analysis.json").write_text(json.dumps(result, indent=2)) (folder/Path(__file__).name).write_bytes(Path(__file__).read_bytes()) HfApi().upload_folder(repo_id=REPO, folder_path=str(folder), path_in_repo="analysis", commit_message="Measure restriction adapter layer likelihood, activations and spectra") print(json.dumps(dict(stage="complete", elapsed=result["elapsed_seconds"])), flush=True) if __name__ == "__main__": main()