Bf16は重みの更新を凍結させている可能性がある
Bf16 might be freezing your weight updates (lacuna.tiptreesystems.com)
要約
bf16フォーマットで重みを保存し、AdamWオプティマイザで直接更新すると、学習率によっては非常に小さな重みしか更新されず、他の重みが事実上凍結されるという現象が観察されました。これは、bf16の精度限界により、学習率に対して小さすぎる重みの変化が丸められて失われるために発生します。この現象は、意図せずモデルの安定性を高める一方で、本来のファインチューニングの挙動とは異なる結果をもたらす可能性があります。
全文翻訳
bf16は重みの更新を凍結させている可能性がある
モデルをbf16で重みを保存してファインチューニングし、torch.optim.AdamWに直接書き込ませると、興味深い挙動が見られます。bf16は、重み自体の約1/512以上の変化でなければ変化を表現できないため、学習率が小さい場合、最も小さな重みしか更新できません。
その結果、小さな重みをトレーニングし、残りを凍結するスパースなファインチューナーとなり、カットオフは学習率によって設定されます。これは実際に存在するもので、私はこの方法でファインチューニングした査読付き論文を半ダース（数千の引用があるもの）見つけましたが、そのどれもこの挙動に言及していませんでした。
私はQuiet-STaRという論文を掘り下げました。なぜなら、（a）クールな論文であること、（b）重みを公開していること、（c）コードが動作することです！そこで、ベースモデルと比較し、トレーニング中に更新されたのは小さな重みだけであることを確認しました。その後、すべての更新を適用するオプティマイザでそのレシピを再実行し、それが発散するのを見ました。より高精度なオプティマイザは、同じ結果を得るために3〜10倍小さい学習率を必要とし、私はそれ以上の結果を達成できませんでした。
これはBarron, Gnanapragasam & Tschannenの「A Fix for Softmax and Cross-Entropy」（2026年10月）から始まりました。彼らは、クロスエントロピーの標準的な計算方法が、モデルが十分に確信を持てば（bf16で約5.5のマージンを超えると、彼らはこれを精度クリフと呼んでいます）、正しい答えの勾配を正確にゼロに丸めてしまうことを示しており、これがAdamと相互作用して、グロッキング中に見られる「スリングショット」損失スパイクを生成します。丸め処理を修正すると、スパイクは消えます。トレーニングダイナミクスとして研究されていたものが、丸め処理だったのです。これにより、他にこのような奇妙なものが世の中にないか疑問に思いました。これはその一つであり、逆方向に機能します。彼らの丸めバグはトレーニングに不安定さを加えますが、これはそれを隠します。
コード、ログ、評価結果はmweiss17/quietstar-bf16-adamwにあります。
1. bf16が最適化にどう影響するか
浮動小数点数は、ここで重要な2つの部分を持っています。指数部は2のべき乗の間隔（1〜2、2〜4、4〜8など）を選択し、仮数部はそれらの内部にある固定数の等間隔のスロットを選択します。すべての間隔は、どれだけ広くても同じ数のスロットを受け取ります。そのため、値はゼロの近くでは密接し、数値が大きくなるにつれて広がります。隣接する値のギャップは、2のべき乗を横切るたびに倍増します。図1は、間隔あたり4つのスロットしかないおもちゃのフォーマットでこれを示しています。
bf16とfp32はどちらも8ビットの指数部を持つため、同じ範囲の数値をカバーします。違いは仮数部にあります。fp32は23ビットの仮数部を持つため、間隔あたり8,388,608個のスロットがあります。bf16は上位7ビットのみを保持するため、128個です。bf16の数値はfp32の数値の最初の半分であり、丸められます。そのギャップは65,536倍広く、数値自体の約1/128から1/256になります。
次に、オプティマイザが行うように、重みに小さな変更を加えます。結果は、フォーマットが保持できる値に丸める必要があります。変更がギャップの半分未満の場合、元の場所に戻って丸められ、更新は失われます。Adamのステップは学習率のおおよそなので、1e-6では典型的なMistralの重みは更新を失いますが、8倍小さい重みは保持します。
これがトレーニングに何をするかを見るために、図1のおもちゃのフォーマットに戻り、2つの重み、w₁とw₂でAdamを実行します。各ドットはフォーマットが保持できる値です。2つのAdamのコピーは同じ点から開始し、同じ最小値を目指します。1つ（オレンジ）は重みをトイフォーマットで保持し、各ステップ後に丸めます。もう1つ（青）はフルプレシジョンコピーを保持します。
bf16は、間隔あたり128個の値を持つことで、4個ではなく、重み自体の約1/512未満の変更が失われるという点で同じ図になります。
2. bf16がAdamWを特に台無しにする方法
パラメータがbf16で、下にfp32マスターコピーがない場合、AdamWには3つの問題があります。
- Weight decayはノーオペレーションです。 Decayは各重みを(1 − lr·wd)で乗算します。これはこれらの設定で約1e-9の相対的な変化です。これは各重み、各ステップで丸められて失われます。
- Adamの2次モーメントが凍結します。 モーメントはパラメータのdtypeに格納されます。β₂ = 0.999では、0.1%の減衰が丸められて失われるため、vは決して縮小せず、真の値の半分程度まで成長すると、その増分も丸められて失われます。これは数百ステップ後に発生するため、以下の100ステップの実行には影響しませんが、通常のファインチューニングにはすべて影響します。
- ギャップの半分未満の更新はドロップされます。 これはこの記事の主題です。
明確にしておくと、これらは既知の問題です。Collage（ICML 2024）は機能しました。Hugging Faceの投稿[