KTO 训练流程总结

在 7B 参数规模以下，对齐相较于单纯的 SFT 几乎没有增益

DPO 的默认学习率为 5e-7，但我们发现 KTO 使用 5e-6 作为默认学习率效果更佳

一、 数据预处理与超参数初始化

1. 数据格式转换
   • 若原始数据为偏好对 x, y_w, y_l（y_w 优于 y_l），则标记 y_w 为合意样本、y_l 为不合意样本；若为原生二元反馈，直接按标签划分两类样本。
   • 处理数据不平衡：设合意样本数为 n_D、不合意样本数为 n_U，调整权重超参数 \lambda_D（合意权重）、\lambda_U（不合意权重），满足：\frac {\lambda_D n_D}{\lambda_U n_U} \in \left [1, \frac {3}{2}\right]
   • 示例：若 n_D:n_U=1:10，设 \lambda_U=1，则 \lambda_D \in [10,15]。
2. 模型与超参数配置
   • 确定参考模型 \pi_{ref}：可为有监督微调（SFT）模型或预训练模型；待优化模型 \pi_\theta 初始化与 \pi_{ref} 一致。
   • 设置核心超参数：
     • 风险厌恶系数 \beta：大型模型 / 已 SFT 用 0.01\sim0.1；小型模型 / 无 SFT 用 0.1\sim1.0
     • 学习率：推荐 5e-6（搭配 AdamW 优化器）
     • 微批次大小：\ge2（推荐 8\sim128，需满足参考点估计的样本量要求）

二、 单步训练核心计算流程

对每个训练微批次执行以下步骤，迭代至模型收敛：

步骤 1：计算隐含奖励 r_\theta(x,y)

对批次内每个样本 (x,y)，计算当前模型相对于参考模型的概率比对数，即隐含奖励： r_\theta (x, y) = \log \frac {\pi_\theta (y | x)}{\pi_{ref}(y | x)}

• 合意样本期望 r_\theta>0（当前模型生成概率高于参考模型）
• 不合意样本期望 r_\theta<0（当前模型生成概率低于参考模型）

步骤 2：估计参考点 \hat{z}_0

参考点用于衡量全局 “基准奖励”，控制损失函数饱和，采用不匹配对有偏估计：

1. 对批次内 m 个样本，构造不匹配对 (x_1,y_2), (x_2,y_3), ..., (x_m,y_1) (j=(i+1)\mod m）
2. 计算不匹配对的隐含奖励均值，截断为非负值：
   \hat{z}0 = \max\left(0, \frac{1}{m} \sum{i=1}^m \log \frac{\pi_\theta(y_{j} | x_i)}{\pi_{ref}(y_{j} | x_i)}\right)

• 简化场景：若 SFT 数据是 KTO 数据的子集，可直接设 \hat{z}_0=0，无需额外计算。

步骤 3：计算人类感知价值函数 v(x,y)

基于前景理论的逻辑函数（避免指数函数数值不稳定），分样本类型计算价值：

v(x, y) = \begin{cases} \lambda_D \cdot \sigma\left(\beta \cdot (r_\theta(x, y) - \hat{z}0)\right) & \text{y为合意样本} \\ \lambda_U \cdot \sigma\left(\beta \cdot (\hat{z}_0 - r\theta(x, y))\right) & \text{y为不合意样本} \end{cases}

其中 Sigmoid 函数 \sigma (t) = \frac {1}{1+e^{-t}}，确保价值在收益区呈凹形、损失区呈凸形，契合人类决策偏差。

风险厌恶 / 寻求程度由 β 控制；β 越大，收益区域的风险厌恶程度越高，损失区域的风险寻求程度也越高。在实际应用中，对于已进行 SFT 的大型模型，β 在 [0.01, 0.10] 范围内的较小值效果更佳；对于未进行前期 SFT、直接进行 KTO 训练的小型模型，β 在 [0.10, 1.00] 范围内的较大值效果更佳。

步骤 4：计算 KTO 损失函数 \mathcal{L}_{KTO}

损失目标是最大化价值函数，形式为批次样本的期望均值：

\mathcal{L}{KTO} = \frac{1}{m} \sum{(x,y) \in \text{batch}} \left[\lambda_y - v(x, y)\right]

其中 \lambda_y 对应样本类型的权重（合意取 \lambda_D，不合意取 \lambda_U），保证损失非负。

步骤 5：模型参数更新

1. 对损失函数关于模型参数 \theta 求梯度（\hat{z}_0 不参与反向传播，避免训练不稳定），梯度公式为：
   \nabla_\theta \mathcal{L}<em>{KTO} = \frac{1}{m} \sum</em>{(x,y) \in \text{batch}} d(y) \cdot \lambda_y \cdot \sigma(\beta z) \cdot (1-\sigma(\beta z)) \cdot \beta \cdot \nabla_\theta \log \pi_\theta(y | x)

其中 d(y) 为符号因子：合意样本取 -1，不合意样本取 +1；z 为价值函数内的自变量（合意样本为 r_\theta-\hat{z}_0，不合意样本为 \hat{z}_0-r_\theta）。用梯度下降算法（如 AdamW）更新参数 \theta，完成单步训练。

三、 训练收尾与调优

1. 收敛判断：通过验证集的生成质量评估（如 GPT-4 胜率、任务指标）判断模型是否收敛，避免过拟合。
2. 简化训练：大参数量预训练模型（如 Llama-3 8B）可跳过 SFT 阶段，直接用预训练模型作为 \pi_{ref} 启动 KTO 训练。
3. 内存优化：若需降低内存占用，可假设 \pi_{ref} 为均匀分布，此时隐含奖励简化为 r_\theta (x,y)=\log \pi_\theta (y|x) - H (\pi_\theta (y'|x))（H 为熵），无需加载参考模型权重。