From 16f4a057aab1b4aa79052e2b11b5da39cc0a4299 Mon Sep 17 00:00:00 2001 From: Nils Date: Wed, 15 Jul 2026 02:33:14 +0200 Subject: [PATCH] TiedAlphaAdapter: learned per-dim alpha with tied B (anchored by construction); pre-registration item 14 Co-Authored-By: Claude Fable 5 --- results-loop/PROTOCOL_UNIFIED.md | 17 ++++++++++++++++ scripts/eval_loop_code.py | 7 +++++-- scripts/loop_common.py | 34 ++++++++++++++++++++++++++++++++ scripts/train_merge_code.py | 13 +++++++++--- 4 files changed, 66 insertions(+), 5 deletions(-) diff --git a/results-loop/PROTOCOL_UNIFIED.md b/results-loop/PROTOCOL_UNIFIED.md index 3a500bb..81d127f 100644 --- a/results-loop/PROTOCOL_UNIFIED.md +++ b/results-loop/PROTOCOL_UNIFIED.md @@ -163,3 +163,20 @@ merge level (35.7/39.3/42.9% at k=2/4/8 — a one-item-per-depth-doubling crawl that at k=8 reaches what the contractive merge reaches at k=4, never approaching the amortization ceiling from above). 4x parameters bought nothing. Depth-monotone computation did not emerge at this budget. + +14. **Tied-alpha arm (pre-registered 2026-07-15, before training).** + TiedAlphaAdapter: x = (1−a)⊙e + a⊙ŝ + MLP([e;ŝ]), a = σ(â) per-dim + learned, init a=0.3 everywhere (bit-equal to MergeAdapter at step 0, + verified). B tied to (1−a): convex combination keeps the LTI fixed + point on the e–ŝ segment (substrate-anchored by construction), + ρ = max(a) < 1 guaranteed, +d≈1.5K params. Standard curriculum, + s0 = band(e), e400, eval ks 0,2,4,8 on 250 items. This is the one + untested cell combining parcae's learnable decay with the merge's + anchoring. Predictions: (a) substrate fidelity preserved (easy ≈ + merge's 88%, unlike both rec arms' ~70%) because anchoring, not + ρ, controls fidelity; (b) hard-bucket at merge level (no significant + gain — per-dim constant α is not where capability lives, per the + adaptive-α E2B result); (c) learned a drifts slightly DOWN from 0.3 + (as in parcae). If (a) holds while rec arms failed it, the + fixed-point-location dial is causally isolated: same learnable-decay + freedom, only the tie to (1−a) differs from parcae. diff --git a/scripts/eval_loop_code.py b/scripts/eval_loop_code.py index 78b3197..4c1d9c2 100644 --- a/scripts/eval_loop_code.py +++ b/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) diff --git a/scripts/loop_common.py b/scripts/loop_common.py index 523c171..f0d640b 100644 --- a/scripts/loop_common.py +++ b/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 diff --git a/scripts/train_merge_code.py b/scripts/train_merge_code.py index b642e0a..4ec1bcb 100644 --- a/scripts/train_merge_code.py +++ b/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 # @@ -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: