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

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Nils
2026-07-15 02:33:14 +02:00
co-authored by Claude Fable 5
parent f370883e23
commit 16f4a057aa
4 changed files with 66 additions and 5 deletions
+17
View File
@@ -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, crawl that at k=8 reaches what the contractive merge reaches at k=4,
never approaching the amortization ceiling from above). 4x parameters never approaching the amortization ceiling from above). 4x parameters
bought nothing. Depth-monotone computation did not emerge at this budget. bought nothing. Depth-monotone computation did not emerge at this budget.
14. **Tied-alpha arm (pre-registered 2026-07-15, before training).**
TiedAlphaAdapter: x = (1a)⊙e + a⊙ŝ + MLP([e;ŝ]), a = σ(â) per-dim
learned, init a=0.3 everywhere (bit-equal to MergeAdapter at step 0,
verified). B tied to (1a): 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 (1a) differs from parcae.
+5 -2
View File
@@ -91,6 +91,8 @@ def main():
help="ParcaeAdapter (rho(A)<1 by construction)") help="ParcaeAdapter (rho(A)<1 by construction)")
ap.add_argument("--perdepth", action="store_true", ap.add_argument("--perdepth", action="store_true",
help="PerDepthAdapter (Bae-style per-iteration merges)") 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", ap.add_argument("--halt", action="store_true",
help="record per-item convergence depth (free-ACT probe)") help="record per-item convergence depth (free-ACT probe)")
args = ap.parse_args() args = ap.parse_args()
@@ -99,8 +101,9 @@ def main():
model, tok = load_model(dtype=torch.bfloat16) model, tok = load_model(dtype=torch.bfloat16)
tok.padding_side = "left" tok.padding_side = "left"
looper = BandLooper(model) looper = BandLooper(model)
from loop_common import PerDepthAdapter from loop_common import PerDepthAdapter, TiedAlphaAdapter
cls = (PerDepthAdapter if args.perdepth cls = (TiedAlphaAdapter if args.tiedalpha
else PerDepthAdapter if args.perdepth
else ParcaeAdapter if args.parcae else ParcaeAdapter if args.parcae
else RecurrentAdapter if args.rec else RecurrentAdapter if args.rec
else AdaptiveMergeAdapter if args.adaptive else MergeAdapter) else AdaptiveMergeAdapter if args.adaptive else MergeAdapter)
+34
View File
@@ -180,6 +180,40 @@ class ParcaeAdapter(RecurrentAdapter):
return self.A_diag().max().item() 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): class PerDepthAdapter(nn.Module):
"""Depth-wise relaxation (Bae et al. 2024, entrance-level): iteration t """Depth-wise relaxation (Bae et al. 2024, entrance-level): iteration t
gets its OWN merge adapter — breaks time-invariance, so each loop step gets its OWN merge adapter — breaks time-invariance, so each loop step
+10 -3
View File
@@ -21,7 +21,8 @@ import torch
import torch.nn.functional as F import torch.nn.functional as F
from loop_common import (AdaptiveMergeAdapter, BandLooper, MergeAdapter, from loop_common import (AdaptiveMergeAdapter, BandLooper, MergeAdapter,
ParcaeAdapter, PerDepthAdapter, RecurrentAdapter) ParcaeAdapter, PerDepthAdapter, RecurrentAdapter,
TiedAlphaAdapter)
from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt
sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) 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", ap.add_argument("--parcae", action="store_true",
help="rec regime with rho(A)<1 by construction " help="rec regime with rho(A)<1 by construction "
"(diag-negative-exp ZOH parameterization)") "(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, ap.add_argument("--perdepth", type=int, default=0,
help="PerDepthAdapter: one merge per iteration depth " help="PerDepthAdapter: one merge per iteration depth "
"(Bae-style relaxation), standard curriculum") "(Bae-style relaxation), standard curriculum")
@@ -80,7 +84,8 @@ SUFFIX = ((f"_s{SEED}" if SEED else "")
+ ("_warm" if ARGS.warm else "") + ("_warm" if ARGS.warm else "")
+ (("_parcae" if ARGS.parcae else "_rec") + str(ARGS.recmax) + (("_parcae" if ARGS.parcae else "_rec") + str(ARGS.recmax)
if ARGS.rec else "") 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> PAUSE_ID = 6 # <unused0>
@@ -139,7 +144,9 @@ def main():
p.requires_grad_(False) p.requires_grad_(False)
looper = BandLooper(model) looper = BandLooper(model)
d = model.config.get_text_config().hidden_size 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, adapter = PerDepthAdapter(d=d, alpha=ARGS.alpha,
n_depth=ARGS.perdepth).cuda() n_depth=ARGS.perdepth).cuda()
elif ARGS.parcae: elif ARGS.parcae: