大语言模型(LLMs)在各领域表现亮眼,但一碰到长文本就容易“露怯”。问题出在哪?两大硬伤:注意力机制的计算量随序列长度平方增长,生成时缓存键值对的内存占用也线性攀升。说白了,序列一长,模型既算不动也存不下。
更让人头疼的是,检索增强生成(RAG)虽然能补充外部知识,但检索回来的信息如果超出模型能处理的上限,反而会拖后腿——生成质量不升反降(见下图黄色部分)。这时,一个叫 MemLong 的方法站了出来:它像给模型外设了一个“记忆银&行”,专门存储历史上下文,需要时再精准调取。

这项研究的目标很明确:既能高效处理长文本,又不牺牲短文本上的表现;既要扩展上下文长度,又要保证计算效率。MemLong 的检索与记忆机制如下图所示(具体流程见原文配图),简单说就是:把过去的上下文和知识存进一个不可训练的记忆库,然后从中检索块级的键值对,喂给模型。
关键思想
提出 MemLong——一种高效、轻量的 LLM 扩展方法。
核心思路:将过去的上下文和知识存储在一个不可训练的记忆库中,并利用这些存储的嵌入来检索块级键值(K-V)对,供模型输入。
MemLong 适用于任何仅解码器的预训练语言模型:只需添加一个用于记忆和检索的 ret-mem 组件,以及一个融合局部与记忆信息的检索因果注意力模块。
这其实是一个组合拳:把不可微分的检索-记忆模块和部分可训练的解码器语言模型结合起来,从而增强长文本上下文建模能力。外部检索器负责挖历史信息,再通过细粒度、可控的检索注意力机制,将语义相关的信息块整合到模型中。这样做不仅提升了模型处理长文本的能力,还能保持信息分布的一致性,避免训练过程中间出现分布偏移。
来看一个具体示例。底层模型保持静态,对整个数据块 Ci 进行因果语言建模;随后,Ci 被缓存为嵌入和键值对形式。上层则进行微调,协调检索偏好并整合检索到的内容。整个流程清晰而克制,不需要大动干戈地改模型结构。
论文:https://arxiv.org/pdf/2408.16967
代码:https://github.com/Bui1dMySea/MemLong/blob/main/run_clm_no_trainer.py
