人工海马网络用于高效的长上下文建模

于红博

摘要: 长序列建模面临着一个基本的折衷，即在类似RNN的模型中压缩固定大小内存的效率与基于注意力的Transformer中无损增长内存的保真度之间的折衷。受认知科学中的多存储模型的启发，我们引入了一种人工神经网络的记忆框架。我们的方法将Transformer的KV缓存维持为无损短期内存的滑动窗口，同时一个可学习的模块，称为人工海马网络（AHN），循环地将窗口外的信息压缩成固定大小的紧凑长期记忆。为了验证这一框架，我们使用现代RNN-like架构实例化AHNs，包括Mamba2、DeltaNet和Gated DeltaNet。对长上下文基准LV-Eval和InfiniteBench的大量实验表明，AHN增强模型始终优于滑动窗口基线，并且实现了与全注意力模型相媲美甚至更优的性能，同时大幅减少了计算和内存需求。例如，用AHNs增强Qwen2.5-3B-Instruct将推理FLOPs减少了40.5％，内存缓存减少了74.0％，同时将其在LV-Eval（128k序列长度）上的平均分数从4.41提高到5.88。代码可在以下链接找到：https://github.com/ByteDance-Seed/AHN。

更新时间: 2025-10-08 17:59:55

领域: cs.CL,cs.AI,cs.LG

下载: http://arxiv.org/abs/2510.07318v1