ParcaeAdapter: rho(A)<1 by construction (ZOH negative-diag), rho logging in rec arms, pre-registration item 12
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -117,3 +117,19 @@ number for the unified adapter exists at time of writing.
|
|||||||
Known risk, stated in advance: 600 steps may be far too little for
|
Known risk, stated in advance: 600 steps may be far too little for
|
||||||
this regime (McLeish et al. use ~50B tokens); a null here bounds the
|
this regime (McLeish et al. use ~50B tokens); a null here bounds the
|
||||||
cheap-retrofit budget only, not the regime.
|
cheap-retrofit budget only, not the regime.
|
||||||
|
|
||||||
|
12. **Parcae-constrained recurrent arm (pre-registered 2026-07-15, before
|
||||||
|
training; Prairie et al. 2026 parameterization).** Same as item 11 but
|
||||||
|
A = exp(−Δt·exp(a)) diagonal → ρ(A) < 1 by construction; init exactly
|
||||||
|
the α=0.3 merge (verified bit-equal at init). ρ(A) logged every 10
|
||||||
|
steps in BOTH arms. Theory-derived predictions, stated in advance:
|
||||||
|
(a) contraction ⇒ fixed point is a function of e ⇒ the Parcae arm
|
||||||
|
SATURATES in k (no depth-monotone gain) and its converged performance
|
||||||
|
is amortizable — if so, our deflationary result is a corollary of
|
||||||
|
ρ<1, and our observed k≈3–4 convergence is the geometric rate 0.3^k;
|
||||||
|
(b) the UNCONSTRAINED item-11 arm either drifts toward ρ≥1 (watch the
|
||||||
|
ρ log: divergent runs should show ρ≥1 before loss spikes) or, if it
|
||||||
|
gains monotone depth-performance, does so with ρ near 1 — the edge of
|
||||||
|
stability is where genuine iteration must live. Either outcome
|
||||||
|
formalizes "the anchor coefficient is the stability dial" as
|
||||||
|
"the anchor coefficient is the spectral radius".
|
||||||
|
|||||||
@@ -16,7 +16,7 @@ from pathlib import Path
|
|||||||
import torch
|
import torch
|
||||||
|
|
||||||
from loop_common import (AdaptiveMergeAdapter, BandLooper, MergeAdapter,
|
from loop_common import (AdaptiveMergeAdapter, BandLooper, MergeAdapter,
|
||||||
RecurrentAdapter)
|
ParcaeAdapter, RecurrentAdapter)
|
||||||
from prep_mbpp import DIRECT_SUFFIX, extract_code, mbpp_prompt, run_tests
|
from prep_mbpp import DIRECT_SUFFIX, extract_code, mbpp_prompt, run_tests
|
||||||
|
|
||||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||||
@@ -79,13 +79,16 @@ def main():
|
|||||||
ap.add_argument("--adaptive", action="store_true")
|
ap.add_argument("--adaptive", action="store_true")
|
||||||
ap.add_argument("--rec", action="store_true",
|
ap.add_argument("--rec", action="store_true",
|
||||||
help="RecurrentAdapter (noise h0, learned A/B)")
|
help="RecurrentAdapter (noise h0, learned A/B)")
|
||||||
|
ap.add_argument("--parcae", action="store_true",
|
||||||
|
help="ParcaeAdapter (rho(A)<1 by construction)")
|
||||||
args = ap.parse_args()
|
args = ap.parse_args()
|
||||||
ks = [int(x) for x in args.ks.split(",")]
|
ks = [int(x) for x in args.ks.split(",")]
|
||||||
|
|
||||||
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)
|
||||||
cls = (RecurrentAdapter if args.rec
|
cls = (ParcaeAdapter if args.parcae
|
||||||
|
else RecurrentAdapter if args.rec
|
||||||
else AdaptiveMergeAdapter if args.adaptive else MergeAdapter)
|
else AdaptiveMergeAdapter if args.adaptive else MergeAdapter)
|
||||||
kw = ({"alpha0": args.alpha} if args.adaptive else {"alpha": args.alpha})
|
kw = ({"alpha0": args.alpha} if args.adaptive else {"alpha": args.alpha})
|
||||||
adapter = cls(d=model.config.get_text_config().hidden_size, **kw).cuda()
|
adapter = cls(d=model.config.get_text_config().hidden_size, **kw).cuda()
|
||||||
|
|||||||
@@ -127,6 +127,58 @@ class RecurrentAdapter(nn.Module):
|
|||||||
+ self.mlp(torch.cat([e32, s_hat], dim=-1)))
|
+ self.mlp(torch.cat([e32, s_hat], dim=-1)))
|
||||||
return out.to(dt)
|
return out.to(dt)
|
||||||
|
|
||||||
|
def rho(self):
|
||||||
|
"""Spectral radius of A (power iteration) — the stability dial."""
|
||||||
|
with torch.no_grad():
|
||||||
|
v = torch.randn(self.A.weight.shape[0],
|
||||||
|
device=self.A.weight.device)
|
||||||
|
for _ in range(30):
|
||||||
|
v = self.A.weight.T @ (self.A.weight @ v)
|
||||||
|
v = v / (v.norm() + 1e-12)
|
||||||
|
return (self.A.weight @ v).norm().item()
|
||||||
|
|
||||||
|
|
||||||
|
class ParcaeAdapter(RecurrentAdapter):
|
||||||
|
"""Recurrent adapter with ρ(A) < 1 GUARANTEED by construction
|
||||||
|
(Parcae parameterization, Prairie et al. 2026): A is the ZOH
|
||||||
|
discretization of a continuous negative-diagonal system,
|
||||||
|
A = exp(−Δt·exp(a)) with learned per-dim a and scalar log-Δt, so every
|
||||||
|
entry lies in (0,1) regardless of optimizer noise. Init matches the
|
||||||
|
fixed merge exactly: a = log(−log α), Δt = 1 → A = α·I."""
|
||||||
|
|
||||||
|
def __init__(self, d=1536, hidden=512, alpha=ALPHA, sigma=1.0):
|
||||||
|
nn.Module.__init__(self)
|
||||||
|
import math as _m
|
||||||
|
self.log_a = nn.Parameter(
|
||||||
|
torch.full((d,), _m.log(-_m.log(alpha))))
|
||||||
|
self.log_dt = nn.Parameter(torch.zeros(()))
|
||||||
|
self.B = nn.Linear(d, d, bias=False)
|
||||||
|
with torch.no_grad():
|
||||||
|
self.B.weight.copy_((1 - alpha) * torch.eye(d))
|
||||||
|
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)
|
||||||
|
self.sigma = sigma
|
||||||
|
|
||||||
|
def A_diag(self):
|
||||||
|
return torch.exp(-torch.exp(self.log_dt) * torch.exp(self.log_a))
|
||||||
|
|
||||||
|
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)
|
||||||
|
)
|
||||||
|
out = (self.A_diag().to(e32.dtype) * s_hat + self.B(e32)
|
||||||
|
+ self.mlp(torch.cat([e32, s_hat], dim=-1)))
|
||||||
|
return out.to(dt)
|
||||||
|
|
||||||
|
def rho(self):
|
||||||
|
with torch.no_grad():
|
||||||
|
return self.A_diag().max().item()
|
||||||
|
|
||||||
|
|
||||||
class BandLooper:
|
class BandLooper:
|
||||||
"""Capture layer-call kwargs once per forward, then re-run L14-30 manually."""
|
"""Capture layer-call kwargs once per forward, then re-run L14-30 manually."""
|
||||||
|
|||||||
@@ -21,7 +21,7 @@ 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,
|
||||||
RecurrentAdapter)
|
ParcaeAdapter, RecurrentAdapter)
|
||||||
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))
|
||||||
@@ -58,7 +58,12 @@ ap.add_argument("--rec", action="store_true",
|
|||||||
"noise h0) + log-uniform random depth 1..recmax, "
|
"noise h0) + log-uniform random depth 1..recmax, "
|
||||||
"truncated bptt (default 4)")
|
"truncated bptt (default 4)")
|
||||||
ap.add_argument("--recmax", type=int, default=16)
|
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)")
|
||||||
ARGS, _ = ap.parse_known_args()
|
ARGS, _ = ap.parse_known_args()
|
||||||
|
if ARGS.parcae:
|
||||||
|
ARGS.rec = True
|
||||||
if ARGS.rec and not ARGS.bptt:
|
if ARGS.rec and not ARGS.bptt:
|
||||||
ARGS.bptt = 4
|
ARGS.bptt = 4
|
||||||
SEED = ARGS.seed
|
SEED = ARGS.seed
|
||||||
@@ -70,7 +75,8 @@ SUFFIX = ((f"_s{SEED}" if SEED else "")
|
|||||||
+ (f"_dk{ARGS.deepk}" if ARGS.deepk else "")
|
+ (f"_dk{ARGS.deepk}" if ARGS.deepk else "")
|
||||||
+ (f"_lr{ARGS.lr}" if ARGS.lr != 1e-3 else "")
|
+ (f"_lr{ARGS.lr}" if ARGS.lr != 1e-3 else "")
|
||||||
+ ("_warm" if ARGS.warm else "")
|
+ ("_warm" if ARGS.warm else "")
|
||||||
+ (f"_rec{ARGS.recmax}" if ARGS.rec else ""))
|
+ (("_parcae" if ARGS.parcae else "_rec") + str(ARGS.recmax)
|
||||||
|
if ARGS.rec else ""))
|
||||||
PAUSE_ID = 6 # <unused0>
|
PAUSE_ID = 6 # <unused0>
|
||||||
|
|
||||||
|
|
||||||
@@ -129,7 +135,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.rec:
|
if ARGS.parcae:
|
||||||
|
adapter = ParcaeAdapter(d=d, alpha=ARGS.alpha).cuda()
|
||||||
|
elif ARGS.rec:
|
||||||
adapter = RecurrentAdapter(d=d, alpha=ARGS.alpha).cuda()
|
adapter = RecurrentAdapter(d=d, alpha=ARGS.alpha).cuda()
|
||||||
elif ARGS.adaptive:
|
elif ARGS.adaptive:
|
||||||
adapter = AdaptiveMergeAdapter(d=d, alpha0=ARGS.alpha).cuda()
|
adapter = AdaptiveMergeAdapter(d=d, alpha0=ARGS.alpha).cuda()
|
||||||
@@ -187,8 +195,12 @@ def main():
|
|||||||
|
|
||||||
log.append({"step": step, "k": k, "loss": loss.item()})
|
log.append({"step": step, "k": k, "loss": loss.item()})
|
||||||
if step % 10 == 0:
|
if step % 10 == 0:
|
||||||
print(f"step {step:4d} k={k} loss={loss.item():.4f} "
|
rho = (f" rho(A)={adapter.rho():.3f}"
|
||||||
|
if hasattr(adapter, "rho") else "")
|
||||||
|
print(f"step {step:4d} k={k} loss={loss.item():.4f}{rho} "
|
||||||
f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
|
f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
|
||||||
|
if rho:
|
||||||
|
log.append({"step": step, "rho": adapter.rho()})
|
||||||
if step % 100 == 99 or step == STEPS - 1:
|
if step % 100 == 99 or step == STEPS - 1:
|
||||||
vals = {}
|
vals = {}
|
||||||
for kk in ((0, 1, 2, 4, 8, 16) if ARGS.rec else (0, 1, 2, 4)):
|
for kk in ((0, 1, 2, 4, 8, 16) if ARGS.rec else (0, 1, 2, 4)):
|
||||||
|
|||||||
Reference in New Issue
Block a user