论文精选

HOLA 给线性注意力加上海马体记忆,解决长上下文遗忘

A Hippocampus for Linear Attention: An Exact Memory for What the Recurrent State Forgets

精选理由

线性注意力模型容易忘事?HOLA 加了个小型精确缓存,长上下文检索比全注意力还强。

AI 摘要

HOLA 为线性注意力模型引入一个专用的精确 KV 缓存作为海马体补丁,与原有的压缩状态共同形成半参数化记忆。在 340M 参数、15B SlimPajama tokens 训练下,HOLA 将 Wikitext 困惑度从 27.32 降至 22.92(-16.1%),甚至低于完整注意力 Transformer++(26.88);LAMBADA 困惑度从 30.95 降至 30.26。在 RULER 的 needle-in-a-haystack 召回任务中,HOLA 在 32k tokens(16 倍训练长度)下仍保持稳健,优于 GDN 和同类缓存方法。

原文 · arXiv cs.AI

A Hippocampus for Linear Attention: An Exact Memory for What the Recurrent State Forgets

Linear-attention and state-space language models compress the prefix into a fixed-size recurrent state, yielding O(1) memory at the cost of a lossy exact memory: when many key--value associations compete, earlier facts are overwritten and needle recall degrades. Inspired by Complementary Learning Systems, we give linear attention a hippocampal complement. HOLA (Hippocampal Linear Attention) keeps the usual delta-rule state as a compressive memory and adds a bounded exact KV cache, forming a semiparametric test-time memory: the state models linearly compressible structure, while the cache stores associations that should not be forced through that state. The cache writes without a learned eviction module, keeping tokens with large beta * ||e||, the prediction residual actually committed to the state; a decoupled RMSNorm-gamma cache read then turns these exact KV pairs into sharp retrieval rather than soft averaging. At 340M parameters trained on 15B SlimPajama tokens, HOLA lowers Wikitext perplexity from 27.32 to 22.92 (-16.1%), below a full-attention Transformer++ (26.88), and improves LAMBADA perplexity from 30.95 to 30.26. It also achieves the best linear in-context retrieval and remains much more robust than GDN or a matched HOLA+recency cache on RULER needle-in-a-haystack recall out to 32k tokens (16x its training length).