给线性注意力加了个记忆库,让模型记住更久远的关键信息。
A Hippocampus for Linear Attention: An Exact Memory for What the Recurrent State Forgets
- 用固定大小缓存+状态压缩双机制,分离可压缩与需精确保存的信息。
- 在340M参数下,维基文本困惑度降至22.92,优于全注意力模型。
- 适合需要长序列精准回忆的场景,如文档问答或知识检索。
线性注意力与状态空间语言模型将前缀压缩为固定大小的状态,虽实现O(1)内存开销,但损失了精确记忆:当多个键值对竞争时,早期信息被覆盖,关键信息召回性能下降。受互补学习系统启发,我们为线性注意力引入海马体式补充机制。HOLA(海马体线性注意力)保留原有的增量更新状态作为压缩记忆,并新增一个有界精确键值缓存,构成半参数化测试时记忆:状态建模线性可压缩结构,而缓存存储不应通过状态强制处理的关键关联。缓存写入无需学习的淘汰模块,仅保留预测残差较大且实际投入状态的条目(beta × ||e|| 大);解耦的RMSNorm-gamma读取机制将这些精确键值对转化为锐利检索,而非软平均。在340M参数、150亿个SlimPajama训练样本上,HOLA将维基文本困惑度从27.32降至22.92(-16.1%),低于全注意力Transformer++(26.88),并使LAMBADA困惑度从30.95降至30.26。其在线性上下文检索中表现最佳,且在长达32,000词的‘针堆’召回任务中,相比GDN或匹配的带时间戳缓存的HOLA仍保持更强鲁棒性(训练长度的16倍)。
原文摘要 · Abstract (English)
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).
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。