Mixture of Tokens: Continuous MoE through Cross-Example Aggregation

Authors: szymon  antoniak
michał  krutul
maciej  pióro
jakub  krajewski
jan  ludziejewski
kamil  ciebiera
krystian  król
tomasz  odrzygóźdź
marek  cygan
sebastian  jaszczur
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 13:28:00

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*.

Main paper: 
 introduction: 
 Transformer-based Large Language Models (LLMs) make up one of the most active fields in AI, exhibiting human-level performance across a variety of tasks, including translation, language understanding, reasoning, and code generation [1, 2, 3] . The exorbitant sizes of all state-of-the-art language models are integral to their success, with parameter counts reaching tens or even hundreds of billions. This phenomenon aligns with findings from [4, 5] , where the latter suggests that the optimal model size grows proportionally to the available computational budget. Given that hardware efficiency has been steadily increasing over the past decade [6, 7] , these findings imply that scaling will continue to be a vital component in training increasingly capable models.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.However, model scaling invariably comes at a cost. Larger models execute more Floating Point Operations (FLOPs) per token, resulting in both training and inference becoming slower and more expensive [8, 9] . Mixture of Experts [10] architectures offer an attractive alternative to standard Transformers by drastically increasing the number of parameters. The core idea is to have multiple experts, each specializing in a different part of the input space. Currently, state-of-the-art models based on MoE leverage sparsity by activating only a fraction of the parameters for each token [11] . This allows the networks to increase the number of parameters by an order of magnitude while keeping the FLOPs per token roughly constant. In this work, we use MoE to signify sparse Mixture of Experts architectures unless explicitly stated otherwise. The aforementioned sparsity is made possible with a router, a small network that selects the best experts for each token. This makes the output of an MoE layer discontinuous with respect to its parameters, as only a subset of the experts is chosen for each token (this is typically done with a discrete top-k operation). The discontinuity and the resulting fluctuations of the router's decisions have been shown to hurt training efficiency [12, 13] and are hypothesized to be a source of training instability in large MoE models [14, 15] . Conversely, existing continuous MoE architectures involve trade-offs, including the inability to scale [16, 17] , or incompatibility with autoregressive decoding [15] . This paper introduces Mixture of Tokens, a novel, continuous Transformer architecture closely related to sparse Mixture of Experts. Similar to MoE, it can support large parameter counts without significant costs in FLOPs. The core idea behind our design is for each expert to process not individual tokens separately, but their combined representation. This technique results in a continuous model that avoids the top-k operation. It requires no additional techniques commonly required in existing MoE designs (both sparse and continuous), such as load balancing losses, calculating solutions to optimization problems, or non-homogeneous training schedules [17, 18, 12] . It is capable of scaling the parameter counts akin to sparse MoEs and is compatible with autoregressive language modeling and generation. Our analysis demonstrates a 3× speedup over a dense baseline and improved stability over MoE.In summary, our contributions are the following:• Introducing the novel Mixture of Tokens (MoT), a continuous Mixture of Experts architecture that mixes tokens from different examples for joint processing. • An analysis of scaling properties of Mixture of Experts models on multiple scales.• Introducing transition tuning, allowing a pretrained MoT model to be tuned for sparse MoE inference if desired. 
background and related work: 
 In this section, we provide an overview of the approaches related to our work and discuss the differences between various MoE designs. We will introduce the Mixture of Tokens architecture in Section 3 and provide a detailed comparison between MoT and related methods in Section 3.4.Transformer scaling has been shown to be a critical factor in achieving state-of-the-art results in language and vision tasks [4, 19] , with the largest disclosed parameter counts in dense models reaching hundreds of billions of parameters [20, 3, 2] . These large models exhibit impressive abilities not present in their smaller counterparts [21] . [4] and [5] have demonstrated that the final model performance is predictable and correlates directly with the model size and the amount of training data. However, increasing model sizes raises the demand for computational resources during both training and inference [22] .Mixture of Experts (MoE) was first introduced by [10, 23] as an ensemble-like neural network comprised of separate sub-networks called experts. The original design used a gating network to select a soft assignment of experts for each input. In the context of Deep Learning, the notion of an MoE layer was introduced in [24] . [25] combined a sparse version of MoE layers with an LSTM to train a model with over 100 billion parameters, which was unprecedented at the time. The design, similar to state-of-the-art MoE models today, used a small routing network to decide the top-k best experts for each input. By choosing only a subset of the experts, they were able to increase the size of the network while keeping FLOPs per token roughly constant. The Transformer was first combined with MoE layers in [26] , where it replaced the Feed-Forward layer. The design was further simplified in [27] , which trained a model with 1.6 trillion parameters using top-1 routing. Since then, a number of studies have investigated different sparse MoE designs [28, 29, 30, 31, 32] . A comprehensive analysis of scaling properties of sparse MoE architectures can be found in [33] . [34] introduced the notion of granularity, which in spirit is similar to the number of groups in MoT, described in Section 3.2.Continuous architectures serve an important role within the field due to their flexible and efficient nature. [17] were pioneers in introducing them in MoE by presenting continuous techniques for calculating encodings of the choice of an expert. In another approach, [16] proposed a method in which they merge experts based on the weights of the router network. In a recent advancement, [15] proposed a continuous variant of MoE for the Vision Transformer, where patches are mixed only within each image.From the very beginning of the Deep Learning field, there has been a shift from discrete functions toward continuous ones. The first perceptron [35] used "all-or-none" activation, supposedly to align with propositional logic. This was later improved with soft activation functions, enabling gradient descent and multi-layer neural networks. Similarly, soft attention, introduced in [36] , enabled RNNs to look at arbitrary input from the past while maintaining the ability to learn the selection with standard gradient descent. This is in contrast to hard attention, which requires, e.g., reinforcement learning techniques. While hard attention could perform on par with soft attention [37, 38] , soft attention, with its simplicity of training, offered better trade-offs and was later used as the basic building block of the Transformer [39] .Mixture of Experts, introduced into Deep Learning by [10, 23, 25] , appears to be inherently a discrete function-after all, the expert either processes a given token or it does not. However, similar to the transition from hard to soft attention, an expert in MoE can "attend" to a combination of tokens, taken as a weighted average. This results in a smooth, continuous model and facilitates more stable training. 
mixture of tokens: 
 The goal of this work is to develop an efficient, continuous architecture that retains the scalability of Mixture of Experts, while simultaneously omitting the top-k operation, which limits a token's exposure to different experts. An intuitive approach to achieving this is to route all tokens to all experts, but this is computationally infeasible for large-scale pretraining. To overcome this constraint, the method explored in this work considers what happens not to an individual token but to a whole group of tokens instead. The main contribution of this work is the observation that allowing an expert to dynamically produce a continuous representation of the entire group of tokens, which is more lightweight to process than each token individually, yields positive results.Algorithm 1 Mixture of Tokens layer 1: for each E in experts do:2:weights E = Softmax(Linear(tokens))3:EQUATION:output E = E(mix) 5: for each i do 6:for each E do 7:update i = E output E * weights i,EMore specifically, in our design, an input batch is divided into groups of tokens, and each group is processed independently. Given a group and a single expert, a scalar weight is produced for each token. The weights are then normalized and used to compute a linear combination of the tokens, which is used as the expert's input. The experts' outputs are used for token updates as follows: for each input token, its update is a linear combination of expert outputs, with the token's mixing weights for each expert as coefficients 2 . A diagram of our method is presented in Figure 1 .To see why this method is scalable, it is helpful to examine the relationship between the number of tokens in a group and the number of experts. Essentially, if these two quantities are equal, the total computation performed by the experts is the same as in the case of top-1 routing. This allows MoT to benefit from the same parameter scaling as seen in MoE, which we confirm empirically in Section 4.2.As we mix tokens from multiple unrelated sequences, we do not expect the model to meaningfully use the information from one sequence to improve prediction in a different sequence. However, we hypothesize that such mixing (1) provides richer feedback (gradients) to train the model, especially the router, and (2) results in a smoother loss landscape, which is resistant to small perturbations in inputs and weights.Intuitively, for a given expert, from the perspective of each token, the token receives a certain amount of update to its representation (in the residual stream) based on:• Itself, producing a proper signal expected to improve the token representation.• Tokens other than itself, which are essentially random tokens from unrelated sequences.As these sequences are randomly sampled from the dataset, the impact from these tokens will point in random directions and, essentially, just add some amount of noise to the token update.We generally expect neural networks to be resistant to a certain amount of noise added to them. Moreover, while the signal-to-noise ratio worsens for tokens with low expert weight, the expert weight also modulates the magnitude of the update. Therefore, the amount of noise added to the representation is limited.We stipulate that MoT experts will learn to focus on a single token or a small number of tokens, thereby minimizing noise and approximating sparse MoE when optimal. However, other tokens will be assigned nonzero weight, allowing some information to flow to the router for each and every token-expert pair, unlike sparse MoE. Additionally, the output of MoT is more continuous, with small perturbations of input/weights corresponding to small changes in the output rather than large discrete jumps, as may occur in the case of sparse MoE.Building on the design described above, we experiment with feeding more than one mixture into each expert. Without further modifications, this approach would result in a linear increase in computational costs for each additional mixture processed. To circumvent this added cost, MoT uses more experts, but each expert has a proportionally reduced hidden dimension. This way, each mixture is processed by a smaller expert, and the layer's total number of parameters, as well as the number of FLOPs used by all experts, remains approximately the same. We find that this design consistently improves MoT as the number of processed mixtures increases, just as it does in sparse MoE [34] . Likewise, the optimal granularity aligns roughly with the number of mixtures in this work.The question of how token groups are decided within a batch is crucial for compatibility with autoregressive training and inference. The main insight here is that tokens from the same sequence cannot be placed in a single group, as the mixing operation would result in an information leak. Due to this restriction, MoT groups tokens from different examples based on their position in the sequence. Thus, all tokens within a group share the same position in their respective sequences. As previously mentioned, to maintain a constant number of FLOPs per token, an increase in the number of experts means an equal increase in group size. An illustration of how grouping is done within a batch of tokens is shown in Figure 3 . In this example, the group size is 2. Note that the maximum possible group size equals the size of the batch.Scaling The technique featured in [17] is based on a continuously differentiable sparse top-k router, which is a major advantage compared to the common top-k gating. However, this approach requires that all experts are utilized in a portion of training, rendering it computationally prohibitive for models with large numbers of experts. The architecture based on merging experts proposed in [16] also offers an attractive, continuous alternative to top-k gating, yet the cost of merging all experts again scales linearly with the number of experts. To address this, the technique is applied once per sequence, which limits the expressive power of the final model.Training stability [26] reported instabilities during the training of large MoE models, stemming from the inaccuracies when calculating router weights in low precision. To stabilize the training, they resorted to using full precision. [27] made progress in using mixed precision when training MoE by using selective high precision for gating. When comparing [26, 27] to MoT, an advantage of our technique emerges -it is more robust to training in lower precision than other methods. We conjecture, that this is due to the merging mechanism being less susceptible to rounding errors than gating in sparse MoEs.Token dropping Token dropping is a phenomenon where tokens do not receive an update from any expert. This can happen when the expert was selected by too many tokens in a batch [27, 26, 40] or, in the case of routing experts to tokens, when a token is not selected by any expert [30] . Existing techniques to combat this phenomenon offer a partial solution, yet the problem persists. In contrast, tokens in MoT are part of every mixture produced within their group; hence, they always receive an update.Auto-regressive decoding Mixture of Tokens is based on the concept of merging tokens before they are processed by an expert. An encoder-only design of a similar nature is featured in concurrent work [15] . The technique is based on merging patches within an image for vision models, i.e., within a single sample. This should be contrasted with MoT, which merges tokens from different sequences within a batch. This crucial difference allows MoT to be compatible with autoregressive training and inference.The time complexity of our approach is identical to that of the Token Choice and Expert Choice methods. In all cases, the cost of computing routing logits is of orderO(d model • N experts • N tokens ).Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [NA] Justification: The paper doesn't contain theoretical results or theorems, as it relies on experimental data.Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
experiments: 
 The focus of this work is to investigate the efficiency of Tokens on autoregressive language modeling. To measure model quality, we pretrain models for a fixed number of tokens and compare final perplexity in accordance with existing MoE literature [28, 27] . In all experiments, the models are trained on the C4 dataset3 [41] and use the GPT-2 tokenizer. Unless specified otherwise, we use mixed precision, where all heavy computation is done in bfloat16, whereas the optimizer state and weights are kept in full precision. To study the stability of our model, we experiment with training fully in reduced precision.Our main result is a substantial speed-up of MoT models compared to dense Transformers (Figure 7 ) and results comparable to sparse MoEs (Figure 6 ). What follows is the analysis of the scaling properties of the MoT architecture with respect to the number of parameters (Figure 4 ) and the number of mixtures sent to each expert (Figure 5 ). In addition to using the Transformer as a baseline, we also compare against Token Choice [27] and Expert Choice [30] as sparse MoE baselines. Given that Expert Choice is sensitive to the size of the batch, in order to avoid discrepancy between training and inference, we group tokens prior to routing in training Expert Choice models.Mixture of Tokens models demonstrate strong scaling properties with respect to the number of parameters. As seen in Figure 4 , increasing the number of experts in MoT layers while using the same compute budget yields consistent improvements. All MoT models are a strict improvement over the Transformer. The figure also includes an ablation experiment, where the mixing weights are fixed to 1/n, where n is the group size. This corresponds to a uniform mixing strategy; the performance of that model clearly suffers, confirming that MoT layers learn non-trivial mixing strategies.The increased number of token mixtures described in Section 3.2 represents another axis of scaling for MoT models, once again exhibiting consistent improvements. We hypothesize that this phenomenon is due to two mechanisms: First, the model becomes more expressive with a larger number of smaller experts. Second, the model can allocate its focus (the mixing weights) more flexibly to more important tokens while reducing the updates for trivial ones.Crucially, the performance of Mixture of Tokens is comparable to that of the strong Mixture of Experts baselines (Figure 6 ). An increased number of mixtures allows it to compete with both Expert Choice and Token Choice architectures. As the sparse routing is hypothesized to contribute to training instabilities in large sparse models, Mixture of Tokens, being continuous, presents a promising alternative. To investigate training instabilities at the scale we experiment on, we trained models entirely in bfloat16, as opposed to the mixed precision used in all other experiments. The results confirm that MoT is more resistant to lower precision training: as the precision of training decreases, the performance of Expert Choice drops below that of Mixture of Tokens, despite the former attaining better perplexity using mixed precision. We find this to be evidence of the architecture's potential for stable training at higher model scales. See Table 1 for details.Finally, we combine our findings on MoT scaling properties to train our most efficient MoT model and compare it to the Transformer baseline (Figure 7 ). The result is a model that achieves the final loss of the baseline in one-third of the training steps. This represents a 3× improvement in terms of the compute budget. Mixture of Tokens suffers from a drawback common to MoEs, namely, it does not support unbatched inference. This is a direct consequence of its design -in the forward pass, it groups several tokens from different examples in the batch. With the growing adoption of Large Language Models on consumer hardware [43, 44] , this lack of support could hinder the architecture's wider adoption. While a Mixture of Tokens with a group size of one is technically possible, in order to keep FLOPs constant, the layer would need to trivially reduce to a standard Transformer MLP.To address this issue, we demonstrate that the weights learned by the Mixture of Tokens can be used to directly initialize a Token Choice model of the same specifications (number of experts and expert size). The layer responsible for producing mixing weights is utilized to initialize the sparse router. In order to mitigate the difference in performance that is caused by this change in architecture, we train the entire new model (no weights are frozen) for 10% of the total pretraining steps of the original model in order to recover the original model's performance (measured in eval loss). We call this technique transition tuning. This way, it is possible to train with Mixture of Tokens and enjoy unbatched generation at inference time. We hypothesize that this pipeline would be especially attractive in setups where having parts of the model train in higher precision is impossible, e.g., on specialized, low-precision hardware. The results are presented in Figure 8 . Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: The paper comes with a complete training and evaluation code, the settings of experiments are contained in the paper, and the dataset is publicly available. Therefore, the paper provides all the information needed to reproduce the results. 
limitations and future work: 
 With the strong performance of MoT on medium-sized models, an obvious next step is to train larger models. This would present an opportunity to validate the stability results on larger models, where training instabilities are more common.As with most Mixture of Experts models, the memory footprint of MoT layers is substantial. Scaled models require large amounts of RAM on specialized hardware for training, making their adoption expensive. To this end, an attractive future direction would be to investigate model distillation with Mixture of Tokens models.In this work, we experimented only with text modality in an autoregressive manner. Other modalities, such as vision, heavily overlap with the approach presented in work concurrent to ours [15] .This mixing of tokens from different sequences and the requirement of performing batched inference may be undesirable in some use cases. While performing unbatched inference is always inefficient with LLMs, as the memory throughput to access model weights becomes the bottleneck, unbatched inference still finds its uses. Even though transition tuning solves this problem, exploring different inference strategies might bring new insights.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: Justification: The utilized computational resources are summarized in Appendix E.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
conclusions: 
 In this work, we presented the Mixture of Tokens, a novel continuous Mixture of Experts architecture compatible with autoregressive decoding. This architecture scales to model sizes similar to sparse Mixture of Experts models, matches their performance, and is more resistant to training instabilities due to lower precision training. Moreover, we introduced transition tuning, a technique for initializing an MoE model with another pretrained MoE model of a different architecture, and showed that the new model achieves the performance of the original one using a fraction of the compute budget. Guidelines: 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes]Justification: The societal impact of this paper is no different from any other work improving the efficiency of LLMs. This is described in more detail in the Appendix D.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: We do not release any pretrained models and use publicly available data.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 
licenses for existing assets: 
 Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We are the creators of the code used in the paper. We cited the original paper of the dataset and included the URL to a version that was used in the training setup. We also mentioned the license of the dataset, which we properly respected.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: • The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example (a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully. (c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.Justification: Authors confirm adherence to the NeurIPS Code of Ethics in every respect.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).The authors note that an MoT layer admits an efficient vectorized implementation, where all meaningful computations are done with batched matrix multiplications.https://huggingface.co/datasets/c4, dataset licensed under ODC-By. 


Decision: The paper proposes an interesting approach to MoE-style autoregressive modeling using soft combinations of tokens along the batch dimension. The reviewers unanimously liked the idea and found the paper easy to follow. Concerns about computational cost and downstream evaluations were addressed in the rebuttal. I feel that for future work, more extensive downstream evaluation of this architecture would be interesting, however, in accordance with the reviewers I believe that the paper & rebuttal provide sufficient evidence in their current state to justify publication. Congratulations on the acceptance!

Reviews: review: paper_summary: paper_summary: This paper proposes a new routing algorithms for MoEs: the mixture of tokens. The context is the following: routing in MoEs is tricky because one token gets usually assigned to one or a few experts and so the gradient feedback to update the router is not great. Therefore, several recent papers proposed to either average the experts parameters [1,2] or mixing the tokens [3] to overcome this issue. This paper focuses on this research direction. Prior work [3] did it in the case of computer vision where causality is not an issue. This paper tries to adapt mixture of tokens to the case of autoregressive models. They first explain how they do the mixture of tokens in section 3.1 and then detail how they manage to ensure causality in section 3.2. Lastly in section 4, they show the benefits of their method over other routing algorithms: in particular, they show that Mixture of Tokens (MoT) minimizes the eval loss much faster than dense Transformers and the performance is slightly better than standard MoEs with token choice and expert choice routings.






[1] Zhong, Z., Xia, M., Chen, D., & Lewis, M. (2024). Lory: Fully Differentiable Mixture-of-Experts for Autoregressive Language Model Pre-training. arXiv preprint arXiv:2405.03133.
[2] Muqeeth, M., Liu, H., & Raffel, C. (2023). Soft merging of experts with adaptive routing. arXiv preprint arXiv:2306.03745.
[3] Puigcerver, J., Riquelme, C., Mustafa, B., & Houlsby, N. (2023). From sparse to soft mixtures of experts. arXiv preprint arXiv:2308.00951.
strength_weakness: strength_weakness: In general, I am a bit doubtful about the methods that attempt to merge either tokens or parameters in order to solve the ill-posed problem of routing in MoEs. For me, MoEs are primarily introduced for efficiency in that we can increase the size of the model while keeping the same number of FLOPs. In my opinion, the approaches that merge tokens/parameters lose the efficiency aspect. I think the authors tried to alleviate this issue by using smaller experts (as done in [1]) but I don't think this totally solves the problem. [1] also shows that when the granularity is too high, MoEs performance drop. In any case, I am happy that some researchers tried mixing tokens in the context of autoregressive models but I do not fully believe in it. Here are some more precise questions I would like to ask to the authors:

- **Computational cost**:Can the authors analyze the additional number of FLOPs that their method incurs compared to standard token choice? Or to other merging methods like the parameter merging in [2]?

- **Any advantage over expert choice/token choice?**: do you think there is any benefit from using mixture of tokens over other existing routing algorithms? It is not obvious from Figure 6. I understand the primary motivation and it makes sense to me. But in practice, it does not seem that the merging methods yield any significant benefits over token or expert choice.

- **Preserving the causal structure**: This is always a challenge for routing algorithms like expert choice or Lory that do not naturally preserve causality. Have you tried other schemes to preserve the causal structure? Would the one that is used when running expert choice on autoregressive models fail? Can you clarify why you believe there is no information leakage, sorry I may have missed this point?


[1] Krajewski, J., Ludziejewski, J., Adamczewski, K., Pióro, M., Krutul, M., Antoniak, S., ... & Jaszczur, S. (2024). Scaling laws for fine-grained mixture of experts. arXiv preprint arXiv:2402.07871.
[2] Zhong, Z., Xia, M., Chen, D., & Lewis, M. (2024). Lory: Fully Differentiable Mixture-of-Experts for Autoregressive Language Model Pre-training. arXiv preprint arXiv:2405.03133.
questions: questions: I have written my questions above.
limitations: limitations: The authors have clearly mentioned the limitations of their approach.
paper_summary: paper_summary: In this paper, the authors propose a novel method called Mixture of Tokens (MOT), an expert-based architecture that addresses the drawbacks of existing Mixture of Experts (MoE) approaches, such as reduced training efficiency, instability, and the necessity of using load balancing losses and elaborate training schedules. These issues are often caused by the top-k operation used to select the active experts. The authors' approach combines tokens into groups (subsets of several tokens) and uses a weighted sum of tokens within these groups as input to each expert, with each expert using its own weights for the sum. This allows each token within the group to be processed by each expert, resolving the issues associated with top-k routing. The authors validate their method through extensive experiments on an autoregressive language modeling task, demonstrating significant improvements in training stability, reduced training costs (compute budget), and improved scaling with regard to model size.
strength_weakness: strength_weakness: * One of the weaknesses the authors pointed out in the limitations section is the necessity to use batched inference for the approach, which could limit the applicability scope of the method. While there are still important applications for batched inference, I believe that resolving this issue could significantly increase the influence of the method

* From the perspective of the experimental evaluation, I would be curious to see evidence that the behavior demonstrated in the paper would hold in other domains, such as images. Additionally, it seems that the proposed approach could resemble the method in [1] a lot.

* In terms of experiments, it would also be interesting to see comparisons with other MoE approaches mentioned in the related work, beyond the reported Token Choice and Expert Choice, or more novel approaches such as [2] and similar methods.

[1] Puigcerver, Joan, et al. "From sparse to soft mixtures of experts." ICLR 2024
[2] Anthony, Quentin, et al. "Blackmamba: Mixture of experts for state-space models." ICLR 2024
questions: questions: * As stated in the weaknesses section, is it possible to demonstrate the effectiveness of the approach beyond the autoregressive language modeling task? I believe it could significantly strengthen the paper.

* In addition to the previous question, could you elaborate on what the major differences are between the proposed approach and the [1] method? Is it the fact that in MoT mixing happens within different batch inputs and in [1] we mix tokens within the same image? If so, then the novelty of MoT (compared to [1]) comes from the grouping and the application to the autoregressive task, which could undermine the contribution of the paper.  

[1] Puigcerver, Joan, et al. "From sparse to soft mixtures of experts." ICLR 2024
limitations: limitations: No significant limitations of the paper, though the authors' discussion on limitations is much appreciated.
paper_summary: paper_summary: This paper proposes a new MoE architecture called Mixture of Tokens (MoT). The motivation for this architecture is twofold: first it is the training instabilities incurred (among other things) by low precision training in standard sparse MoEs; secondly, it is the discontinuous nature of sparse MoEs that makes them hard to train. This work seem to be highly inspired by [1].

In contrast to MoEs, which select top-k subset of experts for each token, MoT aggregates tokens within groups created by subsetting examples across the batch dimension. The cross-example token aggregation weights are produced by a router conditioned on the tokens in the corresponding group. The aggregated group representations are then passed through all the experts. The output token representation is produced by a weighted average of all experts' outputs for a group, where the weights are the original aggregation weights. 

The paper measures validation perplexity of MoT comparing it to compute matched dense transformer and sparse MoEs demonstrating its superiority over the prior and on par performance with the later. MoT shows better training stability in low precision training regime as compared to sparse MoEs.

[1] Puigcerver, Joan, et al. "From sparse to soft mixtures of experts." arXiv preprint arXiv:2308.00951 (2023).
strength_weakness: strength_weakness: One weakness of this paper is a lack of intuitive explanation of why mixing across examples makes sense? I understand that it results in more stable training due to lack of discreet operations, but why does it make sense from the modelling perspective? 

If sentences in the batches are consecutive interrelated text passages, then MoT effectively increases the context length by letting the model to also rely on cross-example information. Can this have an effect on generalization ability? e.g. in instruction fine-tuning, examples in the batches are usually not related to each other, hence MoT type of pre-training may be detrimental for downstream task adaptation of these models, where downstream tasks examples are completely unrelated. 

Therefore, another potential weakness for me is a lack of  generalization/downstream adaptation analysis. Does the training stability advantage of MoT improves generalization and adaptability of the model to downstream tasks? Or does cross-example information sharing actually weakens downstream task generalization?
questions: questions: - ll. 52 - 61: I am uncertain if points 2 and 4 qualify as contributions of this work. These appear to be properties of the proposed MoT method rather than contributions in themselves.
- ll. 55-57: If I understand correctly, the key comparison should be with other sparse MoE methods rather than dense transformers. Therefore, why emphasize a speedup over dense transformers as a contribution when the more relevant comparison is the speedup over other sparse MoE architectures?
- l. 104: "Mixture of Experts, introduced into deep learning by [24],..." does not seem to be correct. MoEs have a much longer history, please consider works [2,3];
- l. 190: why half of feed-forward layers are replaced with MoT layers?
- ll. 167 - 169: an important difference to [15] that might be mentioned explicitly here, is that MoT aggregates tokens across examples in a batch, while [15] aggregates patches within each example.
- caption Figure 6: if I am not mistaken, MoT here not only matches but outperforms best sMoE
- since one of the main advantages of the proposed MoT is that it does not require additional load balancing tricks etc., it would be interesting to include a empirical comparison (or at least a discussion) to existing sparse MoEs that also do not require such tricks, see e.g. [4]
- how are examples in the batches created? Are these consecutive interrelated text passages from the same underlying text?
- MoT's weighted token grouping is reminiscent of an attention operation, but across examples in a batch. Is there any intuition for why aggregating information across examples can be useful? One potential way it can be useful, is that if sentences in the batches are related (e.g. maybe these are consecutive sentences from the underlying training text) then of course  incorporating information from previous/future examples can help reducing loss (training and validation), since the examples (samples in a batch) are interrelated. 

Overall, I believe this paper has significant potential. However, I have concerns regarding downstream task generalization and the relationship between the examples in the training batches. I am willing to reconsider my rating upon rebuttal.

[2] Robert A Jacobs, Michael I Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79–87, 1991b.

[3] Robert A Jacobs, Michael I Jordan, and Andrew G Barto. Task decomposition through competition in a modular connectionist architecture: The what and where vision tasks. Cognitive science, 15(2):219–250, 1991a

[4] Csordás, Róbert, Kazuki Irie, and Jürgen Schmidhuber. "Approximating two-layer feedforward networks for efficient transformers." arXiv preprint arXiv:2310.10837 (2023).
limitations: limitations: Limitations are addressed in Sec. 5. 
I think one missing limitation of this work is the fact that no downstream task generalization is evaluated.

Comments: comment: We thank the reviewer for their response. We are glad that we were able to clarify the points raised.
We appreciate the reviewer's suggestion and are grateful for the improved score. We will include a clarification on the computational cost in the camera-ready version of the paper.
Thank you for your detailed rebuttal and for addressing the questions raised. I appreciate the clarifications provided, particularly regarding batched inference and the differences between MoT and SoftMoE. Given the solid contributions and your thorough responses, I am maintaining my original rating.
I would like to thank the reviewers for their reply. I appreciate their clarification regarding the computational cost and I believe that this should be added to the final version of the paper. I increase my score by 1 point.
We appreciate the reviewer's engagement and we are thankful for the update of the score.
I want to thank the authors for their elaborate replies! I will increase my score.