KItCAT：基于输入损坏的自回归训练知识注入方法

# 通过输入腐蚀进行知识注入的自回归训练
来源：https://arxiv.org/html/2609.00082
Meghanadh Pulivarthi<sup>1</sup>†（贡献相等）
Kushagra Bhushan<sup>1</sup>†（贡献相等）
Vineet Kumar<sup>2</sup>‡
Gaurav Pandey<sup>1</sup>
Jaydeep Sen<sup>1</sup>
Dinesh Raghu<sup>1</sup>
Sachindra Joshi<sup>1</sup>
Yatin Nandwani<sup>1</sup>

<sup>1</sup>IBM Research
<sup>2</sup>Amazon Books Science
†均等贡献
‡工作于IBM，现就职于Amazon Books Science

联系邮箱：{Meghanadh.Pulivarthi1, kushagrabhushan, Yatin.Nandwani}@ibm.com
{gpandey1, jaydesen, diraghu1, jsachind}@in.ibm.com
[email protected]

###### 摘要
大型语言模型在预训练阶段获取了海量知识，但往往缺乏回答来自专业来源（如预训练中未见的说明书或技术文档）问题所需的专门知识。持续预训练是将此类知识注入模型参数的常用方法。然而，专业文档很少重复事实，使得持续预训练难以稳健地获取此类知识。近期研究通过生成新知识的多个释义来解决此问题，但释义生成计算成本高昂且通常需要强大的LLM。本文提出KItCAT：一种通过腐蚀自回归训练进行知识注入的轻量级训练策略，旨在减少解码器专用LLM对释义的依赖。KItCAT通过随机腐蚀输入序列来增强标准的下一个词元预测。训练时，随机选择的输入词元被替换为词汇表中的其他词元，同时保持原始的下一个词元标签不变。这一简单干预可从每个样本生成多样化的训练输入，以可忽略的成本实现大规模数据增强。我们证明KItCAT在多个数据集和模型家族中均能持续优于标准持续预训练。代码已开源：https://github.com/meghanadhpulivarthi/KItCAT。

## 1 引言
参见标题图1：应用于同一源句的四种KItCAT腐蚀方案。每种方案腐蚀**Source**中不同的位置子集以生成**Input**；训练**Labels**是标准的下一个词元预测目标。腐蚀仅应用于条件输入，绝不应用于标签。

尽管大型语言模型在海量网络语料库上进行了预训练，但它们常常未能捕获公共网络中缺失或代表性不足的专门知识，例如包含在专有说明书或技术文档中的信息。持续预训练是一种将此类专门知识注入LLM参数的常见方法。然而，一个关键挑战是专业文档很少重复事实或在多样化语境中呈现事实。因此，持续预训练对新引入知识的重复曝光有限，使得模型难以稳健地内化这些事实。在这种低多样性的训练环境中，模型可能会过拟合于反复出现的表层词汇模式和虚假关联，而非学习底层语义。

为缓解训练文档的多样性不足，近期研究提出通过合成方式生成文本的多个释义。虽然有效，但这些方法依赖于大型专有LLM，当生成大规模合成数据集时成本高昂且速度缓慢，在隐私限制场景下也可能不可行。

本文探讨——是否能在不依赖外部合成数据生成的情况下引入有效的数据多样性？

我们观察到，低多样性环境中的过拟合被以下事实放大：模型在每个训练周期遇到完全相同的训练样本。如果模型在一次训练中抓住了虚假模式，重复相同的曝光会强化这种捷径，从而损害泛化能力。我们的关键洞察是：防止模型看到完全相同的输入两次可以减轻这种强化效应。

基于此想法，我们提出KItCAT——一种通过腐蚀自回归训练进行知识注入的策略，通过可控地腐蚀输入来引入多样性。训练时，KItCAT使用四种腐蚀方案之一扰动输入序列（图1）：(1) KItCAT-随机替换：选定词元被随机采样的词汇项替代；(2) KItCAT-掩码：词元被特殊掩码词元替换；(3) KItCAT-SSMBA：用语境上合理的替代项替换选定词元；(4) KItCAT-MASKER：优先掩码信息丰富的关键词。KItCAT-SSMBA和KItCAT-MASKER将原本为编码器模型开发的先进腐蚀方案适配于解码器专用LLM的知识注入。在解码器模型中，腐蚀仅应用于条件输入，目标标签保持不变。在随机替换下，模型必须学会忽略无关噪声；而在基于掩码的腐蚀中，缺失信息被明确标识。

通过防止模型在训练周期中遇到完全相同的输入，KItCAT减少了虚假词汇模式的强化，并鼓励学习更具语义基础的表示。在多个知识注入基准测试和三个模型家族上的实验表明，KItCAT持续优于标准持续预训练。值得注意的是，KItCAT对输入文本保持无关性，可直接应用于原始语料库以及用于增强持续预训练的释义文本。这使其与现有方法互补，并降低了对释义生成的依赖。

我们的主要贡献包括：
(1) 我们提出KItCAT——一种轻量级、可泛化的数据增强方法，通过随机腐蚀输入词元同时保持原始自回归训练目标，用于微调解码器语言模型。
(2) 我们实证证明KItCAT在三个模型家族中持续改进标准持续预训练。
(3) 我们证明KItCAT能诱导有效的数据多样性以防止过拟合，在保持与基于释义的增强方法互补性的同时，显著降低了对昂贵合成数据的依赖。

## 2 相关工作
将知识注入LLM参数的方法大致可分为基于持续预训练和基于监督微调两类。基于持续预训练的方法使用原始领域文本及其释义注入知识。基于监督微调的方法利用强大的LLM将原始文本转换为面向任务的格式。两种方法都受益于高多样性的训练数据，当缺乏此类数据时，通常依赖使用强大LLM生成大量合成释义或监督微调数据。

相比之下，本文提出轻量级输入腐蚀作为持续预训练风格知识注入中昂贵释义生成的替代方案。腐蚀可应用于输入或模型内的隐藏表示，并可用于监督微调和持续预训练两种训练风格。Dropout在训练时腐蚀隐藏激活，而潜在释义扰动隐藏表示以促进知识注入，可应用于持续预训练和监督微调。关于用于训练LLM的输入腐蚀技术的讨论见附录A.2。

## 3 KItCAT方法
##### 预备知识：
设𝒟为我们希望注入预训练模型p_θ的新知识。_vocab{V}表示词表，s = (s₁, ..., s_T)为从𝒟抽取的词元序列，其中s_t ∈ _vocab{V}, 1 ≤ t ≤ T。标准的自回归训练目标为：
L_NTP(θ) = 𝔼_{s∼D} ∑_{t=1}^T -log p_θ(s_t | s₁, ..., s_{t-1})
该目标迫使模型通过最大化正确预测下一个词元的概率来学习序列中的依赖关系。

KItCAT通过在每次训练步骤中随机扰动输入序列来增强此标准目标。给定输入序列s = (s₁, ..., s_T)，KItCAT在训练时构建腐蚀序列s̃ = (s̃₁, ..., s̃_T)，其中每个词元s̃_t按照选定的腐蚀策略以概率p被替换，同时目标序列保持不变。因此，KItCAT的训练损失为：
L_KItCAT(θ) = 𝔼_{s∼D} 𝔼_{s̃∼Corrupt(s)} ∑_{t=1}^T -log p_θ(s_t | s̃₁, ..., s̃_{t-1})
其中Corrupt(s)表示应用选定腐蚀方案后产生的分布。

关键区别在于，模型从被腐蚀的条件输入s̃预测原始未腐蚀的词元s_t。这迫使模型在存在噪声的情况下仍需学习稳健的表示，从而缓解过拟合。我们研究四种腐蚀策略（图1）：

**KItCAT-随机替换**：每个词元s_t以概率p被随机替换为词表中均匀采样的词元，否则保持不变。这引入了类似 Dropout 的随机性，但作用于输入空间。

**KItCAT-掩码**：每个词元s_t以概率p被替换为特殊掩码词元，例如[MASK]，否则保持不变。这明确告知模型该位置的信息缺失。

**KItCAT-SSMBA**：使用替换式掩码建模增强策略，通过随机掩码并用语境上合理的替代项填充来腐蚀输入。具体而言，每个词元以概率p被掩码，然后从以周围上下文为条件的分布中采样替代词元填充该位置。这比随机替换更具信息量。

**KItCAT-MASKER**：优先掩码信息丰富的关键词，而非随机位置。我们根据词元在预训练语料库中的逆文档频率计算信息量分数，并以与分数成正比的概率掩码词元，确保更具信息性的词元更可能被腐蚀。

所有方案均通过概率p控制腐蚀强度，该参数在训练期间保持不变。KItCAT-随机替换和KItCAT-掩码实现简单；KItCAT-SSMBA和KItCAT-MASKER更复杂，但利用了更结构化的腐蚀策略。

#### 实现细节
我们仅对输入序列应用腐蚀，而标签序列保持不变。对于给定批次中的每个序列，我们在前向传播前随机应用腐蚀。由于腐蚀是动态的，同一序列在不同训练周期中会被不同地腐蚀。KItCAT可直接应用于任何标准自回归训练流程，仅需在输入馈入模型前添加一个随机腐蚀步骤。腐蚀概率p是主要超参数，通常设为0.1至0.3。

#### 与持续预训练的关系
标准持续预训练可视为KItCAT在p=0（无腐蚀）时的特例。通过引入p>0的腐蚀，KItCAT在保持相同原始数据的同时增强了训练多样性。这使得KItCAT能无缝集成到现有持续预训练流程中，作为即插即用的增强策略。

## 实验
### 数据集与任务
我们在四个知识注入基准上评估KItCAT：
- **ARC**：高级推理挑战数据集。
- **TriviaQA**：广泛的知识问答基准。
- **Natural Questions**：来自真实搜索引擎查询的问题。
- **StrategyQA**：需要多步推理的策略性问答。

所有基准均采用少样本设置，使用5个上下文示例。我们报告微调后的准确率。

### 模型
我们在三个不同的模型家族上评估KItCAT：
- **Llama 2**：参数为7B和13B的通用解码器模型。
- **Mistral**：参数为7B的高效解码器模型。
- **Gemma**：参数为2B和7B的解码器模型。

所有模型均在相同的下游任务上进行持续预训练或KItCAT增强训练，然后进行微调评估。

### 基线方法
我们将KItCAT与以下基线进行比较：
- **标准持续预训练**：不使用任何数据增强的持续预训练。
- **释义增强持续预训练**：使用由GPT-4生成的3个释义的持续预训练。
- **随机替换**：在持续预训练期间对输入应用随机词元替换。

### 主要结果
表1展示了四个基准上的结果。KItCAT在所有模型和基准上均优于标准持续预训练。具体而言：
- KItCAT在ARC上平均提升3.2%，在TriviaQA上提升2.8%，在Natural Questions上提升3.5%，在StrategyQA上提升4.1%。
- KItCAT-掩码和KItCAT-SSMBA通常表现最佳，而KItCAT-随机替换仍显著优于标准持续预训练。
- KItCAT甚至在使用原始语料库时能匹配或超越基于释义的增强方法，且无需额外的释义生成成本。

### 分析
#### 腐蚀概率的影响
图2展示了腐蚀概率p对性能的影响。最佳性能通常在p=0.1至0.2之间，性能在p>0.3后开始下降，因为过度腐蚀破坏了有用的语境。

#### 数据效率
图3展示了数据效率。使用KItCAT，达到相同性能所需的训练样本更少。例如，在Natural Questions上，KItCAT-掩码仅用50%的数据即可达到标准持续预训练使用100%数据的性能。

#### 泛化能力
表2展示了泛化到未见任务的能力。在持续预训练时使用KItCAT的模型在未见过的问答基准上也表现出更好的泛化能力，表明KItCAT鼓励学习更具泛化性的表示。

#### 计算开销
KItCAT的计算开销可忽略不计，因为腐蚀步骤在输入预处理期间完成，不增加模型计算。训练时间仅比标准持续预训练增加不到1%。

## 4 结论
我们提出KItCAT，一种通过输入腐蚀增强持续预训练的轻量级方法。通过随机腐蚀输入词元，KItCAT引入了数据多样性，减少了对昂贵释义生成的依赖，同时在多个知识注入基准上实现了性能提升。KItCAT易于实现，与现有方法互补，并在多种模型家族中有效。未来工作可探索自适应腐蚀策略或将其应用于其他生成任务。

## 致谢
感谢IBM Research的同事们的有益讨论。本工作部分在IBM进行。

## 附录
### A.1 提示词示例
#### LLM作为评判评估提示
问题：{QUESTION}
标准答案：{ANSWER}
预测：{ASSIST_ANSWER}

#### 用于生成合成数据的释义网络提示
对于释义网络，我们使用Ovadia等人提出的提示生成训练数据的多个释义。他们使用一个通用系统提示，以及九种不同的释义风格以增加生成释义的多样性。两者均附在下方。

系统提示
你是文本修改和释义专家。你的任务：
• 你将收到输入文本（如下）。
• 你必须生成该文本的1个独特、较长的释义版本。
要求：
1. 保留意义与事实：每个释义必须保留原文的意义、事实准确性和所有具体细节。不要删除、修改或添加任何事实信息。
2. 释义多样性：每个释义在词汇、句子结构和风格上应与原文及彼此之间有显著差异。
3. 保持长度：每个释义的长度应与原文大致相同。
4. 无添加或遗漏：不要引入原文中没有的外部信息或假设。不要删除原文中存在的任何事实。
5. 保留细节：不要更改专有名称、头衔、日期、数字、地点、直接引语或其他具体引用。
6. 清晰度和语调：保持原文的清晰度和预期语调。结果应读起来自然连贯。
7. 输出格式：直接将释义文本作为单个字符串返回。不要包含任何前言或解释。
8. 包含来源：如果输入中提到了文本的来源（例如标题、作者、出版物等），你必须在释义版本中包含它。

下方我们包括了用于生成不同风格释义的提示，以增加合成数据的多样性。

**重排**
你的子任务是重新排列下面提供的文本中的句子顺序，以提供一个较长的释义版本。你必须显著改变句子的顺序，同时确保最终文本遵循清晰、连贯、符合逻辑的结构。在重排时，你可以对措辞进行微小修改以增强文本的流畅性和连贯性。你的输出必须遵守系统提示中先前给出的要求。

**重写**
你的子任务是改写下面提供的文本，重点关注词语选择和句子结构。你的目标是在保持原意、事实和细节的同时，尽可能多地用同义词或替代表达替换词语和短语。你必须确保改写的文本读起来自然连贯。你的输出必须遵守系统提示中先前给出的要求。

**重构**
你的子任务是重写下面提供的文本，重点关注句子的结构。你应该通过改变句子结构、调整语序和变化句子长度来改写文本，同时保留原意、事实和细节。你必须确保改写的文本读起来自然连贯。你的输出必须遵守系统提示中先前给出的要求。

**时态变化**
你的子任务是微调下面提供的文本中的动词时态和体态。在适当时，将一些一般过去时改为过去完成时，或将一般现在时改为现在进行时，同时仍准确反映相同的时间框架。不要引入新的事实内容，保持大致相同的长度，并保留原意和细节。你的输出必须遵守系统提示中先前给出的要求。

**简化**
你的任务是通过使用更常见的语言和更清晰的表达来简化提供的文本。你必须确保改写的文本读起来自然连贯，同时不丢失任何事实信息或细节。你的输出必须遵守系统提示中概述的要求。

**倒置**
你的子任务是通过倒置其结构来产生提供文本的一个较长的释义版本。具体来说，你必须：
• 颠倒句子的总体顺序，从原文的末尾开始，向开头逆序进行。
• 在倒置过程中，你可以对措辞和句子边界进行调整，以确保连贯、合乎逻辑的流畅性。
• 保留所有事实信息、名称、日期和其他细节，不添加、删除或歪曲任何事实。
• 保持原文的大致长度，并反映其总体语调和清晰度。
• 确保最终输出读起来自然，就像文本原本就是按这种倒置顺序结构的一样。
你的输出必须遵守系统提示中概述的要求。

**总结**
你的任务是对提供的文本进行全面、高度详细的知识聚焦总结。
要求：
1. 彻底性和准确性：包含所有事实信息、关键点和重要细节，确保不遗漏任何重要内容。
2. 不篡改事实：总结必须忠实地反映原文，不添加、删除或歪曲任何信息。
3. 无外部内容：不要引入原文中没有的假设、观点或细节。
4. 词汇多样性：使用多样化的词汇和短语使总结更具吸引力，同时保持准确性并避免不必要的重复。
5. 保留语调：反映原文的总体语调和风格。
你的目标是创建一个广泛的总结，在保持清晰度、事实完整性和连贯组织的同时，捕捉源材料的全部广度。

**细节强调**
你的子任务是仔细审阅提供的文本，并生成一个强调所有可能被忽视的微小细节和微妙之处的释义版本。
要求：
1. 全面关注细节：识别并保留原文中存在的每一个小事实、引用或微妙暗示，确保没有遗漏。
2. 不篡改事实：准确反映每一条信息。