RecurrentAdapter arm: Huginn-regime retrofit (learned A/B, noise h0, randomized depth) + pre-registration item 11

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

scripts/loop_common.py

@@ -84,6 +84,50 @@ class AdaptiveMergeAdapter(nn.Module):
         return out.to(dt)
 
 
+class RecurrentAdapter(nn.Module):
+    """Huginn-style recurrent-state update on the frozen band
+    (arXiv 2502.05171 regime: h_{t+1} = A·h_t + B·e + Transformer(h_t, e)).
+
+    The band's residual stream supplies the "+Transformer" term, so the
+    adapter computes the band input x_t = A·ĥ_t + B·e + MLP([e;ĥ_t]) with
+    LEARNED d×d maps A, B (init A=α·I, B=(1−α)·I: starts exactly at the
+    fixed merge). h_0 is norm-scaled noise via init_state — combined with
+    randomized-depth training this targets depth-monotone iteration rather
+    than our anchor-dominant fixed point.
+    """
+
+    def __init__(self, d=1536, hidden=512, alpha=ALPHA, sigma=1.0):
+        super().__init__()
+        self.A = nn.Linear(d, d, bias=False)
+        self.B = nn.Linear(d, d, bias=False)
+        with torch.no_grad():
+            self.A.weight.copy_(alpha * torch.eye(d))
+            self.B.weight.copy_((1 - alpha) * torch.eye(d))
+        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.sigma = sigma
+
+    def init_state(self, e):
+        """h_0: per-position Gaussian noise scaled to the anchor's norm."""
+        n = torch.randn_like(e.float())
+        n = n * (e.float().norm(dim=-1, keepdim=True)
+                 / (n.norm(dim=-1, keepdim=True) + 1e-6)) * self.sigma
+        return n.to(e.dtype)
+
+    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)
+        )
+        out = (self.A(s_hat) + self.B(e32)
+               + self.mlp(torch.cat([e32, s_hat], dim=-1)))
+        return out.to(dt)
+
+
 class BandLooper:
     """Capture layer-call kwargs once per forward, then re-run L14-30 manually."""
 
@@ -181,7 +225,10 @@ class BandLooper:
             logits = self.suffix_logits(s, calls, last_only=last_only)
             return (logits, [s]) if return_states else logits
         with torch.no_grad():
-            s = self.band(e, calls)  # s_0: no trainable params upstream
+            # s_0: no trainable params upstream (noise state for recurrent
+            # adapters — the band(e) warm start would hide the B·e path)
+            s = (adapter.init_state(e) if hasattr(adapter, "init_state")
+                 else self.band(e, calls))
         states = [s]
         n_nograd = max(0, k - bptt) if bptt else 0
         for i in range(k):
@@ -270,7 +317,8 @@ class BandLooper:
             calls, _ = self.capture(input_ids, attention_mask,
                                     logits_to_keep=1)
             e = self._hin[self.l0]
-            s = self.band(e, calls)
+            s = (adapter.init_state(e) if hasattr(adapter, "init_state")
+                 else self.band(e, calls))
             x_star = e
             for _ in range(k):
                 x_star = adapter(e, s)