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,
|
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 = (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)")
|
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)
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
Reference in New Issue
Block a user