用记忆网络+滑动窗口,让大模型高效处理超长文本。
AllMem: A Memory-centric Recipe for Efficient Long-context Modeling
- 结合滑动窗口与非线性记忆网络,实现长序列高效建模。
- 4k窗口在37k上下文任务中仅损失0.83分,128k场景下超越全注意力。
- 可替换预训练模型的注意力层,适合需要长上下文的部署场景。
大语言模型在长序列任务中因自注意力机制带来的计算复杂度和内存开销而性能受限。为此,我们提出 extsc{AllMem},一种新型混合架构,将滑动窗口注意力(SWA)与非线性测试时训练(TTT)记忆网络结合。该方法使模型能有效扩展至超长上下文,同时缓解灾难性遗忘。不仅突破了线性记忆模型的表征瓶颈,还显著降低长序列推理时的计算与内存消耗。此外,我们设计了内存高效的微调策略,将预训练模型中的标准注意力层替换为带记忆的滑动窗口层,实现任意现成LLM向 extsc{AllMem} 架构的高效转换。实证表明,4k窗口模型在37k长上下文的LongBench上接近无损表现,仅比全注意力下降0.83分;在128k上下文的InfiniteBench上,8k窗口版本优于全注意力,验证了参数化记忆在抑制噪声、保持鲁棒长程建模方面的有效性,且避免了全局注意力的高昂成本。
原文摘要 · Abstract (English)
Large Language Models (LLMs) encounter significant performance bottlenecks in long-sequence tasks due to the computational complexity and memory overhead inherent in the self-attention mechanism. To address these challenges, we introduce \textsc{AllMem}, a novel and efficient hybrid architecture that integrates Sliding Window Attention (SWA) with non-linear Test-Time Training (TTT) memory networks. \textsc{AllMem} enables models to effectively scale to ultra-long contexts while mitigating catastrophic forgetting. This approach not only overcomes the representation constraints typical of linear memory models but also significantly reduces the computational and memory footprint during long-sequence inference. Furthermore, we implement a Memory-Efficient Fine-Tuning strategy to replace standard attention layers in pre-trained models with memory-augmented sliding window layers. This framework facilitates the efficient transformation of any off-the-shelf pre-trained LLM into an \textsc{AllMem}-based architecture. Empirical evaluations confirm that our 4k window model achieves near-lossless performance on 37k LongBench with a marginal 0.83 drop compared to full attention. Furthermore, on InfiniteBench at a 128k context, our 8k window variant outperforms full attention, which validates the effectiveness of our parameterized memory in mitigating noise and maintaining robust long-range modeling without the prohibitive costs of global attention.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。