ZeroPoint v0 research prototype: new from-scratch NumPy framework, memory-slot architecture, audited Transformer comparison
5c6c847 verified | """预注册实验：ZeroPoint 长序列变差的原因是不是槽值无界累积？ema 写入能不能修好？ | |
| python3 experiments/ema_longrun.py | |
| 实验设计在跑之前写定（见 PREREG），结果出来后不改判定标准。结果：experiments/ema_longrun.md + results/ 存档。 | |
| """ | |
| import numpy as np | |
| from common import CFG, ITERS, T, batches, evaluate, prepare_data, save_result, train | |
| from models import ZeroX | |
| import zerocg as zc | |
| SEEDS = [0, 1, 2] | |
| LR = 6e-3 # compare.md 里 ZeroPoint 的最佳学习率 | |
| LONG_T = 4096 | |
| LONG_N = 32 # 外推测试序列条数（compare.py 只有 8 条，这里加到 32） | |
| SEGMENTS = [(0, 64), (64, 256), (256, 1024), (1024, 4096)] | |
| PROBE_POS = [64, 256, 1024, 4000] # 测槽值范围的位置 | |
| PREREG = """# 预注册实验：长序列变差与 ema 写入 | |
| > 写于实验之前（2026-10-18）。跑完后只在下面追加结果，不改这一节。 | |
| **观察**：compare.md 里 ZeroPoint（无 ema）在 1024~4095 字处的 loss（1.70）高于 64~255 字处（1.62）。 | |
| **假设 H1**：原因是长期槽（λ 接近 1）的数值随序列长度无界累积，超出训练时（64 字）见过的范围。 | |
| **预测**： | |
| - ① 直接测：推理时槽值的最大绝对值，无 ema 的模型在 4000 字处明显大于 64 字处（比值 > 2）；有 ema 的模型比值接近 1（< 1.5）。 | |
| - ② 间接测：开 ema 重新训练（其余完全相同），1024~4095 字处的 loss 不再明显高于 64~255 字处：差值 < 2 × 种子间标准差。 | |
| **H1 不成立时会看到**：① 无 ema 的槽值没有明显增长；或 ② 开 ema 后远处 loss 照样升高。 | |
| **判定**：① 和 ② 都成立 → 接受 H1，且 ema 是有效的修法；只有 ② 成立 → ema 有效但原因另有其他； | |
| 都不成立 → 放弃 H1，下一步查学到的 λ 值。另外要求 ema 不损害短距离：64 字处验证 loss 的差别 < 2 × 标准差。 | |
| **对照**：两组用同样的 3 个种子、同样 {iters} 步、学习率 {lr}、同样语料；外推用 {n} 条 {lt} 字的验证序列， | |
| 两组都用分块并行算法算 loss（compare.py 里 ZeroPoint 用的是递推算法，和 Transformer 不对等，这里统一）。 | |
| """ | |
| def long_loss(model, val_data): | |
| """(各段 loss 均值) ；分 4 批算，省内存。""" | |
| nll_all = [] | |
| for x, y in batches(val_data, 4, LONG_T, 7, batch=LONG_N // 4): | |
| with zc.no_grad(): | |
| logits = model(x).data | |
| logp = logits - logits.max(-1, keepdims=True) | |
| logp = logp - np.log(np.exp(logp).sum(-1, keepdims=True)) | |
| nll_all.append(-np.take_along_axis(logp, y[..., None], -1)[..., 0]) | |
| nll = np.concatenate(nll_all) | |
| return np.array([nll[:, a:b].mean() for a, b in SEGMENTS]) | |
| def slot_magnitude(model, val_data): | |
| """递推跑 4096 字，记录各位置处所有层槽值的最大绝对值。返回 {位置: 最大|S|}。""" | |
| x, _ = next(batches(val_data, 1, LONG_T, 7, batch=4)) | |
| states = model.init_state(4) | |
| out = {} | |
| for t in range(LONG_T): | |
| _, states = model.step(x[:, t], states, t) | |
| if t + 1 in PROBE_POS: | |
| out[t + 1] = max(float(np.abs(S).max()) for S in states) | |
| return out | |
| def main(): | |
| text, V, train_data, val_data = prepare_data() | |
| print(PREREG.format(iters=ITERS, lr=LR, n=LONG_N, lt=LONG_T), flush=True) | |
| res = {} | |
| for ema in (False, True): | |
| name = "ema" if ema else "no-ema" | |
| vals, longs, mags = [], [], [] | |
| for seed in SEEDS: | |
| np.random.seed(seed) | |
| model = ZeroX(V, n_slot=16, ema=ema, **CFG) | |
| vals.append(train(name, model, train_data, val_data, seed, LR)) | |
| model.eval() | |
| longs.append(long_loss(model, val_data)) | |
| mags.append(slot_magnitude(model, val_data)) | |
| print(f" => {name} seed={seed}: val {vals[-1]:.4f}; long {np.round(longs[-1], 4)}; " | |
| f"|S| {[round(mags[-1][p], 2) for p in PROBE_POS]}", flush=True) | |
| res[name] = dict(val=np.array(vals), long=np.array(longs), | |
| mag=np.array([[m[p] for p in PROBE_POS] for m in mags])) | |
| # ---- 结果 + 按预注册的标准判定 ---- | |
| lines = [PREREG.format(iters=ITERS, lr=LR, n=LONG_N, lt=LONG_T), "\n## 结果（实验之后写）\n", | |
| "### 64 字处的验证 loss（均值 ± 标准差，3 个种子）\n", "| 模型 | 验证 loss |", "|---|---|"] | |
| for name, r in res.items(): | |
| lines.append(f"| {name} | {r['val'].mean():.4f} ± {r['val'].std():.4f} |") | |
| lines += ["", f"### 长度外推（{LONG_N} 条 × {LONG_T} 字，分块并行算法；均值 ± 标准差）\n", | |
| "| 位置 | " + " | ".join(res) + " |", "|---|" + "---|" * len(res)] | |
| for k, (a, b) in enumerate(SEGMENTS): | |
| lines.append(f"| {a} ~ {b - 1} | " + " | ".join(f"{r['long'][:, k].mean():.4f} ± {r['long'][:, k].std():.4f}" | |
| for r in res.values()) + " |") | |
| lines += ["", "### 槽值的最大绝对值（递推推理，3 个种子平均）\n", | |
| "| 位置 | " + " | ".join(res) + " |", "|---|" + "---|" * len(res)] | |
| for k, p in enumerate(PROBE_POS): | |
| lines.append(f"| {p} | " + " | ".join(f"{r['mag'][:, k].mean():.2f}" for r in res.values()) + " |") | |
| no, em = res["no-ema"], res["ema"] | |
| ratio_no = (no["mag"][:, 3] / no["mag"][:, 0]).mean() | |
| ratio_em = (em["mag"][:, 3] / em["mag"][:, 0]).mean() | |
| p1 = ratio_no > 2 and ratio_em < 1.5 | |
| far, near = em["long"][:, 3], em["long"][:, 1] | |
| gap, noise = (far - near).mean(), 2 * max(far.std(), near.std()) | |
| p2 = gap < noise | |
| gap_no = (no["long"][:, 3] - no["long"][:, 1]).mean() | |
| short_ok = abs(em["val"].mean() - no["val"].mean()) < 2 * max(em["val"].std(), no["val"].std()) | |
| if p1 and p2: | |
| verdict = "① ② 都成立：**接受 H1**，ema 是有效的修法" | |
| elif p2: | |
| verdict = "只有 ② 成立：ema 有效，但长序列变差的原因不是（或不只是）槽值无界" | |
| elif p1: | |
| verdict = "只有 ① 成立：槽值确实无界，但 ema 没有修好远处的 loss，H1 不足以解释变差" | |
| else: | |
| verdict = "① ② 都不成立：**放弃 H1**，下一步查学到的 λ 值" | |
| lines += ["", "## 判定（按预注册的标准）\n", | |
| f"- 预测 ①（槽值比值 4000 字 / 64 字）：无 ema {ratio_no:.2f}，有 ema {ratio_em:.2f} → **{'成立' if p1 else '不成立'}**", | |
| f"- 预测 ②（ema 模型 1024~4095 与 64~255 的 loss 差）：{gap:+.4f}，噪声阈值 {noise:.4f} → **{'成立' if p2 else '不成立'}**" | |
| f"（对照：无 ema 的差是 {gap_no:+.4f}）", | |
| f"- ema 不损害短距离：64 字处 loss 差 {em['val'].mean() - no['val'].mean():+.4f} → **{'成立' if short_ok else '不成立'}**", | |
| f"\n**结论**：{verdict}" + ("" if short_ok else "；但 ema 损害了短距离效果，不能直接采用")] | |
| save_result("ema_longrun", lines) | |
| if __name__ == "__main__": | |
| main() | |