EM-LLM: Human-Inspired Episodic Memory for Infinite Context LLMs

Source | HN Comments

EM-LLM 是一种受人类启发、无需微调的 LLM，旨在处理无限上下文长度。它结合了人类情景记忆和事件认知，通过贝叶斯惊奇和图论进行事件分割，并采用两阶段检索机制。实验表明，EM-LLM 在 LongBench 和 $\infty$-Bench 基准测试中表现优异，超越了 SOTA 检索模型和 RAG 方法，甚至在某些任务中超过了全上下文模型。该模型在 10M tokens 上成功检索，并与人类事件感知有很强的相关性，为研究人类记忆机制提供了新框架。

EM-LLM：面向无限上下文 LLM 的类人情景记忆

此仓库包含 ICLR 2025 上发表的 EM-LLM 代码版本：[openreview link]。

快速链接

概述

虽然典型的 LLM 在处理大量上下文时会遇到困难，但人脑擅长组织和检索跨越一生的经验。在这项工作中，我们引入了 EM-LLM，它将人类情景记忆和事件认知的关键方面集成到 LLM 中，无需微调，使其能够处理实际上无限的上下文长度，同时保持计算效率。 EM-LLM 使用贝叶斯惊奇和图论边界细化的组合，以在线方式将 token 序列组织成连贯的情景事件。当需要时，这些事件通过两阶段记忆过程进行检索，结合基于相似性的检索和时间上连续的检索，以实现高效且类似人类的访问相关信息。在 LongBench 和 $\infty$-Bench 基准测试中的实验表明，EM-LLM 具有卓越的性能，在各种基线 LLM 中始终优于 SOTA 检索模型 InfLLM。此外，EM-LLM 在广泛的任务中优于 RAG，同时需要类似的资源。值得注意的是，EM-LLM 的性能甚至在大多数任务中超过了全上下文模型，同时成功地跨越了 10M tokens 执行检索，对于这些模型来说，这种规模在计算上是不可行的。我们的分析揭示了 EM-LLM 的事件分割与人类感知的事件之间存在很强的相关性，这表明该人工系统与其生物对应物之间存在桥梁，从而为探索人类记忆机制提供了一个新的计算框架。

架构

图 1： 每个 LLM 层中记忆形成和检索的架构。 形成： 输入序列最初通过惊奇度进行分割（①中的紫色虚线），然后基于群论度量细化分割（②中的绿色虚线）。保留初始 token 和局部上下文。 检索： 通过 k-NN 搜索 ③ 和从情景记忆中选择连续事件 ④。

结果

图 2： （左） EM-LLM$_S$ 与 RAG（NV-Embed-v2 检索器）与全上下文的比较，以 LLaMA-3.1-8B 作为基本 LLM，在 LongBench 上进行评估。 （右） 在 $\infty$-Bench 扩展版本的 Retrieve.PassKey 上比较各种长序列方法（根据其上下文窗口长度排序）。

点击此处查看更完整的结果表格。

用法

配置

安装 requirements：

python3 -m pip install --upgrade pip
pip install -r "${base_dir}/requirements.txt"
pip install -e "${base_dir}/."

用于配置的 YAML 文件可以在 config/ 目录中找到。

以下是这些文件中包含的每个参数的细分：

verbose: false # 处理完一个示例后，打印问题/预测/答案
compute_ppl: true # 打印并记录每个示例/块的困惑度
return_block_size: true # 打印并记录每个示例/块的块大小
logging: true # 将日志保存到输出目录，并在多处理期间标记单个worker日志
em_splitter: surprisal # 将块拆分为内存块的方法（surprisal、random、sentence）
max_len: 2147483647 # 使用截断之前的最大序列长度
chunk_size: 512 # 解码期间分块输入的大小
conv_type: mistral-inst # 会话模板类型
extended_passkey: 1024 # 以数千个 token (k) 为单位扩展 infinite-bench 的 passkey 任务的长度
model:
 type: em-llm # 用于推理的模型（此版本中仅提供 em-llm）
 path: mistralai/Mistral-7B-Instruct-v0.2 # HuggingFace 模型路径
 min_block_size: 8 # 最小可能的块大小 - 小于此值的块将扩展到此大小
 max_block_size: 128 # 最大可能的块大小 - 大于此值的块将拆分到此大小
 n_init: 128 # 上下文窗口中包含的初始 token 数
 n_local: 4096 # 上下文窗口中包含的本地 token 数
 n_mem: 2048 # 上下文窗口中包含的检索到的 token 数（包括相似性和连续性缓冲区）
 repr_topk: 4 # 每个内存单元中被认为是代表性元素的得分最高的 token 数
 max_cached_block: 512 # 保留在 GPU 内存中的内存块数 - 必须大于 n_mem/min_block_size
 exc_block_size: 512 # 作为执行块一次查询的 token 数 - 每个执行块执行一次 n_mem token 的检索
 base: 1000000 # RoPE base
 distance_scale: 1.0 # RoPE 距离比例
 surprisal_threshold_gamma: 1.0 # 惊奇度计算中的标准差缩放因子（见论文）
 min_free_cpu_memory: 100 # 分配内存块时要保持空闲的最小 CPU RAM 量 (GB)
 disk_offload_threshold: 300000 # 超过此序列中 token 数量应使用磁盘卸载
 vector_offload_threshold: 50000 # 超过此序列中 token 数量应将代表性 token 卸载到 CPU 内存
 similarity_refinement_kwargs: # 与我们论文的边界细化步骤直接相关的参数
  similarity_refinement: false # 是否使用边界细化
  refine_with_buffer: true # 如果为 True，则邻接矩阵将在其邻接矩阵的计算中包括相邻块的一部分 - 旨在使分割与相邻块更兼容，但也增加了计算时间
  refine_from_layer: 20 # 计算邻接时使用的层
  similarity_metric: modularity # 用作细化过程中目标的度量标准：模块化或电导率（或 intra_inter_sim，但到目前为止效果不佳）
 contiguity_buffer_kwargs: # 与连续性缓冲区直接相关的参数
  use_contiguity_buffer: true # 是否使用连续性缓冲区
  contiguity_buffer_size: 0.3 # 专用于连续性缓冲区的 n_mem token 的比例
 uniform_blocks: false # 忽略 em_splitter（上面）并将块分割为大小为 max_block_size 的固定大小的块（上面）
 random_topk_blocks: false # 检索随机块而不是 topk 最相似的块

评估

数据准备 我们采用 $\infty$-Bench 和 LongBench 进行模型评估。您可以通过运行以下命令下载数据集。

bash scripts/download.sh

响应生成 您可以通过运行以下命令来评估 EM-LLM。您还可以选择传入以下参数以适应您的硬件资源

bash scripts/run.sh
  -m|--model # DEFAULT: mistral; OPTIONS: mistral,llama3,llama31,phi3_mini,phi35_mini - 评估期间要使用的基本 LLM。
  -b|--benchmark # DEFAULT: long-bench; OPTIONS: long-bench,infinite-bench,passkey - 要评估的基准。 Passkey 评估 InfiniteBench 的 passkey 检索任务的扩展版本（有关上下文长度参数，请参见 yaml）。
  -w|--world-size # DEFAULT: number of visible GPUs - 评估期间要使用的 GPU 总数。
  -n|--num_gpus_per_job # DEFAULT: 1 - 要归因于每个作业的 GPU 数量。 如果 >1，模型层将均匀分布在多个 GPU 上。
  -r|--rank_offset # DEFAULT: 0 - 忽略脚本可见的前 n 个 GPU。 在单个节点上运行多个实验时很有用。
  -o|--allow_disk_offload # DEFAULT: False - 是否允许内存块的动态磁盘卸载（有关更多详细信息，请参见我们论文的附录）。 在单 GPU 实例中，这也会将代表性 token 卸载到 CPU 内存。

引用

如果您发现 EM-LLM 有用，请引用以下论文：

@inproceedings{fountas2025humaninspired,
  title={Human-inspired Episodic Memory for Infinite Context {LLM}s},
  author={Zafeirios Fountas and Martin Benfeghoul and Adnan Oomerjee and Fenia Christopoulou and Gerasimos Lampouras and Haitham Bou Ammar and Jun Wang},
  booktitle={The Thirteenth International Conference on Learning Representations},
  year={2025},
  url={https://openreview.net/forum?id=BI2int5SAC}
}