ディープラーニングモデルのユニットスケーリング

ディープラーニングの話をすると、浮動小数点フォーマットの重要性が大事だよね。これらのフォーマットは、コンピュータが数をどう保存して処理するかを決めるんだ。それぞれのフォーマットには特定のレイアウトがあって、指数ビットと仮数ビットって呼ばれる部分が含まれてる。指数は数のスケールを理解する手助けをして、仮数は精度を示してるんだ。

ここで、ディープラーニングで使われる一般的な浮動小数点フォーマットをざっと見てみよう：

#提案されたFP8フォーマット

最近、ディープラーニング用の新しい8ビット浮動小数点フォーマットが提案されたんだ。これらのフォーマットには、4つの指数ビットを持つE4と、5つの指数ビットを持つE5がある。これらのデザインは、従来のIEEE 754標準と比較して特別な値の数を減らすことで調整されてる。このシフトは、表現するための最大値と最小値をわずかに異なるものにしてるんだ。

これらのフォーマットを分析する際の重要な概念の一つはバイアス値。バイアスは指数から引かれるもので、IEEE 754フォーマットでやるのと似てる。このことは、数のスケーリングが、表現できる値の範囲を制御できることを意味してる。スケーリングの観点に好みが傾くかもしれないけど、それはソフトウェアの実践に合ってるからで、指数-バイアスの観点はハードウェアに関係してるんだ。

全体として、新しいFP8フォーマットは既存の8ビットフォーマットにかなり似てるんだ。

#単位標準偏差は正しい基準か？

通常分布のテンソルを単位標準偏差で初期化するのは理にかなってると思う。目的は、浮動小数点テンソルの絶対値を使用するフォーマットによって定義された特定の限界内に保つことなんだ。

通常分布からの絶対値の挙動を見てみると、「折りたたまれた正規分布」と呼ばれる変種を考える。ほとんどの確率はこの分布で特定のポイントの周りに集中してる。

従来の浮動小数点フォーマットでは、値は通常、定義された区間内で変動して、その中点は数値範囲の中心に大体一致してる。エラーを避けるためにも、この中点に値をスケーリングすることが大事なんだ。

最終的に、単位標準偏差を目指すことで、ほとんどの値が数値範囲の中心に近づくと結論づける。ただし、この中心性がディープラーニングモデルのトレーニング中に通常分布のテンソルにとって最良の測定かどうかは議論されてる。

中心の下でスケーリングを支持する考えは、これらのサブノーマル値が余分な範囲を提供することだけど、精度が低くなるんだ。さらに、ディープラーニングのトレーニングではオーバーフローよりもアンダーフローが起こることが少ないんだ。

一方で、中心以上でスケーリングを好むのは、特に勾配がトレーニング中に減少するかもしれないっていう考えに基づいてる。だから、スケーリングがこの予想される減少を相殺するのに役立つかもしれない。

これらの視点から、単位スケーリングを目指すのは合理的な中間地点に見える。もし目標が分布を中心点に完璧に合わせることなら、少し大きめのスケールが必要かもしれない。ただ、その違いは最小限だから、単位スケールが強い選択肢で、さまざまな確立された技術とも一致してるんだ。

#単位スケーリングと新たな外れ値

最近の大規模言語モデルの量子化に関する研究では、外れ値を管理するための技術が紹介されたんだ。これらの外れ値は、データの特定の部分に現れる大きな値で、モデルのサイズが大きくなるにつれて特に目立つんだ。

外れ値の一つの課題は、単一の外れ値が他の値の全体的な量子化精度に悪影響を与えること。外れ値が大規模モデルのINT8量子化精度を低下させることが示されていて、単位スケーリングに対する外れ値の影響についての研究が行われてる。

私たちは、単位スケーリングが外れ値の大きさに大きな変化をもたらすとは考えていない。これは、外れ値が活性化テンソルによく見つかるもので、ノーマライゼーション層がスケールをコントロールすることで、ユニットモデルと非ユニットモデルの間で一貫したスケールを維持するからなんだ。

それでも、外れ値に対しては、単位スケーリングが既存の文献で記録された方法よりも安定すると予想されてる。これは主に、単位スケーリングが浮動小数点フォーマットで扱われるのに対し、既存の研究は整数量子化に焦点を当てていて、十分なダイナミックレンジがないからなんだ。

その結果、FP8フォーマットをINT8の代わりに使うことで、外れ値に関する複雑さを軽減できると考えてる。フォーマット間の信号対ノイズ比の比較も、この期待をさらに支持してるんだ。

例として、外れ値の存在が量子化分布を歪ませ、ほとんどの量子化ビンが空になるシナリオを考えてみよう。逆に、浮動小数点フォーマットはより良い分布を提供し、外れ値に適した少数の大きな値と、非外れ値用の小さい値が多数入ってるんだ。

これらの結果を分析すると、単位スケーリングが外れ値と戦う際に有望であることが明らかになるけど、課題はまだ残るかもしれない。

#理論的結果

#例 - スケーリング分析

多層パーセプトロン（MLP）の例を見てみよう。これは、特定の操作を使って入力が出力に変換されることを理解するとこから始まる。そこでは、エレメントごとの活性化関数が適用されるんだ。

この分析では、1つの層の検討を2つの異なる部分に分けることに注目する。手順は特定のアルゴリズムに従って、出力が独立した入力に基づいてどう分散を生むかを説明する。

この考えを続けて、勾配を計算する際に、入力と活性化の間の相関についていくつかの仮定を認める。これらの仮定のいくつかは問題があるかもしれないけど、私たちの目的には簡単な分析に役立つんだ。

#提案を支持する証明

一般的なオプティマイザ、確率的勾配降下法（SGD）とAdamのトレーニングダイナミクスを掘り下げてみる。スケーリングされていないモデルのトレーニングダイナミクスは、単位スケーリングされたモデルのものと似てる。

トレーニングの軌道は、モデル内のパラメータがトレーニング中にどのように進化するかを示す。SGDを使用すると、特定の定数と勾配に基づいてこれらの軌道を概説できる。

もしスケーリングされたモデルの軌道が、スケーリングされていないモデルのものと一貫した関係を維持するなら、効果的に同じトレーニングダイナミクスを共有していることが結論できる。

Adamの場合、その振る舞いは対角スケーリングに対してある程度不変だ。この特性は、スケーリングされたモデルとスケーリングされていないモデルが等しいトレーニング軌道で一致できることを保証しているんだ。

#例 - スケーリングされた計算グラフ

あるベクトル空間の値を持つスケーリングされていない操作を考えてみよう。スケーリングされた計算グラフを定義するとき、それは同じ操作を反映する必要がある。前向きと後ろ向きの値の一貫した方法を確立することで、スケーリングされていないグラフが正しくオペレーションを表現していることを検証できるんだ。

最終的に、スケーリングされたグラフもまたスケーリングされた操作を正しく表現できることを示すことで、表現が正確であることを保証するんだ。

#他のスキームの制約スケーリングされた計算グラフ

他のスケーリング戦略を比較すると、それが制約スケーリングされた計算グラフのフレームワークにどうフィットするかが見えるよ。

例えば、Glorot初期化を適用する場合、特定のスケールを考慮して、値を組み合わせるために使用することが必要で、ランタイム中に不要な乗算が必要ないようにするんだ。

損失スケーリングでは、スケーリングされたアイデンティティオペレーションを導入して、勾配を同期させつつ、それらの勾配全体で単一のスケールを共有できるようにする。このような戦略は慎重に選ばれるべきなんだ。

#単位スケーリングオペレーション集

単位スケーリングには、各オペレーションのスケーリングファクターを慎重に選ぶことが必要だ。これらのファクターは、さまざまなオペレーションと勾配オペレーションの挙動の分析から派生してるんだ。

プラクティショナーを助けるために、一般的なオペレーションのスケーリングされたバージョンとそれぞれのスケーリングファクターを示すリファレンステーブルをまとめたよ。

#ソフトマックスと他のオペレーション

ソフトマックスのようなオペレーションを見てみると、行列の掛け算の後にどう変わるかがわかる。ソフトマックスを適切にスケールするためには慎重なアプローチが必要で、次のプロセスを通じて一貫した値を伝播させることができるんだ。

同様に、ソフトマックスをクロスエントロピーと組み合わせる時は、それを単一の操作として扱う。後ろ向きのオペレーションに注目することで、トレーニング中に動きをチェックしつつ、必要なスケーリングファクターを導出できるんだ。

#既存モデルとの単位スケーリングの整合

単位スケーリングを考慮してモデルを設計することは、既存のモデルを適応させるのとは異なることがある。そういう場合は、単位スケーリングモデルが基準モデルと同等のパフォーマンスを発揮するために追加の考慮が必要なんだ。

#活性化関数とその重要性

活性化関数はモデルのパフォーマンスに大きく影響するんだ。入力スケールに基づいてその挙動が異なるから、単位スケーリングに移行する際は、活性化関数の前後でスケーリングファクターを導入して、入力が一貫していることを確保する必要があるんだ。

#共有パラメータと重みの勾配

操作間で頻繁に使用される共有パラメータは、モデル内のユニットのスケーリング比率に従う必要がある。必要に応じて調整を行い、単位スケーリングモデルが基準モデルのパフォーマンスに一致するように保証しなきゃいけないんだ。

#例：BERTの整合

単位スケーリングをBERTモデルに整合させる実験では、パフォーマンスレベルが従来のモデルと一致するように細心の注意を払ったよ。

#埋め込み層

BERTは、単語、セグメント、位置の埋め込みを含む複数の埋め込みで動作する。この埋め込みを統合するために、単位スケーリングを実施する際には、埋め込み間で一貫した重み設定を行うんだ。

#マルチヘッド自己注意

マルチヘッド自己注意プロセスでは、ソフトマックスオペレーションの整合を保つために必要な調整を行い、モデルのパフォーマンスの整合性を維持するようにするんだ。

#シーケンス長の考慮

シーケンスに基づいて動作するモデルをスケーリングする際は、シーケンスの長さの影響を考慮することが非常に重要なんだ。だから、トレーニング中のパディングしていないトークンの平均数に基づいて計算を調整する必要がある。

#計算オーバーヘッド

単位スケーリングを取り入れると、通常は最小限の計算オーバーヘッドが発生するんだ。これには、主にモデルのトレーニング中の追加の関数呼び出しとして現れる。浮動小数点オペレーションを分析すると、オーバーヘッドは他のプロセスに比べて小さいままだよ。

#スケーリングファクターの統合

前のオペレーションに固定のスケーリングファクターを統合することで、効率を高めることができる。この方法は、追加のメモリトリップを削減しつつ、通信時間を短縮することができ、場合によっては低精度出力を可能にするんだ。

#追加の実験詳細と結果

#キャラクター言語モデル

Wikipediaからの生データセットを使用して、前処理なしで実験を行ったんだ。結果は、体系的な学習率の調整を通じて収集され、パフォーマンスを最大化することができたよ。

#混合精度トレーニング

FP16トレーニングでは、すべてのコンポーネントがそれぞれのフォーマットで一貫して保存されたため、安定性を確保しつつパフォーマンスの効率を最大化できた。

#追加の結果

結果は、単位スケーリングされたモデルが非スケーリングモデルと比較してパフォーマンスを維持することを示していて、トレーニングシナリオで探求された技術の効果を再確認するものになったんだ。

#結論

単位スケーリングは、特に大規模言語モデルの文脈において、ディープラーニングモデル内で実用的なアプローチを提供するんだ。さまざまな浮動小数点フォーマット、その表現、パフォーマンスを分析することで得られた洞察は、この急速に進化する分野におけるトレーニング戦略の進化を強調しているんだ。