item 22 pre-registered: E2-L rung B internalization ladder (front-first deletion, 10 pauses/step, warm-started d=1..3); trainer gains --drop-steps/--warm-start/--steps/--lr; smoke-tested 2 steps on Spark

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

scripts/train_carry_cot.py

@@ -43,12 +43,35 @@ ap.add_argument("--lensnoise", default=None, metavar="RANK,SCALE",
                      "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)")
+ap.add_argument("--drop-steps", type=int, default=0, metavar="D",
+                help="E2-L rung B: delete the first D scratchpad steps, "
+                     "each replaced by --pause-per-step extra pauses")
+ap.add_argument("--pause-per-step", type=int, default=10,
+                help="pauses per deleted step (median step = 10 tokens)")
+ap.add_argument("--warm-start", default=None, metavar="ADAPTER_PT")
+ap.add_argument("--steps", type=int, default=STEPS)
+ap.add_argument("--lr", type=float, default=LR)
 ARGS = ap.parse_args()
+STEPS, LR = ARGS.steps, ARGS.lr
 TAG = ("carrycot_ff" if ARGS.feedforward else "carrycot") + (
+    f"_b{ARGS.drop_steps}" if ARGS.drop_steps else "") + (
     f"_ln{ARGS.lensnoise.replace(',', '_')}" if ARGS.lensnoise else "") + (
     f"_s{ARGS.seed}" if ARGS.seed else "")
 
 
+def drop_cot_steps(cot, d):
+    """Delete the first d scratchpad lines (front-first: the deleted
+    computation must ride the pause-chain before the visible remainder).
+    Returns (new_cot, n_deleted); unparseable cots pass through intact."""
+    lines = [l for l in cot.split("\n") if l.strip()]
+    ans = [l for l in lines if l.startswith("Answer")]
+    steps = [l for l in lines if not l.startswith("Answer")]
+    if len(ans) != 1 or not steps:
+        return cot, 0
+    n = min(d, len(steps))
+    return "\n".join(steps[n:] + ans), n
+
+
 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.
@@ -88,8 +111,9 @@ def build_batch(tok, items, p, device="cuda"):
                  add_special_tokens=False)["input_ids"]
         a = tok(it["cot"] + "<end_of_turn>",
                 add_special_tokens=False)["input_ids"]
-        seqs.append(pr + [PAUSE_ID] * p + a)
-        labs.append([-100] * (len(pr) + p) + a)
+        pi = p + it.get("extra_pauses", 0)
+        seqs.append(pr + [PAUSE_ID] * pi + a)
+        labs.append([-100] * (len(pr) + pi) + a)
         plens.append(len(pr))
     T = max(len(s) for s in seqs)
     pad = tok.pad_token_id or 0
@@ -125,12 +149,21 @@ def main():
             if it["split"] == "train"}
     cots = json.load(open(OUT / "gsm_cot_data.json"))
     data = []
+    n_dropped = 0
     for r in cots:
         it = star.get(r["idx"])
         if it is None:
             continue
+        cot, ndel = (drop_cot_steps(r["cot"], ARGS.drop_steps)
+                     if ARGS.drop_steps else (r["cot"], 0))
+        n_dropped += ndel
         data.append({"question": it["question"], "label": r["label"],
-                     "cot": r["cot"]})
+                     "cot": cot,
+                     "extra_pauses": ndel * ARGS.pause_per_step})
+    if ARGS.drop_steps:
+        print(f"rung B d={ARGS.drop_steps}: {n_dropped} steps deleted "
+              f"across {len(data)} items "
+              f"({ARGS.pause_per_step} pauses each)", flush=True)
 
     model, tok = load_model(dtype=torch.bfloat16)
     for pp in model.parameters():
@@ -147,11 +180,16 @@ def main():
                                    float(sc)).cuda()
         print(f"lens-noise: rank={r} scale={sc} on jbar L{BAND[0]}",
               flush=True)
+    if ARGS.warm_start:
+        (adapter.base if ARGS.lensnoise else adapter).load_state_dict(
+            torch.load(ARGS.warm_start, map_location="cuda"))
+        print(f"warm-started from {ARGS.warm_start}", flush=True)
     opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01)
 
     keep = [it for it in data
             if len(tok(it["question"])["input_ids"])
-            + len(tok(it["cot"])["input_ids"]) + 16 <= 460]
+            + len(tok(it["cot"])["input_ids"])
+            + it.get("extra_pauses", 0) + 16 <= 460]
     rng.shuffle(keep)
     pool = {l: [it for it in keep if it["label"] == l]
             for l in ("easy", "hard")}