KTO Trainer

概要

TRLは、Kawin Ethayarajh、Winnie Xu、Niklas Muennighoff、Dan Jurafsky、Douwe Kielaによる論文KTO: Model Alignment as Prospect Theoretic Optimizationで説明されているKahneman-Tversky Optimization (KTO) Trainerをサポートしており、言語モデルの学習に使用できます。

論文の要旨は以下の通りです:

Kahneman & Tversky の prospect theory（期待効用理論）は、人間がランダム変数を偏りはあるが明確に定義された方法で知覚することを示しています。例えば、人間は損失回避的であることで知られています。本研究では、LLM を人間のフィードバックと整合させるための目的関数が、これらの偏りの多くを暗黙的に取り込んでいることを示します。これらの目的関数（例えば DPO）が交差エントロピー最小化を上回る成功は、それらが人間を考慮した損失関数（HALOs）であることにも一部起因しています。しかし、これらの手法が人間に帰属させる効用関数は、prospect theory の文献におけるものとは異なります。人間の効用を Kahneman-Tversky モデルで表現し、現在の手法が好みの対数尤度を最大化するのに対し、生成の効用を直接最大化する HALO を提案します。このアプローチを Kahneman-Tversky Optimization（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 x 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にはペアになっていない好みのデータセットが必要です。代わりに、ペアになった好みデータセット（単に好みデータセットとも呼ばれます）を提供することもできます。この場合、トレーナーは、選択された応答と拒否された応答を分離して非ペア形式に変換し、選択された補完に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}

理論上、データセットには少なくとも1つのchosen completionと1つのrejected completionが含まれている必要があります。ただし、only chosenデータ、またはrejectedデータのみを使用してKTOを正常に実行したユーザーもいます。rejectedデータのみを使用する場合は、保守的な学習率を採用することをお勧めします。

KTOメソッドの深掘り

Kahneman-Tversky Optimization (KTO)は、与えられた入力に対して出力が望ましいか望ましくないかという二値のシグナルのみを用いて言語モデルをアラインメントするための学習手法です。好ましい/好ましくない補完のペアではなく、この二値シグナルを使用します。KahnemanとTverskyのProspect Theoryに基づき、human-aware loss (HALO)を定義します。HALOは生成の効用を直接最大化するように設計されており、人間の損失回避を反映するために、望ましい例と望ましくない例を非対称に重み付けします。

このセクションでは、KTOが実際どのように動作するかを分解して説明します。主なステップである前処理と損失の計算について扱います。

前処理とトークン化

学習中、各サンプルにはプロンプト、completion、およびその補完が望ましい（True）か望ましくない（False）かを示すブール値のlabelが含まれている必要があります。期待される形式の詳細については、Dataset formatsを参照してください。 KTOTrainerは、モデルのトークナイザーを使用して各入力をトークン化します。

損失の計算

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$ はシグモイド関数、 $\beta > 0$ は参照モデルからの逸脱を制御し、 $\mathrm{KL}$ は推定されたKLダイバージェンス項です。重み $w(y)$ は、望ましい補完に対して desirable_weight 、望ましくない補完に対して undesirable_weight が使用され、望ましい例と望ましくない例の数の不均衡を補正するために用いられます。

損失の種類

記録されたメトリクス

学習と評価の実行中、以下のメトリクスを記録します：

• global_step: 現在までに実行されたオプティマイザーステップの総数。
• epoch: データセットの反復に基づいた現在のエポック番号。
• num_tokens: 現在までに処理されたトークンの総数。
• loss: 現在のログ区間におけるKTO損失の平均値。
• entropy: マスクされていないトークンにおける、モデルが予測したトークン分布の平均エントロピー。
• aux_loss: Mixture-of-Experts モデルのロードバランス用の補助損失。router_aux_loss_coef によってスケーリングされ、損失に追加される前の値です。モデルが MoE モデルであり、router_aux_loss_coef がゼロでない場合にのみ記録されます。
• kl: ポリシーモデルと参照モデルの間の推定KLダイバージェンスの平均値。KTO損失における参照点として使用されます。
• learning_rate: 現在の学習率。スケジューラを使用している場合、動的に変更される可能性があります。
• grad_norm: グラディエントクリッピングの前に計算される、勾配のL2ノルム。
• logits/chosen: モデルが選択済み（望ましい）補完のトークンに割り当てた平均ロジット値。
• logits/rejected: モデルが拒否（望ましくない）補完のトークンに割り当てたロジット値の平均。
• logps/chosen: モデルが選択済み（望ましい）補完のトークンに割り当てた平均対数確率。
• logps/rejected: モデルが拒否（望ましくない）補完のトークンに割り当てた平均対数確率。
• 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)

model_init_kwargs={"dtype": torch.bfloat16} 引数を KTOConfig に渡すことで、これを行うことができます。

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の学習を継続することもできます。そのためには、まずPeftModelをKTOTrainerの外で読み込み、peft_config引数を渡さずにそのままトレーナーに直接渡します。

アダプターを学習する際は、学習されるのが新しいパラメータのみであるため、フルファインチューニングよりも高い学習率を通常使用します。

KTOによるツール呼び出し

KTOTrainerは、tool calling機能を持つモデルのファインチューニングを完全にサポートしています。この場合、各データセットの例には以下を含める必要があります。

• ツール呼び出し（tool_calls）やツール応答（toolロールのメッセージ）を含む、会話メッセージ（プロンプトと補完）
• tools 列に表示される利用可能なツールのリスト。通常、JSON スキーマとして提供されます

期待されるデータセットの構造の詳細については、Dataset Format — Tool Callingセクションを参照してください。

ビジョン言語モデルの学習

KTOTrainerは、Vision-Language Models (VLMs) の学習を完全にサポートしています。VLMを学習するには、image列（サンプルごとに1枚の画像）または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では、切り詰めを行うと画像トークンが削除され、トレーニング中にエラーが発生する可能性があります。これを回避するには、KTOConfigでmax_length=Noneを設定してください。これにより、モデルは画像トークンを切り詰めずに、シーケンス全体の長さを処理できます。

コピーしました

KTOConfig(max_length=None, ...)

max_lengthは、データセット全体で切り詰めによって画像トークンが削除されないことを確認した場合にのみ使用してください。

スクリプトの例

KTO手法を用いてモデルを学習するためのサンプルスクリプトを提供しています。スクリプトはtrl/scripts/kto.pyで利用できます。

Qwen2 0.5B モデルでUltraFeedback データセットを用いて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

使用のコツ

Mixture of Expertsモデル向け: 補助損失の有効化

MOEは、エキスパート間の負荷がほぼ均等に分配されている場合に最も効率的です。
好適性チューニング時にMOEを同様に訓練するために、ロードバランサーの補助損失を最終損失に追加することが有効です。

このオプションは、モデル設定（例: MixtralConfig）で output_router_logits=True を設定することで有効になります。
補助損失が総損失にどれだけ寄与するかを調整するには、モデル設定でハイパーパラメータ router_aux_loss_coef=...（デフォルト: 0.001）を使用します。

バッチサイズに関する推奨事項

1ステップあたりのバッチサイズは4以上、実効バッチサイズは16から128の範囲に設定してください。実効バッチサイズが大きくても、1ステップあたりのバッチサイズが不適切な場合、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