Hugging Face Transformers 勾配蓄積修正
Hugging Face は、transformers の Trainer において勾配蓄積がフルバッチ学習と同等の結果を出さないという数式上の不整合を修正しました。この修正により、損失は蓄積ステップ内のすべてのバッチにわたる 非パディングトークンの総数 に基づいて計算され、単にバッチごとの損失を平均するだけではなくなります。
勾配蓄積エラーの根本原因
勾配蓄積は、より大きなバッチサイズでの学習と数式上で同一であることが期待されています。しかし、勾配蓄積のオン・オフで損失が一致しないという不整合が見つかりました。
この問題は、transformers モデルの モデリングコード内で提供されている「default」損失関数 に起因しています。これらの関数は labels と input_ids がモデルに渡されたときに自動的に呼び出され、ユーザーにとって API をシンプルにしています。Causal Language Modeling（因果言語モデル）などのトークンレベルタスクでは、正しい損失計算は勾配蓄積ステップ内の すべてのバッチにわたる非パディングトークンの総数 で総損失を割る必要があります。従来の実装はバッチごとの損失を平均しており、特定のユースケースに対して数式的に誤っていました。
修正の技術的実装
精度と計算の誤りを解消するため、Hugging Face は損失計算ロジックを変更しました。主な変更点は、デフォルトの mean リダクションから sum リダクションに切り替え、項目数（num_items）で割ることです。
# Corrected loss calculation
loss = nn.functional.cross_entropy(shift_logits, shift_labels, ignore_index=-100, reduction="sum")
loss = loss / num_items
長期的なアーキテクチャ変更
今後の損失計算問題を防ぎ、柔軟性を高めるために、Hugging Face は 2 つの主要な構造変更を実装しています。
1. 自動損失補正
デフォルト損失関数を使用しているユーザーに対して、ライブラリは勾配蓄積時に自動的に必要な修正を適用し、報告・利用される損失が正確になるようにします。
2. カスタマイズ可能な損失 API
内部ライブラリの問題がユーザーをブロックしないよう、Trainer に独自の損失関数を直接渡せる API を導入します。
PreTrainedModel を継承するすべてのモデルは現在 loss_function プロパティを持ちます。このプロパティは config.loss_type によって決定され、ユーザーは LOSS_MAPPING を変更することで損失をカスタマイズできます。
def my_super_loss(logits, labels):
return nn.functional.cross_entropy(logits, labels, ignore_index=-100)
LOSS_MAPPING["my_loss_type"] = my_super_loss
デプロイと利用可能性
この修正は以下の 2 つのプルリクエストを通じてロールアウトされます。
- PR #34191: 最も人気のあるモデル向けに最初の変更を実装。
- PR #34198: ユーザーが独自の損失関数を提供できるようにし、バッチごとのサンプル数を計算に利用できるようにする第二の変更を実装。
ユーザーは main ブランチから transformers ライブラリをインストールすることで、これらの修正をすぐに利用できます。
pip install git+https://github.com/huggingface/transformers
Sources
- OriginalFixing Gradient Accumulation