item 21: GSM carry-cot 57.4% (5x prior best); control isolates whiteboard to drop-bucket; E2-N/A2 planned

scripts/train_carry_cot.py

@@ -38,11 +38,42 @@ P_BY_LABEL = {"easy": 2, "hard": 6}
 ap = argparse.ArgumentParser()
 ap.add_argument("--feedforward", action="store_true")
 ap.add_argument("--seed", type=int, default=0)
+ap.add_argument("--lensnoise", default=None, metavar="RANK,SCALE",
+                help="E2-N1: inject noise into the carried state during "
+                     "training, shaped by the top-RANK sensitivity "
+                     "directions of jbar at the band entrance; noise norm "
+                     "= SCALE * per-position state norm (e.g. 32,0.05)")
 ARGS = ap.parse_args()
 TAG = ("carrycot_ff" if ARGS.feedforward else "carrycot") + (
+    f"_ln{ARGS.lensnoise.replace(',', '_')}" if ARGS.lensnoise else "") + (
     f"_s{ARGS.seed}" if ARGS.seed else "")
 
 
+class LensNoiseWrapper(torch.nn.Module):
+    """Perturb the carried state s (not the anchor e) before the merge,
+    within the span of the lens's top-r readout-sensitive directions.
+    Train-time only (noise_on flag); eval and checkpoints use .base."""
+
+    def __init__(self, base, jbar_layer, rank, scale):
+        super().__init__()
+        self.base = base
+        self.scale = scale
+        self.noise_on = True
+        J = jbar_layer.float()
+        _, _, Vt = torch.linalg.svd(J, full_matrices=False)
+        self.register_buffer("V", Vt[:rank].T.contiguous())  # (d, r)
+
+    def forward(self, e, s):
+        if self.noise_on and self.scale > 0:
+            z = torch.randn(*s.shape[:-1], self.V.shape[1],
+                            device=s.device, dtype=torch.float32)
+            n = z @ self.V.T
+            n = n * (s.float().norm(dim=-1, keepdim=True) * self.scale
+                     / (n.norm(dim=-1, keepdim=True) + 1e-6))
+            s = (s.float() + n).to(s.dtype)
+        return self.base(e, s)
+
+
 def lr_at(step):
     if step < WARMUP:
         return LR * (step + 1) / WARMUP
@@ -107,6 +138,15 @@ def main():
     looper = BandLooper(model)
     adapter = MergeAdapter(
         d=model.config.get_text_config().hidden_size).cuda()
+    if ARGS.lensnoise:
+        r, sc = ARGS.lensnoise.split(",")
+        jbar = torch.load(Path(__file__).resolve().parent.parent
+                          / "results/jbar.pt", map_location="cpu")["Jbar"]
+        from loop_common import BAND
+        adapter = LensNoiseWrapper(adapter, jbar[BAND[0]], int(r),
+                                   float(sc)).cuda()
+        print(f"lens-noise: rank={r} scale={sc} on jbar L{BAND[0]}",
+              flush=True)
     opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01)
 
     keep = [it for it in data
@@ -146,11 +186,15 @@ def main():
             print(f"step {step:4d}  {lbl:4s} loss={loss.item():.4f}  "
                   f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
         if step % 200 == 199 or step == STEPS - 1:
+            if ARGS.lensnoise:
+                adapter.noise_on = False
             for l in ("easy", "hard"):
                 v = val_loss(looper, adapter, tok, val[l], P_BY_LABEL[l])
                 print(f"  val@{step}: {l}={v:.3f}", flush=True)
-            torch.save(adapter.state_dict(),
-                       OUT / f"adapter_{TAG}_e{step+1}.pt")
+            if ARGS.lensnoise:
+                adapter.noise_on = True
+            sd = (adapter.base if ARGS.lensnoise else adapter).state_dict()
+            torch.save(sd, OUT / f"adapter_{TAG}_e{step+1}.pt")
             json.dump(log, open(OUT / f"train_{TAG}_log.json", "w"))
     print("done", flush=True)