TiedAlphaAdapter: learned per-dim alpha with tied B (anchored by construction); pre-registration item 14

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

scripts/eval_loop_code.py

@@ -91,6 +91,8 @@ def main():
                     help="ParcaeAdapter (rho(A)<1 by construction)")
     ap.add_argument("--perdepth", action="store_true",
                     help="PerDepthAdapter (Bae-style per-iteration merges)")
+    ap.add_argument("--tiedalpha", action="store_true",
+                    help="TiedAlphaAdapter (learned per-dim alpha, tied B)")
     ap.add_argument("--halt", action="store_true",
                     help="record per-item convergence depth (free-ACT probe)")
     args = ap.parse_args()
@@ -99,8 +101,9 @@ def main():
     model, tok = load_model(dtype=torch.bfloat16)
     tok.padding_side = "left"
     looper = BandLooper(model)
-    from loop_common import PerDepthAdapter
-    cls = (PerDepthAdapter if args.perdepth
+    from loop_common import PerDepthAdapter, TiedAlphaAdapter
+    cls = (TiedAlphaAdapter if args.tiedalpha
+           else PerDepthAdapter if args.perdepth
            else ParcaeAdapter if args.parcae
            else RecurrentAdapter if args.rec
            else AdaptiveMergeAdapter if args.adaptive else MergeAdapter)

scripts/loop_common.py

@@ -180,6 +180,40 @@ class ParcaeAdapter(RecurrentAdapter):
             return self.A_diag().max().item()
 
 
+class TiedAlphaAdapter(nn.Module):
+    """Merge with LEARNED per-dim constant alpha, B tied to (1-a):
+        x = (1-a)*e + a*s_hat + MLP([e;s_hat]),  a = sigmoid(a_hat)
+    Convex combination => the LTI fixed point cannot leave the e-s_hat
+    segment (substrate-anchored by construction) and rho = max(a) < 1.
+    Init a = alpha everywhere: bit-identical to MergeAdapter at step 0."""
+
+    def __init__(self, d=1536, hidden=512, alpha=ALPHA):
+        super().__init__()
+        import math as _m
+        self.a_hat = nn.Parameter(
+            torch.full((d,), _m.log(alpha / (1 - alpha))))
+        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)
+
+    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)
+        )
+        a = torch.sigmoid(self.a_hat)
+        out = ((1 - a) * e32 + a * s_hat
+               + self.mlp(torch.cat([e32, s_hat], dim=-1)))
+        return out.to(dt)
+
+    def rho(self):
+        with torch.no_grad():
+            return torch.sigmoid(self.a_hat).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

scripts/train_merge_code.py

@@ -21,7 +21,8 @@ import torch
 import torch.nn.functional as F
 
 from loop_common import (AdaptiveMergeAdapter, BandLooper, MergeAdapter,
-                         ParcaeAdapter, PerDepthAdapter, RecurrentAdapter)
+                         ParcaeAdapter, PerDepthAdapter, RecurrentAdapter,
+                         TiedAlphaAdapter)
 from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt
 
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
@@ -61,6 +62,9 @@ ap.add_argument("--recmax", type=int, default=16)
 ap.add_argument("--parcae", action="store_true",
                 help="rec regime with rho(A)<1 by construction "
                      "(diag-negative-exp ZOH parameterization)")
+ap.add_argument("--tiedalpha", action="store_true",
+                help="merge with learned per-dim constant alpha, B tied "
+                     "to (1-a); standard curriculum")
 ap.add_argument("--perdepth", type=int, default=0,
                 help="PerDepthAdapter: one merge per iteration depth "
                      "(Bae-style relaxation), standard curriculum")
@@ -80,7 +84,8 @@ SUFFIX = ((f"_s{SEED}" if SEED else "")
           + ("_warm" if ARGS.warm else "")
           + (("_parcae" if ARGS.parcae else "_rec") + str(ARGS.recmax)
              if ARGS.rec else "")
-          + (f"_pd{ARGS.perdepth}" if ARGS.perdepth else ""))
+          + (f"_pd{ARGS.perdepth}" if ARGS.perdepth else "")
+          + ("_ta" if ARGS.tiedalpha else ""))
 PAUSE_ID = 6  # <unused0>
 
 
@@ -139,7 +144,9 @@ def main():
         p.requires_grad_(False)
     looper = BandLooper(model)
     d = model.config.get_text_config().hidden_size
-    if ARGS.perdepth:
+    if ARGS.tiedalpha:
+        adapter = TiedAlphaAdapter(d=d, alpha=ARGS.alpha).cuda()
+    elif ARGS.perdepth:
         adapter = PerDepthAdapter(d=d, alpha=ARGS.alpha,
                                   n_depth=ARGS.perdepth).cuda()
     elif ARGS.parcae: