arXiv:2510.07318cs.CLcs.AI2025-10被引 9

用仿海马体网络实现高效长序列建模,兼顾记忆精度与计算效率。

Artificial Hippocampus Networks for Efficient Long-Context Modeling

  • 模仿认知科学的多存储模型,用滑动窗口+可学习压缩模块处理长序列
  • 在128k序列上提升评测得分至5.88,同时减少40.5%算力和74%显存占用
  • 适合需要长上下文但受限于算力的大型语言模型部署场景

长序列建模面临传统RNN类模型固定大小内存的效率与基于注意力的Transformer无损增长内存之间的根本权衡。受认知科学中多存储模型启发,我们提出一种人工神经网络记忆框架:保留Transformer KV缓存的滑动窗口作为无损短期记忆,同时引入可学习的‘人工海马体网络’(AHN)将窗口外信息递归压缩为固定大小的紧凑长期记忆。为验证该框架,我们使用现代RNN类架构(如Mamba2、DeltaNet、GatedDeltaNet)实例化AHN,并用于增强开源大模型。提出一种高效的自蒸馏训练方法,仅优化AHN参数而冻结基础模型全部参数。推理时,默认设置32k的大滑动窗口,仅当序列长度超过此阈值时激活AHN,有效缓解注意力在该尺度下的二次复杂度问题。在LV-Eval与InfiniteBench等长上下文基准上的实验表明,经AHN增强的模型持续优于滑动窗口基线,性能接近甚至超越全注意力模型,同时大幅降低计算与内存开销。例如,对Qwen2.5-3B-Instruct增广后,推理FLOPs减少40.5%,缓存内存下降74.0%,在128k序列长度下平均得分从4.41提升至5.88。代码已开源:https://github.com/ByteDance-Seed/AHN。

原文摘要 · Abstract (English)

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.

长序列建模海马体网络记忆机制大模型优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。