Artificial Hippocampus Networks for Efficient Long-Context Modeling

作者/机构: Yunhao Fang∗, Weihao Yu∗,†, Shu Zhong, Qinghao Ye, Xuehan Xiong¶, Lai Wei (ByteDance Seed)

A1 主要贡献

核心问题: 长序列建模面临一个根本性的权衡：一方面是类RNN模型中压缩性固定大小内存的高效率，另一方面是基于注意力机制的Transformer中无损增长内存的高保真度。Transformer的KV缓存虽然能无损保留所有信息，但其大小随序列长度线性增长，计算成本呈二次方增长，这在处理超长序列时成为巨大挑战。

研究目标: 本文旨在提出一种新的人工神经网络内存框架，以结合上述两种内存范式的优点，实现高效的长上下文建模。该框架的目标是既能保留近期上下文的精确信息，又能以紧凑的方式总结历史信息，从而在大幅降低计算和内存需求的同时，保持甚至超越全注意力模型的性能。

创新点/主要贡献:
本文的主要贡献有两方面：
1. 提出人工海马网络（AHN）概念: 受到认知科学中多重存储模型（Multi-Store Model, MSM）的启发，本文引入了人工海馬網絡（AHN）的概念。该框架将Transformer的KV缓存的滑动窗口作为无损的短期记忆，同时使用一个可学习的AHN模块，循环地将滑出窗口的信息压缩成一个固定大小的长期记忆状态。这种方法使得模型能够同时利用两种记忆来进行高效的长上下文建模。

图1
图1

图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。

  1. AHN的实例化与高效训练: 为了验证AHN框架的有效性，本文使用现代类RNN架构，如Mamba2、DeltaNet (DN) 和 GatedDeltaNet (GDN)，将AHN实例化为AHN-Mamba2、AHN-DN和AHN-GDN。此外，本文提出了一种高效的自蒸馏训练方案，在冻结基础大语言模型权重的情况下，仅训练AHN模块的参数。实验结果表明，这些实例在长序列基准测试中显著提高了模型的效率，同时取得了与全注意力模型相当甚至更优的性能。

A3 背景知识/关键Observation/设计原则

2.1 神经网络中的记忆

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缓存，以应对这一挑战。

2.2 记忆管理

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。

A2 方法细节

3.1 预备知识

自注意力机制: 大多数现代自回归大语言模型基于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）矩阵：

公式1
公式1

其中 $W_Q$, $W_K$ 和 $W_V$ 是可训练的权重矩阵。然后，注意力输出计算为值向量的加权和：
公式2
公式2

其中 $M \in \mathbb{R}^{L \times L}$ 是因果掩码，定义为如果 $j \le i$，则 $M_{ij} = 1$，否则 $M_{ij} = 0$。

3.2 人工海马网络

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})$ 来更新压缩记忆。这个循环记忆更新定义为：

公式3
公式3

其中，$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$ 从压缩记忆和无损记忆中获取信息以产生输出：

公式4
公式4

带有AHN的整体模型机制如图2a所示。此外，附录中的图6展示了带有注意力汇（attention sinks）【【91, Efficient streaming language models with attention sinks, 2024】】的AHN示意图。

3.3 实例化

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
图2

图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】】:

公式5
公式5

与压缩所有过去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)$进一步调制，然后通过一个线性投影进行变换：

公式6
公式6

与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和注意力机制的输出相加：

公式7
公式7

复杂度分析: 表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、归一化和矩阵元素求和。灰色项与其他项相比可进一步省略。

表1
表1

公式8
公式8

3.4 训练框架

自蒸馏训练: 虽然带有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所示。

A4 实验环境

  • 模型架构:

    • 基础模型：开源的Qwen2.5-Instruct系列（3B, 7B, 14B）【【93, Qwen2. 5 technical report, 2024】】。
    • AHN模块：使用三种现代循环模型实现：Mamba2【【18, Transformers are ssms: generalized models and efficient algorithms through structured state space duality, 2024】】、DeltaNet【【70, Linear transformers are secretly fast weight programmers, 2021】, 【97, Parallelizing linear transformers with the delta rule over sequence length, 2024】】和GatedDeltaNet【【96, Gated linear attention transformers with hardware-efficient training, 2024】】。

  • 数据集:

    • 训练数据：ChatQA2数据集【【92, ChatQA 2: Bridging the gap to proprietary LLMs in long context and RAG capabilities, 2025】】，一个包含多样化长上下文任务的开源集合。
    • 评估数据：全面的长上下文基准测试套件，包括LongBench【【7, Longbench: A bilingual, multitask benchmark for long context understanding, 2024】】、InfiniteBench【【105, ınftybench: Extending long context evaluation beyond 100k tokens, 2024】】和LV-Eval【【103, Lv-eval: A balanced long-context benchmark with 5 length levels up to 256k, 2024】】，以及一个来自PG19【【68, Compressive transformers for long-range sequence modelling, 2020】】的说明性示例。

  • 硬件配置: 论文未明确提及具体的GPU型号、数量或CPU等硬件配置。

  • 软件配置:

    • 代码实现：所有AHN实例均在PyTorch【【63, Pytorch: An imperative style, high-performance deep learning library, 2019】】中实现。
    • 依赖库：基于LLaMA-Factory【【108, Llamafactory: Unified efficient fine-tuning of 100+ language models, 2024】】和Flash Linear Attention【【95, Fla: A triton-based library for hardware-efficient implementations of linear attention mechanism, 2024】】。
    • 训练框架：采用自蒸馏损失，冻结基础LLM，仅训练新初始化的AHN模块。
    • 优化器：使用AdamW【【53, Decoupled weight decay regularization, 2019】】优化器，学习率为1e-4，前10%的步骤进行线性预热，然后余弦衰减。
    • 训练设置：所有模型在ChatQA2数据集上训练一个epoch，全局批量大小为128。

A4 实验结果

4.2 一个说明性示例

  • 实验内容: 在一个来自PG19的57K token长篇文章上，对比了基础的Qwen-3B-Instruct模型与其AHN-GDN增强版本。
  • 实验结果与分析:
    • 效率: AHN增强模型显著降低了计算复杂度和内存占用（图3a和3b）。
    • 性能: 标准Qwen模型在超过32K token上下文窗口后，困惑度急剧上升；而AHN-GDN增强模型则保持了持续低的困惑度（图3c）。
    • 内存: 基础模型在使用FlashAttention时内存使用量线性增长，而AHN-GDN则使CUDA内存使用量几乎保持恒定（图3d），显示了其处理长上下文序列的有效性。

图3
图3

图3 Qwen2.5-3B-Instruct模型的复杂性分析和模型困惑度，带与不带AHN。仅当序列长度超过窗口大小（本例中为32K）时激活AHN。(a) 带AHN的模型在序列长度方面具有线性计算复杂性。(b) 带AHN的模型保持一致的内存缓存大小。(c) PG19测试集第一本书（57K token）的困惑度结果。Qwen-3B-Instruct在其预训练上下文长度之外性能下降，而AHN增强模型保持持续低的困惑度。(d) 同一样例下的峰值GPU内存。

4.3 长上下文基准测试

  • 超长上下文评估 (LV-Eval & InfiniteBench)
    • 实验内容: 在LV-Eval和InfiniteBench的128k长度子集上，将AHN增强模型与全注意力、带注意力汇的滑动窗口注意力（SWA）以及压缩Transformer（CT）进行比较。滑动窗口方法（SWA和AHN）使用32k token的无损内存。
    • 实验结果与分析 (表2):
      • AHN增强模型在几乎所有任务上都一致优于SWA基准。
      • 令人瞩目的是，它们的性能也超过了全注意力模型，证明了压缩记忆机制的有效性，同时大幅节省了计算和内存。

表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】】和最大或平均池化作为压缩函数实现。

表2
表2

  • 长上下文评估 (LongBench)
    • 实验内容: 在LongBench上选择了6个平均序列长度超过8192 token的任务进行评估，所有方法的无损内存预算固定为8192 token。
    • 实验结果与分析 (表3):
      • AHN增强模型再次取得了比所有基准（SWA, CT）都更优的准确率。
      • 结果表明，循环隐藏状态有效地捕捉和利用了历史信息，从而在各种场景中提升了性能。

表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】】和最大或平均池化作为压缩函数实现。

表3
表3

4.4 消融研究

  • 训练目标：自蒸馏 vs. 下一词元预测

    • 实验内容: 对比使用KL散度的自蒸馏和使用交叉熵（CE）损失的下一词元预测两种训练目标。
    • 实验结果与分析 (表4):
      • 使用CE损失导致LongBench上的性能显著下降。
      • 原因推测：CE提供的学习信号稀疏，而自蒸馏通过教师模型的完整输出分布提供了更密集的指导，迫使AHN学习更具泛化性的上下文表示。

  • 随机化 vs. 固定窗口

    • 实验内容: 对比使用随机化滑动窗口大小和固定窗口大小进行训练。在推理时测试模型对不同窗口大小的泛化能力。
    • 实验结果与分析 (图4):
      • 使用随机化窗口训练的模型泛化能力更强，对特定配置的过拟合程度较低。
      • 在896到8064 token的不同滑动窗口大小配置下，AHN增强模型均保持了强大的性能。

图4
图4

图4 AHN模块在LongBench上展示了强大的上下文泛化能力。

表4 AHN训练设计选择的消融实验。我们消融了两个因素：（1）训练目标，比较自蒸馏（KL损失）与下一词元预测（无全注意力教师模型，CE损失），以及（2）随机化与固定滑动窗口配置。所有实验均基于带有AHN-GDN的Qwen2.5-7B-Instruct。

表4
表4

4.5 通过梯度可视化探测AHN

  • 实验内容: 通过可视化自蒸馏损失相对于窗口外token嵌入的梯度，来理解AHN如何有效地压缩和利用窗口外信息。梯度幅值小意味着该token的信息已被AHN的压缩记忆很好地捕捉。
  • 实验结果与分析 (图5):
    • 在一个数学问题的例子中，AHN倾向于保留数学符号和数字的信息（梯度小，绿色区域），而忽略不太关键的词元，如代词和特殊符号（梯度大，红色区域）。
    • 这表明AHN是一个有针对性的压缩模块，能有效识别并存储关键信息。

图5
图5

图5 绿色区域标记了L2梯度幅值低的token，表示它们被AHN优先选择存储在压缩记忆中；红色表示相反。

A5 结论

本文介绍了一种名为人工海马网络（AHN）的新型轻量级架构组件，旨在增强Transformer模型以进行高效的长序列处理。AHN通过维护一个滑动窗口的KV缓存作为无损记忆，同时将窗口外的信息转化为一个固定大小的压缩记忆，解决了标准Transformer的效率瓶颈。这种方法使得AHN增强模型在长序列上能够实现每token恒定的内存和计算复杂度。实验证明，AHN可以显著减少内存缓存大小和计算量，同时在长上下文基准测试中保持有竞争力的性能。

局限性与未来工作
* 局限性：AHN的固定大小压缩记忆不可避免地会带来信息损失，可能在需要精确召回的任务上影响性能。此外，由于本研究采用参数高效的自蒸馏设置，性能上限受限于基础模型的能力。
* 未来工作：未来的研究可以探索更强的召回机制和全参数训练，以进一步释放AHN的潜力。在应用场景方面，AHN框架为信息稀疏或资源受限的长上下文领域开辟了机会，例如终身学习、流式视频处理以及在边缘设备上的部署。

A6 附录

A 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实例，压缩记忆的更新规则是：

公式9
公式9

至于AHN-DN，其更新规则可以表示为：
公式10
公式10

AHN-Mamba2和AHN-DN的输出规则与AHN-GDN相同，如公式6所示。

带注意力汇的AHN: 我们还提供了带有注意力汇（attention sinks）【【91, Efficient streaming language models with attention sinks, 2024】】的AHN增强网络示意图，如图6所示。

图6
图6

图6 带有AHN的模型示意图。在此例中，注意力汇的数量为2，滑动窗口长度为3。当输入序列长度小于或等于注意力汇和窗口长度之和时，模型操作与标准Transformer完全相同。对于更长的序列，AHN持续将窗口外的token压缩为紧凑的记忆表示。模型随后利用注意力汇和滑动窗口内的无损信息，以及压缩记忆来生成下一个token。

B 额外的基准测试结果

长上下文场景的进一步检验: 本节进一步检验了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个滑动窗口。

表5
表5

表6 LV-Eval 128k子集中所有21个任务的完整结果。所有基于滑动窗口的方法都使用32768个token的无损内存，包括128个注意力汇和32640个滑动窗口。

表6
表6