2510.07318_AHN: Artificial Hippocampus Networks for Efficient Long-Context Modeling

From Moonlight

三句摘要

  1. 🧠 该论文引入了人工海马网络(AHNs),旨在解决长序列建模中效率与保真度的权衡问题,通过将Transformer的无损短期KV记忆转化为固定大小的压缩长期表示。

  2. 💾 AHNs结合了滑动窗口注意力机制用于保留近期上下文,并利用一个类RNN的模块循环压缩窗口外信息,从而实现了每个token恒定的计算成本和内存缓存大小。

  3. 🚀 在LV-Eval和InfiniteBench等长上下文基准测试上,AHN增强模型显著降低了FLOPs和内存占用,同时性能优于滑动窗口基线,并可媲美甚至超越全注意力模型。

关键词

  • Artificial Hippocampus Networks (AHNs): 是一种用于高效长上下文建模的创新型神经网络模块。它受到认知科学中“多存储模型”(Multi-Store Model, MSM)的启发,旨在将Transformer模型中超出滑动窗口范围的“无损短时记忆”(如Attention机制的KV缓存)递归地压缩成一个固定大小的“压缩长时记忆”。这种设计允许模型在处理超长序列时,既保留近期精确的上下文信息,又能对历史信息形成紧凑的总结,从而显著降低计算和内存成本,同时保持或提升模型性能。AHN可以实例化为RNN类架构,如Mamba2、DeltaNet和GatedDeltaNet。

  • Long-Context Modeling: 指的是处理和理解非常长的序列(如长篇文章、书籍、代码或对话历史)的能力。传统的Transformer模型在处理长序列时,由于注意力机制的二次方复杂度(计算量)和KV缓存的线性增长(内存占用),会面临巨大的计算和内存挑战。长上下文建模的目标是开发能够有效且高效地处理这些超长序列的模型。

  • Lossless Memory: 指的是能够精确存储输入信息而不会丢失任何细节的记忆形式。在Transformer模型中,KV(Key-Value)缓存是典型的无损短时记忆。每个输入token都会生成一个KV对,并被添加到缓存中,使得模型在生成下一个token时可以直接访问过去所有token的信息。尽管这种记忆方式保证了信息的高保真度,但其大小会随着序列长度线性增长,导致在处理长序列时内存消耗巨大。

  • Compressed Memory: 指的是一种将信息压缩成固定大小表示的记忆形式,以实现效率。在循环神经网络(RNN)系列模型中,隐藏状态(hidden state)通常被视为一种压缩记忆。它通过在每个时间步更新来整合历史信息,从而保持恒定的内存占用和每步计算成本,适合处理长序列。然而,这种压缩过程不可避免地会导致信息丢失,尤其是在需要精确回忆长距离细节的任务中。

  • Sliding Window Attention: 是一种注意力机制的变体,它限制了注意力计算的范围,只关注输入序列中一个固定大小的“窗口”内的token。这意味着每个token只关注它之前的若干个token,而不是整个序列。这种方法大大降低了计算复杂度,从二次方降低到接近线性,但代价是模型无法直接访问窗口之外的长期历史信息。AHNs在此基础上,对滑出窗口的token进行压缩,以弥补其不足。

  • KV Cache: 是Transformer模型在自注意力计算中用于加速推理的关键组件。对于每个输入token,模型会生成一个Key(K)向量和一个Value(V)向量。这些KV对被存储起来,形成KV缓存。当模型需要生成下一个token时,查询(Query, Q)向量会与KV缓存中的所有Key向量进行匹配,计算出注意力权重,然后用这些权重对Value向量进行加权求和,从而得到当前token的表示。KV缓存使得模型无需为每个新token重新计算所有先前token的注意力,从而显著加速了生成过程。在AHNs框架中,KV缓存被用作无损短时记忆。

  • Recurrent Neural Networks (RNNs): 是一类能够处理序列数据的神经网络。RNN的核心思想是在处理序列时维护一个内部状态(hidden state),该状态会随着每个新输入的token而更新,从而将历史信息传递到下一个时间步。这种结构使得RNN在理论上能够捕捉长距离依赖关系,并且其计算和内存成本通常与序列长度呈线性关系,因此在处理长序列时具有较高的效率。AHN的设计就借鉴了RNN这种高效的、固定大小的状态更新机制。

  • Transformer: 是一种基于自注意力机制的深度学习模型架构,自2017年提出以来,已成为自然语言处理(NLP)领域的主导模型。Transformer模型摒弃了RNN的循环结构,通过自注意力机制让模型能够并行地关注输入序列中的任意两个位置之间的关系,从而在长距离依赖建模方面表现出色。其核心是“Attention is All You Need”的理念。然而,标准的Transformer在处理极长序列时,其计算量和内存需求会随着序列长度的平方增长,这限制了其应用范围。

  • Multi-Store Model (MSM): 是认知科学中关于人类记忆的理论模型,由Atkinson和Shiffrin于1968年提出。该模型认为人类记忆包含三个主要组成部分:感觉记忆(sensory memory)、短时记忆(short-term memory,STM)和长时记忆(long-term memory,LTM)。信息从感觉记忆进入短时记忆,并在编码、储存和提取过程中被加工。通过复述等过程,短时记忆中的信息可以被转移到长时记忆中,在那里可以长期储存。 the paper’s AHN framework is inspired by the hippocampus’s role in consolidating short-term memories into long-term representations, aligning with this multi-store concept.

  • Self-distillation: 是一种模型训练技术,它利用一个更强大或更复杂的“教师模型”来指导一个“学生模型”的学习。在AHNs的上下文中,作者使用一个全注意力(full attention)的LLM作为教师模型,其输出概率分布被视为“真实”的。然后,他们训练一个AHNs增强的“学生模型”,使其输出概率分布尽可能地接近教师模型的输出。通过最小化KL散度等损失函数,学生模型(即AHNs增强的模型)学会模仿教师模型的行为,从而在保持计算效率的同时,获得接近教师模型的性能。这种方法允许在不从头开始训练的情况下,有效地将AHNs整合到现有的大型语言模型中。

摘要

该论文提出了“人工海马网络”(Artificial Hippocampus Networks, AHNs)框架,旨在解决大型语言模型在处理长序列时面临的效率与信息保真度之间的核心矛盾。传统的Transformer模型依赖于“无损增长记忆”(lossless growing memory),即注意力机制中的Key-Value (KV) 缓存,虽然能保留所有token级别的精确信息,但其KV缓存大小与序列长度呈线性增长,计算成本呈二次方增长,导致长序列处理效率低下。而RNN类模型则使用“压缩固定大小记忆”(compressive fixed-size memory),能够保持恒定的每token计算成本和内存占用,但不可避免地损失信息。

受认知科学中“多存储模型”(Multi-Store Model, MSM)的启发,该研究将Transformer的滑动窗口KV缓存作为“无损短期记忆”(lossless short-term memory),而将窗口外的历史信息通过一个可学习的模块——AHN——周期性地压缩成固定大小的“压缩长期记忆”(compressed long-term memory)。当token移出滑动窗口时,AHNs会被激活,对该token对应的KV对进行压缩,更新其固定大小的压缩记忆状态。

核心方法和技术细节:

  1. 整体框架: 模型在每个时间步处理输入token \(x_t\) 并生成输出 \(y_t\)

    • 滑动窗口注意力 (Sliding Window Attention, SWA): 在预定义的滑动窗口 \((W)\) 内,模型采用标准的因果注意力机制来维持对最近token的“无损短期记忆”。这意味着窗口内的KV缓存是精确且未经压缩的。

    • 人工海马网络 (AHN): 当输入序列长度超过窗口大小时,AHNs开始工作。对于每个刚刚移出滑动窗口的KV对 \((k_{t-W}, v_{t-W})\),AHN模块会对其进行处理,并将其信息递归地压缩到固定大小的压缩记忆 \(h_{t-W}\) 中。这个过程可以表示为: $\( h_{t-W} = \text{AHN}((k_{t-W}, v_{t-W}), h_{t-W-1}) \)\( 其中 \)h_{t-W-1}$ 是前一个时间步的压缩记忆。一旦KV对被压缩,原始的窗口外KV对就可以被安全丢弃,从而控制内存增长。

    • 信息整合: 当前的查询 \(q_t\) 会同时从无损的滑动窗口记忆 \(\{(k_i, v_i)\}_{i=t-W+1}^t\) 和压缩长期记忆 \(h_{t-W}\) 中获取信息,以生成最终的输出: $\( y_t = f(h_{t-W}, \{(k_i, v_i)\}_{i=t-W+1}^t, q_t) \)$ 通过这种方式,模型既能保持对近期上下文的精确理解,又能以紧凑的形式保留历史信息的概要。

  2. AHN实例化: AHN可以由各种RNN类架构实现,以利用其高效的递归特性。论文中使用了三种现代线性循环模型进行实例化:Mamba2、DeltaNet (DN) 和 GatedDeltaNet (GDN)。以AHN-GDN为例,其记忆更新规则基于门控Delta规则: $\( h_{t-W} = \alpha(x_{t-W})(I - \beta(x_{t-W})k_{t-W}^T k_{t-W})h_{t-W-1} + \beta(x_{t-W})k_{t-W}^T v_{t-W} \)\( 其中 \)\alpha(x_{t-W})\( 和 \)\beta(x_{t-W})\( 是门控函数,取决于token \)x_{t-W}\(。与原始GatedDeltaNet压缩所有过去token不同,AHN-GDN仅压缩滑动窗口外的token。输出的整合方式为: \)\( y_{\text{AHN},t} = \gamma(x_t)q_t h_{t-W} W_o \)\( 其中 \)\gamma(x_t)\( 是一个门控函数,\)W_o\( 是可学习的输出线性投影。最终输出是AHNs输出与注意力机制输出的和: \)\( y_t = y_{\text{AHN},t} + \text{Attention}(\{(k_i, v_i)\}_{i=t-W+1}^t, q_t) \)$

  3. 计算与内存复杂度分析: AHNs显著降低了计算复杂度和内存占用。对于序列长度 \(L\)、隐藏维度 \(D\)、注意力头数 \(N_q, N_{kv}\)、头维度 \(H\) 和滑动窗口大小 \(W\),标准全注意力模型的计算复杂度为 \(O(L^2)\),内存缓存为 \(O(L)\)。而AHNs将计算复杂度降至 \(O(WL)\) (对于 \(L>W\)),内存缓存则保持为恒定的 \(O(W)\) (加上AHNs固定大小的记忆)。

  4. 训练框架: AHNs采用一种高效的“自蒸馏”(self-distillation)训练方案。将一个预训练的开源大型语言模型(如Qwen)作为“教师模型”(teacher model),其权重在训练过程中被冻结。学生模型是相同的基础LLM,但其注意力机制被修改为在每个层上操作于有限的滑动窗口,并增强AHNs。通过最小化学生模型输出概率分布 \(p\) 与教师模型输出概率分布 \(p'\) 之间的Kullback-Leibler (KL) 散度来训练学生模型: $\( l = \text{KL}(p'||p) \)$ 此框架只优化AHNs的参数,从而实现了计算效率。训练时会随机化AHNs的起始位置和滑动窗口大小,以增强模型的泛化能力。

实验结果与贡献:

  • 效率提升: AHN-augmented模型在长上下文基准测试(LV-Eval、InfiniteBench)中显著降低了FLOPs和内存缓存(例如,Qwen2.5-3B + AHN在LV-Eval上FLOPs减少40.5%,内存缓存减少74.0%)。在PG19这类超长文本上,AHN-GDN能保持恒定的GPU内存占用和稳定的困惑度,而基础模型在超出预训练上下文长度后困惑度急剧上升。

  • 性能提升: AHN-augmented模型在LV-Eval和InfiniteBench上持续优于仅使用滑动窗口的基线,甚至可以与全注意力模型性能相当或超越。在LongBench的六个平均长度超过8k的任务上,AHNs同样表现出优异的准确性。

  • 消融研究: 验证了自蒸馏训练(提供密集学习信号)和随机化滑动窗口大小(促进泛化能力)是AHNs成功的关键设计选择。

  • 梯度可视化: AHNs倾向于保留数学符号和数字等关键信息,而忽略代词和特殊token等次要信息,表明其作为有针对性压缩模块的有效性。

  • 局限性: 尽管AHN在效率和性能间取得了平衡,但其固定大小的压缩记忆仍不可避免地导致信息损失,可能在需要精确召回(exact recall)的任务上表现不佳(如RULER的NIAH测试所示),这为未来的研究提供了方向,例如探索更强的召回机制或进行全参数训练。

总结来说,该论文通过引入AHNs,有效地融合了RNN的效率和Transformer的表达能力,为长序列建模提供了一个新颖且高效的解决方案。

Abstract

本论文提出了一种受认知科学中“多存储模型”启发的神经网络记忆框架,旨在解决长序列建模中的一个核心问题:在RNN类模型的压缩固定大小记忆(高效但有信息损失)与Transformer注意力机制的无损增长记忆(高保真但低效)之间取得平衡。

核心内容讲解:

  • 提出的方法

    • 该方法结合了两种记忆机制:

      • 无损短时记忆:使用Transformer的KV缓存,维护一个滑动窗口,保留最近的精确上下文信息。

      • 压缩长时记忆:引入一个可学习模块——人工海马网络(Artificial Hippocampus Network, AHN),将滑动窗口之外的信息压缩为固定大小的长期记忆。

    • 这种结构模拟了人脑中短期记忆与长期记忆的协同机制。

  • 模型实现

    • AHN模块使用了现代RNN类架构进行实例化,包括 Mamba2、DeltaNet 和 GatedDeltaNet

  • 实验验证

    • 在两个长上下文基准测试集 LV-EvalInfiniteBench 上进行了大量实验。

    • 结果显示:

      • AHN增强模型在性能上显著优于滑动窗口基线模型

      • 并且其表现与全注意力模型相当甚至更优

      • 同时,在计算量和内存消耗方面大幅降低。

  • 具体效果示例

    • Qwen2.5-3B-Instruct 模型为例:

      • 增加AHN模块仅增加0.4%的参数量。

      • 在128k长度的LV-Eval任务中:

        • 计算量(FLOPs)减少40.5%。

        • 缓存内存减少74.0%。

        • 平均得分从4.41提升至5.88。

总结:

本研究提出的AHN框架,通过结合无损短时记忆和压缩长时记忆,实现了在长序列建模中兼顾效率与准确性的新方法。实验验证表明,该方法在保持高性能的同时,显著降低了资源消耗,具有良好的应用前景。

1 Instruction

本节主要介绍了人工系统中记忆建模的发展历程,并引出了本文提出的一种新型记忆机制——人工海马网络(Artificial Hippocampus Networks, AHNs)


1.1 记忆建模的历史发展

  • 早期方法:研究者尝试通过循环神经网络(RNNs)来建模记忆功能,其核心思想是将序列信息压缩到一个固定大小的隐藏状态中,作为记忆使用。这种方法在处理长序列时效率较高,但存在信息丢失的问题,尤其在需要精确长距离信息回忆的任务中表现不佳。

  • 后续发展:随着注意力机制Transformer架构的提出,记忆建模进入新阶段。其中,键值缓存(KV Cache)被用作记忆,具有无损记忆的特点,能保留所有token级别的信息,因此记忆容量更大。然而,KV缓存的缺点是内存和计算成本随序列长度线性/平方增长,在处理极长序列时效率低下。

  • 效率与精度的权衡:RNNs的压缩记忆效率高但精度低,而Transformer的无损记忆精度高但效率低。这种权衡启发研究者参考人类大脑的记忆机制,尤其是记忆的多存储模型(Multi-Store Model, MSM)


1.2 人类记忆机制的启发

  • MSM模型指出,人类大脑拥有短期记忆(工作记忆)长期记忆。短期记忆容量有限,但信息是无损的;长期记忆则通过海马体不断将短期记忆压缩并存储到皮层中。

  • 受此启发,作者提出了一种人工神经记忆框架,结合了无损短期记忆与压缩长期记忆的优点。


1.3 本文提出的方法:人工海马网络(AHN)

  • 核心思想:使用Transformer的KV缓存作为无损短期记忆,当信息超出滑动窗口范围时,由一个可学习的压缩模块(AHN)进行处理,将其压缩为固定大小的长期记忆状态

  • 结构设计

    • AHN模块可以采用类似RNN的结构。

    • 整体框架如图1(a)所示。


1.4 实验与结果

  • 模型实例化:作者将AHN与Mamba2、DeltaNet(DN)和GatedDeltaNet(GDN)结合,分别构建了AHN-Mamba2、AHN-DN、AHN-GDN

  • 实验数据

    • LV-EvalInfiniteBench等长上下文基准测试中,AHN增强模型显著优于滑动窗口模型,并接近或超过全注意力模型

    • 以Qwen2.5-3B-Instruct为例,加入AHN后:

      • 参数增加0.4%

      • FLOPs减少40.5%

      • 内存缓存减少74.0%

      • 在128k长度的LV-Eval上,平均得分从4.41提升至5.88


1.5 贡献总结

  1. 提出AHN概念:通过将滑动窗口外的无损记忆持续压缩为长期记忆,实现高效长上下文建模

  2. 实现实例与验证:构建了AHN-Mamba2、AHN-DN、AHN-GDN,并通过自蒸馏训练验证其有效性。实验表明,这些模型在长序列任务中显著提升效率,同时保持竞争力。


总结

本节系统回顾了记忆建模的发展路径,从RNN到Transformer,再到结合两者优势的AHN框架。重点在于提出一种模仿人类记忆机制的高效建模方法,并通过实验验证其在长序列任务中的优越性。

3 Method

3.1 预备知识(Preliminary)

本节介绍了现代自回归大语言模型的核心结构——Transformer,其核心机制是自注意力机制(self-attention)
输入序列为 \( X = (x_1, x_2, ..., x_L) \in \mathbb{R}^{L \times D} \),通过线性变换生成查询(Q)、键(K)、值(V)矩阵:

\[ Q = XW_Q, \quad K = XW_K, \quad V = XW_V \]

然后通过以下公式计算注意力输出:

\[ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_{in}}} \odot \mathcal{M} \right)V \]

其中 \(\mathcal{M}\) 是因果掩码(causal mask),确保模型只能关注当前及之前的 token。

重点:这部分是 Transformer 的标准结构,为后续 AHN 的引入提供基础。


3.2 人工海马网络(Artificial Hippocampus Networks, AHNs)

定义

受生物海马体启发,AHNs 用于压缩历史信息,将其整合为固定大小的递归状态。模型使用一个大小为 \(W\) 的滑动窗口,在窗口内的 token 保留完整记忆,窗口外的 token 则通过 AHN 压缩为紧凑表示。

递归更新公式为:

\[ h_{t-W} = \text{AHN}((k_{t-W}, v_{t-W}), h_{t-W-1}) \]

其中 \(h_{t-W}\) 是压缩后的记忆状态。

与无损记忆的结合

  • 滑动窗口内:使用标准因果注意力机制,保留无损记忆。

  • 窗口外:使用 AHN 压缩 KV 对,丢弃原始信息。

  • 最终输出基于压缩记忆和窗口内 KV 缓存共同计算:

\[ y_t = f(h_{t-W}, \{(k_i, v_i)\}_{i=t-W+1}^t, q_t) \]

重点:AHNs 的核心作用是高效压缩长上下文,从而减少内存和计算开销。

图示说明

图 2(a) 展示了 AHN 的整体机制,图 2(b) 展示了训练框架(见下节)。


3.3 实现方式(Instantiation)

架构选择

AHNs 可以用类似 RNN 的结构实现。本文实验中使用了三种结构:

  • Mamba2

  • DeltaNet(DN)

  • GatedDeltaNet(GDN)

并分别命名为 AHN-Mamba2、AHN-DN、AHN-GDN。

AHN-GDN 的实现细节

AHN-GDN 使用**门控增量规则(gated delta rule)**更新记忆:

\[ h_{t-W} = \alpha(x_{t-W})(I - \beta(x_{t-W})k_{t-W}^T k_{t-W})h_{t-W-1} + \beta(x_{t-W})k_{t-W}^T v_{t-W} \]

输出通过门控函数 \(\gamma(x_t)\) 调制,并通过线性投影:

\[ y_{\text{AHN}, t} = \gamma(x_t) q_t h_{t-W} W_o \]

最终输出为 AHN 输出与注意力机制输出的和:

\[ y_t = y_{\text{AHN}, t} + \text{Attention}(\{(k_i, v_i)\}_{i=t-W+1}^t, q_t) \]

重点:AHN-GDN 仅压缩滑动窗口外的 token,保留窗口内信息,兼顾效率与准确性。

复杂度分析

表 1 和图 3 展示了使用 AHN-GDN 后的效率提升:

  • 内存复杂度:从 \(O(L)\) 降低到 \(O(W)\)

  • 计算复杂度:从 \(O(L^2)\) 降低到 \(O(WL)\)

重点:引入 AHN 显著提升了模型在长上下文下的计算效率和内存效率。


3.4 训练框架(Training framework)

自蒸馏训练

为了更高效地训练 AHN 增强模型,采用**自蒸馏(self-distillation)**方法:

  • 教师模型:使用预训练的开源大模型(如 Qwen)

  • 学生模型:在教师模型基础上修改注意力机制,加入 AHN

  • 目标函数:最小化 KL 散度:

\[ l = \text{KL}(p' || p) \]

其中 \(p'\) 是教师模型输出,\(p\) 是学生模型输出。

训练策略

  • 冻结主模型参数,仅训练 AHN 参数

  • 提高训练效率,避免从头训练

重点:该训练方法利用已有大模型的知识,仅训练新增的 AHN 模块,节省资源且效果良好。


总结

本章系统介绍了 AHN 的设计与实现:

  1. 背景:基于 Transformer 的自注意力机制

  2. 核心设计:AHNs 通过递归压缩历史信息,实现高效长上下文建模

  3. 具体实现:使用 Mamba2、DeltaNet、GDN 等结构实现 AHN

  4. 效率提升:显著降低内存和计算复杂度

  5. 训练方法:采用自蒸馏策略,仅训练 AHN 模块,提升训练效率

核心贡献:提出 AHN 结构,解决长上下文建模中的效率瓶颈,同时保持模型性能。

4 Experiments

4.1 实验设置

模型与数据集

本研究基于Qwen2.5-Instruct系列(3B、7B、14B)构建AHN(Artificial Hippocampus Networks),并在三种现代循环模型上实现AHN模块:Mamba2、DeltaNet(DN)和GatedDeltaNet(GDN)。训练数据采用ChatQA2数据集,评估则在多个长上下文基准上进行,包括LongBench、InfiniteBench和LV-Eval,并辅以PG19的示例。

基线方法

对比方法包括滑动窗口注意力(SWA)与Attention Sinks结合、以及Compressive Transformers(CT,使用最大池化和平均池化)。所有方法在相同的无损内存预算下进行比较,确保公平性。

实现细节

AHN模块基于PyTorch实现,使用LLaMA-Factory和Flash Linear Attention框架。训练时冻结基础LLM,仅训练AHN模块,采用自蒸馏损失(KL散度)进行优化。训练过程中对AHN模块的起始位置和滑动窗口大小进行随机化,以增强泛化能力。优化器为AdamW,学习率1e-4,线性预热后余弦衰减,训练一个epoch,全局batch size为128。

表2:性能与效率分析

表2展示了不同模型在LV-Eval和InfiniteBench的128k长度子集上的表现。AHN模块在几乎不增加参数的情况下,显著优于SWA和CT方法,尤其在14B模型中表现突出。AHN-GDN在多个任务中达到最高得分,同时保持较低的计算和内存开销。


4.2 示例分析

通过将滑动窗口外的历史信息压缩为固定大小的记忆,AHN显著降低了计算复杂度和内存占用。以PG19中57K token的文本为例,标准Qwen模型在超过32K token后困惑度急剧上升,而AHN-GDN模型保持低困惑度。此外,AHN-GDN在处理长上下文时CUDA内存使用几乎恒定,而基础模型随上下文增长线性上升,验证了AHN在长上下文处理中的高效性。


4.3 长上下文基准测试

超长上下文设置(128K token)

在LV-Eval和InfiniteBench的128K token子集上,AHN模块在几乎所有任务中均优于SWA和CT方法,甚至超过全注意力模型,说明其压缩机制在保持性能的同时显著节省资源。

长上下文设置(平均长度 > 8K)

在LongBench的六个长任务上,AHN模块在固定8192 token内存预算下仍优于基线方法,尤其在7B和14B模型中表现突出。AHN-DN和AHN-GDN在多个任务中取得最高分,表明其在真实场景中具有良好的泛化能力。

表3:LongBench任务表现

表3展示了不同模型在LongBench六个任务上的准确率。AHN模块在3B、7B、14B模型中均优于SWA和CT,尤其在7B和14B模型中提升显著,说明AHN在不同规模模型中均有效。


4.4 消融实验

训练目标:自蒸馏 vs 下一词预测

使用自蒸馏(KL散度)训练的AHN优于使用交叉熵(CE)损失的下一词预测方法。CE损失提供的学习信号稀疏,容易导致模型走捷径,而自蒸馏提供更密集的指导,帮助模型学习更泛化的上下文表示。

窗口大小:随机 vs 固定

使用随机滑动窗口大小训练的AHN在不同上下文长度下表现更稳定,而固定窗口容易过拟合特定配置。实验表明,随机化训练有助于提升模型的上下文泛化能力。

表4:训练设计消融

表4展示了不同训练策略在LongBench任务上的表现。自蒸馏 + 随机窗口的组合效果最佳,显著优于其他组合。


4.5 梯度可视化分析

通过可视化自蒸馏损失的梯度,研究AHN如何压缩和利用窗口外信息。公式定义了梯度对窗口外token的影响,梯度小的token表示其信息已被压缩记忆有效捕获。

图5:梯度可视化结果

图5显示,AHN倾向于保留数学符号和数字等关键信息,忽略代词和特殊token等次要信息,验证了其作为目标压缩模块的有效性。


总结

本章系统评估了AHN在长上下文建模中的性能与效率。实验表明,AHN在多个长上下文基准上均优于现有方法,尤其在处理超长上下文时表现出色。消融实验验证了自蒸馏训练和随机窗口设计对模型泛化能力的重要性。梯度可视化进一步揭示了AHN在压缩过程中对关键信息的保留机制,证明其作为高效长上下文建模模块的潜力。

5 Conclusion and discussion

本节总结了人工海马网络(Artificial Hippocampus Networks, AHNs)的核心贡献与实验结果,并讨论了其局限性与未来研究方向。

核心贡献与实验结果

作者提出了一种新型、轻量级的模型组件——人工海马网络(AHNs),用于增强Transformer模型在处理长序列时的效率。AHNs通过维护一个滑动窗口的KV缓存作为无损记忆,同时将窗口外的信息压缩为固定大小的记忆表示,从而解决了传统Transformer在长序列处理中的效率瓶颈。

这种设计使得使用AHNs增强的模型在处理长序列时,每个token的内存和计算复杂度都保持为常数级别。实验结果表明,AHNs能够在保持在长上下文任务中良好性能的同时,显著减少内存缓存大小和计算开销。

局限性与未来工作

尽管AHNs在计算效率与记忆保真之间取得了良好平衡,但其固定大小的压缩记忆结构不可避免地带来信息损失,可能影响需要精确回忆的任务表现(详见附录)。此外,由于本研究采用了参数高效的知识蒸馏设置,模型性能受限于基础模型的能力。

未来的研究方向包括探索更强的记忆召回机制,以及采用全参数训练方法,以进一步释放AHNs的潜力。

应用前景

在应用场景方面,AHNs框架适用于信息稀疏或资源受限的长上下文任务,例如终身学习、流视频处理以及边缘设备上的部署,具有良好的应用前景。

Acknowledgement

Acknowledgement(致谢)

本节为论文的致谢部分,作者对Shi Guang、Haoqi Fan、Tianle Cai等多位研究人员表示感谢,感谢他们在论文研究过程中提供的宝贵讨论和建议。这部分内容属于常规致谢,未涉及具体研究细节,因此不做展开。

6 AHN instantiation

本节主要介绍了如何使用 Mamba2DateNet(DN) 来实例化 AHN(Artificial Hippocampus Networks),并给出了相应的记忆更新规则。


AHN-Mamba2 的记忆更新规则

  • AHN-Mamba2 的更新公式如下:

    \[ h_{t-W} = \text{AHN-Mamba2}((k_{t-W}, v_{t-W}), h_{t-W-1}, x_{t-W}) \]
    \[ = \exp(-\Delta(x_{t-W})A) h_{t-W-1} + \Delta(x_{t-W-1}) k_{t-W}^T v_{t-W} \]
  • 重点说明:该公式借鉴了 Mamba2 的状态空间模型(SSM)思想,通过指数衰减项(exp 项)和输入相关的 Δ 函数来控制记忆的衰减与更新,实现对长序列的高效建模。


AHN-DN 的记忆更新规则

  • AHN-DN 的更新公式如下:

    \[ h_{t-W} = \text{AHN-DN}((k_{t-W}, v_{t-W}), h_{t-W-1}, x_t) \]
    \[ = (\mathbf{I} - \beta(x_{t-W}) k_{t-W}^T k_{t-W}) h_{t-W-1} + \beta(x_{t-W}) k_{t-W}^T v_{t-W} \]
  • 重点说明:该公式基于 DateNet 的设计,使用 β 函数控制记忆更新的权重,通过矩阵运算实现对历史信息的选择性保留与压缩。


输出规则(Output Rule)

  • AHN-Mamba2 和 AHN-DN 的输出规则与 AHN-GDN 相同,具体见公式 6

  • 说明:输出部分统一设计,确保不同 AHN 实例在输出层保持一致性。


AHN 增强模型的结构图示

  • 提供了结合 Attention Sinks 的 AHN 增强模型示意图(见图6)。

  • 重点说明

    • Attention Sinks 是一种保留关键历史信息的机制,与滑动窗口配合使用。

    • 当输入序列长度 ≤ (Attention Sinks 数量 + 滑动窗口长度) 时,模型等同于标准 Transformer。

    • 当序列更长时,AHN 会将滑动窗口外的信息压缩为紧凑的记忆表示。

    • 模型在生成下一个 token 时,会同时利用:

      • 注意力锚点(Attention Sinks)中的无损信息,

      • 滑动窗口内的信息,

      • 以及压缩后的记忆表示。


总结

本节详细介绍了 AHN 的两种实例化方式(Mamba2 和 DateNet),分别给出了其记忆更新公式,并强调了它们在长序列建模中的压缩与更新机制。同时,结合 Attention Sinks 的结构设计,展示了 AHN 在处理超长上下文时的高效性与灵活性。

7 Additional benchmark results

本节进一步探讨了AHNs(Artificial Hippocampus Networks,人工海马网络)在长上下文场景中的有效性,并展示了更多的基准测试结果。同时,也指出了由于压缩记忆的有损特性,AHNs在精确回忆任务中存在固有局限性。

LV-Eval

本节展示了在128k上下文设置下,所有11个LV-Eval任务的完整结果。所有模型均配置了32768个token的无损记忆,其中包括128个token的注意力沉(attention sinks)和一个32640个token的滑动窗口。

RULER

RULER是一个综合性基准测试,通过增加任务难度和引入更多类别,扩展了标准的“大海捞针”(NIAH)范式。作者在RULER-128k子集的所有NIAH任务中评估了一个AHN增强模型(AHN-GDN),基础模型为Qwen2.5-7B-Instruct。为了公平比较,AHN-GDN和使用注意力沉的滑动窗口注意力均配置了128个注意力沉和32640个token的滑动窗口。结果显示(见表5),AHN-GDN的表现与滑动窗口注意力相当,但在精确回忆任务上明显逊色于全注意力模型。这反映了有损压缩的固有局限性:虽然AHN增强模型能够实现高效的长上下文推理,但在需要从压缩记忆中精确回忆的任务上存在不足。这一局限性为未来的研究提供了方向,例如开发能够在保留关键信息于无损记忆的同时,利用压缩提高效率的记忆管理机制。

表5:RULER-128k中的“大海捞针”(NIAH)任务表现

方法

single_1

single_2

single_3

multikey_1

multikey_2

multikey_3

multivalue

multiquery

全注意力(Full Attn)

98.60

97.20

98.40

89.20

23.60

23.20

55.40

85.45

注意力沉 + 滑动窗口注意力(Sinks + SWA)

26.80

25.40

28.00

27.80

10.60

9.00

22.95

24.00

AHN-GDN

26.80

25.20

28.20

27.40

11.40

8.60

23.45

23.35

表6:LV-Eval 128k子集的完整结果

表6展示了在LV-Eval的21个任务中,不同模型在128k上下文设置下的表现。所有基于滑动窗口的方法均使用了32768个token的无损记忆,包括128个注意力沉和32640个token的滑动窗口。

模型

数据集

全注意力

注意力沉 + 滑动窗口注意力

CT-Max

CT-Average

AHN-Mamba2

AHN-DN

AHN-GDN

Qwen2.5-3B-Instruct

平均值

4.41

4.59

4.12

4.47

5.13

5.68

5.88

…(略去具体任务数据)

Qwen2.5-14B-Instruct

平均值

4.99

5.69

5.28

5.64

6.43

6.50

6.51

总结来看,AHNs在长上下文任务中表现出一定的有效性,尤其在某些任务上优于传统方法。然而,由于其压缩记忆的有损特性,AHNs在需要精确回忆的任务上表现不佳。这表明未来的研究可以探索如何在保持效率的同时,更好地保留关键信息以提高精确回忆能力。