Mixture-of-Depths: トランスフォーマーベースの言語モデルにおける計算リソースの動的割り当て

Mixture-of-Depths: Dynamically allocating compute in transformer-based language models

April 2, 2024

著者: David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, Adam Santoro

cs.AI

要旨

Transformerベースの言語モデルは、入力シーケンス全体にFLOPsを均一に分散させます。本研究では、Transformerが代わりに、シーケンス内の特定の位置にFLOPs（または計算リソース）を動的に割り当てることを学習し、モデルの深さにわたって異なる層でシーケンスに沿った割り当てを最適化できることを示します。私たちの手法では、特定の層でセルフアテンションとMLP計算に参加できるトークン数（k）を制限することで、総計算量の予算を強制します。処理されるトークンは、ネットワークがトップkルーティングメカニズムを使用して決定します。kは事前に定義されるため、このシンプルな手順は、他の条件付き計算技術とは異なり、既知のテンソルサイズを持つ静的な計算グラフを使用します。しかし、k個のトークンの識別情報は流動的であるため、この手法は時間とモデルの深さの次元にわたってFLOPsを不均一に消費することができます。したがって、計算量の総計は完全に予測可能ですが、トークンレベルでは動的かつ文脈依存です。この方法で訓練されたモデルは、計算リソースを動的に割り当てることを学習するだけでなく、効率的にそれを行います。これらのモデルは、同等のFLOPsと訓練時間でベースライン性能を達成しつつ、フォワードパスごとのFLOPsはわずかで、訓練後のサンプリング中にステップする速度が最大50％向上する可能性があります。

English

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.