目标

该文章的目的是引入EM-LLM框架，通过模拟人类的情景记忆机制，解决大语言模型（LLMs）在处理长上下文时的局限性​；EM-LLM的设计灵感源于人类大脑的事件认知与回忆过程，目标是使LLMs能够有效地处理无限长度的上下文，提高检索和推理性能，同时保持计算效率​；

方法

EM-LLM 的实现方法包括 事件分段 和 记忆检索 两个核心阶段，同时引入了基于贝叶斯突发的初始边界识别与图论边界细化方法，详细步骤如下：

---

1. 事件分段（Memory Formation）

步骤 1：基于“贝叶斯突发”的初始事件边界识别

• 目的：通过惊讶值（Surprise） 识别事件边界。

• 实现原理：

  • 惊讶值是一个概率值，表示当前 token 的不确定性（负对数似然）。当预测 token 的概率低时，惊讶值高，表明事件边界的可能性大。

  $$
  \log P(x_t \mid x_1, ..., x_{t-1}; \theta) > T
  $$

  其中：

  • $ x_t $ 是当前 token，
  • $ P(x_t \mid x_1, ..., x_{t-1}; \theta) $表示模型预测当前 token 的概率，
  • T 是动态阈值，由当前窗口的惊讶值的均值和方差决定，公式为：
    $$
    T = \mu_{t-\tau:t} + \gamma \sigma_{t-\tau:t}
    $$
  • $ \mu $：最近窗口的惊讶值均值，$ \sigma $：标准差，$ \gamma $ 是缩放因子。

具体步骤：

1. 在 token 序列中计算每个 token 的惊讶值。
2. 将惊讶值高于阈值 T 的 token 作为初始事件边界。
3. 输出初始事件边界集合 $ B = {b_1, b_2, ..., b_k} $。

---

步骤 2：基于图论的边界细化

• 目的：优化事件边界，使得每个事件内部的 token 相似度高、事件之间相似度低。
• 方法：
  • 将 Transformer 注意力头中的 key 值 的相似性矩阵视为加权图的邻接矩阵。
  • 通过图论中的模块度（Modularity） 和 导通度（Conductance） 进行边界细化。

关键公式：

1. 相似度矩阵：
   $$
   A_{ij} = \text{sim}(K_i, K_j)
   $$
   其中 $ K_i $ 和 $ K_j $ 是 attention head 的 key 值，相似度采用点积计算。

2. 模块度（衡量事件内部密度）：
   $$
   f_M(A, B) = \frac{1}{4m} \sum_{i,j} \left[ A_{ij} - \frac{\sum_i A_{ij} \cdot \sum_j A_{ij}}{2m} \right] \delta(c_i, c_j)
   $$
   其中 $ \delta $ 是 Kronecker delta，表示两个节点是否属于同一个事件。

3. 导通度（衡量事件之间的分割质量）：
   $$
   f_C(A, B) = \min_{S \in V} \frac{\sum_{i \in S, j \notin S} A_{ij}}{\min(\text{vol}(S), \text{vol}(V \setminus S))}
   $$

具体步骤：

1. 将初始边界 B 输入，基于 A 计算事件内部的密度和跨边界的相似性。
2. 对每个初始边界，在事件边界 $ [\alpha, \beta] $ 区间内寻找最佳边界点 $ \hat{\beta} $，优化模块度或导通度。
3. 输出优化后的边界集合 B' 。

---

2. 记忆检索（Memory Retrieval）

步骤 3：两阶段记忆检索

• 目的：从大量的事件记忆中检索与当前任务相关的信息，并动态添加到 LLM 上下文窗口中。

阶段 1：基于相似性的事件检索

1. 相似性计算：对于当前 query，使用k-近邻（k-NN） 方法检索最相似的事件。
   • 相似性基于事件代表 token（例如注意力分数最高的 token）与 query 的点积相似度。
2. 事件选择：检索出最相关的 $ k_s $ 个事件，作为相似性缓冲区。

阶段 2：基于时间连续性的检索

1. 对于相似性检索出的事件，额外检索其相邻事件（在原始顺序中靠近的事件）。
2. 将这些相邻事件添加到连续性缓冲区。
3. 使用队列结构维持事件顺序，保证时间连续性。

最终上下文窗口组成：

• 初始 token：任务开头的部分 token。
• 局部上下文：当前窗口的最近 token。
• 相似性缓冲区：检索到的最相关事件。
• 连续性缓冲区：相邻事件，增强时间连续性。

---

整体实现流程总结

1. 输入：长上下文的 token 序列。
2. 事件分段：
   • 使用惊讶值识别初始边界。
   • 利用图论方法优化事件边界，生成事件集合。
3. 记忆检索：
   • 基于相似性进行 k-NN 检索获取相关事件。
   • 添加相邻事件维护时间连续性。
4. 输出：构建优化后的上下文窗口（初始 token + 相似性事件 + 连续性事件 + 局部上下文），供 LLM 进行推理。

---

算法复杂度

• 惊讶值计算：线性复杂度 O(n) ，其中 n 是 token 数量。
• 边界细化： O(k*n) ，其中 k 是事件边界数。
• 记忆检索: $O(k_s + k_c)$，检索相关和相邻事件。

此流程实现了人类记忆机制的模拟，使 LLM 能够在 无限上下文 中高效地进行检索和推理。

下游表现

EM-LLM的表现经过了LongBench和∞-Bench两大长上下文任务基准测试，结果如下：

对比表现：

• 优于InfLLM（现有最先进KV检索模型）和RAG方法，在80%的任务组上取得领先
• 在检索（例如KV、Passage、Number）和问答任务（如NarrativeQA、HotpotQA）上表现出高达**40%**的提升​

与全上下文模型对比：

• EM-LLM在大部分任务中表现优于全上下文模型，且能够成功进行10百万个tokens的检索，这对当前全上下文模型来说是不可行的​

与人类认知的对比：

EM-LLM的事件分段与人类感知的事件高度相关，揭示了其与人类情景记忆之间的联系，提供了一种新的框架来研究人类记忆机制​