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
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()