KFPT：用于语言模型训练的可靠性与不确定性过滤自蒸馏方法

《Knowledge-Based Systems》：KFPT: Reliability and Uncertainty Filtered Self-Distillation for Language Model Training

　　穆扎法尔·卡安·尤切 | 梅赫梅特·法蒂赫·阿马西亚利
土耳其伊斯坦布尔埃森勒尔，34220，达武特帕萨校区，耶尔德兹技术大学计算机工程系

摘要
基于Transformer的大语言模型通常通过使用下一个标记的交叉熵损失函数，在大规模文本语料上进行训练。尽管交叉熵损失函数具有可扩展性和稳定性，但在实际应用中可能存在一些问题，如过度自信、在罕见或困难标记上的学习信号较弱，以及训练目标与生成时的行为之间存在差异。在这项工作中，我们提出了知识过滤相位训练方法，这是一种两阶段方案，无需额外的教师模型即可增强训练信号。在第一阶段，该方法在交叉熵损失函数中加入选择性正则化项，并定期通过对同一文本进行二次前向传播、在注意力掩码中屏蔽小块区域，从而平均交叉熵损失值，使更新更加稳定。在第二阶段，该方法通过以跨度丢弃带来的第二视角的分布作为目标，加入单向KL一致性项；该项会根据参考视角的可靠性以及模型的不确定性进行选择性加权，仅在有用位置加强其作用。我们还通过数学定理和证明分析了第一阶段和第二阶段中使用的额外项为何能够发挥效果。在大量实验中，我们将KFPT与多种模型架构及训练方式进行了比较，同时与强大的交叉熵损失函数基线以及先前的无教师模型优化方法进行了对比。结果表明，KFPT通常能提升模型的准确性并降低困惑度，其性能优于现有文献中的无教师模型方法。

引言
基于Transformer的大语言模型是当前自然语言处理研究的核心，它们为包括文本生成、摘要生成和问答在内的多种任务提供了统一的“基础模型”框架[1]、[2]。这类模型最常见的训练方式是使用基于下一个标记预测的自回归目标函数及其标准损失函数——交叉熵损失函数，在大规模文本语料上进行训练[3]、[4]。在训练过程中，模型被引导为每个位置的正确下一个标记赋予更高的概率；由于该目标函数实现简单且能够在大规模场景下稳定优化，因此已成为实际应用中的标准选择[5]。这种“单一模型/多任务”方法的实用性还得到了进一步支持，相关研究要么将不同任务统一到文本到文本的框架下，要么将面向生成的预训练技术应用于各种下游任务[6]、[7]。

该方法成功的另一个原因是Transformer架构凭借其自注意力机制，能够直接建模上下文中的长距离依赖关系，同时还能实现并行计算[2]。事实上，多项研究表明，随着模型规模、数据量和计算资源的增加，模型性能通常会以可预测的方式提升[8]、[9]。因此，下一个标记的交叉熵损失函数被广泛视为现代大语言模型训练的默认起点，它提供了一个定义明确的概率学习框架，并且在实际应用中能够高效地在大规模场景下进行训练[10]。

尽管标准的下一个标记交叉熵损失函数易于扩展且效果显著，但在实际应用中仍存在一些众所周知的缺陷。首先，仅最大化概率并不一定能得到校准良好的模型；因为即使模型为其输出赋予了高概率，仍有可能出错，从而导致“过度自信”问题[11]。其次，由于自然语言的频率分布极不均衡，罕见或困难的标记在交叉熵损失函数下的总损失中所占比例很小，因此获得的训练信号较弱；为此，已有研究提出通过重新加权来解决这种不平衡问题[12]。第三，由于模型在训练时通常以真实的先前标记作为条件，但在推理时必须依靠自身的预测，这就导致了暴露偏差，错误会逐步累积，尤其是在长文本生成过程中，模型质量会下降[13]。最后，虽然交叉熵损失函数优化的是标记级别的概率，但在实际应用中，重复生成以及倾向于使用高频词等不良生成行为可能会变得明显；为此，有研究提出了不可靠性损失函数，通过修改训练目标来直接惩罚这些问题[14]。此外，也有研究探讨了通过抑制模型置信度来缓解交叉熵损失函数缺陷的正则化方法[15]。

尽管标准的下一个标记交叉熵损失函数在实际应用中仍是一种强大且可扩展的起点，但人们也发现，它本身并不足以有效约束模型如何构建其预测分布。因此，近年来，“无教师模型目标优化”这一研究方向逐渐兴起，其目标是无需额外的教师模型即可强化训练目标。这一领域的共同理念是让交叉熵损失函数的信号更侧重于提高正确标记的概率，通过利用模型自身的知识或训练过程中产生的辅助信息，使信号更具信息量。

在这一方向上最自然的例子之一就是自我知识蒸馏技术，该技术将模型自身的输出作为“软目标”反馈给模型自身。在自然语言处理领域，这种方法利用模型预测分布中的相对概率作为正则化信号，无需依赖外部教师模型；通过这种方式，它能促使模型更一致地表示各种可能的输出，而不是过度偏向某一个正确结果，同时也能为理解分布的相对结构提供更丰富的学习信号[16]。

自我蒸馏理念的延伸版本适用于更复杂的场景，比如微调任务，这类动态方法从“近期历史数据”中生成教师模型信号。DynSDPB技术能够利用之前小批量训练中模型输出的软分布，无需在每一步都进行单独的教师模型前向传播，这样既降低了计算成本，又起到了减少不稳定性和微调过程中泛化能力下降风险的正则化作用[17]。

无教师模型优化并不仅限于“重复使用输出分布”；有些方法通过促使模型在不同噪声或视角条件下表现得更稳定，间接增强了基于交叉熵损失函数的训练效果。LR-Drop技术通过在各层之间引入更结构化的一致性约束，减少了由dropout随机性引起的训练与推理之间的差异，旨在提高表示的稳定性并减轻微调过程中的过拟合问题[18]。类似地，还有“损失函数重构”方法，它们在交叉熵损失函数无法产生足够强信号的情况下，重新定义损失目标：MiLe损失函数通过动态估计模型输出分布的熵值来反映标记的难度，从而减少简单/高频标记的主导地位，增强对困难/罕见标记的学习信号[19]。

另一类重要的研究方向则是试图突破下一个标记目标函数的时间局限性。Beyond Multi-Token Prediction方法通过添加能够捕捉更长范围结构的额外信号来增强目标函数，例如未来内容的摘要，从而更直接地解决长上下文生成中的错误累积和规划问题；这一方法表明，无需教师模型，也能让目标函数更接近“实际使用时的目标”[20]。最后，监督微调过程中由于交叉熵损失函数过度依赖单一参考而导致的多样性丧失，也是无教师模型目标优化研究中的一个重要问题：Preserving Diversity in Supervised Fine-Tuning方法致力于通过设计能够减少交叉熵损失函数导致的分布窄化效应、并保留有效替代选项的概率质量的目标函数，来缩小训练目标与期望的生成行为之间的差距[21]。

知识蒸馏技术在其最一般的形式中，不仅让较小的学生模型以“正确标签”（硬目标）作为训练依据，还会让其以较大教师模型生成的软概率分布作为训练依据。在早期的“模型压缩”研究领域，将教师模型预测的概率传递给学生模型，通常能够通过提供关于决策边界的更丰富信号，从而实现更高效的学习目标[22]、[23]。在实践中，这种传递通常是通过在交叉熵损失函数中加入学生模型与教师模型分布之间的KL项来实现的，或者通过基于教师模型中间表示的“提示”类型约束来实现；利用中间层的提示来训练轻量级和深度模型是这一思路的经典示例[24]。在序列生成任务中，知识蒸馏不仅可以在标记级别进行，还可以在序列级别进行；例如在机器翻译任务中，序列级的知识蒸馏可以利用教师模型生成的更为“确定/清晰”的目标序列，从而帮助学生模型学习[25]。这一领域还发展出了多种变体，它们利用中间模型，比如“教师助手”，来弥补教师模型与学生模型之间的能力差距，从而简化学习过程[26]。

在大语言模型知识蒸馏的背景下，由于训练与推理流程之间的差异而产生的暴露偏差，依然是导致学生模型在自由文本生成过程中不断积累错误的一个根本性问题。为了解决这种分布偏移问题，有一类重要的研究将知识蒸馏过程从严格的匹配流程转变为模仿学习问题。具体而言，通过将交互式数据聚合算法，如DaGger和DaD，整合到知识蒸馏流程中，可以让学生模型不仅记住正确的步骤，还能从自身可能出错的生成路径中恢复过来；这种结构上的转变成功克服了传统知识蒸馏方法中存在的累积误差问题，显著提升了生成质量[27]。

在基于Transformer的语言模型中，知识蒸馏技术已经发展到这样的程度：其“通用定义”已经不再局限于标记或序列输出层面；在BERT风格的表示模型中，不仅可以蒸馏输出分布，还可以蒸馏自注意力结构以及中间层的表示。沿着这一思路，Patient KD和TinyBERT等设计通过系统地调整模型各层之间的关系，取得了出色的压缩和适配效果[28]、[29]。同样，MobileBERT以及以注意力机制为核心的MiniLM系列则体现了根据模型架构和计算限制来更直接地调整知识蒸馏目标的做法[30]、[31]。另一方面，也有更多理论研究从更宏观的角度解释了知识蒸馏为何以及何时有效，同时也有全面的综述文章系统总结了该领域的现有研究成果[32]、[33]。总体来看，这些研究既阐明了基于教师模型的知识蒸馏之所以具有“丰富信号”优势的原因，也说明了由于教师模型的额外成本以及教师模型错误可能传播等问题，人们为何会自然而然地转向无教师模型/自我知识蒸馏技术[34]。

基于教师模型的方法长期以来一直是语言模型训练中“丰富目标信号”的主要手段。在最经典的形态中，知识蒸馏旨在将高容量教师模型生成的软概率分布传递给学生模型；这样一来，学生模型不仅能被引导 toward正确答案，还能从教师模型提供的关于“不同选项之间的相对接近程度”的信息中受益[35]。在大型语言模型的应用中，这一理念被广泛用于实现诸如通过教师-学生模型训练来创建更轻量的模型，或是让模型在特定任务上表现出更可控的行为等目标。特别是那些使用教师模型生成的目标进行微调的方法，能够提高学生模型的样本效率，使其输出分布的校准程度更高，同时还能将教师模型所具有的“风格/质量”特征传递给学生模型[36]。

除了传统的序列级匹配之外，将大型语言模型的复杂推理能力转移到小型语言模型中，也已成为克服硬件限制的关键方向。为了解决在模型规模缩小过程中逻辑推理能力的损失问题，人们开发了指令微调技术，旨在让大型教师模型的“思维链”处理方式与小型学生模型相契合[37]。然而，那些仅使用教师模型生成的合成数据进行训练的模型，往往容易陷入记忆化的倾向，其在处理未见过的、属于其他领域的任务时，具备的泛化能力十分有限。为了解决这一泛化瓶颈，最近的研究提出了“自我优化”机制，即让学生模型利用自身的输出，而非仅仅依赖教师模型的输出。通过采用两阶段训练框架，即先进行思维链训练，然后再通过直接偏好优化方法来评估并选择学生模型自身生成的推理路径中最一致的那一条，这些方法打破了模型对静态教师数据的过度依赖，使得模型的推理能力更具动态性和适应性[38]。

尽管如此，基于教师模型的知识蒸馏仍然存在一定的实际成本：维护额外的教师模型（并且通常需要在每一步都进行一次额外的前向传播），会增加整体的计算和内存需求；此外，教师模型本身的容量限制和数据偏差，也可能会直接影响学生模型的学习过程。在我们的研究中，我们明确设定了一个约束条件：不使用额外的教师模型。因此，我们的目标是通过重构目标函数和学习信号本身，在单个模型的训练过程中获得蒸馏技术的一些优势（更丰富的学习信号、更稳定的优化效果以及更好的泛化能力）。在这方面，我们的方法不依赖于依赖教师模型的性能提升，而是属于无需教师模型的范畴，旨在在相同的参数预算和相当的计算成本下实现目标层面的改进。KFPT是一种微调方案，它利用同一模型自身输出生成的可靠“软目标”信号来支持训练的最后阶段（无需使用额外的教师模型），同时保持标准的下一个标记交叉熵训练不变。该方法分为两个阶段：在第一阶段，仅使用交叉熵进行稳定的基础更新；经过指定步数后，第二阶段会引入单向KL项，此时“教师”信号取自通过应用跨度降采样得到的同一模型的另一种视图。KL项并非均匀应用于所有标记：(i) 根据教师视角下正确标记与最强替代标记之间的对数几率差值来判断为可靠的标记会获得更高权重；(ii) 基于学生模型熵值的不确定性过滤器会在学生模型仍存在不确定性的区域进一步强化蒸馏信号。在交叉熵训练阶段，我们使用基于前两大概率差距的聚焦加权RU项来强调“困难”样本；而在KL训练阶段，该项的贡献会被减半，以维持优化的平衡。最后，在过渡到KL阶段时，学习率会按照两阶段调度方案逐渐降低，同时还会应用确定的跨度偏移机制，以提高对长文本的覆盖能力。在这项工作中，我们将这种无需教师模型的目标改进方法命名为KFPT，并做出了以下贡献：•无需教师模型的蒸馏：无需额外教师模型，我们就通过利用同一模型的两种视图生成的单向KL项来丰富交叉熵训练内容。•可控可靠性与不确定性的KL项：我们同时运用基于黄金边际的置信度信号以及基于学生模型熵值的不确定性过滤器，对标记层面的KL项进行加权处理。•针对困难标记的聚焦正则化：在交叉熵训练阶段，我们通过基于边际的聚焦RU项来促进分类分离；而在KL训练阶段则减弱其作用，以保持稳定性。•稳定的过渡策略：我们在训练的最后阶段逐步引入KL信号，并通过分阶段的动态学习率调度来维持优化的平衡。总体而言，KFPT旨在通过引入可控可靠性与不确定性的自我蒸馏机制，解决交叉熵算法因单一且困难的训练目标、稀有/困难标记上较弱的学习信号，以及训练目标与生成时行为之间存在差异等问题而导致的过度自信倾向——而且无需付出引入额外教师模型的代价。在这方面，该方法填补了“无需教师模型的目标改进”领域中的实际空白：它在保留交叉熵算法稳定优化特性的同时，使训练信号更具选择性和情境敏感性。整个流程如图1所示。

KFPT概述
知识过滤阶段训练（KFPT）是一种两阶段训练方案，用于在无需额外教师模型的情况下提升自回归语言模型训练的目标函数。该方法在整个训练过程中都保持标准的下一个标记交叉熵目标，但通过两个附加组件来增强学习信号：(i) 一个选择性计算的正则化项，用于限制模型走向过度自信；(ii) 一个加权的KL一致性项。

KFPT的理论分析
在本节中，我们从每个阶段新增项旨在保障的目标出发，对KFPT的两阶段目标函数进行分析。在后续小节中，我们首先介绍第一阶段的分析，并通过定理和证明说明新增项如何与(i) 小扰动下的决策不变性以及(ii) 基于边际的泛化界限相关联。随后我们分析第二阶段，指出蒸馏KL项在忽略常数项后，可被视为一种……

实验部分
本节介绍了为评估所提出的KFPT方法而进行的实验。实验结果涵盖了不同的模型架构和数据规模。模型性能通过在训练前预留的固定测试集上的下一个标记预测准确率来衡量，该测试集包含总共100份文档。表格中，最佳结果以绿色标出。

为何KFPT是一种有效的方法？
KFPT之所以有效，主要是因为它能够以选择性强且稳定的方式，针对标准下一个标记交叉熵在实际应用中表现较弱的场景进行优化，同时不会干扰交叉熵本身的训练过程。实验表明，KFPT在不同架构（CosmosGpt、Llama、Qwen）和数据规模下都能保持稳定的性能：在Stories数据集上，与交叉熵相比，它在2GB及更大规模（20GB）的设置中都能提升Top-1/Top-5准确率，同时在大多数情况下还能降低困惑度。

结论
在这项工作中，我们提出了知识过滤阶段训练（KFPT），这是一种无需额外教师模型即可强化自回归语言模型训练目标函数的两阶段训练方案。KFPT在整个训练过程中都保持标准的下一个标记交叉熵目标，同时通过选择性正则化项在训练初期抑制模型的过度自信倾向；在后期，则通过引入无需教师模型的KL一致性信号来进一步提升模型性能。

生成式AI使用声明
在准备这项工作时，作者使用了ChatGPT来提升文本的语言表达质量和可读性。在使用该工具/服务之后，作者对内容进行了必要的审阅和编辑，并对最终发表的内容承担全部责任。

CRediT作者贡献声明
Muzaffer Kaan Yuce：撰写——初稿、软件实现、方法设计、概念构思。Mehmet Fatih Amasyali：撰写——审稿与编辑、指导监督、方法设计、概念构思。

利益冲突声明
作者声明自己没有已知的可能影响本文研究结果的财务利益或个人关系。

致谢
本研究得到了土耳其科学技术研究委员会（TUBITAK）的资助，项目编号为124E055。此外，该工作还得到了Yildiz技术大学科学研究项目协调处的支持，项目编号为FDK-2024-6421。

Muzaffer Kaan Yuce于2020年在土耳其萨卡里亚大学的计算机工程专业获得学士学位。自2021年起，他成为Yildiz技术大学的博士研究生，目前在该校计算机工程系担任研究助理。他的当前研究兴趣包括自然语言处理、优化算法以及深度学习领域。