1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
|
"""E1 of the magic-s2000 study: trajectory fingerprint over the redx every-100-step checkpoints.
redx recipe = frozen jr 0.1, NO resreg (predates it) — it rode free, made the golden s2000 (val 3.13),
and blew at step 3300 (CE 2.74 -> 41). Question: does rho(step) show a monotone approach to the edge,
with s2000 sitting in a stable-but-critical sweet window before the ~s3200 crossing? That would make
'edge operator' the mechanism of the magic warm start — and abl_delay the way to manufacture it.
Per ckpt: rho/Re_mu of the forward map at the DEEP state (400-step relax; z_T1=150 readings are
state-contaminated per eig_v2_depth), res at 150 (training protocol) and 400, val CE (nb=4).
"""
import torch
from pathlib import Path
import lt_ep_train as L
from eig_control import lead_rho
T1, DEEP, EPS, B, C = 150, 400, 0.1, 6, 1.0
STEPS = list(range(600, 3700, 200))
def load(path):
torch.manual_seed(0)
blk = L.EQBlock(512, 16, 256, 256, c=C, attn_mode='thick'); blk.qknorm = True
ck = torch.load(path, map_location=L.dev)
with torch.no_grad():
for p, w in zip(blk.allp, ck['allp']):
p.copy_(w.to(L.dev))
return blk
print(f"{'ckpt':>6} {'rho@400':>9} {'Re_mu':>8} {'res@150':>9} {'res@400':>9} {'val':>8}", flush=True)
for s in STEPS:
p = Path(f'runs/redx_traj/s{s}.pt')
if not p.exists():
print(f"s{s:<5} MISSING", flush=True); continue
blk = load(p)
torch.manual_seed(42) # SAME batch for every ckpt
idx, _ = L.get_batch('train', B, 256)
xin = blk.embed(idx).detach()
z150 = L.relax(blk, xin.clone(), xin, T1, EPS)
r150 = (L.relax(blk, z150, xin, 1, EPS) - z150).norm().item()
z400 = L.relax(blk, z150, xin, DEEP - T1, EPS)
r400 = (L.relax(blk, z400, xin, 1, EPS) - z400).norm().item()
_, rho, mu = lead_rho(blk, z400, EPS, C, {}, iters=40)
val = L.evaluate(blk, T1, EPS, nb=4)
print(f"s{s:<5} {rho:>9.5f} {mu:>+8.4f} {r150:>9.2e} {r400:>9.2e} {val:>8.4f}", flush=True)
print("EIG_TRAJ_DONE", flush=True)
|
