# 参照モデルのBEMA

この機能は、DPOトレーニング中に参照モデルを更新するためのBEMAアルゴリズムを実装しています。

## 使用方法

```python
from trl.experimental.bema_for_ref_model import BEMACallback, DPOTrainer
from datasets import load_dataset

dataset = load_dataset("trl-internal-testing/zen", "standard_preference", split="train")

bema_callback = BEMACallback(update_ref_model=True)

trainer = DPOTrainer(
    model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
    train_dataset=dataset,
    callbacks=[bema_callback],
)
trainer.train()
```

## DPOTrainer[[trl.DPOTrainer]]

#### trl.DPOTrainer[[trl.DPOTrainer]]

```python
trl.DPOTrainer(*args, **kwargs)
```

[ソース](https://github.com/huggingface/trl/blob/main/trl/experimental/bema_for_ref_model/dpo_trainer.py#L19)

#### train[[trl.DPOTrainer.train]]

```python
train(resume_from_checkpoint: str | bool | None = None, trial: optuna.Trial | dict[str, Any] | None = None, ignore_keys_for_eval: list[str] | None = None)
```

[ソース](https://github.com/huggingface/trl/blob/main/transformers/trainer.py#L1447)

**パラメータ:**

resume_from_checkpoint (`str` または `bool`、*省略可能*) : `str` の場合、以前の `Trainer` インスタンスによって保存されたチェックポイントのローカルパスです。`bool` で `True` に等しい場合、以前の `Trainer` インスタンスによって *args.output_dir* に保存された最後のチェックポイントを読み込みます。指定されている場合、ここで読み込まれたモデル/オプティマイザ/スケジューラの状態からトレーニングが再開されます。

trial (`optuna.Trial` または `dict[str, Any]`、*任意*) : ハイパーパラメータ探索のためのトライアル実行またはハイパーパラメータ辞書。

ignore_keys_for_eval (`list[str]`、*任意*) : モデルの出力が辞書である場合、学習中の評価で予測を収集する際に無視すべきキーのリスト。

**戻り値:** `~trainer_utils.TrainOutput`

グローバルステップ数、トレーニング損失、およびメトリクスを含むオブジェクト。

メインの学習エントリポイント。

#### save_model[[trl.DPOTrainer.save_model]]

```python
save_model(output_dir: str | None = None, _internal_call: bool = False)
```

[ソース](https://github.com/huggingface/trl/blob/main/transformers/trainer.py#L4036)

モデルを保存するため、`from_pretrained()`を使用して再読み込みできます。

メインプロセスからのみ保存します。

#### push_to_hub[[trl.DPOTrainer.push_to_hub]]

```python
push_to_hub(commit_message: str | None = 'End of training', blocking: bool = True, token: str | None = None, revision: str | None = None, **kwargs)
```

[ソース](https://github.com/huggingface/trl/blob/main/transformers/trainer.py#L4307)

**パラメータ:**

commit_message (`str`、*任意*、デフォルトは`"End of training"`) : push 時にコミットするメッセージ。

blocking (`bool`、*任意*、デフォルトは`True`) : `git push`が完了した後にのみ関数が返すかどうか。

token (`str`、*任意*、デフォルトは`None`) : Trainer の元の引数を上書きするための、書き込み権限を持つトークン。

revision (`str`, *任意*) : コミット元の git リビジョンです。デフォルトでは "main" ブランチの先頭です。

kwargs (`dict[str, Any]`、*省略可能*) : `~Trainer.create_model_card` に渡される追加のキーワード引数。

**戻り値:**

`blocking=False`の場合、モデルがプッシュされたリポジトリのURL、または`blocking=True`の場合、コミットの進捗を追跡する`Future`オブジェクトを返します。

`self.args.hub_model_id` のリポジトリに、`self.model` と `self.processing_class` を 🤗 model hub にアップロードします。

## BEMACallback[[trl.BEMACallback]]

#### trl.BEMACallback[[trl.BEMACallback]]

```python
trl.BEMACallback(update_freq: int = 400, ema_power: float = 0.5, bias_power: float = 0.2, lag: int = 10, update_after: int = 0, multiplier: float = 1.0, min_ema_multiplier: float = 0.0, device: str = 'cpu', update_ref_model: bool = False, ref_model_update_freq: int = 400, ref_model_update_after: int = 0)
```

[ソース](https://github.com/huggingface/trl/blob/main/trl/experimental/bema_for_ref_model/callback.py#L59)

**パラメータ:**

update_freq (`int`, *optional*, デフォルトは `400`) : BEMAの重みをXステップごとに更新します。論文ではこれを \\( \phi \\) と表記しています。

ema_power (`float`、*任意*、デフォルトは`0.5`) : EMA減衰係数のべき乗です。論文では\\( \kappa \\)と表記されています。EMAを無効にするには、この値を`0.0`に設定してください。

bias_power (`float`, *optional*, デフォルトは `0.2`) : BEMA スケーリング係数のべき乗です。論文では \\( \eta \\) と表記されます。大きな値（例: `8.0`）を指定すると \\( \alpha_t \\) は `0` に減衰し、バイアス補正が無効化された状態を近似します。一方、`0.0` を指定すると、すべてのステップで \\( \alpha_t \\) が `1` に固定されます（最大かつ減衰しない補正）。

lag (`int`, *optional*, デフォルトは `10`) : 重み減衰スケジュールの初期オフセットです。更新の仮想開始年齢として機能し、初期段階の滑らかさを制御します。論文では \\( \rho \\) と表記されています。

update_after (`int`, *optional*, デフォルトは `0`) : BEMA重みの更新を開始するまでのバーンイン時間。論文では \\( \tau \\) と表記されています。

multiplier (`float`, *optional*, デフォルトは `1.0`) : EMA減衰係数の初期値です。論文では \\( \gamma \\) と表記されています。

min_ema_multiplier (`float`、*任意*、デフォルトは`0.0`) : EMA減衰係数の最小値。

device (`str`、*任意*、デフォルトは`"cpu"`) : BEMAバッファに使用するデバイス。例えば`"cpu"`や`"cuda"`などです。ほとんどの場合、OOMを回避するために、このデバイスはトレーニングに使用するデバイスとは異なるべきです。

update_ref_model (`bool`、*任意*、デフォルトは`False`) : BEMA重みで参照モデルを更新するかどうか。これにより、メインモデルの遅延・平滑化されたバージョンが参照モデルとして作成されます。

ref_model_update_freq (`int`、*optional*、デフォルトは`400`) : このステップ数ごとに、BEMA重みで参照モデルを更新します。

ref_model_update_after (`int`, *optional*, デフォルトは `0`) : 参照モデルの更新を開始する前に待機するステップ数。

[Adam Block](https://huggingface.co/abblock) と [Cyril Zhang](https://huggingface.co/cyrilzhang) による [BEMA](https://huggingface.co/papers/2508.00180)（バイアス補正済み指数移動平均）を実装した [TrainerCallback](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerCallback) です。コードは MIT ライセンスの下で https://github.com/abblock/bema から提供されています。

BEMAは、次のようにスケールするモデルの重みを計算します:

$$ \theta_t' = \alpha_t \cdot (\theta_t - \theta_0) + \text{EMA}_t $$

ここで、\\( \theta_t \\) は現在のモデルの重み、\\( \theta_0 \\) は最初の `update_after` ステップ時点のモデルの重みのスナップショット、\\( \text{EMA}_t  \\) はモデルの重みの指数移動平均、\\( \alpha_t \\) はステップ数 \\( t \\) に従って減衰するスケーリング係数です。

$$ \alpha_t = (\rho + \gamma \cdot t)^{-\eta}. $$

EMAは以下のように計算されます:

$$ \text{EMA}_t = (1 - \beta_t) \cdot \text{EMA}_{t-1} + \beta_t \cdot \theta_t $$

ここで \\( \beta_t \\) は、ステップ数 \\( t \\) に応じて減衰する減衰係数であり、以下のように定義されます

$$ \beta_t = (\rho + \gamma \cdot t)^{-\kappa}. $$

例:

```python
>>> from trl import BEMACallback

>>> trainer = Trainer(..., callbacks=[BEMACallback()])
```