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,
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 = (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)")
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)
+34
View File
@@ -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
+10 -3
View File
@@ -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: