Add the lens-A (whitened) AR, the frozen whitening basis (mu, W) and olens.py load_whitening / whiten / unwhiten

ar_whitened/export.json

@@ -0,0 +1,8 @@
+{
+ "name": "final_olens_lensB/ar_whitened",
+ "ckpt": "ar.whiten.lr7e-5.decay171155/milestone_step_205386",
+ "run": "ar.whiten.lr7e-5.decay171155",
+ "n_tensors": 994,
+ "n_params": 142947328,
+ "overrides": "kind=ar name=ar.whiten.lr7e-5 ar.space=whiten training.lr=7e-05 data.train_dir=/data/pairs/w0,/data/pairs/w1,/data/pairs/w2 data.val_dir=/data/rulers/ar grad_ckpt=False compile=True training.num_workers=16 training.val_batch_size=64 data.basis_path=/data/basis.safetensors training.train_batch_size=72 training.gradient_accumulation_steps=1 effective_batch=576 training.max_steps=213944 wandb_id=ar.whiten.lr7e-5.r1 name=ar.whiten.lr7e-5.decay171155 wandb_id=ar.whiten.lr7e-5.decay171155 branch_from=/runs/ar.whiten.lr7e-5/milestone_step_171155 decay_frac=0.2"
+}

ar_whitened/model.safetensors

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

olens.py

@@ -6,6 +6,10 @@ Self-contained (transformers + peft + safetensors + huggingface_hub). Three piec
   ``ar_direction``), used to score readouts and as the lens's RL reward;
 - ``residual``: capture the subject model's residual for a chat turn.
 
+Also published: the lens-A AR, trained in WHITENED space (``ar_whitened``, same architecture, so
+``load_ar`` loads it), and the frozen per-layer whitening basis it was trained with
+(``whitening/basis.safetensors``: ``load_whitening``, ``whiten``, ``unwhiten``).
+
     from olens import from_hub, load_lens, load_ar, residual, read, ar_direction
     lens = load_lens(from_hub("agu18dec/final_olens", "lens_rl_step600"))
     h = residual(lens, "Who wrote the opera Carmen?", layer=40)   # [5120]
@@ -13,6 +17,12 @@ Self-contained (transformers + peft + safetensors + huggingface_hub). Three piec
     ar = load_ar(from_hub("agu18dec/final_olens", "ar"))
     print(torch.cosine_similarity(ar_direction(ar, "Bizet", 40), h.float().cuda(), dim=0))
 
+    wb = load_whitening(from_hub("agu18dec/final_olens", "whitening",
+                                 filename="basis.safetensors"), layers=[40])  # z = (x - mu_L) W_L
+    ar_w = load_ar(from_hub("agu18dec/final_olens", "ar_whitened"))
+    z_pred = whiten(ar_direction(ar_w, "Bizet", 40).cpu(), 40, wb)
+    print(torch.cosine_similarity(z_pred, whiten(h.cpu(), 40, wb), dim=0))  # its training metric
+
 Conventions (must match training):
 - A residual at layer L is the output of decoder block L (0..63), before the final norm.
 - Text is read with the chat template and enable_thinking=False.
@@ -38,12 +48,13 @@ PROMPT = (
 TAG = "[Layer {n:02d}]"  # the AR's layer prefix (6 tokens)
 
 
-def from_hub(repo: str, folder: str, revision: str | None = None) -> str:
-    """Download one checkpoint folder (``<folder>/model.safetensors``) from the Hub; returns the
-    local path to pass to ``load_lens`` / ``load_ar``."""
+def from_hub(repo: str, folder: str, revision: str | None = None,
+             filename: str = "model.safetensors") -> str:
+    """Download one file of a checkpoint folder (default ``<folder>/model.safetensors``) from the
+    Hub; returns the local path to pass to ``load_lens`` / ``load_ar`` / ``load_whitening``."""
     from huggingface_hub import hf_hub_download
 
-    return hf_hub_download(repo, f"{folder}/model.safetensors", revision=revision)
+    return hf_hub_download(repo, f"{folder}/{filename}", revision=revision)
 
 
 def _lora() -> LoraConfig:
@@ -172,3 +183,39 @@ def ar_direction(ar: AR, text: str, layer: int, max_tokens: int = 32) -> torch.T
     x = torch.tensor([tag + span], device=next(ar.parameters()).device)
     hidden = ar.backbone(input_ids=x).last_hidden_state
     return ar.head(hidden[0, -1].float())
+
+
+class Whitening:
+    """The frozen per-layer whitening basis: ``mu`` [64, 5120] and ``w[L]`` [5120, 5120] (fp32),
+    ``W_L = (Sigma_L + lam_L I)^(-1/2)`` with ``lam_L = 0.1 * mean(diag Sigma_L)``; ``W_L`` is
+    symmetric. Fit once on 1,310,720 residuals per layer (on-policy chat replies) and never refit.
+    """
+
+    def __init__(self, mu: torch.Tensor, w: dict[int, torch.Tensor]) -> None:
+        self.mu, self.w = mu, w
+        self._inv: dict[int, torch.Tensor] = {}
+
+
+def load_whitening(path: str, layers: list[int] | None = None, device: str = "cpu") -> Whitening:
+    """Load the basis; ``layers`` limits which ``W_L`` are read (105 MB each; all 64: 6.7 GB)."""
+    from safetensors import safe_open
+
+    with safe_open(path, framework="pt") as f:
+        mu = f.get_tensor("mu").float().to(device)
+        sl = f.get_slice("w")
+        w = {lay: sl[lay].float().to(device) for lay in (layers or range(N_LAYERS))}
+    return Whitening(mu, w)
+
+
+def whiten(x: torch.Tensor, layer: int, wb: Whitening) -> torch.Tensor:
+    """``z = (x - mu_L) W_L`` for residual(s) ``x`` [..., 5120] taken at ``layer``."""
+    w = wb.w[layer]
+    return (x.float().to(w.device) - wb.mu[layer]) @ w
+
+
+def unwhiten(z: torch.Tensor, layer: int, wb: Whitening) -> torch.Tensor:
+    """The inverse map, ``x = z W_L^(-1) + mu_L`` (``W_L^(-1)`` computed once in fp64, cached)."""
+    if layer not in wb._inv:
+        wb._inv[layer] = torch.linalg.inv(wb.w[layer].double()).float()
+    inv = wb._inv[layer]
+    return z.float().to(inv.device) @ inv + wb.mu[layer]

whitening/basis.json

@@ -0,0 +1,24 @@
+{
+ "rows": 1310720,
+ "ridge": 0.1,
+ "n_shards": 16,
+ "sha256_16": "8d05c083386b6924",
+ "source": "/data/basis.safetensors",
+ "tensors": {
+  "mu": [
+   [
+    64,
+    5120
+   ],
+   "torch.float32"
+  ],
+  "w": [
+   [
+    64,
+    5120,
+    5120
+   ],
+   "torch.float32"
+  ]
+ }
+}

whitening/basis.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:61db172b0841eebf6124ea025f6f667c2c72d68d60b4504987127261f4866d35
+size 6712197280