We would like to thank the reviewer for their questions and suggestions and appreciate the recognition of our paper's strengths. Below, we address the weaknesses and questions mentioned in the review. If our answers address the reviewer's concerns, we would like to kindly ask for the reconsideration of the rating.

Regarding previous related work. We would like to refer to our reply in the general comment.

Regarding weaknesses:

W1: We agree that, apart from experimental results, an intuition behind mixing tokens across examples should be provided. We will add such an explanation to the camera-ready version.

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 argue that such mixing (1) provides richer feedback (gradients) to train the model, especially the router, and (2) provides a smoother loss landscape, resistant to small perturbations in inputs/weights.

Intuitively, for a given expert, from the perspective of each token, the token gets some amount of update to its representation (in the residual stream) based on:

1. Itself, which produces a proper signal expected to improve the token representation.
2. Tokens other than itself, which are essentially random tokens from unrelated sequences. As those sequences are randomly sampled from the dataset, the impact from those 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 some amount of noise added to them (e.g. dropout doesn't hurt). 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 - minimizing the noise and approximating sparse MoE when optimal. Still, other tokens will be assigned nonzero weight, enabling some information to flow to the router for each and every token-expert pair (in contrast to sparse MoE). Moreover, the output of MoT is more continuous, with small perturbations of input/weights corresponding to small changes in the output instead of large discrete jumps with sparse MoE.

W2: During the training and evaluation, we essentially put random, unrelated examples into the same group. For reasons explained in W1 not much cross-example information transfer occurs. Therefore, it does not have any effect on generalization ability.

W3: We believe that comparing perplexity generally predicts model improvements more reliably, especially when trying to predict how particular changes to the model architecture will impact extremely large-scale models. However, during the short rebuttal period, we measured performance on a few benchmarks relevant at this model scale, comparing MoT-Medium to Transformer-Medium, without fine-tuning, zero-shot. In these evaluations, for MoT, just a single evaluation query is put into a batch of 32, with the rest of the batch being comprised of random sequences from the C4 training dataset - ensuring this is still zero-shot.

Regarding questions:

Q1: After some consideration, we agree with the reviewer, and we will move the results/properties of MoT itself (in particular, speed-up against the dense model and improved stability compared to MoE) from the contributions bullet points to the main introduction part.

Q2: We think it is important to compare MoT to both kinds of models—dense models, where the main benefit we show is a significant speedup (point 2), and MoE models, where the main benefit we focus on is increased stability of the training (point 4). We will make those points clearer while moving them into the main text.

Q3: Thank you for bringing this to our attention. This is an oversight on our part. While we already cite [5] (in line 75), we will change the citation in line 104 as well (to both [5] and [6]), to ensure proper credit is given.

Q4: In the literature on MoE (e.g. [1][7][8]), it is quite standard for MoE layers to replace only half of FFN layers - this keeps the majority of MoE advantages while requiring significantly fewer total parameters.

Q5: We will revise this paragraph to highlight the differences and similarities between SoftMoE and MoT more clearly, as suggested.

Q6: While MoT indeed outperforms sparse MoE here, the differences between the compared approaches are relatively minor. Given these results, we think it would be an exaggeration to claim that our method outperformed the others. Instead, we focus on showing other advantages of MoT, like increased training stability.

Q7: Please note that one of the methods we compare our approach to is the expert-choice variant of MoE that operates without load-balancing. However, it employs a sparse router, which requires higher precision and is not fully differentiable, leading to unstable training. By replacing the sparse router with our fully differentiable counterpart, we achieve stable training and enable training entirely in lower precision.

Q8: Training batches are created in the same manner as for regular training. We sample random passages from the training dataset until the desired number of examples in the batch is reached. Consequently, it is unlikely that any two passages within a batch are related.

Q9: See our explanation of intuition behind MoT in the previous section W1 of this response.

We want to thank the reviewer for their comments and questions. We also appreciate the mention of the simplicity of integrating our method into existing approaches. If the reviewer's concerns have been addressed, we would like to kindly ask for the reconsideration of the rating.

Regarding batched inference and applicability. We agree that the necessity of using batched inference, while fine for many industry applications, is a limitation for others. We would like to highlight, however, that transition tuning (introduced in Section 4.4 of our work) can lift this limitation by enabling many benefits of MoT during training (e.g. increased stability) while converting the model to sparse MoE for unbatched inference.

Regarding other domains and tasks other than autoregressive language modeling. We agree that testing the method on other domains is an interesting avenue of research. With limited time and budget, we focused on the domain with, arguably, the most pressing problem of efficiency because of recent scaling efforts across the whole industry. Apart from text the important domain to test MoT on would be vision - but those experiments would heavily overlap with a concurrent paper of SoftMoE. We will properly mention that in the Limitations section of our work.

Regarding the differences with SoftMoE. As the reviewer noted, the main difference between MoT and SoftMoE is that MoT mixing happens between different sequences and not within a sequence/image, enabling its use on autoregressive tasks. We would like to refer to our reply in the general comment for further information about concurrency and the novelty of our work compared to SoftMoE.

We would like to thank the reviewer for their feedback and questions. We also appreciate the recognition of the value of the research problem and the presentation of our work. We hope that the answers below adequately answered the reviewer's questions. If that is the case, we kindly ask for a reconsideration of the paper score.

Regarding SoftMoE. We would like to refer to our reply in the general comment.

Regarding computational cost. The time complexity of our approach doesn't change compared to standard token choice or expert choice methods. The cost of computing routing logits itself takes $O(d_{model}*N_{experts}*N_{tokens})$ FLOPs, and this time complexity is the same for all MoE variants. The mixing and unmixing operations in MoT have the same time complexity. Regarding granularity, our models use values around compute optimal for this model size as calculated in [3]; therefore, according to their experiments, performance should not drop, as in the case of extreme granularity. We will clarify this in the camera-ready version and add a citation to their work.

Regarding the benefits of MoT over expert choice/token choice. We believe that MoT and techniques developed in the future based on MoT will provide better training stability (shown in our work), and therefore enable improved training performance, especially at scale. Moreover, within a continuous setting, each expert is trained on every token, which might be beneficial once we scale to higher expansion rates.

Regarding the causal structure. The original expert choice routing did not try to preserve the causal structure by default, with the Limitations section in [4] stating, "The expert choice method might not immediately apply to auto-regressive text generation as our current implementation takes in the past and future tokens to perform the top-k selection." The approach used in our paper is, we think, the only natural adaptation of expert choice to autoregressive models preventing information leaks. This method is used both in our expert-choice and MoT models. The approach of preventing the causal information leak used in [2] is not applicable to MoT. In MoT, we can never mix tokens from the same sequence, and routing based on the previous tokens/segment (rather than the current tokens/segment) doesn't mitigate the issue.

Regarding no information leakage in MoT. In MoT, there is no information leakage, as we never mix tokens from the same sequence together. Each mixture is based on a specific position in a sequence. So, a token at position $i$ in a sequence $s$ can, in a MoT layer, see tokens from any sequence at position $i$. Combined with the causal attention layer, token at position $i$ can see any tokens at positions $0:i$ in a sequence $s$ (as in standard causal Transformer), and also any tokens at positions $0:i$ in any sequence in a group/batch (through a combination of MoT layer being able to look at a different sequence, and causal attention being able to look at any previous position). However, looking at different sequences in a batch doesn't constitute an information leak since those other sequences are IID examples randomly drawn from the dataset.