ZeroPoint v0 research prototype: new from-scratch NumPy framework, memory-slot architecture, audited Transformer comparison
5c6c847 verified | """长距离精确回忆（联想检索）实验：专门测 ZeroPoint 的已知弱点。 | |
| 序列格式（每个数字是一个词元）： | |
| k1 v1 k2 v2 ... kN vN [问] kq → 模型要说出 kq 对应的 v | |
| 键从 16 个里不重复地抽 N 个，值从另外 16 个里随机抽。要答对，模型必须把 N 对"键 → 值"原样记住。 | |
| 注意力天生擅长这个（用 kq 当查询去找那个键）；ZeroPoint 只有 M 个槽，对数一多就会混。 | |
| python3 experiments/recall.py | |
| 训练时 N 在 2~16 之间随机，测试时分别报告 N = 2, 4, 8, 12, 16 的准确率。 | |
| 结果：experiments/recall.md（最新）+ experiments/results/<时间>-recall.md（存档） | |
| """ | |
| import os | |
| import numpy as np | |
| from common import CFG, save_result | |
| from models import TransformerLM, ZeroX | |
| import zerocg as zc | |
| from zerocg import functional as F | |
| ITERS = int(os.environ.get("RECALL_ITERS", 3000)) | |
| NK, NV = 16, 16 | |
| QUERY = NK + NV # "问" 这个词元 | |
| VOCAB = NK + NV + 1 | |
| MAX_N = 16 | |
| MAX_LEN = 2 * MAX_N + 2 | |
| TEST_N = [2, 4, 8, 12, 16] | |
| SEEDS = [0, 1] | |
| MODELS = { | |
| "Transformer": lambda: TransformerLM(VOCAB, MAX_LEN, **CFG), | |
| "ZeroPoint (16 槽)": lambda: ZeroX(VOCAB, n_slot=16, **CFG), | |
| "ZeroPoint (32 槽)": lambda: ZeroX(VOCAB, n_slot=32, **CFG), | |
| } | |
| def make_batch(rng, B, n): | |
| keys = np.stack([rng.permutation(NK)[:n] for _ in range(B)]) | |
| vals = rng.integers(NK, NK + NV, size=(B, n)) | |
| seq = np.empty((B, 2 * n + 2), dtype=np.int64) | |
| seq[:, 0:2 * n:2], seq[:, 1:2 * n:2] = keys, vals | |
| q = rng.integers(0, n, B) | |
| seq[:, -2], seq[:, -1] = QUERY, keys[np.arange(B), q] | |
| return seq, vals[np.arange(B), q] | |
| def train_and_eval(make, seed): | |
| np.random.seed(seed) | |
| rng = np.random.default_rng(seed) | |
| model = make() | |
| opt = zc.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01) | |
| for it in range(ITERS): | |
| opt.lr = 3e-3 * min(1, (it + 1) / 100) | |
| x, y = make_batch(rng, 32, int(rng.integers(2, MAX_N + 1))) | |
| loss = F.cross_entropy(model(x)[:, -1], y) | |
| opt.zero_grad() | |
| loss.backward() | |
| zc.optim.clip_grad_norm(model.parameters(), 1.0) | |
| opt.step() | |
| model.eval() | |
| test = np.random.default_rng(1000 + seed) | |
| acc = {} | |
| with zc.no_grad(): | |
| for n in TEST_N: | |
| x, y = make_batch(test, 512, n) | |
| acc[n] = float((model(x).data[:, -1].argmax(-1) == y).mean()) | |
| return acc | |
| def main(): | |
| print(f"联想检索：{len(MODELS)} 个模型 × 种子 {SEEDS}，每次 {ITERS} 步", flush=True) | |
| res = {} | |
| for name, make in MODELS.items(): | |
| runs = [train_and_eval(make, s) for s in SEEDS] | |
| res[name] = {n: np.mean([r[n] for r in runs]) for n in TEST_N} | |
| print(f" => {name}: " + " ".join(f"N={n}:{res[name][n]:.0%}" for n in TEST_N), flush=True) | |
| lines = ["# 长距离精确回忆（联想检索）\n", | |
| "`k1 v1 k2 v2 … kN vN [问] kq` → 回答 kq 对应的值。16 个键、16 个值，随机猜的准确率是 1/16 ≈ 6%。", | |
| f"三个模型宽 {CFG['n_embd']}、{CFG['n_layer']} 层、{CFG['n_head']} 头；训练 {ITERS} 步，N 在 2~{MAX_N} 随机；" | |
| f"每格是 512 道测试题 × 种子 {SEEDS} 的平均准确率。\n", | |
| "| 模型 | " + " | ".join(f"N = {n}" for n in TEST_N) + " |", "|---|" + "---|" * len(TEST_N)] | |
| for name, acc in res.items(): | |
| lines.append(f"| {name} | " + " | ".join(f"{acc[n]:.0%}" for n in TEST_N) + " |") | |
| save_result("recall", lines) | |
| if __name__ == "__main__": | |
| main() | |