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/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