Mixture of Tokens: Continuous MoE through Cross-Example Aggregation

要約

Transformer アーキテクチャに基づく Mixture of Experts (MoE) モデルは、言語と視覚のタスクの境界を押し広げています。
これらのモデルの魅力は、FLOP を増加させることなくパラメータ数を大幅に増加できることにあります。
最も広く採用されている MoE モデルは、パラメータに関して不連続であり、多くの場合、スパースと呼ばれます。
同時に、既存の連続 MoE 設計は、スパースな設計に比べて遅れているか、自己回帰デコードと互換性がありません。
完全連続手法の適応が深層学習の包括的なトレンドであるという観察に動機付けられ、私たちは、スパース MoE モデルと同様にパラメーターの数をスケーリングできるシンプルな連続アーキテクチャである Mixture of Tokens (MoT) を開発しました。
従来の方法とは異なり、MoT はさまざまな例からのトークンを混合して各専門家に割り当てます。
このアーキテクチャは、自己回帰トレーニングおよび生成と完全に互換性があります。
当社の最高のモデルは、言語の事前トレーニングにおいて高密度の Transformer モデルと比較してトレーニング速度の 3 倍の向上を達成するだけでなく、最先端の MoE アーキテクチャのパフォーマンスにも匹敵します。
さらに、MoT と MoE 間の密接な関係は、遷移チューニングと呼ばれる新しい技術を通じて実証されています。

要約(オリジナル)

Mixture of Experts (MoE) models based on Transformer architecture are pushing the boundaries of language and vision tasks. The allure of these models lies in their ability to substantially increase the parameter count without a corresponding increase in FLOPs. Most widely adopted MoE models are discontinuous with respect to their parameters – often referred to as sparse. At the same time, existing continuous MoE designs either lag behind their sparse counterparts or are incompatible with autoregressive decoding. Motivated by the observation that the adaptation of fully continuous methods has been an overarching trend in deep learning, we develop Mixture of Tokens (MoT), a simple, continuous architecture that is capable of scaling the number of parameters similarly to sparse MoE models. Unlike conventional methods, MoT assigns mixtures of tokens from different examples to each expert. This architecture is fully compatible with autoregressive training and generation. Our best models not only achieve a 3x increase in training speed over dense Transformer models in language pretraining but also match the performance of state-of-the-art MoE architectures. Additionally, a close connection between MoT and MoE is demonstrated through a novel technique we call transition tuning.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google