nils/jspace
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

PerDepthAdapter (Bae-style per-iteration merges) + convergence-halting probe (free ACT); pre-registration item 13

Browse Source 
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Nils
2026-07-15 01:07:01 +02:00
co-authored by Claude Fable 5
parent f254b20b62
commit d00f120f27
4 changed files with 88 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
results-loop/PROTOCOL_UNIFIED.md
+16
View File
@@ -133,3 +133,19 @@ number for the unified adapter exists at time of writing.
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".
13. **Per-depth adapter arm + free-ACT probe (pre-registered 2026-07-15,
before training).** (a) PerDepthAdapter: one merge adapter per
iteration (n=4, Bae-style depth-wise relaxation at the entrance;
breaks time-invariance — LTV, no fixed-point guarantee), standard
curriculum, e400, eval ks 0,2,4,8. Prediction: lands at or below the
distill/rung-2 amortization ceiling (~46% hard) because depth-indexed
weights add content, not state-evolution; exceeding it would show
per-iteration expressivity was binding and amend the deflationary
claim. Depths >4 reuse adapter 4 (stated: k=8 cell is then
fixed-point-like by construction). (b) Free-ACT probe on the standard
merge arm: record per-item convergence depth (cos>0.9995) at k=8 cap.
Predictions: accuracy unchanged vs fixed k (post-convergence no-ops);
mean k_conv ≈ 3; hard-labeled items converge SLOWER than easy ones
(adaptive compute allocates like ACT without any learned halting
parameter).
scripts/eval_loop_code.py
+24 -4
View File
@@ -28,8 +28,9 @@ OUT = Path(os.environ.get("LOOP_OUT",
@torch.no_grad()
def pass1_at_k(looper, adapter, tok, items, k, batch=8, max_new=220,
feedforward=False, pause=0):
feedforward=False, pause=0, halt=False):
codes = []
k_convs = []
for i in range(0, len(items), batch):
torch.cuda.empty_cache()
chunk = items[i : i + batch]
@@ -42,10 +43,14 @@ def pass1_at_k(looper, adapter, tok, items, k, batch=8, max_new=220,
enc["input_ids"] = torch.cat([enc["input_ids"], pcol], 1)
enc["attention_mask"] = torch.cat(
[enc["attention_mask"], torch.ones_like(pcol)], 1)
cv = {} if halt else None
gen = looper.generate_frozen_prompt(adapter, tok, enc["input_ids"], k,
max_new_tokens=max_new,
attention_mask=enc["attention_mask"],
feedforward=feedforward)
feedforward=feedforward,
conv_out=cv)
if halt:
k_convs.extend(cv.get("k_conv", []))
for j in range(len(chunk)):
txt = tok.decode(gen[j, enc["input_ids"].shape[1]:],
skip_special_tokens=True)
@@ -61,6 +66,9 @@ def pass1_at_k(looper, adapter, tok, items, k, batch=8, max_new=220,
d[1] += 1
per_item = [{"task_id": it["task_id"], "ok": bool(ok)}
for it, ok in zip(items, oks)]
if halt and k_convs:
for r, kc in zip(per_item, k_convs):
r["k_conv"] = kc
return (hits / len(items),
{l: c / n for l, (c, n) in per_label.items()}, per_item)
@@ -81,13 +89,19 @@ def main():
help="RecurrentAdapter (noise h0, learned A/B)")
ap.add_argument("--parcae", action="store_true",
help="ParcaeAdapter (rho(A)<1 by construction)")
ap.add_argument("--perdepth", action="store_true",
help="PerDepthAdapter (Bae-style per-iteration merges)")
ap.add_argument("--halt", action="store_true",
help="record per-item convergence depth (free-ACT probe)")
args = ap.parse_args()
ks = [int(x) for x in args.ks.split(",")]
model, tok = load_model(dtype=torch.bfloat16)
tok.padding_side = "left"
looper = BandLooper(model)
cls = (ParcaeAdapter if args.parcae
from loop_common import PerDepthAdapter
cls = (PerDepthAdapter if args.perdepth
else ParcaeAdapter if args.parcae
else RecurrentAdapter if args.rec
else AdaptiveMergeAdapter if args.adaptive else MergeAdapter)
kw = ({"alpha0": args.alpha} if args.adaptive else {"alpha": args.alpha})
@@ -106,7 +120,13 @@ def main():
t0 = time.time()
acc, by_label, per_item = pass1_at_k(looper, adapter, tok, items, k,
feedforward=args.feedforward,
pause=args.pause)
pause=args.pause, halt=args.halt)
if args.halt:
by_lbl_k = {}
for it, r in zip(items, per_item):
by_lbl_k.setdefault(it["label"], []).append(r.get("k_conv", k))
print(" mean k_conv:", {l: round(sum(v)/len(v), 2)
for l, v in by_lbl_k.items()}, flush=True)
res["ks"][k] = {"acc": acc, "by_label": by_label,
"per_item": per_item}
print(f"k={k}: pass@1={acc:.3f} "
scripts/loop_common.py
+38 -6
View File
@@ -180,6 +180,25 @@ class ParcaeAdapter(RecurrentAdapter):
return self.A_diag().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
can perform a different computation phase instead of converging to a
fixed point. Depths beyond n_depth reuse the last adapter."""
def __init__(self, d=1536, hidden=512, alpha=ALPHA, n_depth=4):
super().__init__()
self.steps = nn.ModuleList(
MergeAdapter(d=d, hidden=hidden, alpha=alpha)
for _ in range(n_depth))
def at(self, t):
return self.steps[min(t, len(self.steps) - 1)]
def forward(self, e, s): # fallback: first-depth adapter
return self.steps[0](e, s)
class BandLooper:
"""Capture layer-call kwargs once per forward, then re-run L14-30 manually."""
@@ -286,14 +305,15 @@ class BandLooper:
for i in range(k):
if i < n_nograd:
with torch.no_grad():
x = adapter(e, s)
x = (adapter.at(i) if hasattr(adapter, "at")
else adapter)(e, s)
if loop_mask is not None:
x = torch.where(loop_mask[..., None], x, e)
s = self.band(x, calls)
s = s.detach()
states.append(s)
continue
x = adapter(e, s)
x = (adapter.at(i) if hasattr(adapter, "at") else adapter)(e, s)
if loop_mask is not None:
x = torch.where(loop_mask[..., None], x, e)
if use_checkpoint:
@@ -351,7 +371,8 @@ class BandLooper:
@torch.no_grad()
def generate_frozen_prompt(self, adapter, tok, input_ids, k,
max_new_tokens=220, attention_mask=None,
stop_strs=(), feedforward=False):
stop_strs=(), feedforward=False,
conv_out=None):
"""Fast equivalent of loop_generate(loop_prompt_only=True).
The looped prompt states are constant across token steps (causality),
@@ -372,9 +393,20 @@ class BandLooper:
s = (adapter.init_state(e) if hasattr(adapter, "init_state")
else self.band(e, calls))
x_star = e
for _ in range(k):
x_star = adapter(e, s)
s = self.band(x_star, calls)
conv = (torch.full((e.shape[0],), -1, dtype=torch.long)
if conv_out is not None else None)
for _i in range(k):
x_star = (adapter.at(_i) if hasattr(adapter, "at")
else adapter)(e, s)
s_new = self.band(x_star, calls)
if conv is not None:
c = nn.functional.cosine_similarity(
s_new.float().flatten(1), s.float().flatten(1), dim=1)
conv[((c > 0.9995).cpu()) & (conv < 0)] = _i + 1
s = s_new
if conv is not None:
conv[conv < 0] = k
conv_out["k_conv"] = conv.tolist()
del calls # prefill re-runs band(x_star) -> same final s as slow path
hook = None
scripts/train_merge_code.py
+10 -3
View File
@@ -21,7 +21,7 @@ import torch
import torch.nn.functional as F
from loop_common import (AdaptiveMergeAdapter, BandLooper, MergeAdapter,
ParcaeAdapter, RecurrentAdapter)
ParcaeAdapter, PerDepthAdapter, RecurrentAdapter)
from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
@@ -61,6 +61,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("--perdepth", type=int, default=0,
help="PerDepthAdapter: one merge per iteration depth "
"(Bae-style relaxation), standard curriculum")
ARGS, _ = ap.parse_known_args()
if ARGS.parcae:
ARGS.rec = True
@@ -76,7 +79,8 @@ SUFFIX = ((f"_s{SEED}" if SEED else "")
+ (f"_lr{ARGS.lr}" if ARGS.lr != 1e-3 else "")
+ ("_warm" if ARGS.warm else "")
+ (("_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 ""))
PAUSE_ID = 6 # <unused0>
@@ -135,7 +139,10 @@ def main():
p.requires_grad_(False)
looper = BandLooper(model)
d = model.config.get_text_config().hidden_size
if ARGS.parcae:
if ARGS.perdepth:
adapter = PerDepthAdapter(d=d, alpha=ARGS.alpha,
n_depth=ARGS.perdepth).cuda()
elif ARGS.parcae:
adapter = ParcaeAdapter(d=d, alpha=ARGS.alpha).cuda()
elif ARGS.rec:
adapter = RecurrentAdapter(d=d, alpha=ARGS.alpha).cuda()