大语言模型何时需要检索？UCLA提出全新自监督选择性检索策略

来源：PaperWeekly

作者：吴迪

单位：UCLA

研究方向 ：自然语言处理

前言

最近一年开始做知识增强语言模型的方向。暑假去到了 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 上扩展是因为代码补全往往在需要补全的位置左侧和右侧都有信息。

训练阶段，我们设计了一种 自监督 的方式来标注选择性检索的标签：

1. 在公开的 repository 中随机采样需要补全的部分 y，把当前文件其余信息记为 X。

2. 检索相关的跨文件信息 CC。
3. 比较模型基于 X vs. 基于 [X: CC] 的输出，通过跟 y 进行比较判断 CC 是否有用。我们使用了文本匹配来判断第二个输出是否离 y 更近，这一指标对于较短的生成长度来说非常接近正确答案。

我们用 多任务学习 来教会模型自我评估+鲁棒地使用检索到的信息。

3.3 实验结果

我们的 Repoformer 模型大幅提高了代码补全的准确率，Repoformer-3B 可以达到 StarCoder 16B 的准确度，Repoformer-16B 创造了新的 SOTA。

和总是检索相比，我们的选择性检索策略可以同时提高代码补全质量和推理速度。