作者/机构: Yunhao Fang∗, Weihao Yu∗,†, Shu Zhong, Qinghao Ye, Xuehan Xiong¶, Lai Wei (ByteDance Seed)
核心问题: 长序列建模面临一个根本性的权衡:一方面是类RNN模型中压缩性固定大小内存的高效率,另一方面是基于注意力机制的Transformer中无损增长内存的高保真度。Transformer的KV缓存虽然能无损保留所有信息,但其大小随序列长度线性增长,计算成本呈二次方增长,这在处理超长序列时成为巨大挑战。
研究目标: 本文旨在提出一种新的人工神经网络内存框架,以结合上述两种内存范式的优点,实现高效的长上下文建模。该框架的目标是既能保留近期上下文的精确信息,又能以紧凑的方式总结历史信息,从而在大幅降低计算和内存需求的同时,保持甚至超越全注意力模型的性能。
创新点/主要贡献:
本文的主要贡献有两方面:
1. 提出人工海马网络(AHN)概念: 受到认知科学中多重存储模型(Multi-Store Model, MSM)的启发,本文引入了人工海馬網絡(AHN)的概念。该框架将Transformer的KV缓存的滑动窗口作为无损的短期记忆,同时使用一个可学习的AHN模块,循环地将滑出窗口的信息压缩成一个固定大小的长期记忆状态。这种方法使得模型能够同时利用两种记忆来进行高效的长上下文建模。
图1 (a) 人工海马网络(AHN)将无损记忆转化为固定大小的压缩表示,以实现高效的长上下文建模。无损记忆(如注意力机制的KV缓存)存储精确的输入信息,但随序列长度增长,导致长序列处理成本高昂。相比之下,压缩记忆(如RNN的隐藏状态)为每个输入token维持恒定的缓存大小和计算成本,但不可避免地会丢失细节。在我们的框架中,滑动窗口注意力将近期的精确上下文作为无损短期记忆,而AHN则循环地将窗口外的信息压缩成固定大小的状态作为压缩长期记忆。这使得模型能够高效处理长序列,同时保留精确的短期信息和历史的紧凑摘要。(b) 在长上下文基准LV-Eval(128k序列长度)上,为Qwen2.5-3B-Instruct增加AHN(参数增加0.4%)可将FLOPs减少40.5%,内存缓存减少74.0%,同时平均分从4.41提高到5.88。
RNNs作为记忆机制的早期尝试: 传统的全连接神经网络缺乏跨时间步维持信息的能力,限制了它们在序列任务中的应用。为了解决这一问题,循环神经网络(RNNs)【【23, Finding structure in time, 1990】, 【35, Neural networks and physical systems with emergent collective computational abilities, 1982】, 【36, Neurons with graded response have collective computational properties like those of two-state neurons, 1984】】被提出,通过在每个时间步更新隐藏状态来持久化信息。然而,普通RNN存在梯度消失和爆炸的问题,难以捕捉长期依赖【【10, Learning long-term dependencies with gradient descent is difficult, 1994】】。为了缓解这些问题,更先进的架构如长短期记忆(LSTM)网络【【33, Long short-term memory, 1997】】和门控循环单元(GRU)【【14, Learning phrase representations using rnn encoder-decoder for statistical machine translation, 2014】】被提出,它们通过门控机制调节信息流,从而更有效地学习长期依赖。由于这类RNN模型为每个输入token维持固定大小的记忆和一致的记忆更新成本,因此它们在处理长序列时非常高效。因此,本文提出的AHN在RNN范式内进行设计,以继承这一优势。
记忆增强网络与注意力机制的演进: 除了基于RNN的架构,研究者还开发了记忆增强神经网络以进一步提升模型的记忆容量。例如,神经图灵机(NTM)【【27, Neural turing machines, 2014】】和可微神经计算机(DNC)【【28, Hybrid computing using a neural network with dynamic external memory, 2016】】引入了外部记忆模块,网络可以从中读写信息。在过去十年中,注意力机制【【6, Neural machine translation by jointly learning to align and translate, 2015】】彻底改变了神经网络处理记忆的方式。完全依赖自注意力机制的Transformer架构【【80, Attention is all you need, 2017】】允许直接访问序列中所有先前的状态,提供了一种无损且可扩展的记忆形式,这在多个领域带来了显著进步【【20, Bert: Pre-training of deep bidirectional transformers for language understanding, 2019】, 【22, An image is worth 16x16 words: Transformers for image recognition at scale, 2021】, 【66, Improving language understanding by generative pre-training, 2018】, 【67, Language models are unsupervised multitask learners, 2019】】,并催生了上下文学习(In-Context Learning)【【11, Language models are few-shot learners, 2020】】和思维链(Chain-of-Thought)推理【【84, Chain-of-thought prompting elicits reasoning in large language models, 2022】】等新技术范式。然而,处理长序列会加剧注意力机制的二次计算复杂度成本【【13, Generating long sequences with sparse transformers, 2019】】。本文提出的AHN通过采用类RNN网络来压缩历史KV缓存,以应对这一挑战。
Transformer中KV缓存的挑战: 类RNN模型【【9, xLSTM: Extended long short-term memory, 2024】, 【14, Learning phrase representations using rnn encoder-decoder for statistical machine translation, 2014】, 【18, Transformers are ssms: generalized models and efficient algorithms through structured state space duality, 2024】, 【23, Finding structure in time, 1990】, 【29, Mamba: Linear-time sequence modeling with selective state spaces, 2024】, 【33, Long short-term memory, 1997】, 【41, Transformers are rnns: Fast autoregressive transformers with linear attention, 2020】, 【64, Rwkv: Reinventing rnns for the transformer era, 2023】, 【76, Retentive network: A successor to transformer for large language models, 2023】, 【96, Gated linear attention transformers with hardware-efficient training, 2024】, 【97, Parallelizing linear transformers with the delta rule over sequence length, 2024】, 【98, Gated delta networks: Improving mamba2 with delta rule, 2025】】通过一个固定大小的隐藏状态来维持记忆,因此记忆缓存不是主要问题。相反,Transformer为输入序列中的每个token存储键值(KV)对,导致KV缓存随序列长度线性增长,消耗大量内存,这对处理长序列构成了重大挑战。为缓解此问题,研究者提出了多种方法【【45, A survey on large language model acceleration based on kv cache management, 2024】】,包括KV缓存选择【【1, Keyformer: Kv cache reduction through key tokens selection for efficient generative inference, 2024】, 【26, Model tells you what to discard: Adaptive KV cache compression for LLMs, 2024】, 【31, Lm-infinite: Zero-shot extreme length generalization for large language models, 2024】, 【47, Snapkv: Llm knows what you are looking for before generation, 2024】, 【52, Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time, 2023】, 【79, Quest: Query-aware sparsity for efficient long-context llm inference, 2024】, 【88, InfLLM: Training-free long-context extrapolation for LLMs with an efficient context memory, 2024】, 【91, Efficient streaming language models with attention sinks, 2024】, 【107, H2o: Heavy-hitter oracle for efficient generative inference of large language models, 2023】】、预算分配【【12, Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling, 2024】, 【24, Ada-kv: Optimizing kv cache eviction by adaptive budget allocation for efficient llm inference, 2024】, 【90, Duoattention: Efficient long-context llm inference with retrieval and streaming heads, 2024】, 【94, Pyramidinfer: Pyramid kv cache compression for high-throughput llm inference, 2024】】、合并【【51, Minicache: Kv cache compression in depth dimension for large language models, 2024】, 【58, Dynamic memory compression: retrofitting llms for accelerated inference, 2024】, 【81, Look-m: Look-once optimization in kv cache for efficient multimodal long-context inference, 2024】, 【83, Model tells you where to merge: Adaptive kv cache merging for llms on long-context tasks, 2024】】、量化【【34, Kvquant: Towards 10 million context length llm inference with kv cache quantization, 2024】, 【49, Awq: Activation-aware weight quantization for on-device llm compression and acceleration, 2024】, 【72, Omniquant: Omnidirectionally calibrated quantization for large language models, 2024】, 【74, Flexgen: High-throughput generative inference of large language models with a single gpu, 2023】, 【89, Smoothquant: Accurate and efficient post-training quantization for large language models, 2023】, 【99, Zeroquant: Efficient and affordable post-training quantization for large-scale transformers, 2022】】、低秩分解【【21, Get more with less: Synthesizing recurrence with kv cache compression for efficient llm inference, 2024】, 【101, Effectively compress kv heads for llm, 2024】】、外部记忆【【62, Memgpt: Towards llms as operating systems, 2023】, 【82, M+: Extending memoryLLM with scalable long-term memory, 2025】】以及神经架构设计【【2, Gqa: Training generalized multi-query transformer models from multi-head checkpoints, 2023】, 【38, Transformer quality in linear time, 2022】, 【50, Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model, 2024】, 【57, Leave no context behind: Efficient infinite context transformers with infini-attention, 2024】, 【73, Fast transformer decoding: One write-head is all you need, 2019】, 【77, You only cache once: Decoder-decoder architectures for language models, 2024】, 【87, Memorizing transformers, 2022】, 【100, Long-context language modeling with parallel context encoding, 2024】】。
与现有记忆管理方法的对比: 在众多策略中,一个直接的方法是使用滑动窗口注意力【【80, Attention is all you need, 2017】】,但这种方法会丢弃窗口外的KV对,从而丢失长程上下文。稀疏Transformer【【13, Generating long sequences with sparse transformers, 2019】】通过保留特定模式位置的KV对来捕捉长程依赖,但仍然会丢弃部分KV缓存。Transformer-XL【【17, Transformer-xl: Attentive language models beyond a fixed-length context, 2019】】通过缓存上一段的隐藏状态作为先进先出(FIFO)记忆,引入了段级循环机制。Compressive Transformer【【68, Compressive transformers for long-range sequence modelling, 2020】】对此进行了扩展,将更早的记忆压缩到次级FIFO记忆中,但当插槽满时仍会丢弃记忆。相比之下,AHN采用类RNN范式,持续将滑动窗口外的KV对压缩到一个终身的压缩记忆中,而不是直接丢弃它们【【48, Jamba: A hybrid transformer-mamba language model, 2024】, 【57, Leave no context behind: Efficient infinite context transformers with infini-attention, 2024】, 【69, Samba: Simple hybrid state space models for efficient unlimited context language modeling, 2025】】。AHN(如AHN-GDN【【98, Gated delta networks: Improving mamba2 with delta rule, 2025】】)还可以动态控制记忆衰减【【18, Transformers are ssms: generalized models and efficient algorithms through structured state space duality, 2024】, 【70, Linear transformers are secretly fast weight programmers, 2021】, 【97, Parallelizing linear transformers with the delta rule over sequence length, 2024】, 【98, Gated delta networks: Improving mamba2 with delta rule, 2025】】。最近的研究将RNN和注意力机制结合在交错的层中【【18, Transformers are ssms: generalized models and efficient algorithms through structured state space duality, 2024】, 【44, Minimax-01: Scaling foundation models with lightning attention, 2025】, 【48, Jamba: A hybrid transformer-mamba language model, 2024】, 【69, Samba: Simple hybrid state space models for efficient unlimited context language modeling, 2025】, 【98, Gated delta networks: Improving mamba2 with delta rule, 2025】】或单个层内【【46, Transmamba: Flexibly switching between transformer and mamba, 2025】, 【57, Leave no context behind: Efficient infinite context transformers with infini-attention, 2024】】。与这些工作不同,本文将压缩模块抽象为AHN概念,从而产生一个更通用的记忆框架。我们采用滑动窗口注意力机制,在token离开窗口时激活AHN。此外,我们引入了一个简单的自蒸馏方案来高效地训练AHN。
自注意力机制: 大多数现代自回归大语言模型基于Transformer架构【【80, Attention is all you need, 2017】】,该架构使用自注意力作为token混合的核心机制。对于一个包含L个token的输入序列 $X = (x_1, x_2, ..., x_L) \in \mathbb{R}^{L \times D}$,自注意力首先通过学习到的线性变换将token投影为查询(Q)、键(K)和值(V)矩阵:
其中 $W_Q$, $W_K$ 和 $W_V$ 是可训练的权重矩阵。然后,注意力输出计算为值向量的加权和:
其中 $M \in \mathbb{R}^{L \times L}$ 是因果掩码,定义为如果 $j \le i$,则 $M_{ij} = 1$,否则 $M_{ij} = 0$。
AHN的定义: 受到多重存储模型(MSM)【【4, Human memory: A proposed system and its control processes, 1968】】和海马体【【71, Loss of recent memory after bilateral hippocampal lesions, 1957】】将无损短期记忆巩固为紧凑长期表示的启发,我们引入了人工海马网络(AHN),通过将历史信息压缩到一个固定大小的循环状态来模拟这一生物功能。AHN与一个大小为W的滑动注意力窗口协同工作。对于在步骤 $t > W$ 处的token,AHN通过处理刚刚离开滑动窗口的键值(KV)对 $(k_{t-W}, v_{t-W})$ 来更新压缩记忆。这个循环记忆更新定义为:
其中,$h_{t-W}$ 是更新后的压缩记忆,总结了截至并包括位置 $t-W$ 的上下文。$h_{t-W}$ 可以是一个向量或矩阵。由于公式3的循环形式,AHN可以用类RNN架构实现,从而实现对长上下文历史的可学习和高效压缩。
与无损记忆的集成: 在预定义的滑动窗口内,应用标准的因果注意力来保留近期token的无损记忆。一旦输入序列长度超过窗口大小,AHN被激活,将窗口外的KV对,即 $(k_{t-W}, v_{t-W})$,压缩到一个固定大小的压缩记忆 $h_{t-W}$ 中。压缩后,窗口外的原始KV对可以被安全地丢弃,只保留窗口内的KV缓存 $\{(k_i, v_i)\}_{i=t-W+1}^t$。最后,当前的查询 $q_t$ 从压缩记忆和无损记忆中获取信息以产生输出:
带有AHN的整体模型机制如图2a所示。此外,附录中的图6展示了带有注意力汇(attention sinks)【【91, Efficient streaming language models with attention sinks, 2024】】的AHN示意图。
AHN的实现: 如上所述,AHN可以使用类RNN架构进行实例化。在我们的实验中,我们专注于现代线性循环模型,因为它们具有高效的并行训练能力。具体来说,我们利用了三种架构,包括Mamba2【【18, Transformers are ssms: generalized models and efficient algorithms through structured state space duality, 2024】】、DeltaNet (DN)【【70, Linear transformers are secretly fast weight programmers, 2021】, 【97, Parallelizing linear transformers with the delta rule over sequence length, 2024】】及其增强版GatedDeltaNet (GDN)【【96, Gated linear attention transformers with hardware-efficient training, 2024】】,分别将AHN实例化为AHN-Mamba2、AHN-DN和AHN-GDN。下面,我们以AHN-GDN的实现为例进行介绍,另外两种AHN实例的描述见附录A。
图2 (a) 带有AHN的模型示意图。在此例中,滑动窗口长度为3。当输入序列长度小于或等于窗口长度时,模型操作与标准Transformer完全相同。对于更长的序列,AHN持续将窗口外的token压缩为紧凑的记忆表示。模型随后利用窗口内的无损信息和压缩记忆来生成下一个token。(b) 基于开源LLM的AHN自蒸馏训练框架。训练期间,基础LLM的权重被冻结,仅训练AHN的参数。
AHN-GDN的记忆更新: 具体来说,AHN-GDN通过门控delta规则更新记忆【【70, Linear transformers are secretly fast weight programmers, 2021】, 【96, Gated linear attention transformers with hardware-efficient training, 2024】, 【97, Parallelizing linear transformers with the delta rule over sequence length, 2024】】:
与压缩所有过去token的GatedDeltaNet【【98, Gated delta networks: Improving mamba2 with delta rule, 2025】】不同,AHN-GDN仅压缩滑动窗口之外的token。
AHN-GDN的输出计算: 对于每个位置t,从$x_t$导出的查询$q_t$用于访问压缩记忆$h_{t-W}$。输出通过一个门函数$\gamma(x_t)$进一步调制,然后通过一个线性投影进行变换:
与GatedDeltaNet【【96, Gated linear attention transformers with hardware-efficient training, 2024】】不同,$\gamma(x_t)$的输出是跨头通道共享的标量,并且输出线性层按头进行分组【【39, Convbert: Improving bert with span-based dynamic convolution, 2020】, 【42, Imagenet classification with deep convolutional neural networks, 2012】】,其可学习权重为 $W_o \in \mathbb{R}^{H \times H}$ (H表示头维度)。
最终输出: 最后,我们简单地将来自AHN和注意力机制的输出相加:
复杂度分析: 表1总结了带与不带AHN-GDN的注意力token混合器的计算和内存复杂度,图3比较了Qwen2.5-3B带与不带AHN-GDN的复杂度。如图所示,集成AHN在内存使用和FLOPs方面都显著优于标准全注意力。特别是,AHN-GDN将注意力的计算复杂度降低到序列长度的线性级别,同时保持内存缓存大小恒定。相比之下,普通的全注意力会产生二次方的计算成本和随序列长度线性增长的内存使用。
表1 带与不带AHN-GDN的因果注意力复杂度。此处L:输入序列长度;D:隐藏维度;Nq/Nkv:查询/键值头的数量;H:头维度;W:滑动窗口大小。仅当L>W时激活AHN。FLOPs仅计算矩阵乘法;省略了softmax、归一化和矩阵元素求和。灰色项与其他项相比可进一步省略。
自蒸馏训练: 虽然带有AHN的模型可以从头开始训练,但我们采用了一种计算效率更高的方法,即自蒸馏【【32, Distilling the knowledge in a neural network, 2015】, 【104, Be your own teacher: Improve the performance of convolutional neural networks via self distillation, 2019】, 【106, Deep mutual learning, 2018】】。这使我们能够利用强大的预训练模型。我们的训练框架使用一个开源的LLM(例如,Qwen【【93, Qwen2. 5 technical report, 2024】】)作为教师模型,其输出概率表示为$p'$。学生模型是同一个LLM,但我们修改了其注意力机制,使其在每一层都在一个有限的滑动窗口感受野上操作。这些窗口注意力层随后被AHN增强。学生的输出概率表示为$p$。我们通过最小化Kullback-Leibler(KL)散度来训练学生模仿教师的输出分布:
训练效率: 为了最大化效率,在训练期间基础模型的权重被冻结,只有AHN的参数被优化。该框架如图2b所示。
模型架构:
数据集:
硬件配置: 论文未明确提及具体的GPU型号、数量或CPU等硬件配置。
软件配置:
图3 Qwen2.5-3B-Instruct模型的复杂性分析和模型困惑度,带与不带AHN。仅当序列长度超过窗口大小(本例中为32K)时激活AHN。(a) 带AHN的模型在序列长度方面具有线性计算复杂性。(b) 带AHN的模型保持一致的内存缓存大小。(c) PG19测试集第一本书(57K token)的困惑度结果。Qwen-3B-Instruct在其预训练上下文长度之外性能下降,而AHN增强模型保持持续低的困惑度。(d) 同一样例下的峰值GPU内存。
表2 LV-Eval和InfiniteBench的128k长度子集的性能和效率分析。混合/模型FLOPs比率衡量了token混合器或整个模型相对于全注意力基准的计算成本。除全注意力外,所有方法的注意力汇【【91, Efficient streaming language models with attention sinks, 2024】】和滑动窗口注意力(SWA)的无损内存为32k token。压缩Transformer(CT)【【68, Compressive transformers for long-range sequence modelling, 2020】】使用注意力汇【【91, Efficient streaming language models with attention sinks, 2024】】和最大或平均池化作为压缩函数实现。
表3 基于Qwen2.5的模型在六个LongBench任务上的性能(平均序列长度>8k)。对于所有方法,注意力汇【【91, Efficient streaming language models with attention sinks, 2024】】和滑动窗口注意力(SWA)的无损内存为8192 token。压缩Transformer(CT)【【68, Compressive transformers for long-range sequence modelling, 2020】】使用注意力汇【【91, Efficient streaming language models with attention sinks, 2024】】和最大或平均池化作为压缩函数实现。
训练目标:自蒸馏 vs. 下一词元预测
随机化 vs. 固定窗口
图4 AHN模块在LongBench上展示了强大的上下文泛化能力。
表4 AHN训练设计选择的消融实验。我们消融了两个因素:(1)训练目标,比较自蒸馏(KL损失)与下一词元预测(无全注意力教师模型,CE损失),以及(2)随机化与固定滑动窗口配置。所有实验均基于带有AHN-GDN的Qwen2.5-7B-Instruct。
图5 绿色区域标记了L2梯度幅值低的token,表示它们被AHN优先选择存储在压缩记忆中;红色表示相反。
本文介绍了一种名为人工海马网络(AHN)的新型轻量级架构组件,旨在增强Transformer模型以进行高效的长序列处理。AHN通过维护一个滑动窗口的KV缓存作为无损记忆,同时将窗口外的信息转化为一个固定大小的压缩记忆,解决了标准Transformer的效率瓶颈。这种方法使得AHN增强模型在长序列上能够实现每token恒定的内存和计算复杂度。实验证明,AHN可以显著减少内存缓存大小和计算量,同时在长上下文基准测试中保持有竞争力的性能。
局限性与未来工作
* 局限性:AHN的固定大小压缩记忆不可避免地会带来信息损失,可能在需要精确召回的任务上影响性能。此外,由于本研究采用参数高效的自蒸馏设置,性能上限受限于基础模型的能力。
* 未来工作:未来的研究可以探索更强的召回机制和全参数训练,以进一步释放AHN的潜力。在应用场景方面,AHN框架为信息稀疏或资源受限的长上下文领域开辟了机会,例如终身学习、流式视频处理以及在边缘设备上的部署。
AHN-Mamba2 和 AHN-DN 的更新规则: 本节描述了如何使用Mamba2【【18, Transformers are ssms: generalized models and efficient algorithms through structured state space duality, 2024】】和DeltaNet (DN)【【70, Linear transformers are secretly fast weight programmers, 2021】, 【97, Parallelizing linear transformers with the delta rule over sequence length, 2024】】来实例化AHN。对于AHN-Mamba2实例,压缩记忆的更新规则是:
至于AHN-DN,其更新规则可以表示为:
AHN-Mamba2和AHN-DN的输出规则与AHN-GDN相同,如公式6所示。
带注意力汇的AHN: 我们还提供了带有注意力汇(attention sinks)【【91, Efficient streaming language models with attention sinks, 2024】】的AHN增强网络示意图,如图6所示。
图6 带有AHN的模型示意图。在此例中,注意力汇的数量为2,滑动窗口长度为3。当输入序列长度小于或等于注意力汇和窗口长度之和时,模型操作与标准Transformer完全相同。对于更长的序列,AHN持续将窗口外的token压缩为紧凑的记忆表示。模型随后利用注意力汇和滑动窗口内的无损信息,以及压缩记忆来生成下一个token。
长上下文场景的进一步检验: 本节进一步检验了AHN在长上下文场景中的有效性,展示了额外的基准测试结果,同时也承认了由于压缩记忆的有损性质,其在精确召回任务上的固有限制。
LV-Eval[103]的完整结果: 我们展示了在128k上下文设置下所有11个LV-Eval任务的完整结果。所有模型都配置了32768个token的无损内存,包括128个token的注意力汇和32640个token的滑动窗口。见表6。
RULER[37]基准测试: RULER是一个全面的基准测试,它通过增加任务难度和额外的类别扩展了标准的海底捞针(NIAH)【【40, Needle in a haystack - pressure testing llms, 2023】】范式。我们在RULER-128k子集的所有NIAH任务上评估了一个AHN增强模型(AHN-GDN),使用Qwen2.5-7B-Instruct作为基础模型。为了公平比较,AHN-GDN和带注意力汇的滑动窗口注意力都配置了128个注意力汇和32640个token的滑动窗口。如表5所示,AHN-GDN在精确召回任务上的表现与滑动窗口注意力相当,但明显差于全注意力。这反映了有损压缩的内在权衡:虽然AHN增强模型能够实现高效的长上下文推理,但它们不可避免地在需要从压缩记忆中精确召回的任务上表现不佳。这一局限性为未来的研究提供了机会,例如开发能够在利用压缩效率的同时在无损记忆中保留关键信息的记忆管理策略。
表5 RULER-128k中高级海底捞针(NIAH)任务的性能。两种滑动窗口方法都使用128个注意力汇和32640个滑动窗口。
表6 LV-Eval 128k子集中所有21个任务的完整结果。所有基于滑动窗口的方法都使用32768个token的无损内存,包括128个注意力汇和32640个滑动窗口。