[论文笔记]RA-DIT: RETRIEVAL-AUGMENTED DUAL INSTRUCTION TUNING

引言

今天带来一篇RAG微调的论文笔记——RA-DIT: RETRIEVAL-AUGMENTED DUAL INSTRUCTION TUNING。

为了简单，下文中以翻译的口吻记录，比如替换"作者"为"我们"。

我们引入了检索增强双指令微调(Retrieval-Agumented Dual Instruction Tuning,RA-DIT)：(1)更新预训练的LM以更好地利用检索到的信息； (2) 更新检索器以返回LM所偏好更相关的结果。通过在需要知识利用和上下文感知的任务上进行微调。

1. 总体介绍

现有的RALM(Retrival-Augumented Language Modeling)架构侧重于两个挑战： (i) 增强LLM结合检索知识的能力 (ii) 改进检索组件以返回更相关的内容。

在这项工作中，我们展示了轻量级指令微调本身就能显著提升RALM的性能。提出了RA-DIT，通过在一组任务上进行微调来为任何LLM装备检索能力的方法，这些任务用于训练语言模型预测中的知识利用和上下文感知。使用预训练的LLAMA和最先进的基于双编码器的密集检索器DRAGON+初始化框架。

图 1：RA-DIT 方法分别微调 LLM 和检索器。对于给定示例，LM-ft 组件更新 LLM 以最大化给定检索增强指令的正确答案的可能性；R-ft 组件更新检索器以最小化检索器得分分布与 LLM 偏好之间的 KL 散度。

分两个步骤进行指令微调。对于语言模型微调(LM-ft)，采用标签损失目标，并在每个微调提示中添加一个检索到的背景字段。该字段至于指令之前。通过在微调过程中加入背景文本，引导LLM最佳地利用检索到的信息并忽略干扰内容。

对于检索微调(R-ft)，使用一个基于监督任务和无监督文本完成的组合计算广义LM监督搜索训练目标来更新查询编码器。使检索器能生成与LLM偏好一致的、更具上下文相关的内容。

2. 方法

2.1 架构

语言模型 检索增强型预训练自回归语言模型。特别是，使用LLAMA；

检索器 采用基于双编码器的检索器架构。给定一个语料库$\mathcal{C}$和一个查询$q$，文档编码器将每个文本片段$c\in \mathcal{C}$映射到一个嵌入$E_d(c)$，而查询编码器将$q$映射到一个嵌入$E_q(q)$。基于查询-文档嵌入相似度检索$q$中排名前k的相关文本片段，相似度通过点积计算：
$$s(q,c)=E_q(q)\cdot E_d(c)\text{\hspace{1em}(1)}$$
使用DRAGON +初始化检索器。

并行上下文检索增强 对于给定的语言模型提示$x$，检索出排名靠前的k个相关文本片段${\mathcal{C}}^{\prime }\subset \mathcal{C},|{\mathcal{C}}^{\prime }|=k$。为了保持上下文窗口大小限制，每个检索到的片段都预先添加到提示中，并且对来自多个增强提示的语言模型预测执行并行计算。最终输出概率是每个增强提示的概率混合，并根据片段相关性得分进行加权。
$$p_{LM}(y|x,{\mathcal{C}}^{\prime })=\sum _{c\in {\mathcal{C}}^{\prime }}{p}_{LM}(y|c\circ x)\cdot p_R(c|x)\text{\hspace{1em}(2)}$$
其中$\circ$表示序列拼接；$p_R(c|x)=\frac{\exp s(x,c)}{{\sum }_{c^{\prime }\in {\mathcal{C}}^{\prime }}\exp s(x,c^{\prime })}$ 是在top-k相关片段中重新归一化的检索分数。

使用一对起始(Background)和结束(\n\n)标记来划定增强提示中的检索段。

2.2 微调数据集

语言模型微调数据集(${\mathcal{D}}_L$)包含20个数据集。每个数据集的示例使用手动编译的模板进行序列化。

2.3 检索增强语言模型微调

将每个微调序列分词一个指令段($x$)和一个输出段($y$)。对于每个示例$(x_i,y_i)\in {\mathcal{D}}_L$，根据$x_i$检索$\text{top-}\tilde{k}$相关的文本片段${\mathcal{C}}_i\subset \mathcal{C}$。对于每个检索到的片段$c_{ij}\in {\mathcal{C}}_i$，通过将其作为背景字段添加到指令前面，创建一个单独的微调示例，从而为每个原始示例生成$\tilde{k}$个独立的微调示例： $\{(c_{ij}\circ x_i,y_i|j=1,\cdots ,\tilde{k}\}$。

使用下一词预测目标对语言模型进行微调，最小化每个实例输出段中词的损失：
$$\mathcal{L}({\mathcal{D}}_L)=-\sum _i\sum _j\log p_{LM}(y_i|c_{ij}\circ x_i)\text{\hspace{1em}(3)}$$
在微调过程中集成上下文检索增强带来了双重益处。首先，它使 LLM 能够更好地利用相关背景知识进行预测。其次，即使是最先进的检索器也可能出错并返回不准确的结果。通过训练 LLM 在给出错误检索片段时做出正确预测，我们使 LLM 能够忽略误导性的检索内容，并在这种情况下依赖其参数化知识。

2.4 检索器微调

采用了一种广义的LSR((LM-Supervised Retrieval)训练方法，利用语言模型本身为检索器微调提供检索。

对于检索微调数据集${\mathcal{D}}_R$中的训练样本$(x,y)$，我们定义检索到的片段$c$的LSR分数如下：
$$p_{LSR}(c|x,y)=\frac{\exp (p_{LM}(y|c\circ x)/\tau )}{{\sum }_{c^{\prime }\in \mathcal{C}}\exp (p_{LM}(y|c^{\prime }\circ x)/\tau )}\approx \frac{\exp (p_{LM}(y|c\circ x)/\tau )}{{\sum }_{c^{\prime }\in {\mathcal{C}}^{\prime }}\exp (p_{LM}(y|c^{\prime }\circ x)/\tau )}\text{\hspace{1em}(4)}$$
其中$\tau$是一个温度参数；${\mathcal{C}}^{\prime }\in \mathcal{C}$表示$x$的top-k检索片段。更高的LSR分数表明$c$在提高语言模型预测正确答案的概率方面更有效。LSR训练的目标是让检索器为能够提高LLM生成正确答案的可能性更高的片段分配更高的分数。

为了实现这一点，我们最小化了$p_{LSR}$与公式2中定义的检索器分数$p_R$之间的KL散度：
$$\mathcal{L}({\mathcal{D}}_R)={\mathbb{E}}_{(x,y)\in {\mathcal{D}}_R}KL(p_R(c|x)||p_{LSR}(c|x,y)))\text{\hspace{1em}(5)}$$
在实践中，只更新检索器的查询编码器，因为微调两个编码器会损害性能。

结论

本文提出了一种轻量级的检索增强双指令微调框架 RA-DIT，它可以有效地为任何预训练的 LLM 添加检索能力。RA-DIT 通过检索增强指令微调更新 LLM，以更好地利用检索到的知识并忽略无关或干扰信息。它还通过来自 LLM 的监督对检索器进行微调，以检索能够更好地帮助 LLM 生成正确输出的文本。

总结

⭐ 作者提出了一种检索增强微调的方法，为语言模型和检索器进行微调。引导LLM最佳地利用检索到的信息并忽略干扰内容。