ICLR2025，来自Google DeepMind团队的工作

https://arxiv.org/abs/2410.04343v2

0.目标

先前对于RAG推理扩展的研究主要集中于提供更多的知识，但只增加知识的数量是不够的。当前的LLM在处理长上下文时仍存在挑战。比如，在超长序列中定位有效信息的能力有限、最佳性能往往是在没有充分利用上下文的情况下实现的、超过一定阈值（文档数量）的检索会使性能停滞甚至下降。

目标是找到上下文长度与最优配置之间的关系，能够预测最佳推理参数，最大限度提高RAG性能。其中DRAG的参数为检索到的文档数量和示例数量。IterDRAG的参数为生成次数。

1.贡献

  • 提出两种RAG方法：DRAG（基于演示的RAG，为LLM提供多个RAG示例）和IterDRAG（基于迭代演示的RAG，将输入查询分解为更简单的子查询，迭代检索）。并证明了这两种方法优于仅提供知识的RAG。
  • 得到了RAG的推理缩放定律：在最佳配置下，RAG性能随有效上下文长度线性变化。
  • 根据定律对RAG性能与不同推理参数建模，推导出计算分配模型，为长上下文RAG的优化提供了指导。

2.相关工作

2.1长上下文LLMs

早期采用稀疏/低秩核来减少内存需求。

I. Beltagy, M. E. Peters, and A. Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020.

K. Choromanski, V. Likhosherstov, D. Dohan, X. Song, A. Gane, T. Sarlos, P. Hawkins, J. Davis, A. Mohiuddin, L. Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009.14794,2020.

N. Kitaev, L. Kaiser, and A. Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2019.

M. Zaheer, G. Guruganesh, K. A. Dubey, J. Ainslie, C. Alberti, S. Ontanon, P. Pham, A. Ravula, Q. Wang, L. Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283–17297, 2020

此外递归和状态空间模型（SSMs）被提出，作为基于transformer模型的有效替代方案。

M. Beck, K. Pöppel, M. Spanring, A. Auer, O. Prudnikova, M. Kopp, G. Klambauer, J. Brandstetter, and S. Hochreiter. xLSTM: Extended long short-term memory. arXiv preprint arXiv:2405.04517, 2024

A. Gu and T. Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023.

B. Peng, E. Alcaide, Q. Anthony, A. Albalak, S. Arcadinho, S. Biderman, H. Cao, X. Cheng, M. Chung, L. Derczynski, et al. RWKV: Reinventing rnns for the transformer era. In Findings of the Association for Computational Linguistics: EMNLP 2023, pages 14048–14077, 2023a.

对于因果LLMs，外推和插值方法已被证明在扩展上下文窗口长度方面非常有效。

S. Chen, S. Wong, L. Chen, and Y. Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306.15595, 2023.

B. Peng, J. Quesnelle, H. Fan, and E. Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint arXiv:2309.00071, 2023b.

O. Press, N. A. Smith, and M. Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108.12409, 2021.

Y. Sun, L. Dong, B. Patra, S. Ma, S. Huang, A. Benhaim, V. Chaudhary, X. Song, and F. Wei. A length-extrapolatable transformer. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 14590–14604, 2023.

最近在高效注意力机制方面的进展，使得LLMs能够训练和推理包含数百万个标记的输入序列。

T. Dao, D. Fu, S. Ermon, A. Rudra, and C. Ré. Flashattention: Fast and memory-efficient exact attention with IO-awareness. Advances in Neural Information Processing Systems, 35:16344–16359, 2022

S. A. Jacobs, M. Tanaka, C. Zhang, M. Zhang, L. Song, S. Rajbhandari, and Y. He. DeepSpeed Ulysses:System optimizations for enabling training of extreme long sequence transformer models. arXiv preprint arXiv:2309.14509, 2023.

H. Liu, M. Zaharia, and P. Abbeel. Ring attention with blockwise transformers for near-infinite context. arXiv preprint arXiv:2310.01889, 2023.

J. Achiam, S. Adler, S. Agarwal, L. Ahmad, I. Akkaya, F. L. Aleman, D. Almeida, J. Altenschmidt, S. Altman, S. Anadkat, et al. GPT-4 technical report. arXiv preprint arXiv:2303.08774, 2023.

M. Reid, N. Savinov, D. Teplyashin, D. Lepikhin, T. Lillicrap, J.-b. Alayrac, R. Soricut, A. Lazaridou, O. Firat, J. Schrittwieser, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arXiv:2403.05530, 2024.

G. Team, R. Anil, S. Borgeaud, Y. Wu, J.-B. Alayrac, J. Yu, R. Soricut, J. Schalkwyk, A. M. Dai, A. Hauth, et al. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805, 2023.

2.2上下文学习

上下文学习（ICL）提供了一种计算效率高的方法，通过依赖少量任务示例来提高模型在推理时的表现。

为了进一步提升ICL性能，现有研究集中在预训练策略上，优化语言模型以实现上下文学习。

此外，选择性使用少量示例也被证明有助于提高下游任务的表现。

值得注意的是，重新格式化或找到最佳顺序的场景示例也能提高ICL的有效性。

随着长情境语言模型的出现，在ICL中扩展示例数量成为可能。

例如，Agarwal等人表明，多样本ICL可以减轻LLM中的预训练偏差，从而提高各种任务的ICL性能。

2.3RAG

RAG通过引入外部知识来提高语言模型的表现。

和naiveRAG相比，优化检索阶段可以有效提升上下文相关性，提高生成质量。

REPLUG使用语言模型作为监督来学习一个密集检索器模型。

此外，编码文档可以增加知识检索并提高生成能力。

Izacard和Grave（2021）利用融合解码器架构来编码多个问题-段落对，同时保持模型效率。

或者，有选择地利用文档中的知识可以提高语言模型对无关上下文的鲁棒性。例如，RAFT提出使用负文档训练语言模型以提高生成质量和相关性（Zhang等人，2024）。与我们的工作同时，提出了长文档检索和数据集扩展以优化RAG性能（Jiang等人，2024；Shao等人，2024）。

3.基于RAG的推理扩展策略

我们使用有效上下文长度来测量推理计算，有效上下文长度定义为 LLM 输出最终答案之前，所有迭代的输入token总数。每次的输入受到LLM上下文窗口的限制。忽略输出的token和检索的开销。

3.1DRAG

DRAG基于NaiveRAG构建，将检索到的topk文档和上下文示例集成到输入中。颠倒文档顺序，将排名较高的文档放在离查询近的位置。

3.2IterDRAG

将查询分解为更简单的子查询，每个子查询执行检索并生成中间答案。解决完所有子查询后，将所有上下文、子查询、中间答案组合在一起，生成最终答案。

由于现有的数据集不带有子查询和中间答案，所以让大模型使用约束解码（constrained decoding）并遵循self-ask格式生成example。

生成example的步骤：

在每一轮中，生成子查询，然后将查询到的文档交错放入prompt，再生成中间答案。

最终答案生成。或者达到最大轮数后，强制生成最终答案。

文档、子查询-中间答案、最终答案一起构成example。

推理过程中，上下文示例添加到初始文档之前。

IterRAG还将学习：1.将问题分解为简单可控的子问题 2.提取子问题的相关信息

这一方法有助于提高RAG回答复杂问题的能力。

具体操作见附录H。

4.RAG性能和推理计算规模

4.1给定预算（最大有效上下文长度）下的最佳性能

限制最大输入token即$L_{max}$时，可以通过调整推理参数 $\theta$ 来优化计算资源的使用。

  • 在 DRAG 中，可以调整 检索文档数量 ( k ) 和 上下文示例数量 ( m )；
  • 在 IterDRAG 中，额外引入了 检索与生成的迭代次数 ( n )。

对于每个输入查询及其真实答案 $ (x_i, y_i) \in \mathcal{X} $，我们可以应用参数为 $ \theta $ 的 RAG 推理策略 $ f $，得到预测结果 $ \hat{y}_i = f(x_i; \theta) $，并计算评估指标 $ P(y_i, \hat{y}_i) $。

为了研究 RAG 性能与推理计算量之间的关系，我们在不同的计算预算 $ L_{\text{max}} $ 下采样，并通过枚举不同的 $ \theta \in \Theta $ 来寻找该预算下的最优平均性能 $ P^*(L_{\text{max}}) $：

$P^*(L_{\text{max}}) := \max_{\theta \in \Theta} \left\{ \frac{1}{|\mathcal{X}|} \sum_i P(y_i, f(x_i; \theta)) \ \Bigg| \ \forall i, l(x_i; \theta) \leq L_{\text{max}} \right\}. \quad$

  • $ \mathcal{X} $：测试集，包含输入查询 $ x_i $ 和真实答案 $ y_i $ 的配对 $ (x_i, y_i) $。
  • $ \theta $：RAG 推理参数，包括：
    • $ k $：检索的文档数量，
    • $ m $：上下文示例（in-context examples）数量，
    • $ n $：生成迭代次数（DRAG 中 $ n=1 $，IterDRAG 中 $ n \geq 1 $）。
  • $ f(x_i; \theta) $：使用参数 $ \theta$ 的 RAG 策略对查询 $ x_i $ 的预测结果 $ \hat{y}_i $。
  • $ P(y_i, \hat{y}_i) $：评估指标（如准确率、F1 分数等），衡量预测答案 $ \hat{y}_i $ 与真实答案 $ y_i $ 的匹配程度。
  • $ l(x_i; \theta) $：对查询 $ x_i $ 使用参数 $ \theta $ 时的实际上下文长度（即所有输入 tokens 的总和）。
  • $ L_{\text{max}} $：预算约束，即允许的最大上下文长度。

在所有满足约束的 $ \theta $ 中，选择使平均性能 $ \frac{1}{|\mathcal{X}|} \sum_i P(\cdot) $ 最大的参数组合。

实验：

评估Genmini 1.5 Flash（上下文窗口最高1M）在知识密集型问答数据集（Bamboogle、HotpotQA、MuSiQue 和 2WikiMultiHopQA）上的性能。

评估指标为exact match（EM）、F1、Acc

$L_{max} \in \{16k, 32k, 128k, 1M, 5M \} $tokens

对于DRAG，检索文档数量 $k \in \{0,1,2,5,10,20,50,100,200,500,1000\}$，示例数量 $m \in \{0,2^0,2^1,...,2^8\}$

对于IterRAG，$n \in \{1,2,3,4,5\}$

模型：

  1. zero-shot QA（ZS QA）纯使用LLM自身知识
  2. many-shots QA（MS QA）只加入m个示例
  3. RAG 只使用k个文档
  4. DRAG
  5. IterRAG

4.2总体性能

结论：在任意上下文长度、任意数据集中，都是DRAG、IterRAG效果较好；达不到指定的上下文长度时会被省略。而且最好的P值随上下文长度线性变化。

DRAG、IterdDRAG的性能随上下文长度的扩展而不断提高，而其余方法很快达到峰值。DRAG 在较短的最大长度下表现出色，而 IterDRAG 在更长的有效上下文长度下更有效。

4.3RAG推理缩放定律

  1. 随着有效上下文长度的扩大，最佳性能表现出线性增长。所以可以通过增加计算来提高 RAG 性能，从而在给定可用计算资源的情况下更准确地预测性能。
  2. 对于 $L_{max}$ 高于 $10^5$​，IterDRAG 继续通过交错检索和迭代生成进行有效扩展。说明其更适合长上下文推理。
  3. 上下文超过1M后，性能提升不明显。

4.4对于特定参数的缩放

  1. 文档和示例并非起到同样作用。对于固定配置，增加检索到的文档数量通常会带来更好的性能提升。但k和m也是有阈值的。
  2. 增加示例对IterRAG更有帮助，如示例从0-1时，IterRAG的性能明显提升；而示例对DRAG不明显。

5.长上下文 RAG 的推理计算分配

构建计算分配模型，目的是为了能根据 $L_{max}$ 求 $\theta$。

  1. 性能指标（𝑃）

    • 表示在数据集 $ X $ 上的表现（如准确率），建模为参数 $ \theta $ 的函数。

  2. 资源参数（𝜃）

    • 定义为三维向量 $ \theta := (k, m, n)^T $，包含：
      • ( k )：使用的文档数量
      • ( m )：上下文示例（demonstrations/shots）的数量
      • ( n )：最大迭代/生成步数

  3. 信息量参数（𝑖）

    • 衡量输入内容的信息价值，定义为 $ i := (i_{\text{doc}}, i_{\text{shot}}, 0)^T $：
      • $ i_{\text{doc}} $：文档的信息量
        • 通过“添加1篇文档 vs 零样本（zero-shot）的性能差异”计算。
      • $ i_{\text{shot}} $：示例的信息量
        • 通过“添加1个示例 vs 零样本的性能差异”计算。
      • $ i_{\text{iter}} $ 被忽略（设为0），因实验中发现增加生成步数对性能无显著提升。

  4. 性能模型公式

    $P(\theta) \approx \sigma((a + b \odot i)^T \log(\theta) + c)$

    • 符号说明：
      • $ \odot $：逐元素相乘（Hadamard积）。
      • $ a, b \in \mathbb{R}^3 $：待估计参数，分别表示资源的基础效应和与信息量的交互效应。
      • $ c $：常数偏置项。 a、b、c都是在特定情况下拟合的。
      • $ \log(\theta) $：对资源向量逐元素取对数。
      • $ \sigma $：Sigmoid函数。

$L_{max}$ 相当于 $\theta$ 的限制条件；拟合出abc；目标是最大化 $P(\theta)$ ，就可以得到最好的 $\theta$。