处理无限上下文的记忆增强型大模型

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

虽然典型的 LLM 在处理长序列上下文时表现欠佳，但人脑擅长组织和检索跨越一生的经历。在这项工作中，我们引入了 EM-LLM，一种无需微调就能将人类 episodic memory 和事件认知的关键特性集成到 LLM 中的架构，使其能够处理实际上无限的上下文长度，同时保持计算效率。EM-LLM 使用 Bayesian surprise 和图论边界优化的组合，以在线的方式将 token 序列组织成有连贯性的 episodic 事件。当需要时，这些事件通过两阶段的记忆过程进行检索，结合基于相似度和时间连续性的检索，实现高效的、人类般的相关信息访问。在 LongBench 和 ∞-Bench 基准上的实验表明，EM-LLM 具有优越的性能，在各种基础 LLM 上持续超越最先进的检索模型 InfLLM。此外，EM-LLM 在各种任务上的表现都优于 RAG，同时消耗的资源类似。值得注意的是，EM-LLM 的性能甚至在大多数任务上超过了全上下文模型，同时成功地在 10M tokens 的规模上执行检索——这对这类模型在计算上是不可行的。我们的分析揭示了 EM-LLM 的事件分割与人类感知事件之间的强相关性，为这个人工系统与其生物学对应物之间搭建了桥梁，从而为探索人类记忆机制提供了新的计算框架。

图 1：每个 LLM 层中记忆形成和检索的架构

形成：输入序列最初通过 surprise 进行分割（①中的紫色虚线），然后基于群论指标对分割进行优化（②中的绿色虚线）。初始 token 和局部上下文被保留。检索：通过 k-NN 搜索③和从 episodic memory 中选择连续的事件④。

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

图 2

（左）EM-LLM_S vs. RAG（NV-Embed-v2 retriever）vs. 全上下文，以 LLaMA-3.1-8B 作为基础 LLM，在 LongBench 上评估。（右）各种长序列方法的对比（按其上下文窗口长度排序）在 ∞-Bench 的 Retrieve.PassKey 扩展版本上。

安装要求

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

用于配置的 YAML 文件位于 config/ 目录中。

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

verbose: false  # print the question/prediction/answer after an example has been processed 
compute_ppl: true  # print and log perplexity for each example/chunk
return_block_size: true  # print and log block size for each example/chunk
logging: true  # save logs to output directory and label individual worker logs during multiprocessing
em_splitter: surprisal  # method by which to split chunks into memory blocks (surprisal, random, sentence)

max_len: 2147483647  # maximum sequence length before truncation is used
chunk_size: 512  # size of chunked input during decoding
conv_type: mistral-inst  # conversation template type

extended_passkey: 1024  # length to extend infinite-bench's passkey task to in terms of thousands of tokens (k)

model:
  type: em-llm  # Which model to use for inference (only em-llm is made available in this version)
  path: mistralai/Mistral-7B-Instruct-v0.2  # HuggingFace model path
  min_block_size: 8  # the smallest possible block size - blocks smaller than this will be expanded to this size
  max_block_size: 128  # the biggest possible block size - blocks bigger than this will be split to this size
  n_init: 128  # number of initial tokens to include in context window
  n_local: 4096  # number of local tokens to include in context window
  n_mem: 2048  # number of retrieved tokens to include in context window (includes both the similarity and contiguity buffers)
  repr_topk: 4  # number of top-scoring tokens per memory unit considered as representative elements
  max_cached_block: 512  # number of memory blocks to keep in GPU memory - must be greater than n_mem/min_block_size
  exc_block_size: 512  # number of tokens queried at a time as an execution block - each execution block performs retrieval of n_mem tokens once
  base: 1000000  # RoPE base
  distance_scale: 1.0  # RoPE distance scale
  surprisal_threshold_gamma: 1.0  # the standard-deviation scaling factor in the surprisal calculation (see paper)

  min_free_cpu_memory: 100  # minimum amount CPU RAM (GB) to keep free when allocating memory blocks
  disk_offload_threshold: 300000  # number of tokens in a sequence past which disk offloading should be used
  vector_offload_threshold: 50000  # number of tokens in a sequence past which representative tokens should be offloaded to CPU memory

  similarity_refinement_kwargs:  # parameters relating directly to the boundary refinement step of our paper
    similarity_refinement: false  # whether to use boundary refinement or not
    refine_with_buffer: true  # if True, the adjacency matrix will include part of the neighbouring chunks in its calculation of the adjacency matrix - designed to make segmentations more compatible with neighbouring chunks, but also increases computation time
    refine_from_layer: 20  # which layers to use when calculating the adjacency 
    similarity_metric: modularity  # the metric to use as the objective during refinement: modularity or conductance (or intra_inter_sim but this doesn't work well so far)

  contiguity_buffer_kwargs:  # parameters relating directly to the contiguity buffer
    use_contiguity_buffer: true  # whether to use a contiguity buffer
    contiguity_buffer_size: 0.3  # proportion of n_mem tokens to dedicate to the contiguity buffer

  uniform_blocks: false  # ignore em_splitter (above) and segment chunks into fixed-sized blocks of size max_block_size (above)
  random_topk_blocks: false  # retrieve random blocks rather than the topk most similar blocks

数据准备

我们采用 ∞-Bench 和 LongBench 进行模型评估。你可以通过运行以下命令下载数据集。

bash scripts/download.sh

响应生成

你可以通过运行以下命令评估 EM-LLM。你也可以选择性地传入以下参数以适应你的硬件资源。

bash scripts/run.sh

    -m|--model  # DEFAULT: mistral; OPTIONS: mistral,llama3,llama31,phi3_mini,phi35_mini - Which base LLM to use during evaluation.
    -b|--benchmark  # DEFAULT: long-bench; OPTIONS: long-bench,infinite-bench,passkey - Which benchmark to evaluate. Passkey evaluates an extended version of InfiniteBench's passkey retrieval task (see yaml for context length parameter). 
    -w|--world-size  # DEFAULT: number of visible GPUs - Total number of GPUs to be used during evaluation. 
    -n|--num_gpus_per_job  # DEFAULT: 1 - How many GPUs to attribute to each job. If >1, model layers will be evenly spread over multiple GPUs. 
    -r|--rank_offset  # DEFAULT: 0 - Ignores the first n GPUs visible to the script. Useful when running multiple experiments on a single node.
    -o|--allow_disk_offload  # DEFAULT: False - Whether to allow dynamic disk offloading of memory blocks or not (see the our paper's Appendix for more details). In single-GPU instances this will offload the representative tokens to CPU memory as well.

如果你发现 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}
}