File size: 3,552 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 | """长距离精确回忆（联想检索）实验：专门测 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()
|