Long-sequence modeling faces a fundamental trade-off between the efficiency of compressive fixed-size memory in RNN-like models and the fidelity of lossless growing memory in attention-based Transformers. Inspired by the Multi-Store Model in cognitive science, we introduce a memory framework of artificial neural networks. Our method maintains a sliding window of the Transformer's KV cache as lossless short-term memory, while a learnable module termed Artificial Hippocampus Network (AHN) recurrently compresses out-of-window information into a fixed-size compact long-term memory. To validate this framework, we instantiate AHNs using modern RNN-like architectures, including Mamba2, DeltaNet, and GatedDeltaNet to augment open-weight LLMs. We also propose an efficient self-distillation training method where the base model's all parameters are frozen and only the parameters from AHNs are optimized. For inference, our method sets a default large sliding window size of 32k for attention, and AHNs activate only when the sequence length exceeds the 32k window, addressing the quadratic-complexity issue of attention that emerges at that scale. Extensive experiments on long-context benchmarks LV-Eval and InfiniteBench demonstrate that AHN-augmented models consistently outperform sliding window baselines and achieve performance comparable or even superior to full-attention models, while substantially reducing computational and memory requirements. For instance, augmenting the Qwen2.5-3B-Instruct with AHNs reduces inference FLOPs by 40.5% and memory cache by 74.0%, while improving its average score on LV-Eval (128k sequence length) from 4.41 to 5.88. Code is available at: https://github.com/ByteDance-Seed/AHN.


翻译:长序列建模面临一个根本性的权衡:一方面,类似RNN的模型采用压缩的固定大小记忆以实现高效性;另一方面,基于注意力的Transformer模型则采用无损增长的记忆以保证保真度。受认知科学中多存储模型的启发,我们提出了一种人工神经网络记忆框架。该方法将Transformer的KV缓存作为一个滑动窗口,作为无损的短期记忆;同时,一个称为人工海马网络(AHN)的可学习模块,以循环方式将窗口外的信息压缩为固定大小的紧凑长期记忆。为验证该框架,我们使用包括Mamba2、DeltaNet和GatedDeltaNet在内的现代类RNN架构实例化AHN,以增强开源大语言模型(LLM)。我们还提出了一种高效的自我蒸馏训练方法,其中基础模型的所有参数被冻结,仅优化AHN的参数。在推理阶段,我们的方法为注意力机制设置默认的32k大滑动窗口,仅当序列长度超过32k窗口时,AHN才被激活,从而解决了在该规模下注意力机制二次复杂度的问题。在长上下文基准测试LV-Eval和InfiniteBench上的大量实验表明,增强AHN的模型在性能上持续优于滑动窗口基线,并达到与全注意力模型相当甚至更优的水平,同时显著降低了计算和内存需求。例如,为Qwen2.5-3B-Instruct模型增强AHN后,推理FLOPs减少了40.5%,内存缓存减少了74.0%,同时其在LV-Eval(128k序列长度)上的平均得分从4.41提升至5.88。代码发布于:https://github.com/ByteDance-Seed/AHN。

0
下载
关闭预览

相关内容

UTC: 用于视觉对话的任务间对比学习的统一Transformer
专知会员服务
14+阅读 · 2022年5月4日
【CVPR2020-旷视】DPGN:分布传播图网络的小样本学习
专知会员服务
28+阅读 · 2020年4月1日
【NeurIPS2019】图变换网络:Graph Transformer Network
专知会员服务
112+阅读 · 2019年11月25日
论文浅尝 | GEOM-GCN: Geometric Graph Convolutional Networks
开放知识图谱
14+阅读 · 2020年4月8日
图机器学习 2.2-2.4 Properties of Networks, Random Graph
图与推荐
10+阅读 · 2020年3月28日
【NeurIPS2019】图变换网络:Graph Transformer Network
MNIST入门:贝叶斯方法
Python程序员
23+阅读 · 2017年7月3日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
46+阅读 · 2015年12月31日
国家自然科学基金
8+阅读 · 2015年12月31日
国家自然科学基金
3+阅读 · 2015年12月31日
国家自然科学基金
17+阅读 · 2008年12月31日
VIP会员
相关资讯
论文浅尝 | GEOM-GCN: Geometric Graph Convolutional Networks
开放知识图谱
14+阅读 · 2020年4月8日
图机器学习 2.2-2.4 Properties of Networks, Random Graph
图与推荐
10+阅读 · 2020年3月28日
【NeurIPS2019】图变换网络:Graph Transformer Network
MNIST入门:贝叶斯方法
Python程序员
23+阅读 · 2017年7月3日
相关基金
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
46+阅读 · 2015年12月31日
国家自然科学基金
8+阅读 · 2015年12月31日
国家自然科学基金
3+阅读 · 2015年12月31日
国家自然科学基金
17+阅读 · 2008年12月31日
Top
微信扫码咨询专知VIP会员