ほぼゼロのトークン埋め込みがゲート型モデルの学習を止める理由

導入

ほぼゼロのベクトルは、単なる小さな数値誤差に見えるかもしれない。しかし、ゲート型系列モデルでは学習全体を不安定化させる起点になり得る。論文「Near-Zero Token Embeddings Silently Stall Training in Gated Sequence Models」によると、非常にまれなトークンやテンプレート用トークンは、長期間にわたって入力側の勾配を受け取らないことがある。入力と出力の埋め込みが共有されていない場合、AdamWの重み減衰だけが働き、該当行の値がさらに縮小する。

主なポイント

• 研究チームは291のオープンウェイト言語モデルリポジトリと、280種類の埋め込みテーブルを調査した。重みを共有しないゲート型・ハイブリッドのチェックポイントでは、38件中20件にノルム10^-6未満の到達可能な入力行があった。一方、重み共有モデルでは0件だった。
• 影響はモデル構造に依存する。検証したTransformerでは、系列の先頭以外の位置やBOSトークンによって、この現象が回避された。対してMamba、Gated DeltaNet、xLSTMでは、ほぼゼロのトークンが入力埋め込み勾配を30桁以上、最大35桁増幅する可能性が示された。
• ほぼゼロの入力では、SiLUゲートや線形クエリなど、現在トークンを読み出す要素が消失しやすい。その結果、過去状態の読み出しが抑えられる。複数層で残差のRMSがRMSNormの下限の平方根を下回ると、Pre-RMSNormが逆伝播勾配を大きく増幅する。
• 勾配の一要素が約1.84×10^19を超えると、PyTorch、Hugging Face Trainer、FSDP2のFP32グローバルノルムクリッピングが無限大へオーバーフローする可能性がある。その結果、当該バッチの勾配がゼロになる。Nemotron-3-Nano-4Bのテルグ語継続事前学習では、約20%のステップで更新が失われ、該当行が凍結した。

意義と影響

この結果は、損失や全体勾配ノルムだけでなく、到達可能な各入力埋め込みのノルムも監視すべきだと示している。特に、埋め込みを共有しないモデルとゲート型アーキテクチャでは、希少トークンがゲート、正規化、クリッピングの連鎖を通じて全体の学習に影響する可能性がある。

公開済みチェックポイント向けに、研究はGroupClipを提案する。入力埋め込みのクリッピングを他のパラメータから分離し、64ビットまたは再スケールしたノルムを使う方法だ。報告された実験では、再スケール版が更新の消失をなくし、問題の行を1ステップで修復した。事前学習向けのGateFloorはSiLU(z)+δを用い、δを0.01から0.1に設定する。これによりゼロ入力でもゲートが正になり、最大3.7億パラメータの実験でほぼゼロ勾配を28桁低減し、ベースラインと同程度のパープレキシティを維持した。

希少トークンの埋め込みは、使われないから安全とは限らない。ゲート型モデルでは、極小の一行が数値不安定性の入口になるため、モデルの検査項目に加える価値がある。

出典：Hugging Face Daily Papers