本文へスキップ
AI-Papers

レイヤードロップアウトとは?LLM訓練FLOPsを25%削減し推論を1.5倍高速化

レイヤードロップアウトとは?LLM訓練FLOPsを25%削減し推論を1.5倍高速化
  • Cerebras らのチームが、現代のLLM事前学習でも層単位のドロップアウトを使うべきだと2,400回超の実験で主張
  • 同じ訓練FLOPsなら検証損失が下がり、固定ステップ数では最大25%の訓練FLOPs削減を確認
  • 訓練済みモデルは early exit や層スキップに強くなり、自己投機的デコーディングで最大1.55倍の高速化

研究の背景

ドロップアウトは、学習中にネットワークの一部をランダムに無効化して過学習を抑える古典的な手法です。その層単位版がレイヤードロップアウト(stochastic depth、Transformerブロックまるごとを確率的に飛ばす仕組み)で、画像認識モデルや初期の言語モデルでは正則化と学習安定化の両面で効果が報告されてきました。

ところが現代のLLM事前学習では、この手法はほぼ使われていません。大量のデータで1エポック程度しか学習しない設定では過学習が起きにくく、正則化の必要性が薄いと考えられてきたこと、そして層を飛ばすと精度が落ちるという報告があったことが理由です。

Mostafa Elhoushi らによる本論文「Don't Drop Dropout: Optimizing Layer Sparsity for Efficient LLM Training and Inference」(ICML 2026採択)は、この通説に正面から反論します。適切な設定を選べばレイヤードロップアウトは損失を悪化させず、むしろ計算量あたりの効率を高め、さらに推論時の高速化まで引き出せると主張しています。実験は Cerebras CS-3 システム上で実施されました。

提案する2つのスケジュール

著者らが探索の末に推奨するのは、層方向と時間方向の2つのスケジュールを組み合わせた設定です。層方向にはIncreasing Linear Distribution(ILD)を使い、最初の層のドロップ確率を0、最終層を最大値として線形に増やしていきます。入力に近い浅い層は表現の土台を作るため保護し、深い層ほど積極的に落とすという発想です。

時間方向にはDecreasing Time Schedule(DTS)を採用し、学習開始時に最大の確率でドロップし、訓練が進むにつれて線形に0まで下げます。学習の終盤では通常の密なモデルと同じ状態になるため、最終的な品質を損ないにくくなります。ドロップの単位はバッチ内のシーケンスごとに独立してマスクを生成する方式で、残った層の出力は訓練時に1/ρ倍にスケーリングして期待値を揃えます。

図1: レイヤードロップアウトが訓練の高速化、early exit と層スキップへの頑健性、推論の高速化を一つの仕組みでもたらす
図1: レイヤードロップアウトが訓練の高速化、early exit と層スキップへの頑健性、推論の高速化を一つの仕組みでもたらす(論文 Figure 1)

図1に示すように、この仕組みは訓練と推論の両方に効きます。左は事前学習中に層を確率的に飛ばして訓練を速める部分、中央は訓練済みモデルが early exit(途中の層で出力を打ち切る)や層スキップに対して緩やかにしか劣化しない性質、右はその頑健性がアダプタや自己投機的デコーディングに引き継がれる流れを表しています。著者らはこの性質をゼロショットの「弾力的な深さ」と呼んでいます。

スケーリング則での検証

主張の裏づけとして、271M から 8.2B パラメータまでの6段階のモデルサイズで、最大160Bトークン、2,400回を超える訓練を実施しています。アーキテクチャは ALiBi 位置埋め込みと squared ReLU を用いた decoder-only Transformer で、パラメータあたり20トークン(20 TPP)という計算最適な予算を基準としています。

ILD で約10%の訓練FLOPsを削減した場合、検証損失の悪化はごくわずかにとどまりました。しかもモデルが大きくなるほど差は縮まり、スケールに対して不利にならない傾向が読み取れます。

モデルサイズ

20 TPP・約10%FLOPs削減時の検証損失の変化

271M

-0.39%

503M

-0.24%

906M

-0.17%

さらに 8.2B モデルで最大ドロップ率を0.99まで引き上げる強気の設定を試したところ、ベースラインと競合する精度を保ったまま訓練FLOPsを25%削減できたと報告されています。同じ訓練FLOPs予算で比べた場合には、レイヤードロップアウトを使ったモデルのほうが検証損失が低くなりました。

推論での恩恵

訓練時に層を飛ばす経験を積ませておくと、推論時に層を減らしても壊れにくいモデルができます。全層の75%の地点で early exit させても精度低下はごくわずかで、1層おきに層をスキップする実験では、906M モデルの損失がベースラインの6.446に対して2.129まで改善しました。層を落とした密なモデルが事実上崩壊するのに対し、ドロップアウトで訓練したモデルは動作を保てるということです。

この頑健性は自己投機的デコーディング(浅い層だけで下書きを作り、全層で検証する高速化手法)とも相性がよく、3.9B モデルで最大1.55倍の速度向上を出力品質を落とさずに達成しています。推論高速化の方向性としては、KVキャッシュを削減するアプローチなど計算の別の部分を削る手法もありますが、本手法は深さ方向の計算を削る点で組み合わせやすいと考えられます。

まとめと今後の展望

実装面での負担が小さい点も実用性を後押しします。既存の学習ループに層のスキップ処理とスケジュールを追加するだけで済むため、大規模な書き換えは不要です。手元の環境で追試しやすい手法と言えるでしょう。

一方で著者らは限界も明示しています。極端に高いドロップ率ではハイパーパラメータの転移が難しいこと、Transformerブロック単位以外の粒度を十分に探索していないこと、Mixture-of-Depths のような学習ベースの深さ最適化手法との直接比較がないことなどです。MoE や非Transformer系アーキテクチャへの一般化、適用可能な最大ドロップ率を予測するスケーリング則も未解決の課題として残されています。

とはいえ、訓練コストと推論コストの双方を同じ仕組みで削れるという結果は、大規模モデルを扱う現場にとって検討する価値のある選択肢です。「もうドロップアウトは要らない」という前提を、実験規模で問い直した研究と位置づけられます。

論文情報: "Don't Drop Dropout: Optimizing Layer Sparsity for Efficient LLM Training and Inference"(Mostafa Elhoushi et al., 2026) arXiv:2609.05275, CC BY 4.0

本記事の図(図1)とサムネイルは上記論文より引用しています(縮小・形式変換あり)。

シェア:

投稿には GitHub アカウントが必要です。投稿内容は公開され、利用規約に反するものは予告なく削除します。