File size: 4,865 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 | """预注册分析：长期槽到底在干什么？（用 diagnose 保存的 9 个模型，不重训）
python3 experiments/slots.py
背景：diagnose 发现无 ema 的模型只有 6% 的写入落在半衰期 > 16 字的槽上，但消融说明"分层初始化"有用（去掉变差 0.051）。
结果：experiments/slots.md + results/ 存档。
"""
import os
import numpy as np
from common import HERE, batches, prepare_data, save_result
from models import ZeroX
OUT = os.path.join(HERE, "out")
SEEDS = [0, 1, 2]
T_RUN = 2048
BUCKETS = [(0, 4), (4, 16), (16, 64), (64, 1e9)] # 半衰期分桶（字）
PREREG = """# 预注册分析：长期槽在干什么
> 写于分析之前（2026-11-01）。跑完只追加结果，不改这一节。用 experiments/out/diag_*.npz（T64 / T64-ema / T256 各 3 种子）。
**矛盾**：T64 模型写入只有 6.1% 落在半衰期 > 16 字的槽；但消融里去掉分层初始化变差 0.051。
**三个可区分的假设**
- H9 读取端在用：长期槽虽然写得少，但**读地址**落在长期槽上的占比明显高于写占比（读占比 ≥ 2 × 写占比）。
→ 长期槽像"慢变的背景"，少量写入、持续读取。
- H10 中期槽才是关键：分层初始化的作用主要来自半衰期 4~64 字的槽，不是 > 64 的；读写都集中在 4~64 桶（该桶读 + 写占比合计 > 50%）。
- H11 训练长度决定长期槽的使用：T256 模型在 > 16 字桶上的写占比 ≥ 2 × T64 的。
来源：diagnose 里 T256 远处 loss 明显更好，推测它学会了用更长的槽。
**判定**：3 个种子平均；各桶占比按所有层、所有头的平均。H9 / H11 的"2 倍"是无依据的整数阈值，仅作记录；
同时报告原始占比，让读者自己看。
"""
def main():
text, V, train_data, val_data = prepare_data()
x, _ = next(batches(val_data, 1, T_RUN, 7, batch=4))
cfg = dict(n_embd=64, n_head=4, n_layer=2, n_slot=16)
rows = {}
for cond in ("T64", "T64-ema", "T256"):
reads, writes = np.zeros(len(BUCKETS)), np.zeros(len(BUCKETS))
for seed in SEEDS:
model = ZeroX(V, ema=(cond == "T64-ema"), **cfg)
model.load(os.path.join(OUT, f"diag_{cond}_s{seed}.npz")).eval()
states = model.init_state(4)
r_acc = [np.zeros((b.mix.H, b.mix.M)) for b in model.blocks]
w_acc = [np.zeros((b.mix.H, b.mix.M)) for b in model.blocks]
for t in range(T_RUN):
_, states = model.step(x[:, t], states, t)
for l, b in enumerate(model.blocks):
a, w = b.mix.last_addr
r_acc[l] += a.mean(0)
w_acc[l] += w.mean(0)
for l, b in enumerate(model.blocks):
hl = np.log(0.5) / np.log(b.mix.retention().data) # (H, M)
r = r_acc[l] / r_acc[l].sum(-1, keepdims=True)
w = w_acc[l] / w_acc[l].sum(-1, keepdims=True)
for k, (lo, hi) in enumerate(BUCKETS):
m = (hl >= lo) & (hl < hi)
reads[k] += (r * m).sum() / (r.shape[0]) / len(model.blocks)
writes[k] += (w * m).sum() / (w.shape[0]) / len(model.blocks)
rows[cond] = (reads / len(SEEDS), writes / len(SEEDS))
lines = [PREREG, "\n## 结果（分析之后写）\n",
f"验证集 4 条 × {T_RUN} 字递推；占比 = 该桶内所有槽的地址权重之和（读 / 写各自归一化），层和头平均，3 个种子平均。\n",
"| 条件 | 量 | " + " | ".join(f"半衰期 {lo}~{'∞' if hi > 1e8 else int(hi)}" for lo, hi in BUCKETS) + " |",
"|---|---|" + "---|" * len(BUCKETS)]
for cond, (r, w) in rows.items():
lines.append(f"| {cond} | 读 | " + " | ".join(f"{v:.3f}" for v in r) + " |")
lines.append(f"| {cond} | 写 | " + " | ".join(f"{v:.3f}" for v in w) + " |")
r64, w64 = rows["T64"]
r256, w256 = rows["T256"]
long_r, long_w = r64[2] + r64[3], w64[2] + w64[3]
h9 = long_r >= 2 * long_w
h10 = (r64[1] + r64[2] + w64[1] + w64[2]) / 2 > 0.5
h11 = (w256[2] + w256[3]) >= 2 * long_w
lines += ["", "## 判定（按预注册的标准）\n",
f"- H9 读取端在用长期槽：> 16 字桶 读 {long_r:.3f} vs 写 {long_w:.3f} → **{'成立' if h9 else '不成立'}**（2 倍阈值无依据，仅作记录）",
f"- H10 中期槽是关键：4~64 字桶 读写平均占比 {(r64[1] + r64[2] + w64[1] + w64[2]) / 2:.3f} > 0.5 → **{'成立' if h10 else '不成立'}**",
f"- H11 训练长度 256 更用长期槽：> 16 字桶写占比 T256 {w256[2] + w256[3]:.3f} vs T64 {long_w:.3f} → **{'成立' if h11 else '不成立'}**（2 倍阈值无依据，仅作记录）"]
save_result("slots", lines)
if __name__ == "__main__":
main()
|