Mixture of Tokens: Continuous MoE through Cross-Example Aggregation

TL;DR

Mixture of Tokens (MoT) is developed, a simple, continuous architecture that is capable of scaling the number of parameters similarly to sparse MoE models and is fully compatible with autoregressive training and generation.

Abstract

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.

Figures (8)

• Figure 1: Mixture of Tokens: Each expert receives a unique mixture of tokens in the group. Mixing weights are determined by the controller, which is a fully connected layer (omitted for clarity). For a given token, its update is a linear combination of expert outputs, with the coefficients equal to the token’s original mixing weights for each expert.
• Figure 2: (Left) Diagram of a standard feed-forward layer featured in the Transformer architecture: each token is processed with the same MLP, independent of other tokens. (Right) Diagram of a Token Choice layer, where each token decides which expert to choose. This way, different experts process a different number of tokens. If one expert is chosen by too many tokens, a portion of the tokens are dropped --- they receive no update.
• Figure 3: Each group consists of tokens with the same position in a sequence. In this example, the group size is $2$. Note that the maximal possible group size is equal to the batch size.
• Figure 4: Scaling with respect to the number of parameters. Also featured are the Transformer baseline and an MoT model with a non-learnable, uniform routing strategy.
• Figure 5: Scaling with respect to the number of token mixtures.