KV 缓存是长上下文推理绕不开的两难:上下文越长,softmax 注意力要读的键值越多,生成越慢;把它们全塞进显存,显存又先撑不住。线性注意力和状态空间模型把上下文压进固定大小的状态,速度和显存都恒定了,但固定状态的容量上限让召回能力随信息量增长而崩掉。德国图宾根大学数学系的 Moritz Brösamle 在 arXiv 提出的 LEMA(Latest Exact Match Attention)选了第三条路:状态可以无限增长,但不放显存、也不逐条扫描——直接当字典查。(https://arxiv.org/abs/2609.25802)
注意力退化成一次哈希查找
LEMA 的规则极简:查询和键先二值化,每个查询只取最近一次精确匹配的键,匹配不到就返回零。论文给出的生成步伪代码只有几行——查一次字典、写一次字典,同键覆盖旧值。于是每个注意力头的 KV 缓存就是一个字典:状态大小随内容增长,理论上限是 2 的头维度次方个条目,实际上只存出现过的键。
理论:与 word-RAM 双向等价
论文证明,带思维链的 LEMA transformer 可以模拟 word-RAM(现代计算机的抽象模型);反过来,word-RAM 也能以每 token 与上下文长度无关的开销模拟 LEMA transformer。作者称据其所知,这种双向对应在此前的注意力变体里没有建立过——硬注意力家族的表达力结果通常只证单向。
训练:从软注意力退火到硬规则
精确匹配不可导,训练是最大难点。论文用直通估计器传二值化梯度,再用 stick-breaking 注意力做软代理,逐步退火到 LEMA。作者坦承这只是第一次尝试:退火对学习率敏感,且训练仍需与序列长度平方成正比的计算量。
实测:召回超 GDN,整体仍落后 softmax
在合成联想召回任务上(词表 4096),只在 8 对关联上训练的 LEMA 几乎完美外推到 4096 对,固定状态的 gated DeltaNet(GDN)则在关联数变大后失效。在 FineWeb-Edu 上训练的 29M 到 834M 参数语言模型里,LEMA 的验证损失大约追平参数量为其 55%-57% 的 softmax 模型。两个长程召回代理测试更能说明差异:重复稀有 bigram 的最远距离桶里,LEMA 损失比自身基线低 2.1 nats,GDN 只低 0.7;RULER 单针检索(S-NIAH-1)上,LEMA 一旦检索成功,重复填充句不再改变状态,检索可以无限持续而不增长状态。
工程:哈希表进内存,生成速度恒定
推理实现把所有头的 KV 缓存做成主内存里一张开放寻址哈希表(线性探测),显存只留权重和激活。在 RTX 3090 上与 vLLM 基准对比,只要哈希表不接近容量上限,LEMA 生成速度恒定、与 GDN 相当。834M 训练模型每头字典在 16k token 后平均 1.7k 条,256k 后 18k 条——论文测试用了 50 GB 的主内存哈希表。代价也有:约 10% 的头从未找到匹配。代码已开源(github.com/moritzbroe/latest_exact_match_attention)。
把"更像计算机的注意力"和"能训练的注意力"接起来,LEMA 给出了一条少见的路线:召回靠无限状态,速度靠 O(1) 查找,显存压力靠内存转嫁。834M 规模的差距说明训练方法远未成熟,但 KV 缓存焦虑的解法未必是更聪明的压缩——也可能是干脆换一种数据结构。