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/jobs/zzz_m_gsm_rungb.sh

@@ -0,0 +1,13 @@
+# gpuq-in:  results-loop/star_data.json results-loop/gsm_cot_data.json results-loop/adapter_carrycot_e400.pt
+# gpuq-out: results-loop/eval_gsm_carrycot_b*.json results-loop/train_carrycot_b*_log.json results-loop/adapter_carrycot_b*.pt
+git pull origin main -q 2>/dev/null
+P=/home/nils/jspace/.venv/bin/python
+export JLENS_MODEL=google/gemma-4-E2B-it LOOP_OUT=/home/nils/jspace/results-loop
+cd /home/nils/jspace/scripts
+WARM=$LOOP_OUT/adapter_carrycot_e400.pt
+for D in 1 2 3; do
+  $P train_carry_cot.py --drop-steps $D --warm-start $WARM --steps 200 --lr 3e-4
+  WARM=$LOOP_OUT/adapter_carrycot_b${D}_e200.pt
+  $P eval_carry_cot.py --adapter $WARM --tag gsm_carrycot_b$D \
+     --grid 0:0,2:$((2+10*D)),2:$((6+10*D)) --n 256
+done

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