ZeroPoint v0 research prototype: new from-scratch NumPy framework, memory-slot architecture, audited Transformer comparison
5c6c847 verified | """预注册实验：查 ZeroPoint 的三个问题 —— 长序列变差的原因、ema 为什么更差、槽为什么塌缩。 | |
| python3 experiments/diagnose.py | |
| 设计在跑之前写定（PREREG），跑完只追加结果。训练好的模型保存在 experiments/out/，供事后再分析。 | |
| 结果：experiments/diagnose.md + results/ 存档。 | |
| """ | |
| import os | |
| import numpy as np | |
| from common import CFG, HERE, ITERS, T, B, batches, prepare_data, save_result, train | |
| from models import ZeroX | |
| import zerocg as zc | |
| SEEDS = [0, 1, 2] | |
| LR = 6e-3 | |
| LONG_T, LONG_N = 4096, 32 | |
| SEGMENTS = [(0, 64), (64, 256), (256, 1024), (1024, 4096)] | |
| OUT = os.path.join(HERE, "out") | |
| CONDITIONS = { # 名字 -> (ema, 训练长度, batch)；每步的字数都是 2048 | |
| "T64": (False, 64, 32), | |
| "T64-ema": (True, 64, 32), | |
| "T256": (False, 256, 8), | |
| } | |
| PREREG = """# 预注册实验：查三个问题 | |
| > 写于实验之前（2026-10-18）。跑完只在下面追加结果，不改这一节。 | |
| **背景**：ema_longrun.md 否定了"槽值无界累积"假设。无 ema 的模型在 1024~4095 字处 loss 比 64~255 字处高约 0.22； | |
| 有 ema 的高约 1.77。compare.md 里 16 个槽只有 3~5 个被写入。 | |
| ## 假设和预测 | |
| **H2（训练长度）**：长序列变差是因为训练只见过 64 字；模型在更长的位置上遇到的槽状态分布是训练时没见过的。 | |
| - 预测：训练长度改为 256（每步字数不变、其余相同）后，1024~4095 与 64~255 的 loss 差 **缩小到 T64 的一半以下**。 | |
| - 不成立时会看到：差距没有缩小（< 2 × 标准差的变化）。 | |
| **H3（ema 为什么更差）**：ema 把长期槽的写入乘上 (1-λ)，训练时（64 字）长期槽几乎不起作用，模型学会只用短期槽； | |
| 到了长序列，读地址仍指向短期槽，但短期槽里的内容和训练时的统计不同。 | |
| - 预测 ③a：ema 模型的**写地址**集中在半衰期短的槽上：半衰期 > 16 字的槽，写入占比 < T64 模型的一半。 | |
| - 预测 ③b：ema 模型学到的长期槽 λ 更接近 1（半衰期更长），或者长期槽的读占比在远处位置上升。 | |
| - 不成立时会看到：两组的写入分布没有明显差别。 | |
| **H4（槽塌缩）**：塌缩不是训练失败，而是写地址的 softmax 竞争让少数槽赢者通吃；在不同序列位置上"赢的槽"不同， | |
| 所以整体统计（compare.md 的 3~5 个有效槽）低估了实际用到的槽数。 | |
| - 预测 ④：按位置分段（0~63 / 64~255 / 256~1023 / 1024~4095）分别统计有效写槽数，各段的并集 > 整体有效槽数 × 1.5； | |
| 同时各段的读/写分布之间有明显差别（任意两段的 JS 散度 > 0.1）。 | |
| - 不成立时会看到：各段的分布几乎相同，用的始终是同几个槽。 | |
| **判定**：每条按上面的数值阈值判，3 个种子均值；差别和 2 × 种子间标准差比。 | |
| """ | |
| def long_nll(model, val_data): | |
| out = [] | |
| 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)) | |
| out.append(-np.take_along_axis(logp, y[..., None], -1)[..., 0]) | |
| nll = np.concatenate(out) | |
| return np.array([nll[:, a:b].mean() for a, b in SEGMENTS]) | |
| def address_stats(model, val_data): | |
| """递推跑 4096 字，按位置分段统计每层每头的平均读/写地址分布。返回 (层, 段, H, M) 两个数组。""" | |
| x, _ = next(batches(val_data, 1, LONG_T, 7, batch=4)) | |
| states = model.init_state(4) | |
| L = len(model.blocks) | |
| H, M = model.blocks[0].mix.H, model.blocks[0].mix.M | |
| reads = np.zeros((L, len(SEGMENTS), H, M)) | |
| writes = np.zeros((L, len(SEGMENTS), H, M)) | |
| for t in range(LONG_T): | |
| _, states = model.step(x[:, t], states, t) | |
| seg = next(i for i, (a, b) in enumerate(SEGMENTS) if a <= t < b) | |
| for l, blk in enumerate(model.blocks): | |
| a, w = blk.mix.last_addr | |
| reads[l, seg] += a.mean(0) | |
| writes[l, seg] += w.mean(0) | |
| for i, (a, b) in enumerate(SEGMENTS): | |
| reads[:, i] /= b - a | |
| writes[:, i] /= b - a | |
| return reads, writes | |
| def eff(p): | |
| """exp(熵)：有效槽数。p 最后一维是槽。""" | |
| p = p / p.sum(-1, keepdims=True) | |
| return np.exp(-(p * np.log(p + 1e-12)).sum(-1)) | |
| def js(p, q): | |
| p, q = p / p.sum(), q / q.sum() | |
| m = (p + q) / 2 | |
| kl = lambda a, b: (a * np.log((a + 1e-12) / (b + 1e-12))).sum() # noqa: E731 | |
| return 0.5 * kl(p, m) + 0.5 * kl(q, m) | |
| def half_lives(model): | |
| """每层每头每槽的半衰期（字）。""" | |
| return np.array([np.log(0.5) / np.log(b.mix.retention().data) for b in model.blocks]) | |
| def main(): | |
| text, V, train_data, val_data = prepare_data() | |
| os.makedirs(OUT, exist_ok=True) | |
| print(PREREG, flush=True) | |
| res = {} | |
| for name, (ema, seq_len, batch) in CONDITIONS.items(): | |
| rows = [] | |
| for seed in SEEDS: | |
| np.random.seed(seed) | |
| model = ZeroX(V, n_slot=16, ema=ema, **CFG) | |
| val = train(name, model, train_data, val_data, seed, LR, seq_len=seq_len, batch=batch) | |
| model.eval() | |
| model.save(os.path.join(OUT, f"diag_{name}_s{seed}.npz")) | |
| ll = long_nll(model, val_data) | |
| reads, writes = address_stats(model, val_data) | |
| hl = half_lives(model) | |
| rows.append(dict(val=val, long=ll, reads=reads, writes=writes, hl=hl)) | |
| print(f" => {name} seed={seed}: val {val:.4f}; long {np.round(ll, 4)}", flush=True) | |
| res[name] = rows | |
| def mean(name, key): | |
| return np.mean([r[key] for r in res[name]], axis=0) | |
| def std(name, key): | |
| return np.std([r[key] for r in res[name]], axis=0) | |
| lines = [PREREG, "\n## 结果（实验之后写）\n", | |
| "### 验证 loss（64 字）和长度外推（32 条 × 4096 字；均值 ± 标准差，3 个种子）\n", | |
| "| 条件 | 64 字验证 | " + " | ".join(f"{a}~{b - 1}" for a, b in SEGMENTS) + " | 远−近 差 |", | |
| "|---|---|" + "---|" * (len(SEGMENTS) + 1)] | |
| gaps = {} | |
| for name in res: | |
| lm, ls = mean(name, "long"), std(name, "long") | |
| g = np.array([r["long"][3] - r["long"][1] for r in res[name]]) | |
| gaps[name] = g | |
| lines.append(f"| {name} | {mean(name, 'val'):.4f} ± {std(name, 'val'):.4f} | " + | |
| " | ".join(f"{m:.4f} ± {s:.4f}" for m, s in zip(lm, ls)) + f" | {g.mean():+.4f} ± {g.std():.4f} |") | |
| # ---- H2 ---- | |
| g64, g256 = gaps["T64"], gaps["T256"] | |
| # T64 的远−近差必须本来就是正的（确实变差了），H2 才有意义 | |
| h2 = g64.mean() > 0 and g256.mean() < 0.5 * g64.mean() and (g64.mean() - g256.mean()) > 2 * max(g64.std(), g256.std()) | |
| # ---- H3：写入在长期槽上的占比；学到的半衰期 ---- | |
| def long_slot_write_share(name): | |
| out = [] | |
| for r in res[name]: | |
| share = [] | |
| for l in range(r["hl"].shape[0]): | |
| mask = r["hl"][l] > 16 # (H, M) | |
| w = r["writes"][l, 0] # 训练长度内（0~63 段）的写分布 (H, M) | |
| share.append((w * mask).sum() / w.sum()) | |
| out.append(np.mean(share)) | |
| return np.array(out) | |
| s64, sema = long_slot_write_share("T64"), long_slot_write_share("T64-ema") | |
| h3a = sema.mean() < 0.5 * s64.mean() | |
| hl64, hlema = mean("T64", "hl"), mean("T64-ema", "hl") | |
| far_read_long = {} | |
| for name in ("T64", "T64-ema"): | |
| vals = [] | |
| for r in res[name]: | |
| v = [] | |
| for l in range(r["hl"].shape[0]): | |
| mask = r["hl"][l] > 16 | |
| v.append((r["reads"][l, 3] * mask).sum() / r["reads"][l, 3].sum() - | |
| (r["reads"][l, 0] * mask).sum() / r["reads"][l, 0].sum()) | |
| vals.append(np.mean(v)) | |
| far_read_long[name] = np.array(vals) | |
| h3b = np.max(hlema) > np.max(hl64) * 1.5 or far_read_long["T64-ema"].mean() > far_read_long["T64"].mean() + 0.1 | |
| lines += ["", "### H3：写入落在半衰期 > 16 字的槽上的占比（0~63 段），和学到的最长半衰期\n", | |
| "| 条件 | 长期槽写入占比 | 最长半衰期（字） | 远处读长期槽占比的变化 |", "|---|---|---|---|"] | |
| for name in ("T64", "T64-ema"): | |
| sh = long_slot_write_share(name) | |
| lines.append(f"| {name} | {sh.mean():.3f} ± {sh.std():.3f} | {np.max(mean(name, 'hl')):.0f} | " | |
| f"{far_read_long[name].mean():+.3f} |") | |
| # ---- H4：分段的有效写槽数、并集、JS 散度 ---- | |
| lines += ["", "### H4：槽的使用（T64，写地址；各层各头平均）\n", | |
| "| 段 | 有效写槽数 |", "|---|---|"] | |
| union_ratio, js_min = [], [] | |
| for r in res["T64"]: | |
| w = r["writes"] # (L, seg, H, M) | |
| overall = eff(w.mean(1)) # (L, H) | |
| per_seg = eff(w) # (L, seg, H) | |
| used = (w / w.sum(-1, keepdims=True)) > 1.0 / w.shape[-1] # 用得比平均多的槽 (L, seg, H, M) | |
| union = used.any(1).sum(-1) # (L, H) | |
| union_ratio.append((union / overall).mean()) | |
| js_vals = [js(w[l, i, h], w[l, j, h]) for l in range(w.shape[0]) for h in range(w.shape[2]) | |
| for i in range(len(SEGMENTS)) for j in range(i + 1, len(SEGMENTS))] | |
| js_min.append(np.max(js_vals)) | |
| w_mean = mean("T64", "writes") | |
| for i, (a, b) in enumerate(SEGMENTS): | |
| lines.append(f"| {a}~{b - 1} | {eff(w_mean[:, i]).mean():.1f} |") | |
| lines.append(f"| 整体 | {eff(w_mean.mean(1)).mean():.1f} |") | |
| union_ratio, js_min = np.array(union_ratio), np.array(js_min) | |
| h4 = union_ratio.mean() > 1.5 and js_min.mean() > 0.1 | |
| lines += [f"\n各段用到的槽的并集 ÷ 整体有效槽数 = {union_ratio.mean():.2f}；各段分布之间最大的 JS 散度 = {js_min.mean():.3f}"] | |
| lines += ["", "## 判定（按预注册的标准）\n", | |
| f"- H2 训练长度：远−近差 T64 {g64.mean():+.4f} → T256 {g256.mean():+.4f} → **{'成立' if h2 else '不成立'}**", | |
| f"- H3a ema 写入避开长期槽：占比 {s64.mean():.3f} → {sema.mean():.3f} → **{'成立' if h3a else '不成立'}**", | |
| f"- H3b ema 的长期槽更长或远处更依赖长期槽 → **{'成立' if h3b else '不成立'}**", | |
| f"- H4 槽塌缩是分段使用不同的槽：并集比 {union_ratio.mean():.2f}，JS {js_min.mean():.3f} → **{'成立' if h4 else '不成立'}**", | |
| "", "模型已保存在 experiments/out/diag_*.npz，可以事后再分析。"] | |
| save_result("diagnose", lines) | |
| if __name__ == "__main__": | |
| main() | |