Exploiting Activation Sparsity with Dense to Dynamic-k Mixture-of-Experts Conversion

Authors: filip  szatkowski
bartosz  wójcik
mikołaj  piórczy
simone  scardapane
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-10 19:15:47

Abstract: Transformer models can face practical limitations due to their high computational requirements. At the same time, such models exhibit significant activation sparsity, which can be leveraged to reduce the inference cost by converting parts of the network into equivalent Mixture-of-Experts (MoE) layers. Despite the crucial role played by activation sparsity, its impact on this process remains unexplored. We demonstrate that the efficiency of the conversion can be significantly enhanced by a proper regularization of the activation sparsity of the base model. Moreover, motivated by the high variance of the number of activated neurons for different inputs, we introduce a more effective dynamic-$k$ expert selection rule that adjusts the number of executed experts on a per-token basis. To achieve further savings, we extend this approach to multi-head attention projections. Finally, we develop an efficient implementation that translates these computational savings into actual wall-clock speedup. The proposed method, Dense to Dynamic-$k$ Mixture-of-Experts (D2DMoE), outperforms existing approaches on common NLP and vision tasks, reducing inference cost by up to 60\% without significantly impacting performance.

Main paper: 
 introduction: 
 Transformers have become a predominant model architecture in various domains of deep learning such as machine translation [47] , language modeling [6, 31] , and computer vision [7, 21] . The widespread effectiveness of Transformer models in various applications is closely related to their ability to scale efficiently with the number of model parameters [20] , prompting researchers to train progressively larger and larger models [45, 19] . However, the considerable computational demands of these models often restrict their deployment in practical settings with limited resources.At the same time, Transformer models exhibit considerable activation sparsity in their intermediate representations [24] , which suggests that most of their computations are redundant. Conditional computation methods can reduce these unnecessary costs by using only a subset of the model parameters for any given input [14] . In particular, Mixture-of-Experts (MoE) layers [38] , consisting of multiple experts that are sparsely executed for any input token, are an effective way to decouple the number of parameters of the model from its computational cost [3] . As shown by [52] , many pre-trained dense Transformer models can be made more efficient by converting their FFN blocks into MoE layers, a process they call MoEfication. We consider the following research question: what is the optimal way to convert a generic Transformer model into an equivalent sparse variant? We identify a series of weaknesses of the MoEfication process limiting the resulting accuracy-sparsity tradeoff, and propose corresponding mitigations as follows. We call the resulting algorithm Dense to Dynamic-k Mixture-of-Experts (D2DMoE) and outline it in Figure 1 .1. First, we analyze the relationship between the activation sparsity of the starting model and the efficiency of the final MoE model. We show that computational savings are directly related to sparsity levels, and we correspondingly enforce higher activation sparsity levels before conversion through a lightweight fine-tuning process, which leads to a substantially improved cost-to-performance trade-off.2. We identify the router training scheme in the original MoEfication algorithm as a limitation of the conversion process. We propose to frame the router training as a regression problem instead, hence our routers directly predict the norm of the output of each expert.3. We show that Transformer models exhibit significant variance of the number of activated neurons, and standard top-k expert selection in the MoE layers is inefficient. We propose an alternative dynamic-k expert selection scheme that adjusts the number of activated experts on a per-token basis. This approach enables the model to efficiently allocate compute between easy and hard inputs, increasing the overall efficiency.4. We generalize the conversion method to any standalone linear layer including gated MLP variants commonly found in modern LLMs [45, 42] and projections in Multi-Head Attention (MHA) layers (which often account for over 30% of total computations in Transformer models [39] ).For MHA, we propose a replacement procedure in which every dense layer is substituted by a two-layer MLP module trained to imitate the output of the original layer.We evaluate D2DMoE across benchmarks in text classification, image classification, and language modeling, demonstrating significant improvements in cost-performance trade-offs in all cases. D2DMoE is particularly well-suited for contemporary hardware, as evidenced by our efficient GPU implementation, which we contribute alongside our proposed method. 
motivation: 
 MoE models have gained a lot of traction over the last years as an effective architecture to decouple the parameter count from the computational cost of the models [56] . In a MoE layer, hard sparsity is usually enforced explicitly by applying a top-k operation on the outputs of a trainable gating layer. However, many recent works [53, 2, 30] have shown that most Transformers, when trained at scale, build intrinsically sparse and modular representations. Zhang et al. [52] proposed to leverage this naturally emerging modularity with MoEfication -a method that converts dense transformer models into MoE models by grouping FFN weights into experts and subsequently learning small routers that determine which experts to activate. Models converted with MoEfication are able to preserve the performance of the original dense models while using only a fraction of their computational cost. However, we believe that the MoEfication procedure is not optimal, and therefore aim to obtain dense-to-sparse conversion schemes that obtain a better cost-performance trade-off. Intuitively, a MoE converted from a sparser base model would be able to perform the original function using a smaller number of experts. To validate this hypothesis, we perform MoEfication on different variants of GPT2-base 4 with varying activation sparsity levels and show the results in Figure 2a . As expected, MoEfication performs better with sparser models. We further investigate the per-token mean and the variance of non-zero neurons in the base and sparsified model, and show the results in Figure 2b . Observe that different layers use a different number of neurons on average. Moreover, the variance of the number of activated neurons is quite high and becomes even more significant in the sparsified model. This means that static top-k gating as used in MoEfication is not optimal for dense-to-MoE converted models, and a more flexible expert assignment rule that would be able to handle the high per-token and per-layer variance could be beneficial to the efficiency of such models, as illustrated at Figure 2c . Such dynamic-k gating requires routers that reliably predict the contribution of each expert. We observe that routers obtained through MoEfication do not accurately reflect this contribution. Moreover, their router training procedure depends on the strict sparsity of the model guaranteed by the ReLU activation function. Therefore, we design a novel router training scheme that directly predicts the contribution of each expert and generalizes to the broader family of activation functions. We combine the proposed components (sparsity enforcement, expert contribution routing, and dynamic-k gating) into a single method that we call Dense to Dynamic-k Mixture-of-Experts (D2DMoE), which we describe in detail in the next Section.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: We include a separate "Limitations" section.Guidelines:• 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. 
method: 
 D2DMoE reduces the computational cost of the model by splitting every MLP module into a MoE layer. In this section, we describe all of its components in detail. A high-level overview of the entire procedure is presented in Figure 1 . The conversion process can be optionally preceded by MHA projection layer replacement (Sec. 3.5), which allows us to apply the same transformation pipeline on all replacement modules.We expect that enforcing higher levels of activation sparsity may allow for the execution of an even smaller number of experts, resulting in overall computational savings. To this end, we induce activation sparsity by fine-tuning the model with an additional loss term that induces activation sparsity [11] . We apply the square Hoyer regularization [22, 17] on the activations of the model:EQUATIONwhere a l is the activation vector from the middle layer of the l-th MLP for input x, and L is the total number of MLPs in the model. Overall, the model is trained with the following cost function:EQUATIONwhere L CE is cross-entropy loss, and α is the hyperparameter that controls the strength of sparsity enforcement. We find that the pre-trained models recover the original performance with only a fraction of the original training budget (eg. 1B tokens for GPT2-base or Gemma-2B, which is less than 1% of the tokens used for pretraining).We split the two-layer MLP modules into experts using the parameter clustering method proposed by Zhang et al. [52] . Assuming the MLP layers are composed of weights W 1 , W 2 and corresponding biases b 1 , b 2 , we treat the weights of each neuron from W 1 as features and feed them into the balanced k-means algorithm [26] that groups neurons with similar weights together. Then, we use the resulting cluster indices to split the first linear layer W 1 , the first bias vector b 1 , and the second linear layer W 2 into n experts of the same size. The second bias b 2 is not affected by this procedure.MoEfication process was designed for standard two-layered MLPs [52] . Recent LLMs [45, 42] have shifted towards gated FFNs, where the activation is realized through a Gated Linear Unit (GLU) [37] , which contains an additional weight matrix for the gate projections. To adapt the expert clustering procedure described above to gated FFN layers, we cluster the weights of the gating matrix W g instead of W 1 , and use the obtained indices to divide the weights of the two other layers. We provide more intuition and details on our method for gated FFNs in Appendix G.In a standard MoE-based model, the gating networks are trained in an end-to-end manner. Contrary to this, we train each gating network independently. We propose to frame the problem of training the router as a regression task and directly predict the ℓ 2 -norm of the output of each expert with the router. Formally, given an input token z, we train D2DMoE router R to minimize the following loss:EQUATIONwhere E i is the i-th expert. We use a small two-layer neural network as the router R and apply an absolute value activation function to ensure non-negative output. This regression-based formulation is still compatible with the commonly used top-k expert selection, but enables more precise attribution of the contribution of each expert, as we show later in the experimental section.Note that Zhang et al. [52] also trains each routing network independently, but their method constructs artificial labels for each input, and then subsequently trains the router as a classifier. We discuss the differences in detail in Appendix A.Commonly used MoE layers always execute top-k experts for each token, where k is a predefined hyperparameter. This means that, regardless of the difficulty of the input, the model spends the same amount of compute on each batch [54] or token [38] . While this may be appropriate if the model is trained with the same restriction, it is suboptimal for a model that was converted from a dense model, as we show in Section 2.Since our router directly predicts the ℓ 2 -norm of the output of each expert, we propose a dynamic-k expert selection method that skips experts for whom the router predicts relatively small output norms. Given a router output vector R(z), we select a hyperparameter τ ∈ [0, 1] and define the expert selection rule G for the i-th element as:EQUATIONNote that as τ increases, the number of executed experts and the overall computational cost decrease.We emphasize that after model deployment τ can be adjusted without retraining. Figure 4 : FLOPs-performance tradeoff comparison of our method and MoEfication [52] on CV and NLP benchmarks. We also include early-exit (ZTW, [49] ) and token dropping baselines (A-ViT, [51] ) for classification. Our method outperforms these baselines across multiple computational budgets.A significant amount of computing in deep neural networks is often spent on dense layers that are not followed by any activation function. Dense-to-sparse-MoE conversion methods cannot reduce the costs of such layers due to a lack of activation sparsity. This determines an upper bound on the possible computational savings. To overcome it, we substitute dense layers with small MLPs with approximately the same computational cost and number of parameters. Each MLP is trained to imitate the output of the original dense layer given the same input by minimizing the mean squared error between the two (akin to a distillation loss).Figure 3 : Multi-Head Attention projection conversion scheme.In our case, for Transformer architectures, we substitute projection matrices along with their biases in every MHA layer, as depicted in Figure 3 . This means that the final model has four MoE layers in the MHA layer and one MoE layer in the FFN layer (either plain or gated) for each Transformer block. Note that we do not modify the computation of the scaled dot-product attention itself and this scheme can be applied to any standalone dense layer.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [NA]Answer: [Yes] Justification: We provide the source code along with a README. Code is easy to run once the environment is set up, and generates the plots automatically. Guidelines: • 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). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: We discuss the impact of our work, which we believe is positive. 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). 
experiments: 
 To analyze the impact of our method, we evaluate its performance on language modeling, text classification, and image classification. We obtain performance versus computational cost characteristics for each method by evaluating the methods with different inference hyperparameters (either τ described in Section 3.4 for D2DMoE or number of experts k for MoEfication; we mark them on the plots with dot markers). We report the computational cost of each method in FLOPs, as it is a device-independent metric that has been shown to correlate well with latency [27] . In addition, we measure the wall-clock execution time of an efficient implementation of our method.For MoEfication, we follow the procedure described by Zhang et al. [52] by converting the activation functions of the pre-trained model to ReLU and then fine-tuning the model. In the case of D2DMoE, we also replace activation functions with ReLU, except for Section 5.4, where we demonstrate that our method performs well also with GELU. To provide a fair comparison, the total training data budget is always the same between different methods. See Appendix J for a detailed description of our setup. The source code for our experiments is available at: https://github.com/bartwojcik/D2DMoE.Vision Transfomer [7] is one of the most popular architectures in computer vision. Since our method applies to any Transformer model, we evaluate it on the popular ImageNet-1k [35] dataset. We use a pre-trained ViT-B checkpoint as the base model and compare D2DMoE with MoEfication in terms of the computational cost versus accuracy trade-off. For broader comparison, we also evaluate the state-of-the-art early-exit method Zero-time Waste (ZTW) [49] , as well as our re-implementation of A-ViT, an efficient token dropping method proposed by Yin et al. [51] . Our results, presented in Figure 4a , demonstrate the significant gains from applying our method over MoEfication. We evaluate our method with BERT-base [6] on the CARER dataset [36] that contains text samples categorized into 6 different emotion categories. We compare the accuracy versus FLOPs trade-off for D2DMoE, MoEfication, and ZTW. We show the results in Figure 4b . The performance of MoEfication gradually deteriorates and completely collapses when the number of executed experts approaches zero. In comparison, D2DMoE maintains the original performance for a wide range of computational budgets, and the performance drop starts at a significantly lower budget. While early-exiting performs well for the lowest budgets, it obtains worse results than D2DMoE at medium budgets and suffers from a gradual performance decline.We evaluate our method on language modeling and compare it with MoEfication using GPT-2base [31] and Gemma-2B [42] . We initialize GPT-2 models from a publicly available OpenAI checkpoint pre-trained on a closed-source WebText dataset and use OpenWebText [12] in all of our experiments. For Gemma-2B, we also start from the publicly available pretrained model and evaluate its language capabilities on the C4 dataset [32] after finetuning. For both models, we use around 1B tokens for the finetuning phase (less than 1% of the cost of original pretraining) and 8-16M tokens for router training. We report the results in this section without the MHA projection replacement, as this task is highly sensitive to changes in attention layers, leading to noticeable loss degradation. For more training details, see Appendix J.3We present test losses for D2DMoE and MoEfication at different compute budgets for GPT-2-base and Gemma-2B in Figures 4c and 4d respectively. Our method outperforms the baseline at every computational budget. The loss of D2DMoE plateaus for higher budget levels, while the baseline displays consistently worse results whenever we lower the computational budget. Notably, for the larger Gemma-2B model our method performs well for most compute budgets, while the performance of MoEfication collapses. The failure of MoEfication can be explained by the emergence of massive activations in large models [40] , which makes it unable to learn reliable routing, as we analyze in more detail in Appendix E.We also provide a downstream evaluation of our Gemma models on the BoolQ dataset. We take the base model, which achieves 68.40% zero-shot evaluation accuracy, and convert it to MoE with D2DMoE and MoEfication. In Table 1 , we report the relative accuracy of the models at different compute budgets. Our method largely retains the performance across multiple compute budgets, while the performance of MoEfication decreases significantly. This shows that the loss-vs-FLOPs results for D2DMoE and MoEfication directly translate to downstream performance on language tasks. For any model acceleration method to be practically useful, it must reduce end-to-end inference execution time on modern GPU hardware. To achieve this, we implement the forward pass of our MoE-layer in the Triton intermediate language [43] , and employ several optimizations for our implementation, including an efficient memory access pattern, kernel fusion, and configuration auto-tuning. As suggested by Tan et al. [41] , our implementation also avoids unnecessary copies when grouping tokens.We verify the performance of our implementation for a single D2DMoE layer (24 experts with expert dimensionality 128) layer in isolation by comparing it with the corresponding MLP module (inner dimensionality 3072) on an NVIDIA A100 GPU. We fill a tensor of size [256 × 197 × 768] (batch size, sequence length, and hidden dimension, respectively) filled with Gaussian noise and use it as input to both modules. The gating network of D2DMoE is included in measurements, but the decisions are overridden with samples from a Bernoulli distribution, and we control how many experts are executed on average by changing the Bernoulli probability. The results, presented in Figure 5 , show that our implementation scales linearly with the number of executed experts, and has negligible overhead. Our method can be almost three times as fast as standard MLP while preserving 99% of the original accuracy.In Appendix C we provide additional wall-clock measurement results along with a more detailed description of our implementation. To accelerate inference D2DMoE leverages input-dependent activation sparsity, a property inherent to almost every Transformer model. However, interaction between D2DMoE and other popular network acceleration techniques, such as pruning [16] or quantization [13, 28] , is unclear. We evaluate D2DMoE in combination with such techniques to demonstrate their complementarity.First, we evaluate D2DMoE applied on top of networks pruned with CoFi, a structured pruning technique introduced by Xia et al. [50] . CoFi removes redundant neurons, attention heads, and sublayers to achieve the desired sparsity ratio, and then subsequently fine-tunes the reduced network. We first prune the base model with CoFi to the desired sparsity level, apply D2DMoE to it, and then evaluate both models on QNLI [48] . In Figure 6 , we show that D2DMoE successfully accelerates inference even on networks pruned to high sparsity levels.In Figure 7 , we also investigate the applicability of D2DMoE to quantized models using dynamic post-training quantization from PyTorch5 on BERT trained on the CARER dataset. Our method is robust to 8-and 16-bit quantization and exhibits only slight variations in performance after quantization. As FLOPs do not take bit width into account, we show quantized models in the same FLOPs range as the original model. In Appendix C, we also present wall-clock time measurements for quantized D2DMoE. 
analysis: 
 In this section, we present in detail additional experiments that provide insights into the performance of our method. Additionally, in Appendix E we analyze the performance of MoEfication with Gemma, in Appendix F we provide the results of router architecture analysis, in Appendix H we conduct experiments corresponding to the ones in Section 5.5 with GELU function, and in Appendix I we show additional visualizations for expert activation patterns.The dynamic-k rule introduces variability in the allocation of the computational budget along the model depth. To explore its scale, we investigate the distribution of the number of executed experts, with and without the activation sparsification phase. In Figure 8a , we show the histograms of the number of activated experts for each FFN layer of the BERT-base model trained on the CARER dataset (additional results are available in the appendix in Appendix I). As expected, the model with enforced activation sparsity requires fewer experts for a given threshold. Both base and sparsified models exhibit significant variance in the number of activated neurons across different layers, which justifies the dynamic-k selection and indicates that computational adaptability mechanisms are crucial for efficient inference in Transformer-based models. D2DMoE also allows the model to allocate different computational resources to various layers. We expect the model to allocate more compute to tokens containing information relevant to the current task. Since each token position in a ViT model corresponds to a separate and non-overlapping part of the input image, we can easily plot a heatmap to indicate the regions of the image where the model spends its computational budget. In Figure 8b we present such an analysis for our converted ViT-B model. As expected, the dynamic-k routing enables the model to minimize the computational effort spent on regions that contain insignificant information.Since our method consists of several steps, the positive impact of each one of them may not be evident. To show the significance of every component, we perform an ablation study by incrementally adding each component to the baseline method. We take a BERT-base model and evaluate the ablated variants in the same setting as described in Section 4.2. The results of this experiment are presented in Figure 9a . As expected, each ablated version of the method improves upon the previous one.The sparsity enforcement phase leads to enhanced performance compared to plain MoEfication. Alternative router training objective and dynamic-k expert assignment further improve the results, but -as the method only operates on the FFN layer -the computational cost cannot go below the cost of the remaining part of the model. Extending D2DMoE to MHA projection layers allows our method to reduce the computational cost further, and the resulting full method retains the accuracy of the original model at around twice fewer FLOPs than MoEfication.To justify our proposed activation sparsity phase, we investigate the impact of the activation sparsity of the base dense model on the performance MoE obtained with our method. We conduct a study similar to the one presented in Figure 2a : we train multiple base models with different activation sparsity enforcement loss weights α and convert them to Mixture-of-Experts models with our method.The results, shown in Figure 9b , highlight the positive correlation between the activation sparsity and the performance of the converted MoE, as higher sparsity in the base model always translates to better performance for D2DMoE. This is consistent with results previously observed for MoEfication. However, our method achieves better results for every base model in all cases, proving that regression routing and dynamic-k selection better utilize the induced sparsity.Activation sparsity works focus their analysis on networks with ReLU activation, as other functions (such as GELU or SiLU) do not guarantee exact sparsity. When analyzing non-ReLU models, such works require fine-tuning with the activation function changed to ReLU (relufication) [52, 27] , which limits their practical applicability. We hypothesize that relufication is not necessary and the models with many near-zero activations effectively function similarly to standard ReLU-based models. To evaluate this hypothesis, we extend the sparsity enforcement scheme to the commonly used GELU activation by penalizing the model for pre-activation values larger than a certain threshold. We first transform the pre-activation values as z ′ = max(0, z -d), where z is the pre-activation value and d is a displacement hyperparameter. Then, we apply the loss from Equation (1) on z ′ . This transformation penalizes only pre-activation values larger than d, and as a result, the model learns to produce values that effectively become negligible post-activation. We empirically find that d = -10 works well with GELU as the output below this value is near zero.To validate our hypothesis, we follow the methodology from Section 4.3 and we train ReLU-and GELU-based GPT-2 with and without sparsity enforcement loss. We show the results in Figure 9c . D2DMoE with a sparsified GELU-based model performs similarly to a sparsified ReLU-based model, while the performance of the non-sparsified GELU-based variant collapses. Within ReLU-based models, the sparsification still enhances the performance of D2DMoE, but the improvements are less drastic, and the behavior of our method does not significantly change as in the case of GELU. This shows sufficient activation sparsity enforcement relieves the model from the dependence on ReLU.A crucial hyperparameter in D2DMoE is the selection of expert size. Smaller experts may allow a more granular selection of executed neurons, likely resulting in a lower computational cost. However, decreasing the expert size increases the number of experts, which translates to a larger router, potentially negating any computational gains. To study the impact of this hyperparameter on our method, we evaluate D2DMoE on GPT-2 with different expert sizes, and show the results in Figure 9d .We observe that our method generally performs better with smaller experts. Those results differ from the ones presented in [52] , where the expert size is significantly higher. The positive correlation between granularity and performance can be explained by the increased levels of activation sparsity in our model, which requires significantly fewer activated neurons (experts). As expected, the performance decreases for the extreme choice of expert size equal to 1 due to significantly higher routing costs. We include additional results for expert granularity with GELU activation in Appendix H. 
related work: 
 Mixture-of-Experts. MoE layers were introduced as an efficient way to further increase the capacity of deep neural networks applied for NLP tasks, initially in LSTM models [38] , and later in Transformers [23] . Since then, they have also been applied to computer vision [33, 5] . MoE layers have gained significant popularity primarily due to their excellent scaling properties [8, 3] . Nonetheless, training such models is challenging, primarily because gating decisions must be discrete to ensure sparse expert selection. Various methods of training were proposed, some of which include reinforcement learning [1] , weighting the expert output by the probability to allow computation of the gradient of the router [38] , or using the Sinkhorn algorithm [3] . Some of those approaches also suffer from the possibility of load imbalance and therefore require auxiliary losses or alternative expert selection methods [9, 54] . Interestingly, in many cases fixed routing functions perform similarly to trainable routers [34] , which suggests that current solutions are largely suboptimal. MoE models can also be derived from pre-trained dense models by splitting the model weights into experts and independently training the routers for each layer [52, 57] , which avoids most of the problems present in end-to-end training.Activation sparsity in Transformers. Li et al. [24] show that ReLU-based Transformer models produce sparse activations in their intermediate representations, an effect that is prevalent across architectures, layers, and modalities. They propose a simple rule for keeping only top-k activations in each MLP layer, which results in a model with comparable performance. Similarly, Mirzadeh et al. [27] demonstrate that ReLU activation function in LLMs encourages ensuing activation sparsity that can be leveraged to skip redundant computations. Tuli and Jha [46] take a step further and design a dedicated Transformer architecture accelerator that also exploits activation sparsity, while Liu et al. [25] proposes to predict activation sparsity structure in LLMs and reduce the model latency by skipping redundant computations. Jaszczur et al. [18] demonstrate that it is possible to train Transformer models from scratch with a fixed level of activation sparsity and obtain similar performance. Finally, a related line of works focuses on sparsity in the attention distributions instead of intermediate representations [4] . None of the above-mentioned methods explore induced activation sparsity as a way to increase computational gains, nor do they address variance of the number of sparse activations on a per-token basis. 
conclusion: 
 We introduce Dense to Dynamic-k Mixture-of-Experts (D2DMoE), a novel approach that induces activation sparsity to improve the efficiency of Transformer-based models by converting their layers to Mixture-of-Experts (MoE). We demonstrate the interplay between the activation sparsity of dense models and the efficiency of converted MoEs. Moreover, we introduce regression-based router training and dynamic-k routing, which enable our method to efficiently utilize the induced sparsity. Finally, we show how dense-to-sparse-MoE conversion approaches can be extended to MHA projections and gated MLPs. Our approach is compatible with the existing Transformer architectures and significantly improves upon existing MoE conversion schemes. Our findings contribute to the ongoing efforts to make Transformer models more efficient and accessible for a wider range of applications, especially in resource-constrained environments. 
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 new data or new models.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.12. 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 always cite or credit the original authors. We always check the license of the assets we use.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. 
 
 Appendix: While D2DMoE displays promising results in reducing the computational cost of inference in Transformer models, a few limitations should be acknowledged. Our proposed sparsity enforcement and router training phases require additional training time. This overhead, while small, must be considered when evaluating the benefits of our approach. Moreover, we demonstrate improved performance over existing approaches on common NLP and CV tasks, but the scope of our experiments is restricted due to limited access to computational resources. Further research is needed to explore its applicability to extremely large models.Our work focuses primarily on fundamental machine learning research and we do not see any specific risks or ethical issues associated with our method. Nevertheless, we recognize the potential for misuse of machine learning technology and advocate for responsible AI practices to mitigate such risks.Our router training procedure is similar to the one proposed in MoEfication [52] , but the source code of the method provided by the authors6 contains a different routing objective than the one reported in the paper. While the paper describes their router training objective as a prediction of the sum of ReLU activation values in each expert, the source code uses prediction labels created from the sum of the activations in the intermediate layer divided by the maximum value in the entire batch and minimize the binary cross-entropy loss. Assuming that a k,j is the activation vector in the hidden layer of expert j for sample k, the label generation for their router can be expressed as:EQUATIONIn comparison to their approach, the router training procedure in D2DMoE differs in multiple aspects:• Our router considers the output of each expert instead of looking at the activations in the intermediate layers. • Instead of using artificially created labels based on the sums of activation values, we predict the ℓ 2 -norm of the output. This has the additional benefit that our router can work with alternative activation functions. • Our router is trained with the mean-squared error instead of the binary cross-entropy loss.The output of our router is constrained to positive values, while the MoEfication router is constrained to outputs in [0, 1].We find that the above differences are responsible for the improved performance of our router (see Figure 9a ).To compare the efficiency of our method with a standard MLP layer in Transformer, we estimate FLOPs in both modules. We assume the layer is composed of two linear transformations, with input and output size d m and hidden dimension ed m , where e is the expansion factor, which is usually equal to 4 in standard Transformer models. We skip the negligible cost of the biases and activation functions for simplicity.One can estimate the cost of the MLP layer in FLOPs, C MLP , as:EQUATIONFor dynamic-k expert selection with n total experts and k experts selected for a given input, the cost of the forward pass is composed of the cost of a forward pass through k experts and the cost of the 2-layer router with input dimension d m , hidden dimension d h and output dimension n. The cost of the single expert pass can be expressed as:EQUATIONand the routing cost can be estimated as:EQUATIONTherefore, the full cost of dynamic-k C dynk can be estimated as:EQUATIONand the cost of our method compared to the cost of standard MLP can be expressed as: For different sizes of router, dynamic-k uses fewer FLOPs than standard MLP as long as the total number of experts is sufficiently large and the number of selected experts is not equal to the total number of experts. For the clarity of presentation, we plot discrete values of k and n as continuous.EQUATIONAs long as the number of selected experts k does not approach the total number of experts n and the hidden dimension of the router does not approach the size of hidden dimension d m , the ratio is significantly below one.Assuming the worst case for second term (n = ed m ), we can estimate the cost ratio as:EQUATIONwhich shows that dynamic-k expert selection only exceeds the FLOPs cost of the standard network when the dynamic-k rule selects almost all experts or the number of experts becomes very high. For an even more detailed analysis, we refer to Figure 10 where the cost ratio between our method and standard MLP is shown, assuming different router sizes and e = 4 as standard for most Transformer models. In practice, we use d h = 128, so in all our experiments d m = 6d h . In Listing 1 we present the pseudocode for our efficient implementation of the forward pass of the D2DMoE module. We skip the pseudocode of the kernel of the second layer as it is similar, but provide the full source code in our code repository. Note that our implementation has multiple points where it could be improved for further performance gains: 1) metadata that is required for the kernels could also be computed with a dedicated kernel to reduce overhead; 2) atomic operations are currently used in the second layer to merge the results from different experts, an alternative implementation that does not use atomic operations could be faster; 3) it could be rewritten in CUDA to make use of dynamic parallelism. We leave those improvements for future work.In the main paper, we have presented wall-clock time measurements of a single D2DMoE layer. Below, we also ensure that our implementation works and performs well when used for the ViT-B model in which each FFN is replaced with a D2DMoE module. In Figure 11 , we measure the averaged processing time and the accuracy of our model. We perform the experiments on an NVIDIA A100 GPU using a batch size of 256. Each point on the x-axis corresponds to a single τ threshold and shows the wall-clock time of processing a single input averaged over the entire ImageNet-1k test set. Dynamic inference with D2DMoE offers up to 30% reduction in processing time without affecting the accuracy.To show that D2DMoE also reduces the execution latency of quantized models, we modify our kernels to handle float16 and int8 data types. In Table 2 we perform a similar experiment to the one from Figure 5 . We sample gating decisions from the Bernoulli distribution with probability p and measure the execution time of our experts for the three data type variants. The results show that both the higher activation sparsity (lower p) of our method and lower-precision data types are complementary in terms of wall-clock time reduction. While we see a smaller improvement from using int8 over float16 on A100, we attribute this to differences between GPU architectures and software support for low-precision arithmetic. In Section 4.5 we have demonstrated that our method is compatible with two popular model compression methods: quantization and pruning. A natural question is whether our method can be effectively applied to models compressed via knowledge distillation. Since distilled models also exhibit activation sparsity that our method relies on, D2DMoE should be applicable to such models. In Figure 12 we demonstrate the results of D2DMoE when applied on a ViT-S model, which has been trained via knowledge distillation [15] with the torchvision ViT-B being used as the teacher model. We see that D2DMoE is also able to reduce the cost of this smaller model.E Routing analysis for large models As presented in Figure 4d , in comparison to other considered benchmarks MoEfication visibly underperforms on language modeling with Gemma-2B. We attribute this to the emergence of massive activations in LLMs that reach a specific scale [40] .Massive activations are outliers along certain feature dimensions whose magnitudes are thousands of times larger than the magnitudes of other activations. The training objective of MoEfication described in Equation ( 5) uses maximum activation over the entire batch to normalize the target label for each expert. Upon encountering large outlier values, those labels become effectively meaningless, as the values for most of the experts become very close to zero. In this case, the router effectively learns to output zero labels for most of the experts aside from the ones corresponding to the outlier values.In comparison to MoEfication, our router training scheme does not make use of such normalization, and should therefore be robust to the emergence of massive activations. To validate this, we apply MoEfication on Gemma-2B, but with our regression routing instead of the original router training strategy. We compare the resulting model with vanilla MoEfication in Figure 13 and notice that replacing the routing scheme is enough for the model to learn effective expert assignment, as even though the expert choice is static and the base model is not sparsified, the cost-loss trade-off has significantly improved. This simple experiment shows that our regression routing objective is more robust than MoEfication when scaling to larger models.In comparison to standard linear routers used in MoE models trained from scratch, routers in MoEfication are 2-layer MLPs. To obtain the best performance with D2DMoE, we compare the linear design with MLPs with different hidden sizes for BERT-base and GPT-2-base on Figures 14a and 14b respectively. Linear routers do not perform well with our method, and overall a 2-layer MLP with a hidden dimension of 128 results in the best performance for both models. Note how for BERT-base, the accuracy curve for a model with the hidden dimension of 128 is slightly worse than for smaller routers, but for harder task with GPT-2 a more complex router is required. Following this analysis, we use 2-layer MLP with a hidden dimension of 128 for most of our experiments in the paper, with the only exception being the larger Gemma-2B model where we scale the hidden dimension accordingly to 512 to match the increase in model dimensionality. To provide better intuition behind the extension of our method to GLU-based gated MLPs mentioned in Section 3.2, we visualize the differences between standard FFN and Gated FFN and the application of our method in Figure 15 . Standard Transformer MLP realizes the following function:EQUATIONwhere W 1 , W 2 are the weights for the upscale and downscale projections 7 and A stands for the activation function. In comparison, gated MLP can be written down as:EQUATIONwhere W g is the weight for the added gate projection.The intuition behind MoEfication, which our method also follows for standard FFNs, is that the sparsity of the intermediate, post-activation representations determines the sparsity of the output representation. Therefore, the expert split is performed based on the weights of the upscale projection, as zeroed neurons in the upscale activations will also result in zeroed outputs of the downscale projection. When extending D2DMoE to Gated MLPs, our intuition is that the gating projections determine the sparsity of all the later representations, as both upscale and downscale are multiplied with the gating values. Therefore, we propose to build the experts through clustering performed on the gating weights W g and use the indices obtained through expert split on gating weights to construct experts from W 1 and W 2 . Following similar reasoning, for GLU-based models, we also perform activation sparsity enforcement on the gating projections instead of upscale projections as described originally in Section 3.1.In addition to experiments in Section 5.5, we present the results of similar ablation carried on the sparsified GPT-2 model with GELU activation. The results, presented in Figure 14c , follow the same pattern as before, which supports our claim that the sparsification enables the GELU-based models to function similarly to ReLU-based ones.Following the analysis for MoE-converted FFN layers in Section 5.1, we present full results for FFN in Figure 16 , and investigate the activation patterns in MHA projections modified with our method in Figures 17 to 20. The projection modules display lower levels of sparsity than FFNs, which is to be expected as our projection layers have lower intermediate dimensionality. Expert selection distribution patterns in Q and K show significant similarity, and the patterns in V and output projections are also similar to a lesser degree. The variance of the number of selected experts in MHA projections is higher than in FFN layers, but it still exists and the distribution in some of the layers seems to be bimodal, which provides further justification for the dynamic-k selection rule.In this Section, we describe the technical details used in the D2DMoE conversion procedure. For full reproducibility, we share the source code that we used for conducting the experiments. All experiments were performed using the PyTorch library [29] on the NVIDIA A100 and V100 GPUs on internal clusters. We utilize the fvcore library to count model FLOPs8 .All methods start with the same pre-trained ViT-B from the torchvision9 library and are trained on ImageNet-1k using the augmentation proposed by Touvron et al. [44] . We use mixup (0.8), cutmix, label smoothing (0.1), gradient clipping (1.0) and the Adam optimizer with a cosine learning rate schedule without warm-up. For D2DMoE, we replace the MHA projections and train the replacements for 3 epochs with the initial learning rate 0.001 and batch size 128, and then finetune the model for 90 epochs with sparsity enforcement weight α = 0.2, initial learning rate 2 • 10 -5 and batch size 512. We then convert the modules into MoE layers, and train the gating networks for 7 epochs with the initial learning rate set to 0.001 and batch size 128. We train ZTW for 100 epochs in total, allocating 5 epochs for ensemble training, while keeping the rest of the original hyperparameters unchanged. For MoEfication, we first convert the pre-trained model to ReLU-based one and finetune for 90 epochs with an initial learning rate of 0.0001 and batch size 256. We then split the weights and train the routers for 10 epochs with the initial learning rate 0.001 and batch size 256.All experiments start from the same pre-trained BERT-base checkpoint. For methods requiring ReLU activation function, we replaced GELU with ReLU and continue model pretraining on concatenated wikipedia [10] and books [55] corpora for 5000 steps on 8 GPUs using main setup from https://github.com/huggingface/transformers/blob/main/examples/ pytorch/language-modeling/run_mlm.py, per device batch size 96 and learning rate 5 • 10 -4 .For MHA projections replacement we use the same corpus and train replaced MLP modules on a single GPU with batch size 128 and learning rate 0.001 for 3000 steps. We finetuned base dense models on CARER dataset for 5 epochs with 2 • 10 -5 learning rate. For sparsity enforcement in D2DMoE we use α linearly increasing from zero to 0.0001 over training. For both MoEfication and D2DMoE we train routers with batch size 64 and initial learning rate 0.001 for 5 epochs. In all experiments, we use Adam optimizer with linear learning rate decay. For MoEfication we use expert size 32, for D2DMoE we use more granular expert size equals 6. For ZTW we trained ICs for 5 epochs with batch size 32 and learning rate 0.01.We base our code and hyperparameters for GPT2-base on the nanoGPT repository provided at https: //github.com/karpathy/nanoGPT. We initialize the model from https://huggingface.co/ openai-community/gpt2. In all pretraining experiments, we initialize models from a publicly available OpenAI checkpoint pre-trained on a closed-source WebText dataset and finetune for the fixed number of 1000 steps with the effective batch size equal to the value in the repository through gradient accumulation. The alpha values for sparsity enforcement can be found at Figure 9b . We train the routers for D2DMoE and MoEfication for 2000 steps using one GPU and tuning the learning rates for a given expert size from the range between 0.002 -0.005. For router training, we use Adam optimizer and cosine warmup scheduler.For Gemma-2B, we start from the checkpoint at https://huggingface.co/google/gemma-2b.We also finetune the model for 1k steps with an effective batch size of 1024, sequence length of 1024 and Adam optimizer with a learning rate of 1e-4. As Gemma's hidden dimension is much larger than the other considered models, we change the hidden dimensionality of the routers to 512 for both our method and MoEfication, but keep the other hyperparameters the same as in the rest of the experiments. For MoEfication, Gemma, we use 512 experts to obtain an expert size comparable to the one in their paper. For our method, we use 2048 experts. In D2DMoE, we set sparsity enforcement weight to 0.00003. We train the routers for 500 steps with Adam and effective batch size of 16 and use a learning rate of 0.001.We report the results for language modeling without the MHA projection replacement step, as we find that it is especially sensitive to changes in the attention layers, which always result in visible loss degradation.Filip integrated the codebase and ran the experiments for GPT-2 and Gemma, performed the activation sparsity analysis, and all the analyses on language modeling models. He contributed to the design of dynamic-k gating and played a primary role in designing the experiments and writing the article.Bartosz set the research direction of the project and proposed the alternative routing scheme, dynamick expert selection, the additional activation sparsity enforcement phase for ReLU and GELU, and the replacement of MHA projection layers. He wrote the shared codebase for the experiments, carried out the ViT-B experiments, implemented the custom Triton kernels for the efficient implementation of the method, and also played a primary role in the writing and editing of the article.Mikołaj made this paper possible by performing all of the experiments at the initial stages of the project and implementing MoEfication and numerous variants of our method. He carried out the BERT experiments, performed weight sparsity compatibility analysis, the ablation study, and contributed to the crafting of the paper.Simone significantly improved the paper's readability and provided invaluable advice for revising it.# compute the necessary metadata Justification: We put a lot of effort to accurately present our contributions.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.We provide the experimental details for this analysis in Section 4.3 and Appendix J.https://pytorch.org/tutorials/recipes/recipes/dynamic_quantization.htmlMoEfication source code for router training is publicly available at: https://github.com/thunlp/ MoEfication/blob/c50bb850307a36f8a0add6123f56ba309a156d13/moefication/utils.py# L188-L260We omit biases for simplicity.https://github.com/facebookresearch/fvcorehttps://pytorch.org/vision/stable/models.html 


Decision: The submission proposes a dynamic-k routing mechanism to select the number of experts in a MoE setting.  The submission initially received split reviews, which improved during the rebuttal process.  All but one reviewer now recommend acceptance.  The remaining reviewer has concerns about novelty, noting other papers that have used a dynamic-k routing mechanism in some way.  The submission also provides an interesting practical implementation and benchmarking results that appear to be of value in themselves.  The authors are strongly recommended to include citations to and discussion about the references pointed out by Reviewer 85BP - this will strengthen the correctness and relevance of the paper.

Reviews: review: paper_summary: paper_summary: The paper introduces a method called D2DMoE aimed at enhancing the efficiency of transformer models. D2DMoE implements a dynamic-k routing mechanism that allows the model to select a variable number of experts based on the input. The method leverages the inherent activation sparsity in transformer models to reduce the number of active parameters during inference, leading to significant computational savings—up to 60%—without compromising performance. The approach demonstrates that by converting dense layers into Mixture-of-Experts (MoE) layers, transformer models can achieve better accuracy-sparsity trade-offs, making them more efficient for various NLP and vision tasks.
strength_weakness: strength_weakness: In my view, this paper appears to be primarily a repackaged sparse-activated pruning technique stemming from the MoE concept. Several concerns arise:

1. Limited comparison with alternative sparsification methods: The paper predominantly contrasts with MoEfication, neglecting a thorough analysis against other compression or sparsification strategies beyond early-exit techniques. Consequently, the comparative experiments presented lack depth and fail to be fully convincing.

2. Lack of practical acceleration results: The paper only presents theoretical reductions in FLOPs and parameter counts. Without actual inference acceleration results on real hardware (e.g., V100, H100, GTX-4090Ti), it's impossible to assess the practical benefits of the method. The additional overhead from the gating mechanism could potentially negate some of the theoretical gains.

3. Questionable novelty of Dynamic-k gating: The proposed expert selection based on ℓ2-norm of output is indeed more reminiscent of pruning techniques than traditional MoE approaches. This calls into question the novelty of the method when viewed in the context of pruning literature, and makes the comparisons with MoE methods potentially unfair or irrelevant.

4. Limited novelty and inadequate discussion of related work: Many of the proposed operations bear similarities to existing sparse activation pruning methods. The paper fails to adequately discuss these connections, instead focusing on less relevant work. This omission of crucial related work in pruning literature significantly undermines the claimed novelty of the approach.

5. Lack of discussion of some SOTA sparse methods: such as Pruner-Zero: Evolving Symbolic Pruning Metric From Scratch for Large Language Models. ICML2024.

6. Theoretical foundation: The paper lacks a strong theoretical justification for why the proposed methods work better, which becomes even more critical given the concerns about its relationship to existing pruning techniques.
questions: questions: See weaknesses.
limitations: limitations: NA
paper_summary: paper_summary: This paper improves over MoEfication with four innovations: (i) enforcing higher activation sparsity; (ii) directly predict the norm of the output of each expert; (iii) dynamic-k expert selection scheme; and (iv) generalization to any standalone linear layer. The resulting method achieves significant improvements in terms of cost-vs-performance trade-offs in text classification, image classification, and language modeling.
strength_weakness: strength_weakness: - The baseline seems weak. ZTW only receives a few citations and is not published at a top machine learning conference/journal.

- Only FLOPs are reported and it is a poor indicator of practical latency/throughput, especially in the LLM era.

- It is not clear if the proposed method could be combined with other methods to show its broader applicability. In particular, quantization is a very general method and is almost adopted as a default in many practical applications.

- The writing could be improved. For example, more technical details of MoEfication should be included.

- In all figures, the lines are continuous. However, these lines should have been a result of interpolation. The original point values are missing.
questions: questions: NA
limitations: limitations: Please see weakness.
paper_summary: paper_summary: This paper proposes a method to convert a dense pre-trained transformer network into its sparse counterpart. The approach begins by fine-tuning a pre-trained network to sparsify its activation values. Subsequently, the method clusters the neurons in the MLP layers to form distinct experts and introduces a router network. Finally, the router network is fine-tuned to predict the norm of the output of each expert. Additionally, this work explores the possibility of modifying attention layers and clustering MLP layers using the GLU activation function. The results demonstrate that the proposed method is robust to high sparsity and outperforms baseline model
strength_weakness: strength_weakness: 1. This method introduces several hyperparameters during training and inference, which could potentially make it difficult to reproduce and deploy.
2. I find that Figure 2 does not justify dynamic-k very well, as most layers have a similar percentage of non-zero activations. However, this concern is alleviated after reviewing Figures 12 through 16.
3. Although this work sparsifies at a coarse granularity, which is beneficial for reducing latency, reporting wall-time would be more convincing than solely reporting FLOPs.
questions: questions: After converting the attention weight matrices, is any router network introduced to determine which head is activated? In the main paper, I only read that the attention matrices are replaced with two-layer networks with comparable parameters and cost. I didn't find any rationale provided for this conversion.
limitations: limitations: 1. Currently, this method is only evaluated on small networks, so it is unknown whether the proposed method can be effectively applied to larger networks.
2. Although this method leverages sparsity in activations to accelerate inference, it does not reduce the number of parameters, and may even increase them, unlike other methods such as pruning that simultaneously reduce the parameter count.
paper_summary: paper_summary: The paper presents Dense to Dynamic-k Mixture-of-Experts (D2DMoE), a method to convert a dense model into a MoE model, that exploits the fact that activations in Transformer models are typically very sparse. The sparsity can be further improved by the addition of square Hoyer regularization during a light fine-tuning phase. Since the sparsity pattern changes over examples, the use of a dynamic selection rule is preferred. The strategy used to convert the dense model into a MLP is almost identical to that of Zhang et al. from 2022, but the loss function used to train the router is different. The router in this paper is trained to predict the norm of the output of the corresponding expert (and the experts with the an expected l2 norms over a certain dynamic threshold are selected). Finally, the present work also MoEfies not only the FFN layers as usual, but also all the linear layers in the MHA (for that, each linear layer is converted into a MLP of the same cost, minimizing the l2 error, similar to distillation). 
The proposed approach is evaluated on different architectures (ViT, BERT, GPT2, Gemma) and benchmarks (including both vision and language). In all cases the proposed approach offeres a better FLOP-quality trade-off than that of Zhang et al. (2022), and achieves a quality similar to the original dense baseline but at significantly fewer FLOPs.
strength_weakness: strength_weakness: - The use of MoEs from scratch is criticized because the training difficulties that these have. However, the training recipes for these have improved significantly in the past two or three years, and are now a component of many state-of-the-art models (e.g. [Mixtral](https://mistral.ai/news/mixtral-of-experts/), [DeepSeekMoE](https://arxiv.org/abs/2401.06066), [Gemini 1.5](https://storage.googleapis.com/deepmind-media/gemini/gemini_v1_5_report.pdf)). Thus, I believe that these should definitely be considered as a baseline, otherwise the potential impact of this work is more limited.
- The proposed method involves many steps. They are well documented (see strengths), but makes the whole approach a bit cumbersome, in my opinion. Alternativelly, I wonder how the dynamic routing, the sparsity-indicuing regularization and other of the components used would perform when training a model from scratch.
- As usual, one must be carefull when doing only FLOP vs quality comparisons, since the proposed method, especially the dynamic selection of experts, may be hard to efficiently implement on modern hardware (GPUs, TPUs). I couldn't find this discussed anywhere in the paper, so I would appreciate more clarity on this (perhaps even additional plots showing Runtime vs quality, even if it's only on the appendix).
questions: questions: See the comments mentioned in the weaknesses.
limitations: limitations: No particular ethical considerations affecting this work, in my opinion.
paper_summary: paper_summary: This paper proposes a more effective dynamic-k expert selection rule that adjusts the number of executed experts on a per-token basis to reduce the high computational cost of the transformer models. They claim their D2DMoE model outperforms existing approaches.
strength_weakness: strength_weakness: 1. The most important contribution is the dynamic-k routing algorithms. However, the main work is a small modification to MoEfication. The difference of the proposed method with MoE-based methods is limited.
2. The experiments are not conducted well. This paper lacks a comprehensive comparison with other methods. There exist many MoE-based VIT methods, such as [1, 2]. Furthermore, why does the authors does not provide any performance number in the table? This paper claims they outperforms existing approaches on common NLP and vision tasks, however, only two methods are compared.
3. The major advantage of D2DMoE compared with MoEfication is saving inference cost. However, the saved cost is quite limited observed from fig. 4.

[1] Mod-squad: Designing mixtures of experts as modular multi-task learners
[2] M3 vit: Mixture-of-experts vision transformer for efficient multi-task learning with model-accelerator co-design
questions: questions: 1. Training Overhead: Can you provide more details on the additional training time required for sparsity enforcement and router training? How does this overhead compare to the overall training time of the original dense models?
2. Real-world Deployment: What steps or considerations are necessary for deploying D2DMoE models in real-world applications, particularly in terms of latency, hardware compatibility, and integration with existing pipelines?
3. Adversarial Robustness: Have you evaluated the robustness of D2DMoE models against adversarial attacks? If not, do you have any plans to explore this aspect in future work?
4. Comparison with Pruning Techniques: How does D2DMoE compare with other model pruning or compression techniques in terms of performance, computational savings, and implementation complexity?
5. Scalability: Can the proposed method be effectively scaled to extremely large models like GPT-3? Are there any specific challenges or considerations when applying D2DMoE to such models?
limitations: limitations: The paper does not sufficiently address potential challenges in real-world deployment, such as latency variations, hardware compatibility, and integration with existing systems.

Comments: title: Further clarification on differences between D2DMoE and multi-task methods provided by the reviewer
Response by Reviewer
About dynamic inference method
comment: We thank all the reviewers for the discussion. 

We have addressed each reviewer's concerns in detail, including the additional experiments and clarifications provided in our initial rebuttal regarding FLOPs vs latency, the A-ViT baseline, and quantization. During the discussion period, we also presented new results with quantization for reviewer uQnr. We also added downstream evaluations for reviewer 85BP, despite receiving their feedback just a day before the end of the discussion period.

Our general impression was that the reviewers’ main concerns revolved around the latency of our method and its compatibility with other model acceleration techniques. We have provided an efficient GPU kernel for our method, which we consider a significant technical contribution. We demonstrated that our method provides a ​​63% wall-clock time reduction for FFN modules and around 30% reduction when evaluated end-to-end. We believe that we have adequately addressed the above-mentioned issues through the discussion. 

We are pleased to note that the three reviewers who recommended acceptance remain positive about our work after the discussion and were satisfied with the responses we provided. We would like to express our gratitude to all the reviewers for their engagement during the discussion period and constructive feedback that helped us improve our paper.
To further emphasize the key differences between multi-task methods referenced in [7,8] and our work, we outline them below:

* Our goal is to reduce inference cost by skipping the redundant computation during the model inference through the MoE framework. In contrast, the referenced works employ MoE primarily to solve the multi-task problem, with **efficiency being only a secondary concern**.

* Both referenced methods use a classic formulation of MoE adapted for the multi-task setting (e.g. [7] optimizes mutual information between experts and tasks). **None of these works mentions activation sparsity or anything similar to our router training scheme**.

* The cited works do always spend the **same amount of compute for each sample** (classic formulation of MoE), with the differences being present only between tasks for [7]. **In contrast, our dynamic-k gating adjusts the computational effort of the model on a per-sample and per-token basis**.

* Our method allows for the **adjustment of computational cost after training** and during deployment through the change of the $tau$ hyperparameter. In contrast, [7,8] must retrain the model each time to target a different budget. While this may be not initially obvious from the FLOPs vs accuracy plots in ours and referenced works, we consider this a vital property of dynamic inference methods.

* Finally, **we target scenarios with limited computational resources** for the adaptation of pre-trained model, while the references provided by the reviewer require a **significant compute budget to train models from scratch** (e.g. [7] reports training for 80 hours with 240 NVIDIA V100 GPUs).
We thank the reviewer for the feedback and the discussion.
We thank the reviewer for the response. We understand the reviewer's concerns with distillation as a baseline, but - similarly to pruning - we see it more as a complementary model acceleration technique [1]. However, due to time constraints, we cannot provide any experiments that could demonstrate it to the reviewer and will only be able to include them in the camera-ready version. 

Additionally, we would like to point out that the CoFi method [2] evaluated in combination with our method in Section 5.6 does use distillation as a part of its training scheme. Our method still allows for computational saving when applied on top of CoFi. Input-dependent activation sparsity is an inherent property in almost all models, and therefore our method should be universally applicable. 

Finally, while we appreciate the thoroughness of the reviewer, to the best of our knowledge so far no papers have examined the compatibility of dynamic inference with all three major model compression methods—pruning, quantization, and distillation. In this context, our comparisons with pruning and quantization are already particularly comprehensive.

We thank the reviewer for the discussion and his suggestions that helped us improve the paper.

#### References

[1] Han, Yizeng, et al. "Dynamic neural networks: A survey." IEEE Transactions on Pattern Analysis and Machine Intelligence 44.11 (2021): 7436-7456.

[2] Xia, Mengzhou, Zexuan Zhong, and Danqi Chen. "Structured pruning learns compact and accurate models." arXiv preprint arXiv:2204.00408 (2022).
We appreciate the reviewer's response, active participation in the discussion, and support for our work. Their feedback has been valuable in helping us improve our paper.
#### References
[1] Hendrycks, Dan, and Kevin Gimpel. "Gaussian error linear units (gelus)." arXiv preprint arXiv:1606.08415 (2016).

[2] Zhang, Hongyi, et al. "mixup: Beyond empirical risk minimization." arXiv preprint arXiv:1710.09412 (2017).

[3] Srivastava, Nitish, et al. "Dropout: a simple way to prevent neural networks from overfitting." The journal of machine learning research 15.1 (2014): 1929-1958.

[4] Zhou, Yanqi, et al. "Mixture-of-experts with expert choice routing." Advances in Neural Information Processing Systems 35 (2022): 7103-7114.

[5] Kaplan, Jared, et al. "Scaling laws for neural language models." arXiv preprint arXiv:2001.08361 (2020).

[6] Clark, Aidan, et al. "Unified scaling laws for routed language models." International conference on machine learning. PMLR, 2022.

[7] Chen, Zitian, et al. "Mod-squad: Designing mixtures of experts as modular multi-task learners." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

[8] Fan, Zhiwen, et al. "M^3vit: Mixture-of-experts vision transformer for efficient multi-task learning with model-accelerator co-design." Advances in Neural Information Processing Systems 35 (2022): 28441-28457.
We thank the reviewer for the response. While the brevity of their comment suggests that they may have already formed a definitive opinion about our work, we believe it is important to clarify our view to other participants. Apart from addressing the raised concerns, we also want to highlight how they have been presented. In particular, we find it challenging to understand why the **reviewer's response took so long, given that it is brief, somewhat vague, and largely reiterates previously stated points**. We point out other issues in our responses below. **We respectfully ask other participants to carefully consider both the reviewer's comments and our discussion to arrive at their own conclusions regarding the validity of the raised concerns.** 
> I still think the dynamic-k routing is a slight modification of MOE's top-k routering.

**The reviewer seems to disregard our discussion, where we clearly state that dynamic-k is not our only contribution.** We also emphasize that reviewer uQnr praised the novelty of our router training scheme.

We have provided a clear justification for the introduction of dynamic-k. This modification is evidently not that obvious given that dynamic-k was not adopted in earlier works. We demonstrate that static top-k gating is inappropriate for MoEs converted from dense models (Figure 2b) and that dynamic-k **enables consistent improvements** (Figure 4, Figure 5a). 
We also highlight that many influential ideas, such as the GELU activation function [1], mixup augmentation [2], dropout [3], or expert choice routing for MoEs [4] could be seen as __slight modifications__. Straightforward, yet well-motivated methods that provide consistent improvements are crucial for advancing the field, and simplicity often facilitates broader adoption. 

> Furthuremore, from fig. 4, we cannot observe the performance directly. Using loss value to represent the performance is not a good choice. 

**The reviewer initially described the cost savings shown in Fig. 4 as minimal. However, following our rebuttal, the reviewer apparently no longer seems to contest the improvement shown in the figure.** Instead, they raise a new concern shortly before the conclusion of the discussion period.

We emphasize that the **reviewer’s statement is inaccurate, as we do report accuracy for classification tasks**. We use test loss only for LLM evaluation, as done in multiple other works. In particular, a significant portion of the **industry and research community heavily invests in research on scaling laws [5, 6]. These studies focus on improving model loss**, which indicates that it is widely considered as an important indicator of model performance.

To further support our claim, **we provide a downstream evaluation of our Gemma models on the BoolQ dataset**. We take the base model, which achieves 68.40% accuracy, and convert it to MoE with D2DMoE and MoEfication and perform zero-shot evaluation. In the table below, we report **relative accuracy** of the models at different compute budgets. Similarly to the loss in Figure 4d, **our method largely retains the performance** across multiple compute budgets, while the **performance of MoEfication decreases significantly**. 

| Compute budget | 100%    | 90%    | 80%    | 70%    | 60%    | 50%    | 25%    | 10%    |
|----------------|---------|--------|--------|--------|--------|--------|--------|--------|
| D2DMoE         | 100.00% | 99.68% | 99.37% | 98.69% | 97.60% | 94.34% | 92.75% | 90.89% |
| MoEfication    | 100.00% | 92.24% | 92.19% | 92.15% | 88.79% | 75.40% | 86.70% | 77.53% |

> Although [1,2] are multi-task methods, they all provide pretrained models. The authors should attempt to compare them because the proposed method is quite similar to [1,2].

We still consider the fact that the referenced works are multi-task learning solutions as a fundamental issue that renders any fair comparison impossible. The reviewer's response and his original review **fail to clarify how the referenced works are similar or relevant to our work**, with the only superficial connection being that they also use the MoE architecture. [7] targets Taskonomy and PASCAL-Context datasets, and [8] targets NYUD-v2 and PASCAL-Context datasets (all of them exclusively multi-task datasets). Neither of these works includes evaluations on baselines that are even remotely similar to ours, nor does our method target multi-task learning. Even though [7,8] might provide pre-trained models, comparison of models trained on different datasets **would mostly reflect performance differences stemming from the use of different training data** instead of those inherent to the method. 

We include the references in a separate response.
Thank the authors for providing further clarification. 

I have read the rebuttal and other reviews. I decide to maintain my original score of 6.
Thank you for conducting the additional experiments. I don’t see any major concerns that would result in a rejection of this paper. However, after reading the comments from the other reviewers, I think there is still room for improvement in experiments. As a result, I decide to maintain my current score.
Thank you very much for addressing my concern, as well as those from other reviewers.

I am still reluctant to increase my score, since I think that the method is quite cumbersome (as I mentioned) given the goal and what it achieves. For instance, the authors have mentioned during the rebuttal that the comparison with MoEs trained from-scratch is not fair because the goal of this work is "to reduce the inference cost of existing pre-trained models". If so, a really imporant and very popular strategy is missing as a baseline: distillation to a smaller model (dense or MoE).
I thank the reivewer's time and effort to give me a detailed response. However, I still think the dynamic-k routing is a slight modification of MOE's top-k routering.  Furthuremore, from fig. 4, we cannot observe the performance directly. Using loss value to represent the performance is not a good choice.  Although [1,2] are multi-task methods, they all provide pretrained models. The authors should attempt to compare them because the proposed method is quite similar to [1,2].
We thank the reviewer for the comments. We understand we have addressed all the other concerns.

>There are also more advanced approaches to dynamic inference, such as dynamic token sparsity, etc. This work seems to have no advantage against these approaches

We refer the reviewer to Figure 3 of the rebuttal PDF, where we add A-ViT [1], a dynamic token sparsity method, as a baseline for the response to reviewer 85BP. The results show that we outperform A-ViT.

>The authors highlight that their work is a dynamic inference approach, so are there any comparisons with previous works on dynamic pruning?

According to a recent survey [2], at least two groups of works use the term “dynamic pruning”.

Dynamic pruning approaches such as [3, 4, 5, 6] differ from static pruning because they allow the network structure to change by pruning or reactivating connections during training. The resulting model is static and therefore should complement dynamic inference methods such as ours (as shown in Section 5.6).

A small subset of dynamic inference works do use the “dynamic pruning” term for methods that dynamically select channels for execution in convolutional layers [7, 8, 9]. However, they all target exclusively CNNs, while our work focuses on Transformer models.

##### References
[1] Yin, Hongxu, et al. "A-vit: Adaptive tokens for efficient vision transformer." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

[2] He, Yang, and Lingao Xiao. "Structured pruning for deep convolutional neural networks: A survey." IEEE transactions on pattern analysis and machine intelligence (2023).
[3] Lin, Tao, et al. "Dynamic Model Pruning with Feedback." ICLR-International Conference on Learning Representations. 2020.

[4] Ruan, Xiaofeng, et al. "DPFPS: Dynamic and progressive filter pruning for compressing convolutional neural networks from scratch." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 35. No. 3. 2021.

[5] Lin, Shaohui, et al. "Accelerating Convolutional Networks via Global & Dynamic Filter Pruning." IJCAI. Vol. 2. No. 7. 2018.

[6] Chen, Zhiqiang, et al. "Dynamical channel pruning by conditional accuracy change for deep neural networks." IEEE transactions on neural networks and learning systems 32.2 (2020): 799-813.

[7] Elkerdawy, Sara, et al. "Fire together wire together: A dynamic pruning approach with self-supervised mask prediction." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.

[8] Gao, Xitong, et al. "Dynamic Channel Pruning: Feature Boosting and Suppression." International Conference on Learning Representations.

[9] Lin, Ji, et al. "Runtime neural pruning." Advances in neural information processing systems 30 (2017).
The authors highlight that their work is a dynamic inference approach, so are there any comparisons with previous works on dynamic pruning?

There are also more advanced approaches to dynamic inference, such as dynamic token sparsity, etc. This work seems to have no advantage against these approaches
We appreciate the reviewer’s involvement and are grateful for a very swift response.

> Although the baseline is from a top conference, it seems too old. (...)

The extension paper from November 2023 demonstrates that ZTW remains a SOTA early exit method and outperforms other baselines such as GPF [1] (ACL2021) and L2W [2] (ECCV2022). 

In the response to reviewer 85BP we have also added A-ViT [3] (CVPR2022), with the results available in Figure 3 of the rebuttal PDF. 

Overall, we compare our method against four baselines (ZTW, MoEfication[4], CoFi[5], A-ViT) and conduct experiments on four datasets from different modalities. This surpasses the evaluation scope of other established works on dynamic architectures published at top conferences, which often evaluate on fewer datasets and baselines or compare against static models instead ([3], CVPR2022; [4], ACL 2022; [5], CVPR2023; [6], EMNLP2023; [7], NeurIPS2021; [8], NeurIPS2022; [9] ICLR2023).

> (...) it's better to report both accuracy and latency, when combining your method with quantization.

To show that quantization reduces the latency of D2DMoE, we modify our kernels to handle `float16` and `int8` data types. We perform a similar experiment to the one from Figure 1 of the rebuttal PDF: 1) we sample gating decisions from the Bernoulli distribution with probability $p$ 2) measure the execution time of our experts for the three data type variants. We present the latency (ms) in the tables below:

### RTX 4090
| p  	 | 0.0 	 | 0.1 	 | 0.2 	 | 0.3 	 | 0.4 	 | 0.5 	 | 0.6 	 | 0.7 	 | 0.8 	 | 0.9 	 | 1.0 	 |
|---------|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|
| float32 | 0.005 | 0.009 | 0.013 | 0.018 | 0.023 | 0.028 | 0.033 | 0.038 | 0.042 | 0.047 | 0.052 |
| float16 | 0.004 | 0.005 | 0.007 | 0.009 | 0.011 | 0.014 | 0.016 | 0.018 | 0.021 | 0.024 | 0.027 |
| int8    | 0.004 | 0.004 | 0.005 | 0.007 | 0.008 | 0.009 | 0.010 | 0.011 | 0.012 | 0.013 | 0.014 |

### A100
| p  	 | 0.0 	 | 0.1 	 | 0.2 	 | 0.3 	 | 0.4 	 | 0.5 	 | 0.6 	 | 0.7 	 | 0.8 	 | 0.9 	 | 1.0 	 |
|---------|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|
| float32 | 0.006 | 0.009 | 0.012 | 0.015 | 0.019 | 0.022 | 0.025 | 0.028 | 0.031 | 0.035 | 0.038 |
| float16 | 0.006 | 0.007  | 0.008 | 0.010 | 0.011 | 0.013 | 0.014 | 0.016 | 0.017 | 0.019 | 0.021 |
| int8    | 0.007 | 0.008 | 0.009 | 0.010 | 0.011 | 0.012 | 0.014 | 0.015 | 0.016 | 0.017 | 0.019 |

The results show that both the higher activation sparsity (lower $p$) of our method and lower-precision data types are complementary in terms of wall-clock time reduction. While we see a smaller improvement from using `int8` over `float16` on A100, we attribute this to differences between GPU architectures and software support for low-precision arithmetic.

##### References
[1] Liao, Kaiyuan, et al. "A global past-future early exit method for accelerating inference of pre-trained language models." Proceedings of the 2021 conference of the north american chapter of the association for computational linguistics: Human language technologies. 2021.

[2] Han, Yizeng, et al. "Learning to weight samples for dynamic early-exiting networks." European conference on computer vision. Cham: Springer Nature Switzerland, 2022.

[3] Yin, Hongxu, et al. "A-vit: Adaptive tokens for efficient vision transformer." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

[4] Zhang, Zhengyan, et al. "MoEfication: Transformer Feed-forward Layers are Mixtures of Experts." Findings of the Association for Computational Linguistics: ACL 2022. 2022.

[5] Chen, Xuanyao, et al. "Sparsevit: Revisiting activation sparsity for efficient high-resolution vision transformer." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

[6] Tan, Shawn, et al. "Sparse Universal Transformer." Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing. 2023.

[7] Rao, Yongming, et al. "Dynamicvit: Efficient vision transformers with dynamic token sparsification." Advances in neural information processing systems 34 (2021): 13937-13949.

[8] Schuster, Tal, et al. "Confident adaptive language modeling." Advances in Neural Information Processing Systems 35 (2022): 17456-17472.

[9] Chataoui, Joud, and Mark Coates. "Jointly-Learned Exit and Inference for a Dynamic Neural Network." The Twelfth International Conference on Learning Representations. 2023.
Thanks for the reply and the additional experiments. Most of my concerns are addressed except two:

- Although the baseline is from a top conference, it seems too old. Additional experiments with newer baselines may not be necessary during this rebuttal period, but I highly recommend trying out more recent ones.

- If your method does not hurt accuracy, and since quantization usually does not hurt accuracy too much, I'm not surprised that when combining your method with quantization, the accuracy will be maintained. However, this may not be true for latency. Therefore, it's better to report both accuracy and latency, when combining your method with quantization.