KTO 训练器

概述

TRL 支持 Kahneman-Tversky 优化（KTO）Trainer，用于训练语言模型，如 Kawin Ethayarajh、Winnie Xu、Niklas Muennighoff、Dan Jurafsky 和 Douwe Kiela 在论文 KTO: Model Alignment as Prospect Theoretic Optimization 中所述。

该论文的摘要如下：

Kahneman & Tversky 的前景理论指出，人类以有偏但定义明确的方式感知随机变量；例如，人类以损失厌恶而闻名。我们表明，用于将 LLM 与人类反馈对齐的目标函数隐式地融入了许多此类偏差——这些目标函数（例如 DPO）相较于交叉熵最小化的成功，部分可归因于它们是人类感知的损失函数（HALOs）。然而，这些方法所赋予人类的效用函数仍与前景理论文献中的效用函数存在差异。利用 Kahneman-Tversky 人类效用模型，我们提出了一种 HALO，它直接最大化生成的效用，而非像当前方法那样最大化偏好的对数似然。我们将这种方法称为 Kahneman-Tversky 优化（KTO），在 1B 到 30B 的规模上，其性能与基于偏好的方法相当或更优。至关重要的是，KTO 不需要偏好数据——对于给定输入，它只需要一个二元信号，表明输出是可取的还是不可取的。这使得它在偏好数据稀缺且昂贵的现实世界中更容易使用。

官方代码可在 ContextualAI/HALOs 中找到。

此训练后方法由 Kashif Rasul、Younes Belkada、Lewis Tunstall、Pablo Vicente 贡献，之后由 Albert Villanova del Moral 进行了重构。

快速入门

本示例演示了如何使用 KTO 方法训练模型。我们使用 Qwen 0.5B 模型 作为基础模型。我们使用来自 KTO Mix 14k 的偏好数据。您可以在以下位置查看该数据集：

以下是用于训练模型的脚本：

# train_kto.py
from datasets import load_dataset
from trl import KTOConfig, KTOTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
train_dataset = load_dataset("trl-lib/kto-mix-14k", split="train")

training_args = KTOConfig(output_dir="Qwen2-0.5B-KTO")
trainer = KTOTrainer(model=model, args=training_args, processing_class=tokenizer, train_dataset=train_dataset)
trainer.train()

使用以下命令执行该脚本：

accelerate launch train_kto.py

在 8 块 H100 GPU 上分布式训练，大约需要 30 分钟。您可以通过查看奖励曲线来验证训练进度。奖励差值呈上升趋势表明模型正在改进，并随时间推移生成更好的响应。

要查看训练后的模型的表现，您可以使用Transformers Chat CLI。

$ transformers chat trl-lib/Qwen2-0.5B-KTO
<quentin_gallouedec>:
What is the best programming language?

<trl-lib/Qwen2-0.5B-KTO>:
The best programming language can vary depending on individual preferences, industry-specific requirements, technical skills, and familiarity with the specific use case or task. Here are some widely-used programming languages that have been noted as popular and widely used:

Here are some other factors to consider when choosing a programming language for a project:

 1 JavaScript: JavaScript is at the heart of the web and can be used for building web applications, APIs, and interactive front-end applications like frameworks like React and Angular. It's similar to C, C++, and F# in syntax structure and is accessible and easy to learn, making it a popular choice for beginners and professionals alike.
 2 Java: Known for its object-oriented programming (OOP) and support for Java 8 and .NET, Java is used for developing enterprise-level software applications, high-performance games, as well as mobile apps, game development, and desktop applications.
 3 C++: Known for its flexibility and scalability, C++ offers comprehensive object-oriented programming and is a popular choice for high-performance computing and other technical fields. It's a powerful platform for building real-world applications and games at scale.
 4 Python: Developed by Guido van Rossum in 1991, Python is a high-level, interpreted, and dynamically typed language known for its simplicity, readability, and versatility.

预期的数据集类型和格式

KTO 需要一个非配对偏好数据集。或者，您也可以提供一个配对偏好数据集（通常简称为偏好数据集）。在这种情况下，训练器会自动将其转换为非配对格式，方法是分离被选中（chosen）和被拒绝（rejected）的回复，并为被选中的补全分配label = True，为被拒绝的补全分配label = False。

KTOTrainer 同时兼容 标准 和 对话 数据集格式。当提供对话数据集时，训练器会自动将聊天模板应用于该数据集。

# Standard format
unpaired_preference_example = {"prompt": "The sky is", "completion": " blue.", "label": True}

# Conversational format
unpaired_preference_example = {"prompt": [{"role": "user", "content": "What color is the sky?"}],
                               "completion": [{"role": "assistant", "content": "It is blue."}],
                               "label": True}

理论上，数据集应至少包含一条被选中的补全和一条被拒绝的补全。不过，一些用户已成功使用仅包含被选中数据或仅包含被拒绝数据的集运行 KTO。如果仅使用被拒绝的数据，建议采用较为保守的学习率。

深入了解 KTO 方法

Kahneman-Tversky 优化（KTO）是一种训练方法，旨在仅使用给定输入的输出是可取的还是不可取的这一二元信号来对齐语言模型，而不是使用偏好/非偏好补全对。该方法借鉴了 Kahneman 和 Tversky 的前景理论，定义了一种人类感知损失（HALO），直接最大化生成的效用，并通过对可取和不可取示例进行非对称加权来反映人类的损失厌恶。

本节详细说明了 KTO 在实际应用中的工作原理，涵盖关键步骤：预处理和损失计算。

预处理和分词

在训练期间，每个样本应包含一个提示（prompt）、一个completion以及一个布尔值label，用于指示该补全是否可取（True）或不可取（False）。有关预期格式的详细信息，请参阅数据集格式。 KTOTrainer 会使用模型的 tokenizer 对每个输入进行分词。

计算损失

KTO 损失函数中使用的 KL 散度项是通过将每个提示词与从批次其他位置抽取的不匹配补全进行配对来估计的。因此，估计 KL 项的损失类型（除 apo_zero_unpaired 外的所有类型）都需要每个设备的训练批次大小大于 1，并采用顺序采样策略，以确保不匹配对在批次内保持稳定。

KTO 中使用的损失（ 论文) 中的公式 7）定义如下： $$\mathcal{L}_{\mathrm{KTO}}(\theta) = \mathbb{E}_{(x,y)}\!\left[ w(y)\Big(1 - v(x, y)\Big) \right]$$

其中值 $v(x, y)$ 为 $$v(x, y) = \begin{cases} \sigma\!\left(\beta\big(\log\frac{\pi_{\theta}(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)} - \mathrm{KL}\big)\right) & \text{if } y \text{ is desirable} \\ \sigma\!\left(\beta\big(\mathrm{KL} - \log\frac{\pi_{\theta}(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}\big)\right) & \text{if } y \text{ is undesirable} \end{cases}$$

此处 $x$ 是提示词， $y$ 是补全， $\pi_{\theta}$ 是正在训练的策略模型， $\pi_{\mathrm{ref}}$ 是参考模型， $\sigma$ 是 sigmoid 函数， $\beta > 0$ 控制与参考模型的偏差，而 $\mathrm{KL}$ 是估计的 KL 散度项。权重 $w(y)$ 对于期望的补全为 desirable_weight ，对于非期望的补全为 undesirable_weight ，用于抵消期望与非期望示例数量之间的不平衡。

损失类型

记录的指标

在训练和评估过程中，我们记录以下指标：

• global_step：迄今为止已执行的优化器步数总数。
• epoch：当前 epoch 编号，基于数据集迭代。
• num_tokens：迄今为止处理的 token 总数。
• loss：当前日志记录区间内的平均 KTO 损失。
• entropy：模型在未屏蔽 token 上的预测 token 分布的平均熵。
• aux_loss：混合专家（Mixture-of-Experts）模型的负载均衡辅助损失，在乘以 router_aux_loss_coef 并加到总损失之前。仅当模型为 MoE 模型且 router_aux_loss_coef 非零时记录。
• kl：策略模型与参考模型之间平均估计的 KL 散度，用作 KTO 损失中的参考点。
• learning_rate：当前学习率，如果使用调度器，该值可能会动态变化。
• grad_norm：梯度裁剪前计算的梯度 L2 范数。
• logits/chosen：模型为所选（理想）补全中的 token 分配的平均 logit 值。
• logits/rejected：模型分配给被拒绝（不理想）补全中各个 token 的平均 logit 值。
• logps/chosen：模型为所选（理想）补全中的 token 分配的平均对数概率。
• logps/rejected：模型为被拒绝（不理想）补全中的 token 分配的平均对数概率。
• rewards/chosen: 为选中的（理想的）补全计算的平均隐式奖励，计算公式为 $\beta \log \frac{\pi_{\theta}(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}$.
• rewards/rejected: 为被拒绝（不理想）的补全计算的平均隐式奖励，其计算方式为 $\beta \log \frac{\pi_{\theta}(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}$.
• rewards/margins：所选完成项与被拒绝完成项之间的平均隐式奖励差值。

自定义

兼容性与约束

当前 KTOTrainer 实现中，部分参数组合被有意限制：

• 针对 lm_head 的 PEFT 适配器不受支持（请改用 modules_to_save=["lm_head"]）。
• sync_ref_model=True 不适用于使用未保留独立 ref_model 的 PEFT 模型进行训练的情况。
• sync_ref_model=True 不能与 precompute_ref_log_probs=True 同时使用。
• precompute_ref_log_probs=True 不支持与 IterableDataset（训练或评估）或视觉数据集一起使用。
• 用于估计 KL 散度项的损失类型（除 "apo_zero_unpaired" 外的所有类型）都需要 train_sampling_strategy="sequential"，并且每个设备的训练批次大小必须大于 1。

模型初始化

您可以直接将from_pretrained()方法的kwargs传递给KTOConfig。例如，如果您想以不同的精度加载模型，类似于

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct", dtype=torch.bfloat16)

您可以通过向 KTOConfig 传递 model_init_kwargs={"dtype": torch.bfloat16} 参数来实现此操作。

from trl import KTOConfig

training_args = KTOConfig(
    model_init_kwargs={"dtype": torch.bfloat16},
)

请注意，from_pretrained() 的所有关键字参数均受支持。

使用 PEFT 训练适配器

我们支持与 🤗 PEFT 库的紧密集成，允许任何用户方便地训练适配器并在 Hub 上共享，而无需训练整个模型。

from datasets import load_dataset
from trl import KTOTrainer
from peft import LoraConfig

dataset = load_dataset("trl-lib/kto-mix-14k", split="train")

trainer = KTOTrainer(
    "Qwen/Qwen2-0.5B-Instruct",
    train_dataset=dataset,
    peft_config=LoraConfig(),
)

trainer.train()

您还可以继续训练您的 PeftModel。为此，首先在 KTOTrainer 外部加载一个 PeftModel，然后直接将其传递给 trainer，而不传入 peft_config 参数。

在训练适配器时，由于只学习新参数，通常会使用比全量微调更高的学习率。

使用 KTO 进行工具调用

KTOTrainer 完全支持对具有 tool calling 能力的模型进行微调。在这种情况下，每个数据集示例应包含：

• 对话消息（提示词和补全），包括任何工具调用（tool_calls）和工具响应（tool 角色消息）
• tools 列中可用工具列表，通常以 JSON 模式提供

有关预期数据集结构的详细信息，请参阅 Dataset Format — Tool Calling 部分。

训练视觉语言模型

KTOTrainer 完全支持训练视觉语言模型（VLM）。要训练 VLM，请提供一个包含 image 列（每个样本一张图像）或 images 列（每个样本一个图像列表）的数据集。有关预期数据集结构的更多信息，请参阅 Dataset Format — Vision Dataset 部分。

from trl import KTOConfig, KTOTrainer
from datasets import load_dataset

trainer = KTOTrainer(
    model="Qwen/Qwen2.5-VL-3B-Instruct",
    args=KTOConfig(max_length=None),
    train_dataset=load_dataset("trl-internal-testing/zen-image", "conversational_unpaired_preference", split="train"),
)
trainer.train()

对于 VLM，截断可能会移除图像 token，从而导致训练过程中出现错误。为避免此问题，请在 KTOConfig 中设置 max_length=None。这允许模型处理完整的序列长度，而不会截断图像 token。

已复制

KTOConfig(max_length=None, ...)

仅在确认截断不会移除整个数据集中的图像标记后，才使用 max_length。

示例脚本

我们提供了一个使用 KTO 方法训练模型的示例脚本。该脚本位于 trl/scripts/kto.py

要在UltraFeedback 数据集上使用Qwen2 0.5B 模型测试 KTO 脚本，请运行以下命令：

accelerate launch trl/scripts/kto.py \
    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \
    --dataset_name trl-lib/kto-mix-14k \
    --num_train_epochs 1 \
    --output_dir Qwen2-0.5B-KTO

使用技巧

对于混合专家模型：启用辅助损失

当负载在专家之间大致均匀分布时，MOE 的效率最高。
为了确保在偏好调优期间以类似的方式训练 MOE，将负载均衡器的辅助损失添加到最终损失中是有益的。

通过在模型配置中设置 output_router_logits=True（例如 MixtralConfig）来启用此选项。
要调整辅助损失对总损失的贡献程度，请在模型配置中使用超参数 router_aux_loss_coef=...（默认值：0.001）。

批量大小建议

请使用至少为 4 的每步批大小（per-step batch size），以及介于 16 到 128 之间的有效批大小（effective batch size）。即使有效批大小很大，如果每步批大小设置不当，KTO 中的 KL 估计也会不准确。

学习率建议

每个 beta 取值都有一个可承受的最大学习率，超过该值后训练性能会下降。对于默认设置 beta = 0.1，大多数模型的学习率通常不应超过 1e-6。随着 beta 减小，学习率也应相应降低。总体而言，我们强烈建议将学习率保持在 5e-7 到 5e-6 之间。即使数据集较小，也不建议使用超出此范围的学习率。相反，应增加训练轮数以获得更好的效果。

数据不平衡

的 desirable_weight 和 undesirable_weight 中的 KTOConfig 分别指分配给理想/正样本和不理想/负样本损失的权重。 默认情况下，两者均为 1。但是，如果其中一种样本的数量明显多于另一种，则应提高数量较少类型的权重，使得（desirable_weight $\times$ 正样本数量）与（undesirable_weight $\times$ 负样本数量）的比例保持在 1:1 到 4:3 之间。

KTOTrainer

>>> from trl import KTOTrainer
>>> from datasets import load_dataset

>>> dataset = load_dataset("trl-lib/kto-mix-14k", split="train")

>>> trainer = KTOTrainer(
...     model="Qwen/Qwen2.5-0.5B-Instruct",
...     train_dataset=dataset,
... )
>>> trainer.train()

KTOConfig