CODING WITH AI · LECTURE 06

返回课堂 · 讲义

已保存可复算的 CPU 教学输出。ICL fixture 为合成评分测试,所有 Lab 均未调用真实模型。

LAB 02 · 按内容读记忆:相似度、干扰与选择性

目标是理解一个读取接口:query 与 key 比较,用匹配程度取回 value。我们先看三个手写记忆,再在明确的随机向量假设下改变数量、维度、相似度和选择性。

证据身份:可复算的 CPU 玩具记忆实验,不是语言模型,不是 ICL 性能评测,也不代表真实长上下文的容量。 用 NumPy,固定种子 20261020,每段均显示具体输入与结果。

你可以先预测:更尖锐的 softmax 总会帮忙吗?当 query 更像错误 key 时,选择性会怎样?

查看与修改计算代码
import numpy as np
np.set_printoptions(precision=4,suppress=True)
SEED=20261020
K=np.array([[1.,0.],[0.,1.],[0.8,0.6]])
V=np.array([10.,20.,30.])
q=np.array([1.,0.])
def softmax(x):
    z=x-np.max(x); e=np.exp(z); return e/e.sum()
def read(K,V,q,beta=4.):
    scores=K@q
    weights=softmax(beta*scores)
    return {'scores':scores,'linear_value':scores@V,'weights':weights,'softmax_value':weights@V}
print('K=',K,'V=',V,'q=',q,sep='\n')
for beta in (1.,4.,16.): print('选择性 beta=',beta,read(K,V,q,beta))

保存的本地 CPU 输出 · associative-memory-01

K=
[[1.  0. ]
 [0.  1. ]
 [0.8 0.6]]
V=
[10. 20. 30.]
q=
[1. 0.]
选择性 beta= 1.0 {'scores': array([1. , 0. , 0.8]), 'linear_value': np.float64(34.0), 'weights': array([0.4573, 0.1682, 0.3744]), 'softmax_value': np.float64(19.17100337591611)}
选择性 beta= 4.0 {'scores': array([1. , 0. , 0.8]), 'linear_value': np.float64(34.0), 'weights': array([0.6814, 0.0125, 0.3062]), 'softmax_value': np.float64(16.247926543286287)}
选择性 beta= 16.0 {'scores': array([1. , 0. , 0.8]), 'linear_value': np.float64(34.0), 'weights': array([0.9608, 0.    , 0.0392]), 'softmax_value': np.float64(10.78331545251376)}

线性读取 Σvᵢ(kᵢᵀq) 会混入第三条相似记忆。softmax 将权重变为非负、和为 1;提高 beta 让读取更集中于得分最高的一项。这里 q 就是第一条 key,所以集中通常有利,但这依赖 query 的正确性。

1. 只改变相似度:两条不同值能否分开?

保持 V、q 和 beta 不变,把第三条 key 改得更接近第一条。相同的 key 配不同的 value 时,接口没有额外信息可判断哪个才是“想要的”。

查看与修改计算代码
for similarity in (0.,0.8,0.99,1.0):
    close=K.copy()
    close[2]=[similarity,np.sqrt(max(0,1-similarity**2))]
    result=read(close,V,q,beta=16.)
    print('第三key与目标余弦=',similarity,'权重=',result['weights'],'读出=',round(float(result['softmax_value']),4))
print('相同key不同value时,两个相同最高分会平分主要权重。')

保存的本地 CPU 输出 · associative-memory-03

第三key与目标余弦= 0.0 权重= [1. 0. 0.] 读出= 10.0
第三key与目标余弦= 0.8 权重= [0.9608 0.     0.0392] 读出= 10.7833
第三key与目标余弦= 0.99 权重= [0.5399 0.     0.4601] 读出= 19.2017
第三key与目标余弦= 1.0 权重= [0.5 0.  0.5] 读出= 20.0
相同key不同value时,两个相同最高分会平分主要权重。

2. 选择性失效:query 偏向另一条记忆

目标仍声明为第一条 value=10,但只把 q 改为更接近第三条的方向。若提高 beta 后结果更远离 10,不是程序坏了,而是接口忠实地放大了错误匹配。

查看与修改计算代码
q_bad=np.array([0.7,np.sqrt(1-0.7**2)])
for beta in (1.,4.,16.):
    r=read(K,V,q_bad,beta)
    print('beta=',beta,'目标=10','weights=',r['weights'],'读出=',round(float(r['softmax_value']),4),'绝对误差=',round(abs(float(r['softmax_value'])-10),4))

保存的本地 CPU 输出 · associative-memory-05

beta= 1.0 目标=10 weights= [0.2986 0.3029 0.3985] 读出= 20.9986 绝对误差= 10.9986
beta= 4.0 目标=10 weights= [0.1912 0.2024 0.6064] 读出= 24.1513 绝对误差= 14.1513
beta= 16.0 目标=10 weights= [0.0097 0.0121 0.9782] 读出= 29.6851 绝对误差= 19.6851

3. 随机向量:数量、维度与读取误差

生成方式:每条 key、value 都从独立标准正态向量采样并逐行单位化。query 等于所查询的 key,或加上固定幅度的独立噪声后单位化。value 维度固定为 16;每次使用三个独立随机流和固定 512×512 key 池,改变记忆数时保留已有 key/value,统一评估前 64 条 query;这样新记忆带来的干扰不会与重新抽样混在一起。

评价分开看:key top-1 检索正确率,与读出向量相对目标 value 的平均平方误差。前者只测排序,后者能显示即便 top-1 正确,其他 value 仍可混入。表中所有单变量对照都以 m=64,d=32,beta=4,noise=0 为基线。

查看与修改计算代码
def normalize(x): return x/(np.linalg.norm(x,axis=-1,keepdims=True)+1e-12)
def experiment(m=64,d=32,beta=4.,noise=0.,seed=SEED):
    if not (1 <= m <= 512 and 1 <= d <= 512): raise ValueError('教学池限制: m,d 必须在 1..512')
    # 独立随机流与固定大池:改变记忆数量时,已有key/value和测试query保持不变。
    keys=normalize(np.random.default_rng(seed).standard_normal((512,512))[:m,:d])
    values=normalize(np.random.default_rng(seed+1).standard_normal((512,16))[:m])
    n_queries=min(64,m)
    noise_vectors=normalize(np.random.default_rng(seed+2).standard_normal((512,512))[:n_queries,:d])
    queries=normalize(keys[:n_queries]+noise*noise_vectors)
    scores=queries@keys.T
    weights=np.exp(beta*(scores-scores.max(axis=1,keepdims=True)))
    weights/=weights.sum(axis=1,keepdims=True)
    out=weights@values
    return {'m':m,'d':d,'beta':beta,'noise':noise,'n_queries':n_queries,
            'key_top1':float(np.mean(scores.argmax(axis=1)==np.arange(n_queries))),
            'read_MSE':float(np.mean((out-values[:n_queries])**2)),
            'mean_target_weight':float(np.mean(weights[np.arange(n_queries),np.arange(n_queries)]))}
base=experiment()
print('基线:',base)
for name,params in [('只增加记忆数量',{'m':256}),('只增加维度',{'d':128}),('只提高选择性',{'beta':16.}),('只增加query噪声',{'noise':2.})]:
    print(name,experiment(**params))
assert experiment()==base

保存的本地 CPU 输出 · associative-memory-07

基线: {'m': 64, 'd': 32, 'beta': 4.0, 'noise': 0.0, 'n_queries': 64, 'key_top1': 1.0, 'read_MSE': 0.02263704699898566, 'mean_target_weight': 0.4063104867825286}
只增加记忆数量 {'m': 256, 'd': 32, 'beta': 4.0, 'noise': 0.0, 'n_queries': 64, 'key_top1': 1.0, 'read_MSE': 0.04617046326457634, 'mean_target_weight': 0.1434130799496341}
只增加维度 {'m': 64, 'd': 128, 'beta': 4.0, 'noise': 0.0, 'n_queries': 64, 'key_top1': 1.0, 'read_MSE': 0.0196146400774959, 'mean_target_weight': 0.4455247343996429}
只提高选择性 {'m': 64, 'd': 32, 'beta': 16.0, 'noise': 0.0, 'n_queries': 64, 'key_top1': 1.0, 'read_MSE': 4.212714514168018e-09, 'mean_target_weight': 0.9997929867384966}
只增加query噪声 {'m': 64, 'd': 32, 'beta': 4.0, 'noise': 2.0, 'n_queries': 64, 'key_top1': 0.671875, 'read_MSE': 0.05399102125984726, 'mean_target_weight': 0.08257851200860024}

怎样解释:随机且足够分散的 key 更容易区分;数量增加会加入更多干扰项;提高 beta 能在匹配正确时压低干扰,但不能修复错误 query 或同 key 冲突。单次种子只是可复算观察,不能当作容量定律。你可以只改 SEED 检查结论是否依赖这一次采样,再增加种子汇总。

4. 接回当前上下文

这本直接构造 K/V,真实 Transformer 的 K/V 由当前 token 通过已训练投影计算。输入示例变化时,本次激活与 K/V 改变,参数未必更新;如何把示例解释为任务规则,还依赖多层已学计算。

MLP 的记忆解释涉及训练形成的参数;这里的当前记忆是显式数组。不能把两者的存储位置、更新方式混为一谈,更不能用本表宣称大模型能可靠读取相同数量的 token。

来源:Alex Wang — You could have invented transformers from memory 提供按内容读取的叙述视角;本本数值与随机代码为课程独立构造,未复刻博客容量实验。继续到 LAB 03,把 K/V 从读出机制转成每个请求的真实字节账本。