来源：PaperWeekly本文约3300字，建议阅读9分钟
选择性检索会是 RALM 重要的基础技术之一，还有很多重要的后续工作可以做。
一、前言最近一年开始做知识增强语言模型的方向。暑假去到了 AWS AI 的 CodeWhisperer 组实习，在 repository-level code completion 这个任务上进行了一些相关研究，有了一些发现，也投中了一篇 ICML Oral，便把感悟写成此文，希望大家多批评指正。
二、背景
2.1 检索增强语言模型（Retrieval Augmented Language Models, RALMs）一般来说，只要语言模型在推理阶段使用了外部知识，就可以被叫作检索增强语言模型。
https://acl2023-retrieval-lm.github.io/slides/3-architecture.pdf
上图截取自 ACL 2023 的 tutorial [1]，我认为可以大致代表 NLP 社区对 RALM 最普遍的认识。可以看到这里的外部知识大多都是非结构化的知识，比如实体或者大段的文本。然而把知识提供给模型的方式和粒度是多种多样的，比如 kNN-LM [2] 是在最终预测的 token 层面上把模型和检索的分布进行插值，RETRO [3] 在中间层设计了特殊的参数来处理检索到的信息，REPLUG [4] 等工作则直接把检索到的信息放置到模型的输入之中。
随着大模型处理信息的能力越来越强，REPLUG 一类的 RALM 系统的表现越来越突出，因为实现更加方便，也不必须微调模型，在现实生活中也更易被广泛应用。
2.2 选择性/适应性检索（Selective/Adaptive Retrieval）在最早的 RALM 系统中，检索以及使用检索到的信息的方式可以被认为是固定的/不可变的。例如 kNN-LM，默认就是每个 token 检索一次，又如 REPLUG，只是根据问题检索一次。这样的设计实现方便，但是在实际使用中就暴露出一些弊端：
时间成本：由于很多知识库的非常大，实际应用中检索甚至可能耗时高过生成回复。必要性：很多生成的任务或者问题并不知识密集（knowledge-intensive），模型自身的能力足够回答，检索并不能带来什么增益。检索质量：受限于知识库的质量和检索器的能力，很多时候检索到的知识并不相关，反而带来过多噪声，对模型是一种干扰。另外，随着模型生成更长的内容，在生成之前检索的内容很可能对当前生成的 token 已经没有很大帮助，需要检索更相关的信息。因此，RALM 领域内出现了一种新的声音：是否可以选择性/适应性地检索？即不需要检索就不检索，需要多次检索就多检索几次。以下两篇早期工作个人认为比较经典：
Efficient Nearest Neighbor Language Models:
https://aclanthology.org/2021.emnlp-main.461/
这篇文章探索了三种提升 kNN-LM 效率的策略，其中第三种策略即是用 language modeling 的目标训练一个适应性的插值权重，在模型本身可以生成得很好的位置关闭检索。
When Not to Trust Language Models: Investigating Effectiveness of Parametric and Non-Parametric Memories
https://aclanthology.org/2023.acl-long.546/
本篇文章发现一个有趣的现象：当问题需要的知识是高频率的热门知识时，检索增强反而会降低模型的性能。根据这一观察，作者提出只对问到长尾的知识的问题进行检索。
2.3 选择性检索的流派选择性检索的工作在 23 年多了很多，个人认为这些工作的不同点主要在于如何回答一个哲学性的问题：到底什么时候需要检索？
以此为据，现有的选择性检索的工作可以大致分为这几个流派。
（一）问题难度
核心思路：根据问题做判断，遇到不需要检索的问题就不检索。这类方法非常直观，也不需要任何模型/检索即可标注标签，但是个人实验发现一般效果最差，主要原因是这类方法并没有参考知识库和模型的任何信息，导致判断比较主观。
代表工作：
When Not to Trust Language Models: Investigating Effectiveness of Parametric and Non-Parametric Memories
https://aclanthology.org/2023.acl-long.546/
The Web Can Be Your Oyster for Improving Language Models
https://aclanthology.org/2023.findings-acl.46/
Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection
https://openreview.net/forum?id=hSyW5go0v8
（二）检索质量
核心思路：尝试检索，但是如果检索到的信息质量不高就不用。这类方法巧妙利用的检索器一般都会输出的相似度分数，过滤掉了大量噪声，一般认为效果很好，只是无论如何都需要进行一次检索，对效率没有什么提高。
代表工作：
You can’t pick your neighbors, or can you? When and How to Rely on Retrieval in the kNN-LM
https://aclanthology.org/2022.findings-emnlp.218/
Corrective Retrieval Augmented Generation（进一步拓展：如果检索到的信息质量不高就再次检索）
https://arxiv.org/pdf/2401.15884
（三）模型信心
核心思路：在模型信心很低时进行检索增强。这类方法让检索的频率和时机更加灵活，也解决了随着模型的生成，之前检索的信息变得越来越无用的问题。这一类方法主要的问题是需要模型是 well-calibrated，否则模型在信息较高时也会犯错。
代表工作：
Active Retrieval Augmented Generation
https://arxiv.org/abs/2305.06983
（四）以结果为导向
核心思路：如果检索能影响/帮助到模型就检索。这类方法一般是效果最好的一类，因为它既参考了知识库和检索器的能力，又参考了模型使用检索信息的能力。唯一的问题是这类方法往往需要收集模型的输出历史，然后进行一定程度的训练。这一类方法也是和 context compression 方法联系最紧密的，例如下面提及的 FILCO 和 RECOMP。
代表工作：
Efficient Nearest Neighbor Language Models
https://aclanthology.org/2021.emnlp-main.461/
Self-Knowledge Guided Retrieval Augmentation for Large Language Models
https://arxiv.org/abs/2310.05002
RECOMP: Improving Retrieval-Augmented LMs with Compression and Selective Augmentation
https://arxiv.org/abs/2310.04408
Learning to Filter Context for Retrieval-Augmented Generation
https://arxiv.org/abs/2311.08377
Repoformer: Selective Retrieval for Repository-Level Code Completion
https://arxiv.org/abs/2403.10059
Adaptive-RAG: Learning to Adapt Retrieval-Augmented Large Language Models through Question Complexity
https://arxiv.org/abs/2403.14403
三、Repoformer
TL; DR：我们设计了一种新的自监督的训练方式，让代码语言模型可以准确判断代码补全是否需要跨文件检索并鲁棒地使用检索到的信息。实验表明我们的方法和基础的 RALM 相比可以同时提高代码补全的精度和推理速度。
论文标题：
Repoformer: Selective Retrieval for Repository-Level Code Completion
论文链接：
https://arxiv.org/abs/2403.10059
Repoformer官网：
https://repoformer.github.io/
3.1 动机对于大型的软件工程项目来说，好的模块化是代码质量的基础，但是模块之间的相互依赖也让利用模型进行代码补全变得更加困难。最近的工作证明跨文件的检索增强非常有效，但是我们发现（1）检索对推理速度的影响非常大；（2）检索带来的准确度提升非常稀疏。如下图所示，大部分的检索对模型代码补全的表现没有影响，而其余检索的一半往往还有负面影响。
3.2 方法我们提出一种选择性检索的方法，让模型在推理过程中先自我评估当前的代码补全问题是否能从检索中获益，然后直接把这一概率以 token 预测的方式展现出来。
推理阶段，我们模型的推理方式可以看做是 fill-in-the-middle 的一种扩展。我们采取在 fill-in-the-middle 上扩展是因为代码补全往往在需要补全的位置左侧和右侧都有信息。
训练阶段，我们设计了一种自监督的方式来标注选择性检索的标签：
在公开的 repository 中随机采样需要补全的部分 y，把当前文件其余信息记为 X。检索相关的跨文件信息 CC。比较模型基于 X vs. 基于 [X: CC] 的输出，通过跟 y 进行比较判断 CC 是否有用。我们使用了文本匹配来判断第二个输出是否离 y 更近，这一指标对于较短的生成长度来说非常接近正确答案。我们用多任务学习来教会模型自我评估+鲁棒地使用检索到的信息。
3.3 实验结果我们的 Repoformer 模型大幅提高了代码补全的准确率，Repoformer-3B 可以达到 StarCoder 16B 的准确度，Repoformer-16B 创造了新的 SOTA。
和总是检索相比，我们的选择性检索策略可以同时提高代码补全质量和推理速度。
同时，Repoformer 可以用来作为即插即用的选择性检索策略，提升一系列其他模型的代码生成质量和精度。
在论文中我们提供了更多分析和结果，欢迎大家查看。
3.4 Repoformer vs. Self-RAG如果对领域的工作比较熟悉的读者，应该会发现 Repoformer 和 Self-RAG 的 formulation 比较相似。这两篇属于同期工作，Repoformer 在 2023 年暑假已经完成，因为公司内部政策和审核原因，2024 年 3 月才和大家见面。
尽管这两篇的 formulation 相似，个人看来它们代表的选择性检索的流派是完全不同的（上文的一 vs 四）。正是因为这两种流派的假设不同，最后得到的模型在选择性检索的精度会相差很多。在 Repoformer 的论文中，我们分析了 Repoformer 放弃检索的例子，发现其中有 85% 以上都准确地对应了无用检索。
ICML Rebuttal 的过程中，我们还将 Repoformer 学习到的检索策略和流派二、三中代表性的策略进行了对比，发现在同样的检索比例下，Repoformer 的策略带来的代码补全精度大大超过了另外两种。
与之相比，个人对 Self-RAG 的一些分析显示 Self-RAG 学习到的选择性检索的策略和随机选择无异。
四、后续工作个人认为选择性检索会是 RALM 重要的基础技术之一，还有很多重要的后续工作可以做，例如对知识来源的批判和选择，对模型自己知识的表达，检索后的归因，迭代式地搜索，针对性地改进知识库的格式等等……也欢迎大家在评论区畅所欲言。
参考文献
[1] https://acl2023-retrieval-lm.github.io/
[2] https://arxiv.org/abs/1911.00172
[3] https://arxiv.org/abs/2112.04426
[4] https://arxiv.org/abs/2301.12652
