[論文レビュー] Mixture-of-Depths: Dynamically allocating compute in transformer-based language models

David Raposo, Sam Ritter+4 more|arXiv (Cornell University)|Apr 2, 2024

Natural Language Processing Techniques被引用数 4

ひとこと要約

この論文は、学習されたtop-kルーティング機構を用いてトークンが層をスキップできるようにすることで、トランスフォーマー言語モデルにおける計算を動的に割り当てるMixture-of-Depths (MoD)を提案する。各層の合計FLOPsを制限する。MoDモデルは、最大50%少ないFLOPsで前向き伝搬を実行し、最大50%速い推論を達成しながら、静的計算グラフにより訓練効率とハードウェア互換性を維持する。

ABSTRACT

Transformer-based language models spread FLOPs uniformly across input sequences. In this work we demonstrate that transformers can instead learn to dynamically allocate FLOPs (or compute) to specific positions in a sequence, optimising the allocation along the sequence for different layers across the model depth. Our method enforces a total compute budget by capping the number of tokens ($k$) that can participate in the self-attention and MLP computations at a given layer. The tokens to be processed are determined by the network using a top-$k$ routing mechanism. Since $k$ is defined a priori, this simple procedure uses a static computation graph with known tensor sizes, unlike other conditional computation techniques. Nevertheless, since the identities of the $k$ tokens are fluid, this method can expend FLOPs non-uniformly across the time and model depth dimensions. Thus, compute expenditure is entirely predictable in sum total, but dynamic and context-sensitive at the token-level. Not only do models trained in this way learn to dynamically allocate compute, they do so efficiently. These models match baseline performance for equivalent FLOPS and wall-clock times to train, but require a fraction of the FLOPs per forward pass, and can be upwards of 50\% faster to step during post-training sampling.

研究の動機と目的

• トランスフォーマーにおける均一な計算割り当ての非効率性、すなわちすべてのトークンが計算ニーズに関係なく同じFLOPsを受ける問題に対処する。
• 動的で文脈に依存する計算割り当てを、トークンおよび層の両次元で実現するが、動的計算グラフを導入しない。
• 知的なルーティングにより、各モデルの異なる深さを通過するトークンを制御することで、前向き伝搬におけるFLOPsを削減しながら、ハードウェア効率と訓練速度を維持する。
• 学習ルーティングが、訓練時間や正確性を犠牲にすることなく、モデルの効率と性能を向上させられるかどうかを検証する。
• 推論時に計算を節約できることを示し、性能の低下なしに高速推論やモデルスケーリングを可能にする。

提案手法

• top-kルーティング機構により、各層で自己注意およびMLP計算に参加するトークンがルーター・ヘッドに基づいて選択され、kは事前に固定される。
• ルーターによって選択されないトークンは残差接続を通って渡され、結果としてトランスフォーマーブロックの計算をスキップし、FLOPsを節約する。
• 計算グラフが静的で、テンソルサイズが事前に分かっているため、ハードウェア効率と既存アクセラレータとの互換性が保証される。
• ルーティング意思決定は、バランス損失を追加で必要とせず、エンドツーエンドで学習され、top-kルーティングによる安定性を活用する。
• 自己回帰的サンプリング中、将来のトークン情報が利用できない状況でも、ルーターは単純な補助分類器または損失関数を用いて蒸留される。
• 各層の計算予算（k）を調整することで、速度と性能のトレードオフを可能にし、より速い推論やモデルスケーリングを実現する。

Figure 1: Mixture-of-Depths Transformer. As in mixture-of-experts (MoE) transformers we use a router to choose among potential computational paths. But unlike in MoE transformers the possible choices are a standard block’s computation (i.e., self-attention and MLP) or a residual connection. Since so

実験結果

リサーチクエスチョン

• RQ1 トランスフォーマーは、静的計算グラフを維持しながら、トークンおよび層の両方で計算を動的に割り当てることができるか？
• RQ2 ルーティングによる動的計算割り当ては、性能の低下なしに顕著なFLOPs削減をもたらすか？
• RQ3 同等の訓練FLOPs予算下で、MoDモデルは標準トランスフォーマーと比較して、より高い性能を達成できるか、またはより速い推論を実現できるか？
• RQ4 推論時に将来のトークン情報が欠落している場合、ルーティング性能にどのような影響が生じるか、またその影響を軽減できるか？
• RQ5 ルーティング機構は、標準ブロックや残差接続を超えた、異なる種類の計算（例：メモリ参照、ツール使用）にも拡張可能か？

主な発見

• MoDモデルは、同じFLOPsおよびウォールクロック時間で訓練された場合、ベースライントランスフォーマーと同等またはそれ以上の性能を達成し、計算効率の向上を示した。
• MoDモデルは、標準トランスフォーマーと比較して、前向き伝搬あたり最大50%のFLOPsを削減し、性能に損なわれることなく高速推論を実現した。
• 訓練FLOPsと時間は同じに保ったまま、微調整後のサンプリング段階で最大50%速い推論ステップを達成した。
• 訓練時のtop-kルーティング機構は、将来のトークン情報が利用できない自己回帰的推論にもうまく一般化され、性能の低下が最小限に抑えられた。
• MoDモデルは知的なルーティングを学習し、簡単な予測に対して不要な計算をスキップしながら、戦略的な注意ルーティングにより長距離依存関係を保持した。
• このフレームワークは拡張可能であり、メモリ参照やツール使用などの多様な計算タイプへのルーティングをサポートでき、制御可能なFLOPsコストを持つ新しいアーキテクチャ設計を可能にする。

Figure 2: Routing schemes. Tokens are funnelled to the computational path of their choice when using token-choice routing (left). If a given path exceeds its capacity (e.g., more than two tokens in this example) then surplus tokens must be dropped (purple token). The exact token that is ultimately d