Mixture of Tokens: Efficient LLMs through Cross-Example Aggregation

要約

Mixture of Experts (MoE) モデルは、トレーニングと推論のコストを維持しながら Transformer モデルのパラメータ数を増加できるという約束にもかかわらず、そのアプリケーションには顕著な欠点があります。
これらのモデルの重要な戦略は、処理されたトークンごとに、最大でも少数のエキスパート、つまり広範なフィードフォワード層のサブセットをアクティブにすることです。
しかし、このアプローチには課題がないわけではありません。
エキスパートとトークンのマッチングの操作は個別に行われるため、MoE モデルはトレーニングの不安定性やエキスパートの利用の不均一などの問題が発生しやすくなります。
補助損失やバランスを考慮したマッチングなど、これらの懸念に対処するために設計された既存の手法は、モデルのパフォーマンスが低下するか、トレーニングがより困難になります。
これらの問題に対応して、私たちは、前述の問題を回避しながら、MoE アーキテクチャの利点を保持する完全に微分可能なモデルであるトークンの混合を提案します。
このアプローチでは、トークンをエキスパートにルーティングするのではなく、さまざまなサンプルのトークンをエキスパートに供給する前に混合し、モデルがトークンとエキスパートのすべての組み合わせから学習できるようにします。
重要なのは、この混合を無効にして、推論中に異なるシーケンスが混合するのを避けることができることです。
重要なことは、この方法は、マスクされた大規模言語モデルと因果的な大規模言語モデルのトレーニングおよび推論の両方と完全に互換性があることです。

要約(オリジナル)

Despite the promise of Mixture of Experts (MoE) models in increasing parameter counts of Transformer models while maintaining training and inference costs, their application carries notable drawbacks. The key strategy of these models is to, for each processed token, activate at most a few experts – subsets of an extensive feed-forward layer. But this approach is not without its challenges. The operation of matching experts and tokens is discrete, which makes MoE models prone to issues like training instability and uneven expert utilization. Existing techniques designed to address these concerns, such as auxiliary losses or balance-aware matching, result either in lower model performance or are more difficult to train. In response to these issues, we propose Mixture of Tokens, a fully-differentiable model that retains the benefits of MoE architectures while avoiding the aforementioned difficulties. Rather than routing tokens to experts, this approach mixes tokens from different examples prior to feeding them to experts, enabling the model to learn from all token-expert combinations. Importantly, this mixing can be disabled to avoid mixing of different sequences during inference. Crucially, this method is fully compatible with both masked and causal Large Language Model training and inference.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google