adaptive-alpha merge (Lys-inspired, trained) + truncated-BPTT deep-k training

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

scripts/eval_loop_code.py

@@ -15,7 +15,7 @@ from pathlib import Path
 
 import torch
 
-from loop_common import BandLooper, MergeAdapter
+from loop_common import AdaptiveMergeAdapter, BandLooper, MergeAdapter
 from prep_mbpp import DIRECT_SUFFIX, extract_code, mbpp_prompt, run_tests
 
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
@@ -75,15 +75,16 @@ def main():
     ap.add_argument("--pause", type=int, default=0,
                     help="append p pause tokens to each prompt")
     ap.add_argument("--alpha", type=float, default=0.3)
+    ap.add_argument("--adaptive", action="store_true")
     args = ap.parse_args()
     ks = [int(x) for x in args.ks.split(",")]
 
     model, tok = load_model(dtype=torch.bfloat16)
     tok.padding_side = "left"
     looper = BandLooper(model)
-    adapter = MergeAdapter(
-        d=model.config.get_text_config().hidden_size,
-        alpha=args.alpha).cuda()
+    cls = AdaptiveMergeAdapter if args.adaptive else MergeAdapter
+    kw = ({"alpha0": args.alpha} if args.adaptive else {"alpha": args.alpha})
+    adapter = cls(d=model.config.get_text_config().hidden_size, **kw).cuda()
     if args.adapter:
         adapter.load_state_dict(torch.load(args.adapter, map_location="cuda"))
         adapter.eval()

scripts/loop_common.py

@@ -53,6 +53,37 @@ class MergeAdapter(nn.Module):
         return out.to(dt)
 
 
+class AdaptiveMergeAdapter(nn.Module):
+    """Merge with state-dependent anchor coefficient (Lys-inspired, trained).
+
+    alpha(e, s) = sigmoid(w·[e;ŝ] + b), per position; w zero-init and
+    b = logit(0.3), so at init this is exactly the fixed alpha=0.3 merge."""
+
+    def __init__(self, d=1536, hidden=512, alpha0=0.3):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(2 * d, hidden), nn.GELU(), nn.Linear(hidden, d)
+        )
+        nn.init.zeros_(self.mlp[2].weight)
+        nn.init.zeros_(self.mlp[2].bias)
+        self.alpha_head = nn.Linear(2 * d, 1)
+        nn.init.zeros_(self.alpha_head.weight)
+        import math as _m
+        nn.init.constant_(self.alpha_head.bias,
+                          _m.log(alpha0 / (1 - alpha0)))
+
+    def forward(self, e, s):
+        dt = e.dtype
+        e32, s32 = e.float(), s.float()
+        s_hat = s32 * (
+            e32.norm(dim=-1, keepdim=True) / (s32.norm(dim=-1, keepdim=True) + 1e-6)
+        )
+        cat = torch.cat([e32, s_hat], dim=-1)
+        a = torch.sigmoid(self.alpha_head(cat))
+        out = (1 - a) * e32 + a * s_hat + self.mlp(cat)
+        return out.to(dt)
+
+
 class BandLooper:
     """Capture layer-call kwargs once per forward, then re-run L14-30 manually."""
 
@@ -117,7 +148,9 @@ class BandLooper:
 
     def loop_logits(self, adapter, input_ids, k, attention_mask=None,
                     use_checkpoint=False, return_states=False, last_only=False,
-                    loop_mask=None, feedforward=False):
+                    loop_mask=None, feedforward=False, bptt=None):
+        """bptt: backprop only through the last `bptt` iterations (McLeish-
+        style truncated BPTT); earlier iterations run under no_grad."""
         """Teacher-forced logits after k merge->band loops. k=0 = plain forward.
 
         loop_mask (B, T) bool: positions where the merge applies; elsewhere the
@@ -143,7 +176,17 @@ class BandLooper:
         with torch.no_grad():
             s = self.band(e, calls)  # s_0: no trainable params upstream
         states = [s]
-        for _ in range(k):
+        n_nograd = max(0, k - bptt) if bptt else 0
+        for i in range(k):
+            if i < n_nograd:
+                with torch.no_grad():
+                    x = adapter(e, s)
+                    if loop_mask is not None:
+                        x = torch.where(loop_mask[..., None], x, e)
+                    s = self.band(x, calls)
+                s = s.detach()
+                states.append(s)
+                continue
             x = adapter(e, s)
             if loop_mask is not None:
                 x = torch.where(loop_mask[..., None], x, e)

scripts/train_merge_code.py

@@ -20,7 +20,7 @@ from pathlib import Path
 import torch
 import torch.nn.functional as F
 
-from loop_common import BandLooper, MergeAdapter
+from loop_common import AdaptiveMergeAdapter, BandLooper, MergeAdapter
 from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt
 
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
@@ -34,6 +34,7 @@ LR = 1e-3
 WARMUP = 20
 MAX_TOK = 512
 K_BUCKETS = [(1, ("easy",)), (2, ("easy", "hard")), (4, ("hard",))]
+# --deepk 16 rescales to [(2, easy), (8, mixed), (16, hard)]
 
 ap = argparse.ArgumentParser()
 ap.add_argument("--seed", type=int, default=0)
@@ -44,11 +45,19 @@ ap.add_argument("--feedforward", action="store_true",
                 help="apply adapter once, no recurrence (pause-FF control)")
 ap.add_argument("--alpha", type=float, default=0.3,
                 help="merge weight (2B-tuned default 0.3; try 0.1-0.15 at 12B)")
+ap.add_argument("--adaptive", action="store_true",
+                help="state-dependent alpha (AdaptiveMergeAdapter)")
+ap.add_argument("--deepk", type=int, default=0,
+                help="scale curriculum depths by deepk/4 (e.g. 16 -> 2/8/16)")
+ap.add_argument("--bptt", type=int, default=0,
+                help="truncated BPTT: grads only through last N iterations")
 ARGS = ap.parse_args()
 SEED = ARGS.seed
 SUFFIX = ((f"_s{SEED}" if SEED else "")
           + (f"_p{ARGS.pause}" if ARGS.pause else "")
-          + (f"_a{ARGS.alpha}" if ARGS.alpha != 0.3 else ""))
+          + (f"_a{ARGS.alpha}" if ARGS.alpha != 0.3 else "")
+          + ("_ad" if ARGS.adaptive else "")
+          + (f"_dk{ARGS.deepk}" if ARGS.deepk else ""))
 PAUSE_ID = 6  # <unused0>
 
 
@@ -107,7 +116,15 @@ def main():
         p.requires_grad_(False)
     looper = BandLooper(model)
     d = model.config.get_text_config().hidden_size
-    adapter = MergeAdapter(d=d, alpha=ARGS.alpha).cuda()
+    if ARGS.adaptive:
+        adapter = AdaptiveMergeAdapter(d=d, alpha0=ARGS.alpha).cuda()
+    else:
+        adapter = MergeAdapter(d=d, alpha=ARGS.alpha).cuda()
+    global K_BUCKETS
+    if ARGS.deepk:
+        f = ARGS.deepk / 4
+        K_BUCKETS = [(max(1, int(k * f)), lbls) for k, lbls in K_BUCKETS]
+        print("K_BUCKETS ->", [(k, l) for k, l in K_BUCKETS], flush=True)
     opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01)
 
     train = [it for it in data if it["split"] == "train"
@@ -135,7 +152,8 @@ def main():
             g["lr"] = lr_at(step)
         logits = looper.loop_logits(adapter, ids, k, attention_mask=msk,
                                     use_checkpoint=True, loop_mask=lmask,
-                                    feedforward=ARGS.feedforward)
+                                    feedforward=ARGS.feedforward,
+                                    bptt=ARGS.bptt or None)
         loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(),
                                lab[:, 1:].flatten(), ignore_index=-100)
         opt.zero_grad(set_to_none=True)