Mixture of Nested Experts: Adaptive Processing of Visual Tokens

Authors: gagan  jain
◇⋆  nidhi hegde
aditya  kusupati
arsha  nagrani
shyamal  buch
prateek  jain
anurag  arnab
sujoy  paul
google  deepmind
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-12 21:35:59

Abstract: The visual medium (images and videos) naturally contains a large amount of information redundancy, thereby providing a great opportunity for leveraging efficiency in processing. While Vision Transformer (ViT) based models scale effectively to large data regimes, they fail to capitalize on this inherent redundancy, leading to higher computational costs. Mixture of Experts (MoE) networks demonstrate scalability while maintaining same inference-time costs, but they come with a larger parameter footprint. We present Mixture of Nested Experts (MoNE), which utilizes a nested structure for experts, wherein individual experts fall on an increasing compute-accuracy curve. Given a compute budget, MoNE learns to dynamically choose tokens in a priority order, and thus redundant tokens are processed through cheaper nested experts. Using this framework, we achieve equivalent performance as the baseline models, while reducing inference time compute by over two-fold. We validate our approach on standard image and video datasets - ImageNet-21K, Kinetics400, and Something-Something-v2. We further highlight MoNE's adaptability by showcasing its ability to maintain strong performance across different inference-time compute budgets on videos, using only a single trained model.

Main paper: 
 introduction: 
 Visual tokens, the fundamental building blocks of image and video representations, often exhibit strong inter-dependencies, spatially in images and spatio-temporally in videos. This offers a potential avenue for optimization in visual processing, as processing every token with equal emphasis may not be necessary for achieving optimal results. Traditional Vision Transformer (ViT) [20] and Video Vision Transformer (ViViT) [2] based models, however, process all tokens with equal emphasis, disregarding this inherent codependency and leading to unnecessary computational burden. This becomes a major bottleneck when deploying these models in real-world scenarios, where computational resources may be limited and real-time processing is required.To this end, conditional computation has become a promising line of research to increase the capacity of a network, while only conditionally activating a part of it during inference. Sparse Mixture of Experts (MoEs) was initially popularized for Natural Language Processing (NLP) [42, 22] ,but it has been gaining attention for furthering conditional computation ideas in vision [39, 1, 35, 50] as well. While MoEs bring in improved performance at a given inference cost, they also increase the overall parameter count, leading to increased storage requirements. Moreover, these works rely on experts that have the same parameter count and compute, limiting their ability to reduce computational costs without resorting to skipping tokens entirely. In this work, we devise the Mixture of Nested Experts (MoNE) framework, which provides a scalable approach to conditional computation, bringing in significant reductions at inference time, while working with the same parameter space as the baseline model. MoNE draws inspiration from nested architectures [47, 31, 53] , particularly MatFormer [19] , that learns multiple representations of the same data with varying levels of details, based on structured slices of the parameter space. MoNE employs these structured nested models as experts in the MoE framework (without increasing parameter count), and learns a network to route tokens to these experts. We explore various design choices and present an effective recipe for allocating compute to experts, assigning tokens to experts, and training the MoNE framework. For the assignment operation, we propose Expert Preferred Routing (EPR), a routing algorithm that greedily assigns tokens to experts under capacity constraints based on router predictions. Figure 1 shows token importance as perceived by MoNE. We propose the following three primary contributions:1. We introduce the novel Mixture of Nested Experts (MoNE) framework to dynamically allocate computational resources for Vision Transformer (ViT) based models.2. Given a fixed parameter count, MoNE offers the flexibility of learning networks at much lower FLOPs (∼ 2.3× on video datasets) and real-time latency (∼ 2×), while being quality neutral.3. Rigorous experiments show that MoNE works well for both image and video transformers, and visualizations depict that tokens routed to larger experts correlate well with regions of interest. 
related work: 
 Transformers [45] have become the de-facto architecture for processing data across multiple modalities spanning language [10, 36] , images [20, 17] , video [2, 49] and audio [23] and combinations thereof [38] . Consequently, there have been numerous efforts to improve the efficiency of transformers to make them more amenable for deployment in real-world applications [44] . These include approaches like efficient approximations of attention [12, 48] , local attention [32, 4, 13] and reducing the number of tokens in the transformer [40, 29, 8] among others. Our work focuses on conditional computation [5, 34] , observing that some input tokens are easier to process than others, and therefore require less computation during inference.Mixtures-of-Experts (MoE) transformers learn to route tokens to one of multiple expert MLPs [42, 22] . Although such models conditionally process input tokens, each expert has the same parameter-and FLOP-count, meaning that the total computation is constant for each input. More relevant to our approach, Mixture of Depths [37] extends the routing logic of MoE to conditionally skip an expert completely, thus total computation for each input varies dynamically. Completely skipping tokens being a hard unretrievable decision, our work chooses from an array of nested network, which effectively process information and help to stabilize training by getting rid of discontinuities.Nested architectures [47, 31, 53] on the other hand, learn hierarchical representations of the input, where the first k hidden dimensions encode the most relevant information. This allows to extract multiple models with varying inference compute from a single trained model, similar to 'Mix-n-Match' in [19] . However, these models do not process tokens adaptively. Our model, in contrast, consists of a learned router which dynamically routes tokens to experts of different hidden dimensions based on the given compute constraints. Therefore, instead of requiring the user to select the hidden dimensions of each transformer layer, our model only needs a single compute constraint input. Moreover, we show experimentally the superior accuracy-efficiency trade-offs achieved by our approach.We note that other conditional computation approaches include "early exiting" [46, 41, 21, 28] such that the processing of "easy inputs" terminates before passing through all layers of the transformer.In addition, the ACT [25] algorithm was proposed for recurrent neural networks, and uses a "ponder cost" to learn a "halting score" for when to stop processing a particular input. This has since been extended to recurrent transformers [15] , and also to each individual token in a transformer [52, 51] , thus adaptively determining which tokens in a transformer to process. In contrast, our approach does not drop tokens, rather processes them with smaller nested models. This allows us to retain most of the information, and hence dampen the effect of irrecoverable decisions. We experimentally verify that our adaptive approach offers strong compute-performance trade-offs. Flextron [11] is a concurrent work, which looks at elastic inference, specified by user latency needs, with a focus on language modeling. Unlike Flextron, MoNE is guaranteed to learn models bounded by the specified latency needs and is able to learn from a single training phase, without using a surrogate model. 
preliminaries: 
 Here, we discuss the concept of nested models, on which we build Mixture of Nested Experts (MoNE), followed by a discussion about Mixture of Nested Experts (MoE), and its differences from MoNE.For the purposes of this work, we use the Vision Transformer (ViT) [20] as an example of a full model, from which nested submodels can be derived. Inspired by MatFormer [19] , we define these submodels for every layer of the network, for both Self-Attention and MLP (see Appendix A.1). The key idea is that in a feature projection operation Wx, whereW = [W [∶ D m ] , W [ D m ∶] ],andW [∶ D m ]denotes "slicing" the first D m dimensions, we can extract a partial projectionW [∶ D m ] x [∶ D m ]. This can be done for any projection in the transformer, and we can extract smaller models from it. We refer to these as nested models, and D m as the nested model dimension. This is shown in Figure 2a . The Extract operation extracts the first D m features and applies the corresponding projection sub-matrix to it, while the Pad operation pads it back to full dimension D before residual connections and LayerNorm. While MatFormer applies the nested structure only to the hidden dimension of the MLP layer, in our approach we extend it to the in-and out-projections of both the Self-Attention (SA) and MLP layer. In the SA block, irrespective of the sub-model used in the in-projections, it is always projected to the model dimension D for the (QK T )V operation. The same thing is performed in MLP, where the hidden dimension is always 4D, as in ViT, irrespective of in/out-projection dimension.We extract E nested models with exponentially-spaced model dimensions. Therefore, for a typical value of E = 4, the model dimension for the nested models are D 8 , D 4 , D 2 , D . Note that while we build upon the idea of nested models from MatFormer, we do not share their training strategy which involves joint optimization through a weighted loss over these submodels. In contrast, we treat these nested models as distinct experts with varying compute requirements. The Mixture of Nested Experts (MoNE) framework (described in detail in Sec. 4.1) then dynamically routes input tokens to these nested experts based on their information content, with the idea that more informative tokens should be processed by larger (and thus more computationally expensive) nested models.A Mixture of Experts (MoE) layer in a transformer can be represented as MoE(x) = ∑ E i=1 g(x) i e i (x), where E is the number of experts, e i () are the expert models each having their own parameters, g ∶ R D → R E is the routing/gating function, which decides the experts which should process x. Note that g is sparse with only k << E non-zero terms. During inference, only those experts are active.MoE strictly increases the parameter count, but maintains the same inference FLOPs by setting k = 1. However, it still needs to process all tokens with the same pre-defined compute. In contrast, in MoNE, we do not extend the parameter count of the model, due to the nesting structure (see Sec. 3.1), and dynamically choose a nested expert during inference. Unlike in MoE, where all experts have the same capacity, in MoNE with k = 1 always, e i ⊂ e i+1 , which allows us to dynamically allocate compute.  
methodology: 
 In this section, we describe the details of our Mixture of Nested Experts (MoNE) framework for efficient inference. We assume a Vision Transformer (ViT) [20] based architecture for our approach, and then extend it to Video ViT (ViViT) [2] as well.Tokenization: In this paper, as our primary focus is images and videos, the model input is in R H×W ×3×T , where T = 1 for images and T > 1 for videos. After tokenization, the input to the transformer is X ∈ R D×N where N is the number of tokens, and D their model dimension. For images, we have N = H p h ⋅ W p w , and for video, N = T p t ⋅ H p h ⋅ W p w , where H, W, T are the input height, width and duration respectively. p h , p w and p t are the patch sizes along these respective dimensions. We use the ViT [20] and ViViT [2] architectures to tokenize images and videos respectively, obtaining a list of tokens X = {x i } N i=1 . MoNE Block: The Mixture of Nested Experts (MoNE) framework is a dynamic routing mechanism that processes visual tokens using nested models with varying computational capacities, instead of processing all tokens with the full model. A pictorial repsentation of the model is presented in Figure 2b . Let B l = {B l 1 , . . . , B l E } denote the nested blocks at a certain layer l with increasing parameter sizes, B l E (.) being the full model block. A router network decides the appropriate nested block to use for every token. Hence information from tokens of different model dimension interact with each other. This is enabled by performing self-attention at the full model dimension D as discussed before. For each token x i , a router produces a probability distribution over the E nested experts, r i = softmax(W r x i + b r ), where W r and b r denote the router weights and bias respectively. These router predictions are sent to an assignment algorithm, which assigns every token to a single appropriate nested expert. Based on the assignments, we update the features for the i th token in the l th layer as follows -EQUATIONwhere the j th nested expert is chosen by the Expert Preferred Router [EPR(.)] algorithm for the i th token as per Eq. 2:EQUATIONNote that the multiplication of the router predictions with the model output in Eq. 1 allows gradient propagation through the router weights. We also introduce a learnable parameter α ∈ [0, 1), initialized to 0, which ensures proper gradient flow during the initial training stages, specifically during finetuning from a pre-trained MatFormer model. Without scaling, a low initial router prediction would dampen the block output, whereas the initial multiplicative factor being 1 ensures a stable starting point.The feature of the last layer x L i is used for downstream applications. For classification tasks, we apply global average pooling on all the token features and apply a linear classifier layer to predict the categories.Within the MoNE framework, the routing strategy is crucial for achieving an optimal balance between performance and computational efficiency. Traditionally there are two primary routing strategiestoken choice [42] and expert choice [39] . In token-choice routing, the router predicts the probability distribution over the available experts, and picks the expert with the highest probability. However, this can suffer from load balancing issues, with most of the tokens being routed to one or few experts. Hence, inference time compute is only bounded by the compute of the full model. On the other hand, in expert choice routing, each expert selects the top-k tokens with the highest preference for that expert. This guarantees perfect bounds on computation. Potential conflicts due to token selection by multiple experts are resolved by prioritizing based on model size.Formally, we consider a given distribution of nested models applied to the tokens, represented as c = {c 1 , . . . , c E }, s.t., ∑ i c i = 1, which we call the capacity distribution over the nested models. The method for obtaining a suitable capacity distribution, given the inference time compute requirements, will be discussed in Sec. 4.3. Given router probabilities r i for N tokens across E experts, we employ an Expert Preferred Routing algorithm (Algorithm 1). This is a greedy assignment approach that gives higher preference to larger nested models, aiming to identify the most important tokens first. We begin by examining the router predictions for the biggest to the smallest model, assigning k j = ⌊c j N ⌋ of the remaining tokens to j th nested model. Any remaining tokens, arising from integer packing constraints, are assigned to the smallest model. Algorithm 1 presents the proposed Expert Preferred Routing (EPR) algorithm.Require: r ∈ R E×N (router predictions), c (capacity distribution, s.t., c T 1 = 1), Ensure: M ∈ {1, . . . , E} N (nested model index)1: M ← 1 N Default assignments to the smallest model 2: for j = E to 1 do 3:k j ← ⌊c j ⋅ N ⌋ 4: I ← Top-k-Index(r[j, . . . ], k i )Returns value and indices of Top-K 5:M [I] ← j 6:r[∶, I] ← 0 Null out assigned ones 7: end for 8: return MThe Expert Preferred Routing (EPR) as described in Section 4.2 needs the individual expert's capacity bounds c i to be specified. To get this, we define a metric called the effective capacity :e c = ∑ E i=1 cidi D,where d i = D 2 E-i is the model dimension of the i th nested model. Given a certain inference FLOP requirement, we can translate that to an equivalent effective capacity e c . Since every token gets processed through exactly one nested expert, this along with the given budget imposes two constraints on the unknown capacity distribution c. However, since the individual expert capacities vary loglinearly, multiple distributions c can lead to the same e c for E > 2 and it is non-trivial to choose one over the other. MoEs generally use auxilliary loss functions [39, 42] to promote equal usage of experts. But in MoNE, that would render a certain fixed capacity, missing out on the flexibility that the framework provides to function with any capacity (as depicted later in Figure 5b ). Hence, we invoke intuitive constraints to solve for c. Specifically, we incentivize the usage of larger models, while also adding an entropy term to ensure uniformity of capacities across experts. Given these constraints, we solve the following optimization problem:EQUATIONIn practice, we set (β, δ) to (10, 2) and use a Sequential Least SQuares Programming (SLSQP) optimizer to solve Eq. 3 for the capacity distribution c, which is then used by EPR (Algorithm 1) to get token to expert mappings. We empirically verify these choices in Section 6.MoNE can be seamlessly adapted for video-based tasks. In videos, there exists another dimension -time -which adds to the significant redundancy in the tokens. Given the large number of tokens that can be obtained from a video, the computational costs grow drastically. To tackle this problem, works in literature factorize computation along space and time [2, 6] , perform local windowed computation [33] , etc. MoNE being a token based approach, directly extends to video encoders.For video processing, we leverage the Factorized Encoder architecture of ViViT [2] . This architecture employs two distinct transformers: spatial and temporal. After tokenization, each temporal index yields a set of tokens representing information from local spatio-temporal neighborhoods. These spatial tokens interact within their temporal index for L s layers, culminating in a single global token per index. Subsequently, a temporal transformer processes these global tokens across L t layers. Given that the spatial transformer significantly dominates computational costs in this model, we integrate MoNE into the spatial component while maintaining full capacity for the temporal transformer. The router predicts expert assignments for all temporal frames independently, which are then consumed by the EPR(.) algorithm to produce frame-wise expert assignments. 
results: 
 In this section, we empirically evaluate MoNE on multiple datasets spanning images and videos for different model sizes, assess its adaptability to stringent FLOP constraints, and depict real-time latency gains achieved by MoNE during inference.We empirically evaluate MoNE on image and video classification. For image classification, we train the network with random initialization. As for video classification, we follow previous literature and start from a pre-trained MatViT [19] model due to the inherent nested structure required in MoNE. We follow the joint training strategy of MatViT, with separate losses an all model granularities. We implement MoNE on JAX [9] using BigVision [7] for image classification and Scenic [16] for video classification. We follow the AugReg [43] training strategy to train all our image classification models. For video classification tasks, we inherit all augmentations and hyperparameter values directly from the ViViT [2] paper.For all experiments in this section, we place a single router at the first transformer layer, and propagate the router decisions to all the layers. We also multiply the router predictions (Eqn 1) to all layers, which ensures differentiable paths through the router network in all layers and allows the more evolved features from later layers to influence router learning. We also perform analysis of router placement in Section 6.Baselines: We first compare with MatViT's nested models. As mentioned in the paper [19] , we perform joint training over all four nested models that we consider in this work -{ D 8 , D 4 , D 2 , D}. MatViT is equivalent to MoNE, with a deterministic router to pass all tokens to the same nested model. We show that adaptively mixing tokens with different model dimensions performs much better across datasets and tasks. We also compare with Mixture of Depths (MoD) [37] , which is also a token routing algorithm, but proposed for language tasks. MoD takes the extreme decision of either processing or skipping for every token in a layer. MoNE, on the other hand, makes fuzzy decisions to choose intermediate-sized models, instead of skipping, which helps to retain significant information at the expense of low compute. We adopt the best reported MoD configuration: processing 12.5% of tokens every other layer while processing all tokens in the remaining layers.We also emphasize that MoNE acts as a complementary framework to traditional MoEs like Sparse VMoE [39] , and inference-time optimization techniques like Token Merging (ToMe) [8] . We present an extended discussion and further results in Appendix A.2, comparing with other adaptive baselines and validating the compounded savings by applying ToMe on MoNE.Images: First, we evaluate MoNE on ImageNet-21k [18] classification using ViT. We experiment with S, B, and L models to showcase the efficacy of MoNE across model sizes. As ImageNet-21k can have multiple labels for an image, we report the commonly used precision@1 metric. Figure 3 shows the results for all the models on ImageNet-21k. MoNE performs much better than MatViT's nested models and MoD, specifically in the low FLOPs regimes. MoNE achieves comparable performance to baselines with around 2× reduction in FLOPs.Following the literature on language models [37, 27] , we experimented with isoFLOPs training, which involves training for the same number of FLOPs as the baseline models. Since MoNE models have fewer FLOPs compared to their ViT counterparts, they require more training epochs to achieve the same total training FLOPs. We conducted this experiment on the S/16 model (see Figure 3a ) and observed additional improvements in MoNE's performance, particularly for the lower FLOPs models.Videos: Since video models rely on heavy pre-training [2] , we first train a baseline model with nested structure on the benchmark datasets -Kinetics-400 [31] and Something-Something-v2 (SSv2) [24] . We use the ViViT Factorized Encoder B/16 model [2] for our experiments and consistently report the 8x1 test accuracy, averaging predictions over 8 temporal clips [2] . Figure 4 FLOPs for MoNE are 162.8 vs 376.3, with similar accuracy -64.6 vs 64.4). We always do isoFLOPs training while fine-tuning these models. We attribute the higher compute gains compared to images due to the greater (spatial and temporal) redundancy in videos, which MoNE exploits well.Inference time capacity adaptation: Capacity adaptation during inference is crucial, as the inference time budget is often dynamic, changing based on user needs. Ideally, a model should adjust with little to no retraining. To evaluate this ability, we test how MoNE, trained at a specific effective capacity (e c ) performs when evaluated at other capacities. Fig. 5 presents the results for image and video classification. We observe that the model adapts well to nearby capacities. However, as expected, its ability declines with extreme shifts in the capacity budget between train and eval. The performance degradation is steeper while adapting a model trained at high capacity to low capacity. We also note that the performance degrades more gracefully in videos than on images, presumably due to the larger temporal redundancy.To enhance model adaptability, we train a model with the capacity sampled uniformly at random from {0. Additionally, the variation of latency and throughput with FLOPs for varying model capacities of MoNE is depicted in Figure 6a to 6d. The plots show that latency and throughput gains scales linearly with FLOPs reductions. It is important to note that inference gains depend heavily on implementation and while a simple high-level efficient implementation of our framework yields gains of this scale, we believe that further improvements can be obtained by optimizing a low-level GPU kernel implementation for MoNE.In addition, it is worth noting that the proposed Expert Preferred Routing (EPR) in Algorithm 1 loops only over the number of experts, which is typically a small number and fixed to 4 in our framework. While the nature of the EPR algorithm does not allow parallelization of the computation any further, the time taken by the algorithm is negligibly small as compared to the total time taken by the model. For comparison, on a V100 GPU, the EPR algorithm adds just 0.5 ms to the forward propagation time of a ViT-B/16 model (190 ms), accounting for < 0.3% of the total computation time.  
router analysis: 
 In this section, we discuss, analyse and visualise the design choices in implementing the router network. We choose the SSv2 dataset for this analysis. We further provide an extended discussion in A.3, detailing on the choice of number of routers, associating router outputs with Task Difficulty, and understanding the implications of having of learnable router.Router Position: As discussed before, we use a single router at the first layer, and propagate its decisions for all layers. While a delayed router might benefit from a more processed feature representation as input, this also diminishes the compute gains, as the initial layers operate at full capacity. We reason this choice by monitoring performance while placing the router at different layers in the network. As Figure 7a suggests, the gains through richer features from the later layers is outweighed by the shift in the curve to the right, and an equivalent capacity with our default router produces higher points on the curve.We vary the number of routers, placing them at different regular intervals in the network in Figure 7b . The decision from one router is carried out until the next router block is encountered. We notice a clear downtrend in performance with increase in number of routers from being present in the first layer to being present in all layers. Intuitively, more routers demand learning more decisions, and the network has to adapt to these decisions, making optimization harder. Comparison with Random Router: We compare our learned router approach to a random router, which maps tokens to nested experts randomly, while still maintaining the capacity limits of each expert (c i ), as computed in Section 4.3. Results in Figure 7c suggests that with lower effective capacities, the random router performance degrades while the learned router still manages to understand relevant patterns from the input, thus upholding performance.Visualizing Important Tokens: The above claim is further backed by visualizing the token importance during inference at a low effective capacity (e c ). We highlight the tokens selected by the largest expert, i.e., the full model on a few images in Figure 8a . It can be easily observed that the tokens sent to the largest model correlate well with the regions of interest in the images. On videos (Figure 8b ) as well, the highlighted regions across temporal stamps consistently track the regions of motion.Capacity Allotment: Given a fixed input capacity e c , we demonstrate the superior performance of our heuristic-based allocation method (Section 4.3) compared to other approaches, as shown in Table 2 . While the Proportionate allocation (assigning capacity inversely proportional to expert compute cost) and Uniform allocation (assigning equal capacity to all experts) show promising results, they lack the flexibility to adapt to varying budgets. Additionally, greedy approaches, such as allocating the entire budget to the largest expert and dropping other tokens (MoD style), or a greedy approach where the largest expert is assigned capacity such that all the remaining tokens are routed through the smallest expert, exhibit inferior performance.  
conclusion: 
 In this work, we presented Mixture of Nested Experts (MoNE), a novel framework for adaptive processing of visual tokens by dynamically allocating computational resources to different tokens. Through a nested structure with shared parameters and the proposed expert-choice routing algorithm, MoNE achieves significant reductions in inference time (over two-fold) without sacrificing accuracy on benchmark image and video datasets. Future works can be centered around extending MoNE to denser tasks like object detection, captioning, etc.Limitations: Extending this to auto-regressive decoding in LLMs is non-trivial, as this is designed primarily with an encoder architecture in mind. We leave this further exploration for future work. Societal Impact: The MoNE framework dynamically allocates computational resources with a given budget, thereby significantly minimizing energy usage and carbon emissions during inference of vision models. MoNE can also play a role in democratization of AI, allowing broader access to trained models without the need for large resources. ). MoNE, in contrast, matches baseline performance with limited inference compute while working with the same parameter space.While generally MoE architectures are designed with the expectation of specialization of experts to certain tasks, this is not always the case. In Mixtral of Experts [30] , the authors do not observe patterns in the assignment of experts based on the topic. In Sparse VMoE [39] , the authors observe very weak correlation of router decisions to categories. In MoNE, overlap between experts allows the largest expert to utilize the full parameter space, meaning complete representation power as enjoyed by the equally-sized vanilla model. Additionally, as shown in Table 5 of [19] , joint optimization of shared experts leads to better performance than having independent experts of the same size.0.6 0.7 0.8 0.9 1.0 1.1 1.2 GFLOPs In Table 3 , we compare MoNE with other baselines, particularly with adaptive computation of dense models. We perform this experiment on ImageNet-1K with a Ti/16 sized model. ACT [25] , PonderNet [3] , DepthAdapt [21] , A-ViT [52] are works with similar motivation of input adaptivity as MoNE, and MoNE shows superior performance. Latency gains on bigger models e.g., ViT-B are even higher, as also observed in literature [8] .We also highlight that MoNE can be utilized as a baseline for further inferencetime optimizations for improving latency.To this end, we apply Token Merging (ToMe) [8] on top of the MoNE style ViT-Ti/16 model trained on ImageNet-1K. For this experiment, we train a model with full capacity till the third layer and then place a router that optimizes for latency for all subsequent layers. ToMe is applied only on the first 3 layers.For fair comparison, we compare the performance drop and quote the same from a ViT-Ti model from the ToMe paper in Figure 10 . Our preliminary results demonstrate that this implementation improves performance compared to ToMe on ViT, and this can be further extended to all MoNE layers, applying it to distinct sets of nested tokens, indicating that ToMe is complementary to MoNE. 
 
 Appendix: Number of Routers: It is important to note that the number of routers in MoNE do not have the same implications as in traditional MoE frameworks. In MoEs, the parameter count increases with the number of layers on which the expert router is placed, and hence we typically see performance gains. Even then, as depicted by Tables 2 and 8 in Sparse VMoE [39] , significantly increasing the parameter count with more routers only marginally improves performance. On the other hand, in MoNE, the parameter size remains fixed irrespective of the number of routers: the only change a router brings is re-assignment of tokens to nested experts while keeping the total compute per layer fixed. We hypothesize that increasing the number of routers leads to slight decrease in performance, as seen in Table 7a due to two reasons:• It brings in additional optimization challenges (also prevalent in the MoE literature [39] )• Reassignment of a token from smaller to larger nested expert limits its information content to the representation power of the smaller expert, therefore not improving performance. The converse case occurs while reassigning from bigger to smaller nested experts, thus losing information.Since MoNE allows flexibility in the placement of routers, an interesting future direction would be to extend MoNE to more challenging task settings, where a higher number of routers might lead to better results.Task Difficulty: To further analyse the decisions made by the MoNE router, we study the visual inputs from the ImageNet-21K dataset deemed most and least compute intensive with respect to the router logits. This analysis experiment is performed in a setting without capacity constraints, in order to understand if the router decisions correlate with task difficulty (i.e. harder to understand inputs are sent to larger experts). Therefore, instead of using the greedy EPR Algorithm 1, we take an argmax over the router logits to make decisions. The results presented in Fig. 11 depict two sets of images, the top-3 images that demand the lowest and the highest compute respectively, according to the router decisions. It can be intuitively observed that the images demanding less compute are visually simple, while the ones demanding highest compute are relatively complex. Learned vs Random Router: Figure 7c shows the performance of the model with a learned vs a random router at different capacities. While for higher capacities, the learned router performs marginally better than the random one, the gap significantly widens as we go to lower capacities, from 0.1% at e c = 0.6 to 1.3% at e c = 0.2. This makes sense: with ample capacity, many tokens can be heavily processed, reducing the need for smart routing. Conversely, in low-capacity scenarios, routing decisions become crucial as only a few tokens can utilize the heavy experts. Interestingly, ViTs inherently shuffle information [14] , potentially even in the "Random" router setting as well, acting as an intrinsic information router. We note that a model trained with a learned router when evaluated with a random router, performs significantly worse (∼ 6% drop in Top1 Accuracy on Ti/16 trained on ImageNet-1K). 


Decision: The paper presents the "Mixture of Nested Experts" (MoNE) concept. It permits to dynamically route tokens in a priority order and thus enables to reduce inference-time compute compared to classic MoE methods. Reviewers appreciated the motivation and results. However, reviewers also raised concerns regarding the choice of baselines (MoE is missing), comparability of results (some methods are fine-tuned and others are trained from scratch) and other minor points. While the rebuttal was able to address many points, some concerns remained (additional baselines like MoE, proper and fair comparison to all methods). AC thinks those remaining concerns are important and strongly encourages the authors to revise the paper to point out/acknowledge those concerns such that future readers are aware. AC trusts that the authors will do so and recommends to accept.

Reviews: review: paper_summary: paper_summary: This paper proposes the Mixture of Nested Experts (MoNE) framework. MoNE is built on top of the MatFormer architecture which utilizes a nested architecture where smaller, less computationally expensive sub-models are nested within larger models. Similar to MatFromer, MoNE uses structured slices of the model weights (the experts) to process information hierarchically. While MatFormer focuses on obtaining many models on the optimal loss-vs-compute curve using the mix'n'match algorithm (during inference), MoNE learns to dynamically route visual tokens to an appropriate expert based on their significance and the computational budget.

The results show that MoNE achieves a favorable accuracy-efficiency trade-off curve for image and video classification tasks compared to the baselines.
strength_weakness: strength_weakness: - Baseline comparisons:  It is conceivable that the dynamic routing approach can outperform the mix'n'match policy of MatFormer during inference as smaller sub-networks can focus on simpler/uninformative tokens while the larger experts focus on more complex/informative tokens. Nonetheless, to fully evaluate the potential of the proposed MoE framework, it would be beneficial to compare it with other MoE architectures, especially those like Sparse Vision-MoE.

- The advantage of having a nested set of experts compared to non-overlapping experts remains unclear. The qualitative results suggest that MoNE utilizes the full model for the most critical tokens, while less informative tokens are processed by smaller sub-networks. Given this, the paper could be strengthened by comparing MoNE to other dynamic token processing methods such as AdaViT [1], SVitt [2], and A-ViT [3], which have significantly enhanced computational efficiency in image and video recognition tasks by learning to dynamically skip tokens.

1. Meng, Lingchen, et al. "Adavit: Adaptive vision transformers for efficient image recognition." CVPR 2022.
2. Li, Yi, et al. "Svitt: Temporal learning of sparse video-text transformers." CVPR 2023.
3. Yin, Hongxu, et al. "A-vit: Adaptive tokens for efficient vision transformer." CVPR 2022.

- Specialization of Experts in MoE Architectures: A primary focus in designing MoE architectures is the specialization of experts to address different aspects of the data distribution. Typically, having non-overlapping diverse experts is desirable. In contrast, MoNE and MatFormer consist of partially overlapping parameters, suggesting that the implicit learning bias in these architectures may differ significantly from that in conventional MoEs. Investigating the differences between MoNE experts and conventional sparse MoE architectures could greatly enhance the reader's understanding of the proposed MoE framework's behavior.

- Worse results when increasing the number of routers: Standard MoE architectures employ MoE layers and routers in every(other) layer, enhancing the model's flexibility and expressiveness. However, it is concerning that increasing the number of routers leads to worse performance in MoNE. This suggests potential optimization challenges in this MoE model compared to standard MoEs. Notably, the best results are obtained when the router is placed at the very first layer, indicating that the dynamic decision-making for all layers is based on relatively simple cues at local feature levels. Further exploration into why the performance deteriorates with more routers compared to traditional MoEs, which typically see improved results, would be insightful.
questions: questions: Please review the weaknesses section.
limitations: limitations: The authors mention one of the limitation of their work regarding extension to sequence modeling tasks.
paper_summary: paper_summary: This paper introduces the concept of Mixture of Nested Experts (MoNE), which utilizes a nested architecture to process visual tokens more efficiently in visual media like images and videos. MoNE aims to leverage redundancy in data, choosing experts in a priority order to process visual tokens, thereby achieving substantial compute savings while maintaining performance. The approach, demonstrated with MoNE's algorithms like MoNE-21K and MoNE-4K, optimizes adaptive processing on standard image and video datasets, significantly reducing computational demands while using a single trained model.
strength_weakness: strength_weakness: 1. My main concerns are in the actual inference speed of your models. The experimental results reported in the paper focus on comparing FLOPs rather than throughput. Could the authors demonstrate some direct comparisons of training/inference speed?

2. The EPR algorithm (Algorithm 1) proposed in the paper seems to implement routing operations through some loops. Is this process actually implemented through loops or parallel computing? Does this operation take up a lot of inference time?

3. I'm also concerning whether the method can benefit dense prediction tasks such as image segmentation. As in those tasks there are less information redundency, will MoNE result in performance degradation?
questions: questions: See weaknesses.
limitations: limitations: Limitations are discussed in the paper.
paper_summary: paper_summary: This paper tried to use Matryoshka mechanism to assign tokens to different experts.
strength_weakness: strength_weakness: 1. Why there is not comparison with FF which owns a single scale? What I mean is, starting for MaViT, and continue finetune the model with the same amount of compute .
2. Why is imagenet accuracy so low compared to a normal ViT? Such as Figure 3.
3. Why is there no experiments on imagenet1k?
questions: questions: 1. In "we place a single router at the first transformer layer, and propagate the router decisions to all the layers.", why doing this? Why can't we learn from MoE papers, and place routers for all transformer layers?
2. How stable is router training? How do you evaluate whether router really can find a suitable compute to match the task difficulty for a specific image/video?
3. In figure 6b, do you mean that there is one router in the whole model? This is suprising. 
4. Seems like random router is only a bit worse than the learned one...
limitations: limitations: 1. Important Tokens: is there a way to measure it?
2. The method description part is not clear to me.
3. Authors should also consider comparing with TokenMerging series of work, since the goal is the same, which is to save compute from the token level.
paper_summary: paper_summary: This paper present a method to select nested portions of a transformer network, using a MoE router-expert assignment method where each expert is a progressively larger slicing of a single underlying model.  A capacity budget determines how many tokens can go to each expert, while a router network scores experts for each token so that the most important tokens (in the sense of benefitting from the larger computations) go to the larger model slices.  Furthermore, when trained using random budget selection, the model can be dynamically scaled to different cost-accuracy tradeoffs at inference time.  The model is evaluated on image and video classification tasks using imnet21k, K400 and SSv2, with excellent flops-accuracy performance compared to baselines, and a set of ablations show the impact of different design choices for router placement.
strength_weakness: strength_weakness: There are a few points that I think were under-explained (see questions below).  In particular the descriptions of the projections around the MLP and SA were a little terse, and I'm still not sure exactly which operations are in the subspaces vs the full dimensional space.

The connection between token redundancy and the operation of the model is also a little tenuous, though likely (see below as well), and while FLOPs are measured explicitly, the impact on both real computation savings and runtime may depend on hardware and distributed computation implementations, which I didn't see mentioned.
questions: questions: * sec 3.1  I don't understand exactly where the projection back out to 4*D or D happens, and if this was a linear projection or padding?  What operations are performed in the full-dimensional space, and which in the subspaces?  Right now it's unclear exactly where the computation savings are, as I'm not sure exactly which operations are in the smaller subspace.

* sec 3.1 l.102:  "it is always projected to the model dimension D for the (QK)V operation."  This could use more explanation on how the projection is done.  Is this linear or padding?  A more explicit equation for the projection from D/m to D could help.

* Which tensors and dimensions are the layernorms performed over?  It appears they are performed within each expert separately (which would make sense since they have different dimensions) but I'm not entirely sure.

* The motivation mentions redundancy in the tokens quite a bit, and while I agree with that in general, I think the link so far is tenuous and wasn't made very explicit.  In particular, if the router network is a linear classifier on the first transformer input, then how can it possibly know which tokens are redundant without comparing between them?  Does such a comparison come from lower conv layers in feature extraction?  Or is it just that some types of tokens tend to be more redundant than others (flat regions, for example, are just by nature of being flat and therefore next to something similar).

* eq 3:  the effective capacity e_c is defined as a linear combination of the dimension ratios, but d-dimensional MLP will use O(d^2) operations, so that a d/2 dimensional MLP will have 1/4 of the multiply-adds compared to a d dimensional one.  Does this matter for the capacity assignments?  Or is this not actually the case because the MLP always has one 4D hidden layer, so it is a linear scaling?

* also eq 3:  the overall intuition makes sense as described, but this particular optimization seems a little arbitrary; in particular, why are the weights delta^i exponential, and why is the entropy term needed in addition to the e_c capacity constraint?  It seems possible the entropy is counteracting the exponential delta weights, when this could have been less aggressively weighted with no entropy?  I don't know that there is necessarily a benefit to equal usage of all experts in this case, either (see my other question on runtime and distributed computation below)

* for each of the different FLOPs levels in the performance curves, what are the expert assignment ratios?

* Another interesting point of comparison would be the accuracy levels of each expert alone --- that is, the 4 single-expert routing assignments (1,0,0,0), (0,1,0,0) etc.

* an interesting extension would be a lowest-capacity expert with 0 in its dimension slice, so no MLP, but still included in self-attention as well.  this one is perhaps beyond the scope of this work but seems natural considering the redundant computation motivations

* something that wasn't touched on was the impact on model runtime and latency, particularly in distributed settings, although that will be hardware-dependent.  Since there are sync points between the experts at all SA layers, how to distribute computation between devices isn't as immediately clear as for many same-capacity experts, since the longest-running computation can be in any of the four expert capacities depending on the assignments

* sec 4.4 videos:  I agree that there is redundant temporal information and that would be a great application of MoNE.  However, the application in the second paragraph doesn't actually seem to leverage or implement that, since it says that it applies the MoNE frame-by-frame.  If that's the case, the temporal dimension wouldn't be used in MoNE to determine important vs redundant computations.  If time and space combinations interleave or alternate, then it could --- but the description indicates that isn't the case.  On the other hand, even without this just looking at a single frame, some regions are more likely to be redundant across frames than others, and so the argument for MoNE may still apply because of that.  Whatever the case, it appears contradictory in the text right now and could use a little more explanation.



Smaller comments:


* l.122:  "with k = 1" isn't very clear, I initially thought it meant "if k = 1" which prompted me to ask, but what if k > 1?  In fact, I think this means to say e_i is a subset of e_i+1 and k is always 1 (i.e. exactly one of the nesting levels is selected)

* meaning of i, j isn't always consistent, e.g. l.114 and 168 sums i over 1..E, but in most other places i is tokens and j is experts, so seeing i index the experts in a few places is a little confusing

* Alg 1 EPR:  Should "T" be "N" here?  Also, the way this is phrased with floor function, there could be a few left-over tokens at the end that are unassigned.  I take it these are assigned to the cheapest expert rather than being dropped?

* l.185:  "flexibility" --- If the flexibility here is dynamic and can be specified at inference time, then c would need to be changed at inference time as well --- does this mean that many different values for c are used a training time?  And how are they randomized or selected?   --- this was answered in adaptive training experiment later in sec 5, but could be good to mention in the method description
limitations: limitations: addressed in the final discussion section
paper_summary: paper_summary: The paper builds on the Mixture-of-Experts paradigm for vision transformer, and adds a hierarchical aspect to it, yielding a Mixture of Nested Experts (**MoNE**). More specifically, experts are defined as subsets of nested channels in the Feed Forward Networks, such that each expert has a different compute cost, and the larger one correspond to the full original model. Nevertheless, some operations are always performed at the full model dimension, after padding the tokens that went through smaller experts as needed: The $(QK^T)V$ operation in the self-attention, the layer norms and the residual connections. In addition, a dynamic budget allocation heuristic is proposed to adapted the classical (uniform) load balancing loss to the scenarios of experts with different capacities.
The proposed method is then evaluated on image and video classification.
strength_weakness: strength_weakness: * **Fixed number of experts:** It is not clear to me why `MoNE` uses a fixed number of nested blocks: It is very natural for `MoEs`, since experts are entirely separate blocks, however in the case of `MoNE`, the router could act directly on the channel and output a binary decision (essentially, having as many experts as dimension $D$). I would expect this to have an impact on the load balancing loss, but it would given much more flexibility to distribute the capacity across experts. More generally, there is a large literature on dynamic sparsity (e.g. *(1)* for recent references) which seems closely related to the design of nested experts and is not mentioned here.

* **Baselines:**  The experiment sections lacks baselines, in particular ones which also consider a for dynamic routing (e.g. there are none for the video classification task). For instance:
	  * Simple mixture of experts. The only baseline considered in the paper is MoD. Not only that the paper compared to *the best reported MoD configuration* (line 233); However MoD experiments seem to have been run on very different datasets, so it is not clear why the best configuration would transpose here.
	  * Token pruning; In particular for  image classification, there is a very large body of literature on token pruning (A-ViT, E-Vit, etc.). Since the only tasks considered in the experiments are classification, pruning tokens should be a valid dynamic computing baseline.
	In addition, some of the metrics reported in the paper could be improved:
	  * As far as I know, ImageNet-21k evaluation is not very standard as there is no official train-validation split.
	  * FLOPs are not enough to show the method's efficiency and they should be accompanied by real latency. For instance, they do not take into account extra padding operations.

* **Relevance of the dynamic budget allocation:** Table 1 shows that a simple uniform allocation for the load balancing loss performs as well as the heuristic proposed in Section 4.3. Therefore it is not clear to me whether the proposed dynamic budget allocation really has a significant impact



#### references
  * (1) Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time, Liu et al, 2023
questions: questions: * In Equation (1), is $r_i$ in fact computed on $x_i$ (the input of the block) or on $z_i$ (the input of the FFN / output of the self-attention) ?
* I assume the padding to the dimension $D$ is necessary to preserve a common dimensionality for the tokens inside a batch ? or is there another motivation behind it ?
limitations: limitations: There is  little discussion about some limitations e.g. (i) why rely on a fixed number of experts or (ii) what are the real latency of the method. The only limitation mentioned is that it would be hard to adapt this scheme to decoder-only LLMs, which seems a bit out of scope with the paper itself.

Comments: title: Reply to rebuttal
responses
comment: We're glad to hear that our rebuttal addressed your concerns.
We would like to clarify some potential misunderstandings by the reviewer, and mention that many of these points are addressed in our [global rebuttal](https://openreview.net/forum?id=HbV5vRJMOY&noteId=m6HR8cjvnf) to all reviewers.

In particular, we have already compared to existing dynamic models on ImageNet-1K in our initial rebuttal. Moreover, our current implementation does not use unnecessary padding, and we have already demonstrated substantial gains in latency (in addition to theoretical FLOPs) in our initial rebuttal. Finally, we also discussed in our initial rebuttal about how our MoNE is complementary to traditional MoEs.

We now detail these points below:

**Baselines**

We compared to existing dynamic computing methods in the second table of our initial [Global Rebuttal](https://openreview.net/forum?id=HbV5vRJMOY&noteId=m6HR8cjvnf). This table shows that we outperform existing approaches in both accuracy and GFLOPs when using the same ViT-Ti backbone. Note that throughput comparisons are difficult as they are hardware and implementation dependent (eg. ViT-B/16 throughput mentioned in [47] is 300 img/sec vs in [38] is 659 img/sec), thus we have only compared our method to a vanilla ViT using the same hardware and code-base.

**ImageNet-1K**

To clarify, our **Baseline Comparisons** table in the [Global Author Rebuttal](https://openreview.net/forum?id=HbV5vRJMOY&noteId=m6HR8cjvnf) is indeed on the **standard ImageNet-1k benchmark**. It indicates the superior performance at lower FLOPs of over other dynamic methods _in addition to_ the extensive ImageNet-21K results shown in Figures 3 and 5 of the paper. Additionally, the same section mentions that Token Merging techniques (ToMe) can be applied on top of MoNE to further achieve latency savings, thus depicting our framework’s flexibility of use and how it is complementary to related work. The performance of MoNE is further shown on **standard** video classification benchmarks (Kinetics400, SSv2), indicating that MoNE is able to achieve baseline performance at significantly lower performance. We hope these clarifications help and are happy to clarify further.

**MoEs**

As detailed in the [Global Author Rebuttal](https://openreview.net/forum?id=HbV5vRJMOY&noteId=v0FuPA3Sk5), MoNE is comparable to a dense model, as it does not increase the parameter space, and we have demonstrated accuracy and efficiency gains over dense models. Extending MoNE to a sparse MoE setup is a promising direction for future work.

Concretely, we can use MoNE in a sparse MoE setup where the experts are multiple MoNE layers, instead of traditional dense layers. This would reduce the compute cost of an MoE whilst still having the same parameter space.

In response to "It would be easy to design a MoE counterpart that would have the same number of inference parameters as MoNE …", it is important to note that the parameter-scaling properties of MoE models are different to dense models, particularly in the "low-parameter" regime. As shown in Table 8 of VMoE [34], MoE B/16 ,models that have the same number of parameters as dense ViT L/16 models underperform the dense models by 2-5% depending on the task.

**Efficient Implementation**

We would like to clarify that while Figure 2a of the paper indicates a padding for bringing feature representations to the same dimension, it is  _only_ for illustration / ease of explanation. In our actual efficient implementation, we handle tokens with different dimensions as different tensors, processing them in parallel without padding, and concatenating only for combined operations that require the whole token set together, which happens to be the case only for the SoftMax operation in SelfAttention. In this way, we match the exact theoretical FLOPs as well as achieve the gains indicated in the **Latency/Throughput** section of the [Global Author Rebuttal](https://openreview.net/forum?id=HbV5vRJMOY&noteId=v0FuPA3Sk5). Consequently, from a memory perspective as well, we do not incur any additional costs as we perform the exact operations _without padding_.
We thank the reviewer’s appreciation. 

The reviewer’s understanding is correct regarding slicing, and we will make it clear in the revised manuscript.

In Figure 2b, we do not show the slicing operation in the intermediate layers and only show in Fig 2a, to avoid an overcomplicated diagram. Figure 2b is just to show the mixture of nestedness across tokens. The slicing at the input is just an indicative of that token processing at a certain nested dimension. The horizontal blue and red bars indicate the dimensions of the MLP and attention hidden dimension respectively, which are always at the full dimension ($H_A, H_V, H_{MLP}$ using the reviewer’s notations).
Thanks a lot for your response. While I appreciate the added throughput results and discussion, I am inclined to keep my rating as I still think this would warrant a more extensive evaluation/discussion/modification of the paper.  In the current state, the paper proposes a technically sound dynamic computing method but does not discuss/compare well to existing related work (e.g. no throughput for dynamic pruning methods), and the core evaluation in the paper is performed on very specific settings/benchmarks (imagenet21k for classification, and video classification), so it is also hard to get an insight on how it compared to existing dynamic methods.

> We benchmark our method on ImageNet-21K for images to depict the large-scale efficacy of our framework. We would like to point out that ImageNet-21K has been extensively benchmarked in the AugReg Paper [38]

While I do agree that ImageNet-21k might be a more robust benchmark than ImageNet-1k, it is much less standard in practice.  With respect to existing baselines, it would be more fair to report ImageNet-21k results *in addition* to 1k rather than omit the 1k benchmark. 

> MoNE, while being a MoE-style framework, differs starkly from the concept behind MoE networks. MoNE does not increase the number of parameters, unlike traditional MoE networks. 

I agree but one may argue that this is the strength of MoE: Increased model capacity / number of training parameters while keeping the number of inference parameters constant (at least for a batch size of 1). It would be easy to design a MoE  counterpart that would have the same number of inference parameters as MoNE, or inversely, the same number of training parameters but faster inference.


> However, in an efficient implementation, padding would not be necessary, since we can maintain features of different dimensions separately, and add only the non-zero regions whenever needed.

I am not 100% convinced it would be so trivial to reach such an efficient implementation (hardware wise), as it would also mean having to constantly handle tokens with different number of dimensions, as opposed to the standard "large tensor multiplication/sums". In that sense do see the advantage of simplicity for padding, but it seems to be a bit of a negative point from a memory usage perspective (though MoE also have a similar issue with batched inference and having to activate many experts)
Thanks for the authors' rebuttal. Most of my concerns are well-addressed and I keep my rating of a weak accept.
Thanks for the responses.  I especially appreciate the new latency and throughput measurements and additional baselines as requested by a couple of the other reviewers.  I will keep my score at 7.

The explanation of slicing above helped, as it explains better that each of the W matrices are sliced on one of their dimensions, but I still think it could be clearer.  I found it helpful to write out the dimensions of all the W matrices and go back over the transformer ops you enumerated, distinguishing between "D" and "H" dims in each:

$W_q : [H_A, D]$,
$W_k : [H_A, D]$,
$W_v : [H_V, D]$,
$W_o : [D, H_V]$,
$W1 : [H_{MLP}, D]$,
$W2 : [D, H_{MLP}]$

in which case all the the $D$ dims are sliced to $D/m$, and applied to correspondingly sliced vectors, and all the $H$ dims are left intact, if I understand correctly.

I also now see that in Fig 2, the Pad operation is not shown in 2b, only 2a, and that the order of slicing and LN is inconsistent between 2a and 2b:  2a applies LN to full D-dim features, which makes sense, but 2b indicates that x is sliced before applying LN, in which case it's unlcear if this means x is padded back up to D and re-sliced, or if it's actually applied as in 2a, with LN before slicing.  I think this is where some of my confusion around these ops came from as well.  Lastly the horizontal blue bar of dim 4D in the MLP (and the horizontal red bar of dim D in SA) are a little confusing as well, as all the other sliced horizontal bars here are W matrices, and so it's not clear what the unlabeled bars correspond to in the figure.
Dear Reviewers and Area Chairs

We wholeheartedly thank you for your thoughtful reviews and the time you've dedicated to improving our work. We've carefully addressed all the points raised and submitted our rebuttal. We would greatly appreciate any further feedback or confirmation on our responses.
Thank you for your valuable feedback and increased score. We are pleased to hear that your main concerns have been addressed.

We appreciate you bringing Flextron to our attention. We have taken note of this concurrent work and agree that it would be beneficial to include a brief discussion in our paper highlighting the differences between the two approaches. We will make sure to incorporate this in our revised version.
Thank you for your detailed responses and the addition of new baseline comparisons. My main concerns have been addressed, and I am pleased to raise my score to 6.

I would like to bring to your attention Flextron published at ICML'24 - "Flextron: Many-in-One Flexible Large Language Model" by Cai, Ruisi, et al. This work presents a very similar idea to your paper. While it is clearly a concurrent work and does not diminish the contributions of your paper, it would be beneficial to include a brief discussion highlighting the differences between the two approaches.