TIS-DPO: Token-level Importance Sampling for Direct Preference Optimization With Estimated Weights

📄 arXiv: 2410.04350v3 📥 PDF

作者: Aiwei Liu, Haoping Bai, Zhiyun Lu, Yanchao Sun, Xiang Kong, Simon Wang, Jiulong Shan, Albin Madappally Jose, Xiaojiang Liu, Lijie Wen, Philip S. Yu, Meng Cao

分类: cs.CL

发布日期: 2024-10-06 (更新: 2025-04-15)

备注: Published in ICLR 2025, code in https://github.com/exlaw/TIS-DPO

💡 一句话要点

TIS-DPO：基于Token重要性采样的直接偏好优化方法

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

关键词: 直接偏好优化 重要性采样 大型语言模型 偏好对齐 Token级别优化

📋 核心要点

  1. DPO忽略了不同token的重要性差异，导致优化效率降低，难以达到最优结果。
  2. TIS-DPO通过token级重要性采样，为每个token分配权重，实现更精细的偏好对齐。
  3. 实验表明，TIS-DPO在多个任务上显著优于基线方法，并能有效识别关键token。

📝 摘要（中文）

直接偏好优化(DPO)因其简单有效而被广泛应用于大型语言模型(LLM)的偏好对齐。然而，DPO被认为是一个bandit问题，其中整个响应被视为单个臂，忽略了token之间的重要性差异，这可能会影响优化效率并难以达到最佳结果。本文提出，DPO的最佳数据集中，获胜和失败响应中的每个token应具有相等的预期奖励，因为token重要性没有差异。由于最佳数据集在实践中不可用，我们建议使用原始数据集进行重要性采样，以实现无偏优化。因此，我们提出了一种token级重要性采样DPO目标，名为TIS-DPO，它根据每个token的奖励为其分配重要性权重。受先前工作的启发，我们使用一对对比LLM的预测概率差异来估计token的重要性权重。我们探索了三种构建这些对比LLM的方法：（1）使用对比提示引导原始LLM，（2）使用获胜和失败响应训练两个独立的LLM，以及（3）使用获胜和失败响应执行正向和反向DPO训练。实验表明，TIS-DPO在无害性和有用性对齐以及摘要任务上显著优于各种基线方法。我们还可视化了估计的权重，证明了它们识别关键token位置的能力。

🔬 方法详解

问题定义：DPO将整个回复视为一个整体，忽略了回复中不同token的重要性差异。这种处理方式可能导致优化效率低下，并且难以达到最优的偏好对齐效果。现有方法未能充分利用token级别的细粒度信息，导致模型无法区分关键token和非关键token，从而影响了模型的学习效果。

核心思路：TIS-DPO的核心思路是为每个token分配一个重要性权重，从而在DPO的优化过程中考虑token级别的差异。论文假设，在理想的偏好数据集中，获胜回复和失败回复中的每个token应该具有相同的预期奖励。由于理想数据集不可得，因此采用重要性采样的方法，利用原始数据集来近似理想数据集，从而实现无偏的优化。

技术框架：TIS-DPO的整体框架是在DPO的基础上引入了token级别的重要性采样。首先，使用对比LLM估计每个token的重要性权重。然后，将这些权重应用到DPO的损失函数中，从而实现token级别的偏好对齐。框架主要包含三个阶段：1) 使用对比LLM估计token重要性权重；2) 将权重融入DPO损失函数；3) 使用加权后的DPO目标进行模型训练。

关键创新：TIS-DPO的关键创新在于引入了token级别的重要性采样机制，从而能够更精细地控制DPO的优化过程。与传统的DPO方法相比，TIS-DPO能够区分不同token的重要性，并根据其重要性调整优化方向。此外，论文还提出了三种构建对比LLM的方法，用于估计token的重要性权重。

关键设计：论文提出了三种构建对比LLM的方法：(1) 使用对比提示引导原始LLM；(2) 使用获胜和失败回复训练两个独立的LLM；(3) 使用获胜和失败回复执行正向和反向DPO训练。Token重要性权重通过对比LLM的预测概率差异来估计。TIS-DPO的损失函数是在DPO损失函数的基础上，对每个token的损失项乘以其对应的权重。具体参数设置和网络结构细节未在摘要中详细说明，需要参考原文。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，TIS-DPO在无害性和有用性对齐以及摘要任务上显著优于各种基线方法。论文还通过可视化展示了TIS-DPO估计的token权重，验证了其能够有效识别关键token位置的能力。具体的性能提升数据需要在原文中查找。

🎯 应用场景

TIS-DPO可应用于各种需要偏好对齐的大型语言模型应用场景，例如对话系统、文本摘要、代码生成等。通过更精细的token级别优化，可以提升模型在安全性、有用性和对齐性方面的表现。该方法有助于构建更安全、更可靠、更符合人类价值观的AI系统，并能提升用户体验。

📄 摘要（原文）

Direct Preference Optimization (DPO) has been widely adopted for preference alignment of Large Language Models (LLMs) due to its simplicity and effectiveness. However, DPO is derived as a bandit problem in which the whole response is treated as a single arm, ignoring the importance differences between tokens, which may affect optimization efficiency and make it difficult to achieve optimal results. In this work, we propose that the optimal data for DPO has equal expected rewards for each token in winning and losing responses, as there is no difference in token importance. However, since the optimal dataset is unavailable in practice, we propose using the original dataset for importance sampling to achieve unbiased optimization. Accordingly, we propose a token-level importance sampling DPO objective named TIS-DPO that assigns importance weights to each token based on its reward. Inspired by previous works, we estimate the token importance weights using the difference in prediction probabilities from a pair of contrastive LLMs. We explore three methods to construct these contrastive LLMs: (1) guiding the original LLM with contrastive prompts, (2) training two separate LLMs using winning and losing responses, and (3) performing forward and reverse DPO training with winning and losing responses. Experiments show that TIS-DPO significantly outperforms various baseline methods on harmlessness and helpfulness alignment and summarization tasks. We also visualize the estimated weights, demonstrating their ability to identify key token positions.