File size: 7,095 Bytes
5c6c847 | 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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | """预注册实验：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()
|