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

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

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)