大语言模型 (LLM) 已经改变了我们与技术交互的方式，但它们存在一个关键弱点: 记忆。虽然它们能够处理并生成类似人类的文本，但在处理超长信息序列 (例如，一整本书、一份冗长的法律文件，或一个复杂的代码库) 时能力有限。这是因为主流架构 Transformer 面临一个根本性的权衡: 它要么保持完美、无损的记忆，但随着上下文增长会变得极其缓慢且昂贵；要么使用压缩的、固定大小的记忆，速度快但不可避免地遗失重要细节。

如果我们能两全其美呢？如果 AI 能像人脑一样，既能完美记忆近期事件，又能将较早的信息高效压缩到一个紧凑的长期存储中，会怎样？

这正是那篇引人注目的新论文——《人工海马体网络: 实现高效长上下文建模》的核心思想。研究人员提出了一种受认知科学启发的创新框架，为 Transformer 提供更复杂的记忆系统。他们的方法称为 人工海马体网络 (AHN) ，取得了显著成果: 在处理长序列时极大降低计算成本与内存使用，同时在挑战性的长上下文基准测试中还能提升性能。

一张图表展示了 AHN 记忆架构，以及一个柱状图显示其效率和性能提升。

图 1: (a) AHN 的核心思想是将无损、不断增长的记忆 (如注意力 KV 缓存) 转化为固定大小的压缩记忆。(b) 将 AHN 添加到一个参数量为 3B 的模型中，可使计算成本 (TFLOPs) 降低 40.5%，内存缓存减少 74.0%，同时提升 LV-Eval 长上下文基准测试评分。

在本文中，我们将深入解析这篇论文——探索现代 AI 面临的记忆困境，剖析这个受大脑启发的解决方案的工作原理，并审视其架构、训练方法和实验结果，这些都表明 AHN 或许是构建更高效、可扩展 LLM 的关键一步。

记忆困境: Transformer vs. RNN

要理解 AHN 的意义，我们首先需要了解神经网络中两种主流的记忆方式: Transformer 所采用的方法，以及其前身循环神经网络 (RNN) 的方法。

RNN: 高效但健忘的书记官

早期的序列模型，如 LSTM 和 GRU，都属于 RNN 范畴。RNN 逐步处理信息，并维护一个隐藏状态——一个固定大小的向量，作为它的记忆。在每一步中，模型接收当前输入与前一个隐藏状态，生成输出并更新记忆。

这种设计效率极高: 每个 token 的计算和内存需求与序列长度无关。然而，这种高效带来了信息瓶颈。将整个历史上下文压缩到一个固定大小的状态意味着必须丢弃信息。在超长序列中，它不可避免地丢失重要细节，使得长程回忆变得困难。

Transformer: 完美但昂贵的图书管理员

Transformer 架构 (发表于 Attention Is All You Need) 彻底改变了序列建模模式。它不使用压缩的隐藏状态，而是采用自注意力机制: 每个 token 的查询 (Query) 向量与所有先前 token 的键 (Key) 向量进行比较，生成注意力得分，用于加权值 (Value) 向量。

这些信息存储在键值 (KV) 缓存中，这是一种无损记忆，保留精确的 token 级信息，使强大的上下文学习和跨距离检索成为可能。

缺点在于: KV 缓存随序列长度线性增长 (\(O(L)\)) ，而为每个新 token 计算它与所有前序 token 的注意力分数，则带来二次方的计算成本 (\(O(L^2)\)) 。对于超长序列，这会造成极高的延迟与内存消耗。

这就是根本性的权衡: RNN 的压缩记忆带来高效性，而 Transformer 的无损记忆带来高保真。

受大脑启发的解决方案: 人工海马体网络

为了解决这一困境，研究人员将目光投向已知最复杂的记忆系统——人脑。认知科学中的多重存储模型将短期 (工作) 记忆与长期记忆区分开来。海马体被认为负责将短期记忆转化并巩固为长期存储。

受此启发，他们提出了人工海马体网络 (AHN) 框架，将无损的短期记忆与压缩的长期记忆融合。

AHN 的工作原理如下:

  1. 无损短期记忆: 模型采用一个滑动注意力窗口，为最近的 \(W\) 个 token (如最后 4096 个) 维护完美的 KV 缓存。
  2. 压缩长期记忆: 当某个 token 的 KV 对即将滑出窗口时，它会传给 AHN，而非直接丢弃。
  3. AHN 模块: 这是一个可学习的、类似 RNN 的组件，使用被移出窗口的 KV 对来循环更新固定大小的隐藏状态，从而生成对久远信息的紧凑摘要。

生成新 token 时，模型会同时利用无损的近期上下文和压缩的长期历史。

一张分步图，展示了 token 如何被滑动窗口和 AHN 处理。

图 2: (a) AHN 工作流程: 序列长度若短于窗口，模型表现与标准 Transformer 相同。随着序列延长，滑出窗口的 token 会被压缩至记忆状态 \(h\)。(b) 自蒸馏训练方案中，AHN 从固定的教师模型学习。

数学上，位于 \(t-W\) 位置的记忆更新为:

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

token \(t\) 的输出 \((y_t)\) 同时依赖压缩记忆与无损窗口:

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

用现代 RNN 构建 AHN

AHN 是一个通用概念，可由不同的循环架构实现。作者探索了 Mamba2、DeltaNet 以及 GatedDeltaNet (GDN) 等变种。

在基于 GDN 的 AHN (AHN-GDN) 中，可学习的门控单元控制新 KV 信息注入记忆状态的方式，实现更细致的更新。查询会从压缩记忆中取出信息，与滑动窗口注意力输出融合。

AHN-GDN 记忆更新的方程。

组合输出遵循以下公式:

结合 AHN 和注意力输出的方程。

此设计通过固定滑动窗口与 AHN 状态大小来保持高效，避免了二次方扩展。

一张表格，比较了全注意力和带有 AHN-GDN 的注意力的计算和内存复杂性。

表 1: 复杂度比较: 全注意力需要 \(O(L^2)\) FLOPs 和 \(O(L)\) 内存；经 AHN 增强的注意力则为 \(O(WL)\) FLOPs 和 \(O(W)\) 内存占用。

用自蒸馏训练 AHN

从零训练大型 LLM 成本高昂。作者采用自蒸馏来高效训练 AHN:

  • 教师模型: 预训练的全注意力 LLM (如 Qwen2.5) 。
  • 学生模型: 将其全注意力模块替换为滑动窗口 + AHN。

教师权重被冻结，仅训练 AHN 参数。训练目标是最小化 KL 散度，使学生模型的预测概率接近教师:

\[ l = \mathrm{KL}(p' \parallel p) \]

这样便将教师模型的长距离依赖处理能力迁移至学生模型的 AHN。

AHN 实测结果

示例: 阅读长篇书籍

研究团队将标准 Qwen2.5-3B LLM 与加入 AHN 的版本在 57K-token 的 PG19 段落上对比。基线模型预训练的上下文长度为 32K。

图表比较了标准模型与带有 AHN 的模型在 FLOPs、内存使用和困惑度上的差异。

图 3: (a, b) AHN 使 FLOPs 线性增长，内存保持恒定。(c) 基线模型困惑度在超过 32K token 后急剧上升，AHN 则保持低且稳定。(d) AHN 占用更少 GPU 内存。

长上下文基准表现

在 LV-Eval 与 InfiniteBench 的 128K-token 序列测试中，AHN 始终优于滑动窗口和 Compressive Transformer 基线，甚至超过全注意力模型，同时 FLOPs 降至后者一半、内存仅为四分之一。

表格展示了不同模型大小和方法在 LV-Eval 和 InfiniteBench 上的性能。

表 2: 从 3B 到 14B 参数的模型中，AHN 变体均比基线、甚至部分全注意力模型表现更佳。

在平均序列长度超过 8K-token 的 LongBench 任务中，AHN 再次展现更高准确率。

表格展示了在 LongBench 任务上的性能。

表 3: AHN 在多种长上下文任务中均优于基线表现。

消融研究洞见

两项设计选择被证明格外重要:

  • 自蒸馏 vs. 下一词预测: 自蒸馏泛化效果更好；下一词预测训练则降低性能。
  • 随机窗口 vs. 固定窗口: 训练时随机化窗口大小，可提升对不同上下文长度的适应性。

一组图表显示，AHN-GDN 在各种无损记忆大小下都能保持高性能。

图 4: AHN-GDN 在不同记忆预算下始终保持最佳性能。

AHN 学会存储的内容

梯度可视化表明，AHN 会优先保留语义重要的 token (如数字与运算符) ，而弱化不相关 token。

一个文本示例，其中的 token 根据梯度大小着色，显示了 AHN 优先处理的内容。

图 5: 在数学任务中，AHN 保留关键数字和符号内容 (绿色) ，弱化填充性 token (红色) 。

结论与未来展望

人工海马体网络是一种巧妙、受生物学启发的方案，解决了长上下文高效处理的难题:

  • 高效性: 每个 token 计算时间呈线性，内存恒定。
  • 高性能: 长上下文基准测试比分优于基线甚至部分全注意力模型。
  • 灵活性: 兼容多种循环架构。
  • 实用性: 可通过轻量训练快速增强现有预训练 LLM。

局限性仍在——有损压缩可能影响需精确回忆的任务。未来或可探索混合策略，在无损保留关键信息的同时，从压缩中获益。

融合神经科学洞见与 AI 工程技术，AHN 或将帮助 LLM 读完整本书、无限期处理数据流，并在受限硬件上运行——向更具可扩展性的终身学习系统迈进。