PerDepthAdapter (Bae-style per-iteration merges) + convergence-halting probe (free ACT); pre-registration item 13

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

scripts/loop_common.py

@@ -180,6 +180,25 @@ class ParcaeAdapter(RecurrentAdapter):
             return self.A_diag().max().item()
 
 
+class PerDepthAdapter(nn.Module):
+    """Depth-wise relaxation (Bae et al. 2024, entrance-level): iteration t
+    gets its OWN merge adapter — breaks time-invariance, so each loop step
+    can perform a different computation phase instead of converging to a
+    fixed point. Depths beyond n_depth reuse the last adapter."""
+
+    def __init__(self, d=1536, hidden=512, alpha=ALPHA, n_depth=4):
+        super().__init__()
+        self.steps = nn.ModuleList(
+            MergeAdapter(d=d, hidden=hidden, alpha=alpha)
+            for _ in range(n_depth))
+
+    def at(self, t):
+        return self.steps[min(t, len(self.steps) - 1)]
+
+    def forward(self, e, s):  # fallback: first-depth adapter
+        return self.steps[0](e, s)
+
+
 class BandLooper:
     """Capture layer-call kwargs once per forward, then re-run L14-30 manually."""
 
@@ -286,14 +305,15 @@ class BandLooper:
         for i in range(k):
             if i < n_nograd:
                 with torch.no_grad():
-                    x = adapter(e, s)
+                    x = (adapter.at(i) if hasattr(adapter, "at")
+                         else 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)
+            x = (adapter.at(i) if hasattr(adapter, "at") else adapter)(e, s)
             if loop_mask is not None:
                 x = torch.where(loop_mask[..., None], x, e)
             if use_checkpoint:
@@ -351,7 +371,8 @@ class BandLooper:
     @torch.no_grad()
     def generate_frozen_prompt(self, adapter, tok, input_ids, k,
                                max_new_tokens=220, attention_mask=None,
-                               stop_strs=(), feedforward=False):
+                               stop_strs=(), feedforward=False,
+                               conv_out=None):
         """Fast equivalent of loop_generate(loop_prompt_only=True).
 
         The looped prompt states are constant across token steps (causality),
@@ -372,9 +393,20 @@ class BandLooper:
             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)
-                s = self.band(x_star, calls)
+            conv = (torch.full((e.shape[0],), -1, dtype=torch.long)
+                    if conv_out is not None else None)
+            for _i in range(k):
+                x_star = (adapter.at(_i) if hasattr(adapter, "at")
+                          else adapter)(e, s)
+                s_new = self.band(x_star, calls)
+                if conv is not None:
+                    c = nn.functional.cosine_similarity(
+                        s_new.float().flatten(1), s.float().flatten(1), dim=1)
+                    conv[((c > 0.9995).cpu()) & (conv < 0)] = _i + 1
+                s = s_new
+            if conv is not None:
+                conv[conv < 0] = k
+                conv_out["k_conv"] = conv.tolist()
             del calls  # prefill re-runs band(x_star) -> same final s as slow path
 
         hook = None