Blame for experiments/diagnose.py

"""预注册实验：查 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()