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:
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user