ZeroPoint v0 research prototype: new from-scratch NumPy framework, memory-slot architecture, audited Transformer comparison
5c6c847 verified | """ZeroPoint vs Transformer 对照实验（字符级语言模型）。 | |
| python3 experiments/compare.py # 默认语料：Python 标准库源码（离线可用） | |
| python3 experiments/compare.py my.txt # 换成自己的文本 | |
| ITERS=2000 python3 experiments/compare.py | |
| 三个模型宽度、层数、头数、MLP 完全一样，唯一区别是"token 混合层 + 位置编码"： | |
| Transformer 自注意力 + 可学习位置编码 | |
| Transformer (sin) 自注意力 + 论文 §3.5 的正弦位置编码（理论上能外推到更长序列） | |
| ZeroPoint 记忆槽，无位置编码 | |
| 流程： | |
| 1. 学习率搜索：每个模型在 LRS 里各试一次（种子 0），取验证 loss 最低的 | |
| 2. 用各自最好的学习率再跑种子 1、2，报告 3 个种子的 均值 ± 标准差 | |
| 3. 长度外推：训练只见过 64 个字，测到 4096 个字 | |
| 结果：experiments/compare.md（最新）+ experiments/results/<时间>-compare.md（存档） | |
| """ | |
| import sys | |
| import numpy as np | |
| from common import (CFG, ITERS, T, batches, evaluate, forward_flops_per_char, is_zeropoint, prepare_data, | |
| save_result, slot_usage, train) | |
| from models import TransformerLM, ZeroX | |
| import zerocg as zc | |
| LRS = [1e-3, 3e-3, 6e-3] | |
| SEEDS = [0, 1, 2] | |
| MODELS = { | |
| "Transformer": lambda V: TransformerLM(V, T, **CFG), | |
| "Transformer (sin)": lambda V: TransformerLM(V, T, pos="sin", **CFG), | |
| "ZeroPoint": lambda V: ZeroX(V, n_slot=16, **CFG), | |
| } | |
| # 长度外推的分段；Transformer 的注意力矩阵是 长度² 的，4096 字放不进内存，只测到 1024 | |
| SEGMENTS = [(0, 64), (64, 256), (256, 1024), (1024, 4096)] | |
| LONG_MAX = {"Transformer (sin)": 1024, "ZeroPoint": 4096} | |
| LONG_N = 8 # 测试序列条数 | |
| def long_context_loss(model, val_data, max_len): | |
| x, y = next(batches(val_data, 1, max_len, 7, batch=LONG_N)) | |
| 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 = -np.take_along_axis(logp, y[..., None], -1)[..., 0] | |
| return np.array([nll[:, a:b].mean() if b <= max_len else np.nan for a, b in SEGMENTS]) | |
| def main(): | |
| path = sys.argv[1] if len(sys.argv) > 1 else None | |
| text, V, train_data, val_data = prepare_data(path) | |
| print(f"语料 {len(text):,} 字符，词表 {V}，每次训练 {ITERS} 步", flush=True) | |
| sweep = {name: {} for name in MODELS} # name -> lr -> 验证 loss（种子 0） | |
| finals = {name: [] for name in MODELS} # name -> [各种子的验证 loss] | |
| long_loss = {name: [] for name in LONG_MAX} | |
| usage = [] | |
| def run(name, seed, lr): | |
| np.random.seed(seed) | |
| model = MODELS[name](V) | |
| val = train(name, model, train_data, val_data, seed, lr) | |
| print(f" => {name} seed={seed} lr={lr:g}: val {val:.4f}", flush=True) | |
| return model, val | |
| for name in MODELS: | |
| best = None | |
| for lr in LRS: | |
| model, val = run(name, 0, lr) | |
| sweep[name][lr] = val | |
| if best is None or val < best[1]: | |
| best = (model, val, lr) | |
| best_lr = best[2] | |
| trained = [best[0]] + [run(name, s, best_lr)[0] for s in SEEDS[1:]] | |
| finals[name] = [best[1]] + [None] * (len(SEEDS) - 1) | |
| for k, model in enumerate(trained): | |
| if k > 0: | |
| finals[name][k] = evaluate(model, val_data) | |
| if name in LONG_MAX: | |
| long_loss[name].append(long_context_loss(model, val_data, LONG_MAX[name])) | |
| if is_zeropoint(model): | |
| usage.append(slot_usage(model)) | |
| # ---- 写结果 ---- | |
| m0 = {name: make(V) for name, make in MODELS.items()} | |
| lines = ["# ZeroPoint vs Transformer 对照实验\n", | |
| f"- 语料：{'Python 标准库源码' if not path else path}，{len(text):,} 字符，词表 {V}；90% 训练 / 10% 验证", | |
| f"- 三个模型：宽 {CFG['n_embd']}、{CFG['n_layer']} 层、{CFG['n_head']} 头、Pre-LN、GELU MLP 4 倍；" | |
| "ZeroPoint 每头 16 个槽、可学习初始状态、分块并行（块长 32）、无位置编码", | |
| f"- 训练：上下文 {T}、batch 32、{ITERS} 步、AdamW（权重衰减 0.01）、预热 100 步 + 余弦、梯度裁剪 1.0", | |
| f"- 学习率：每个模型在 {LRS} 里各试一次（种子 0），取最好的再跑种子 {SEEDS[1:]}", | |
| "- 验证 loss = 40 个固定验证 batch 的平均交叉熵（nats/字符，越低越好），均值 ± 标准差（3 个种子）", | |
| "- 计算量 = 前向时每个字的矩阵乘浮点运算次数（不受机器负载影响；训练一步约 3 倍）\n", | |
| "| 模型 | 参数量 | 最佳学习率 | 验证 loss | 前向计算量（每字） | 推理状态（第 64 字 / 第 4096 字，浮点数） |", | |
| "|---|---|---|---|---|---|"] | |
| for name, vals in finals.items(): | |
| m = m0[name] | |
| best_lr = min(sweep[name], key=sweep[name].get) | |
| state = (f"{m.state_floats():,} / {m.state_floats():,}（固定）" if is_zeropoint(m) | |
| else f"{m.state_floats(64):,} / {m.state_floats(4096):,}（KV 缓存）") | |
| lines.append(f"| {name} | {m.num_parameters():,} | {best_lr:g} | {np.mean(vals):.4f} ± {np.std(vals):.4f} | " | |
| f"{forward_flops_per_char(m, V) / 1e3:.0f} K | {state} |") | |
| lines += ["", "## 学习率搜索（种子 0 的验证 loss）\n", | |
| "| 模型 | " + " | ".join(f"lr={lr:g}" for lr in LRS) + " |", | |
| "|---|" + "---|" * len(LRS)] | |
| for name, d in sweep.items(): | |
| lines.append(f"| {name} | " + " | ".join(f"{d[lr]:.4f}" for lr in LRS) + " |") | |
| lines += ["", f"## 长度外推（训练上下文 {T}；{LONG_N} 条验证序列 × 3 个种子平均）\n", | |
| "| 位置 | ZeroPoint | Transformer (sin) | Transformer（可学习位置编码） |", "|---|---|---|---|"] | |
| lz, ls = np.nanmean(long_loss["ZeroPoint"], axis=0), np.nanmean(long_loss["Transformer (sin)"], axis=0) | |
| for k, (a, b) in enumerate(SEGMENTS): | |
| tag = "（训练见过的长度）" if b <= T else "" | |
| s = "—（注意力矩阵太大，没测）" if np.isnan(ls[k]) else f"{ls[k]:.4f}" | |
| lines.append(f"| {a} ~ {b - 1}{tag} | {lz[k]:.4f} | {s} | {'能跑' if b <= T else f'跑不了（只有 {T} 个位置）'} |") | |
| u = np.mean(usage, axis=0) # (层, 2) | |
| lines += ["", "## ZeroPoint 槽的使用情况（有效槽数，满分 16）\n", | |
| "有效槽数 = 平均读/写地址分布的 exp(熵)：16 个槽用得一样多 = 16，全挤在一个槽 = 1。\n", | |
| "| 层 | 读 | 写 |", "|---|---|---|"] | |
| lines += [f"| {i} | {r:.1f} | {w:.1f} |" for i, (r, w) in enumerate(u)] | |
| save_result("compare", lines) | |
| if __name__ == "__main__": | |
| main() | |