AlignDistil: Token-Level Language Model Alignment as Adaptive Policy Distillation

📄 arXiv: 2503.02832v3 📥 PDF

作者: Songming Zhang, Xue Zhang, Tong Zhang, Bojie Hu, Yufeng Chen, Jinan Xu

分类: cs.CL, cs.AI, cs.LG

发布日期: 2025-03-04 (更新: 2025-07-23)

备注: ACL 2025 Main Conference, code available at: https://github.com/songmzhang/AlignDistil

💡 一句话要点

AlignDistil：提出一种基于Token级别语言模型对齐的自适应策略蒸馏方法

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型对齐 强化学习 人类反馈 直接偏好优化 蒸馏学习 Token级别优化 自适应策略 对比学习

📋 核心要点

  1. 现有LLM对齐方法使用响应级别的奖励，忽略了token级别的细粒度信息，导致优化不充分。
  2. AlignDistil提出了一种token级别的蒸馏方法，利用DPO学习到的奖励来指导token级别的优化。
  3. 实验结果表明，AlignDistil优于现有方法，并实现了更快的收敛速度，提升了模型性能。

📝 摘要（中文）

在大语言模型（LLMs）中，LLM对齐至关重要，通常通过诸如基于人类反馈的强化学习（RLHF）和直接偏好优化（DPO）等方法实现。然而，在大多数现有的LLM对齐方法中，响应中的所有token都使用稀疏的、响应级别的奖励或偏好标注进行优化。忽略token级别的奖励可能会错误地惩罚高质量的token或鼓励低质量的token，从而导致次优的性能和缓慢的收敛速度。为了解决这个问题，我们提出AlignDistil，一种等效于RLHF的token级别奖励优化的蒸馏方法。具体来说，我们将DPO学习到的奖励引入到RLHF目标中，并在理论上证明了该目标与token级别蒸馏过程之间的等价性，其中教师分布线性地组合了来自DPO模型和参考模型的logits。在此基础上，我们通过构建具有正向和反向DPO模型的对比DPO奖励，进一步弥合了DPO模型奖励与纯奖励模型之间的准确性差距。此外，为了避免对不同token的欠优化和过度优化，我们设计了一种token自适应logit外推机制，为每个token构建合适的教师分布。实验结果表明，我们的AlignDistil优于现有方法，并展示了由于其token级别分布奖励优化而带来的快速收敛。

🔬 方法详解

问题定义：现有的大语言模型对齐方法，如RLHF和DPO，通常使用response-level的奖励信号来优化模型。这种方法忽略了token级别的细粒度信息，可能导致高质量的token被错误惩罚，而低质量的token被鼓励，最终影响模型的性能和收敛速度。因此，如何利用token级别的奖励信号进行更有效的模型对齐是一个关键问题。

核心思路：AlignDistil的核心思路是将DPO学习到的奖励信息融入到RLHF的框架中，从而实现token级别的优化。具体来说，论文证明了将DPO奖励引入RLHF目标等价于一个token级别的蒸馏过程。在这个过程中，教师模型由DPO模型和参考模型线性组合而成，从而为每个token提供更精确的优化目标。

技术框架：AlignDistil的整体框架包含以下几个主要模块：1) DPO模型：使用DPO方法训练一个偏好模型，用于学习response-level的奖励信号。2) 参考模型：一个未经对齐的原始模型，用于提供logits的基线。3) 对比DPO奖励：通过构建正向和反向的DPO模型，弥合DPO模型奖励与纯奖励模型之间的差距，提高奖励的准确性。4) Token自适应logit外推机制：根据每个token的特性，自适应地调整教师模型的logits，避免对不同token的欠优化和过度优化。5) 蒸馏训练：使用token级别的蒸馏损失函数，将教师模型的知识迁移到学生模型，实现模型对齐。

关键创新：AlignDistil最重要的创新在于提出了token级别的蒸馏方法，将response-level的奖励信号分解为token级别的优化目标。通过理论证明，将DPO奖励引入RLHF目标等价于token级别的蒸馏过程，从而实现了更细粒度的模型对齐。此外，token自适应logit外推机制能够根据每个token的特性调整优化策略，进一步提升了模型的性能。

关键设计：AlignDistil的关键设计包括：1) 对比DPO奖励：通过正向和反向DPO模型，提高奖励的准确性。2) Token自适应logit外推机制：使用一个可学习的参数来控制DPO模型和参考模型logits的组合权重，从而为每个token构建合适的教师分布。3) 蒸馏损失函数：使用交叉熵损失函数来衡量学生模型和教师模型之间的差异，并使用AdamW优化器进行训练。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，AlignDistil在多个数据集上优于现有的RLHF和DPO方法。例如，在XXX数据集上，AlignDistil的性能提升了X%，并且收敛速度更快。此外，消融实验验证了对比DPO奖励和token自适应logit外推机制的有效性。

🎯 应用场景

AlignDistil可应用于各种需要对齐的大语言模型场景，例如对话系统、文本生成和代码生成等。通过token级别的优化，可以提高生成文本的质量、一致性和安全性，从而提升用户体验。该方法还可以用于模型的持续学习和微调，使其更好地适应新的任务和数据。

📄 摘要（原文）

In modern large language models (LLMs), LLM alignment is of crucial importance and is typically achieved through methods such as reinforcement learning from human feedback (RLHF) and direct preference optimization (DPO). However, in most existing methods for LLM alignment, all tokens in the response are optimized using a sparse, response-level reward or preference annotation. The ignorance of token-level rewards may erroneously punish high-quality tokens or encourage low-quality tokens, resulting in suboptimal performance and slow convergence speed. To address this issue, we propose AlignDistil, an RLHF-equivalent distillation method for token-level reward optimization. Specifically, we introduce the reward learned by DPO into the RLHF objective and theoretically prove the equivalence between this objective and a token-level distillation process, where the teacher distribution linearly combines the logits from the DPO model and a reference model. On this basis, we further bridge the accuracy gap between the reward from the DPO model and the pure reward model, by building a contrastive DPO reward with a normal and a reverse DPO model. Moreover, to avoid under- and over-optimization on different tokens, we design a token adaptive logit extrapolation mechanism to construct an appropriate teacher distribution for each token. Experimental results demonstrate the superiority of our AlignDistil over existing methods and showcase fast convergence due to its token-level distributional reward optimization.