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
Venue: NeurIPS 2024 poster (2024)
Keywords: Mixture of Experts
Matryoshka Representation Learning
Information Compression
Efficient Inference
Primary area: machine_vision
TL;DR: We devise a framework that learns to dynamically choose experts at a token level, from a mixture of nested models falling on a increasing compute-accuracy curve, resulting in superior performance at a provided compute budget.

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.

Mixture of Nested Experts: Adaptive Processing of Visual Tokens
Gagan J
Aditya Kusupati Shyamal Buch
Nidhi Hegde
Arsha Nagrani
Prateek Jain Anurag Arnab Sujoy Paul
Google DeepMind
'University of Washington
Abstract
The visual medium (images and videos) naturally contains a large amount of infor- mation redundancy, thereby providing great opportunity for leveraging efficiency in processing. While Vision Transformer (ViT) based models scale effectively to large data regimes, they fail to capitalize this inherent redundancy, leading higher computational costs. Mixture Experts (MoE) networks demonstrate scal- ability while maintaining same inference- time costs, but they come with larger parameter footprint. We present Mixture of Nested Experts (MoNE), which utilizes nested structure for experts, wherein individual experts fall compute-accuracy curve. Given a compute budget, MoNE learns to dynamically choose tokens in priority order, and thus redundant tokens are processed through cheaper nested experts. Using this framework, we achieve equivalent performance the baseline models, while reducing inference time compute by over two-fold. We validate our approach on standard image and video datasets ImageNet-22 NetgeNet-2111 21 Kinetics400, and Sowethinn-Somttiing--- ability by showcasing its ability to maintain strong performance across different inference-ting compute budgets on videos, using only single trained model.
increasing
We further highlight MoNE's adapt-
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 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 be- comes a major bottleneck when deploying these models real-world scenarios, where computational resources may be limited and real-time processing is required.
To this end, conditional computation has become promising line of research to increase the capacity of a network, while only conditionally activating part of it during inference. Sparse Mixture of Experts (MoEs) was initially popularized for Natural Language Processing (NIPP) [42, 22] has been gaining attention for furthering conditional computation ideas in vision [39, 1, 35, 50] well. given inference cost, they also increase the overall
While MoEs bring in improved performance parameter count, leading to increased storage requirements. Moreover, these works rely experts that have the same parameter count and compute, limiting their ability to reduce computational costs without resorting to skipping tokens entirely.
"quul contribution
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

BC-07
Figure I: MoNE's learned token importance: From left to right, fewer image tokens are processed using the full model fit a compute budget by an increasing threshold on MoNE's router logits.
In this work, we devise the Mixture of Nested Experts (MoNE) framework, which provides scalable approach to conditional computation, bringing in significant reductions at inference time, while working with the same parameter space 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 structured slices of the parameter space. MoNE employs these structured nested models as experts 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 router predictions. Figure 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 fixed parameter count, MoNE offers the flexibility of learning networks at much lower FLOPs (~ 2.3x on video datasets) and real-time latency (~ 2x), while being quality neutral.
3. Rigorous experiments show that MoNE works well for both image and video transformers, and visualizations depict that tokens routed larger experts correlate well with regions of interest.
2 Related Work
Transformers [45] have become the de-facto architecture for processing data across multiple modal- ities 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 transform- ers 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 easier to process than others, and therefore require less computation during inference.
Mixtures-of-EExeere (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 FI.OP-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 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, leam hierarchical representations of the input, where the first & 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 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 sccuracy-efficiene 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 "halting score" for when 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 udaptive approach offers strong compute- perppotaepproomnncee 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 learn from single training phase, without using surrogate model.
3
Preliminaries
Here, we discuss the concept of nested models, on which we build Mixture of Nested Experts (MoNE), followed by discussion about Mixture of Nested Experts (MoE), and its differences from MoNE.
3.1 Nested Models
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 feature projection operation Wx, where W
[ W [2g]]. and W W denotes "slicing" the first D dimensions, we can extract a partial projection W (.2 This m 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 P/m as the nested model dimension. This is shown in Figure 2a. The Extract operation extracts the first P/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-proeections of both the Self-Autention (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 (QKT) operation. The same thing is performed in MLP, where the hidden dimension is always 4D, as in ViT, irrespective of proou--pooecction 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 _ 4. 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 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.
3.2 Mixture of Experts
Mixture of Experts (MoE) layer in a transformer can be represented as MoE(x) ΣE g(x) (x))),x()) where F. is the number of experts, () are the expert models each having their own parameters, g: RDD R is the routing/gating function, which decides the experts which should process x. Note that g is sparse with only l << 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. contrast, MoNE, we do not extend the parameter count of the model, due to the nesting structure (see Sec. 3),, and dynamically choose a nested expert during inference. Unlike in MoF, where all experts have the same capacity, in MoNE with & 1 always, & 6,-1, which allows us to dynamically allocate compute.
3

---

XL
Rod
MLP
MLP
R7-
MLP
Layer Norm
Layer Norm
0/
LN
Mutti-ited Self-Attention
Paa
Layer Norm
Layer Norm
D/-
SA
D14
Router
Router
LN
D
(a)
(b)
Figure 2: (a) Nested model: Partial in- and out-projections in the SA and MLP layers create nested models. m controls the parameter count and the FLOPs of nested models. The self-attention information exchange happens the full model dimension D, MLP dimension is 4D in VIT. (b) Mixture of Nested Experts (MoNE): Each token is routed nested network, denoted by different model dimension in the diagram. Here x; gets routed to nested model with model dimension P/4, whereas X gets to the full model. The information exchange between these tokens of different dimension happens in the self attention block, where they are always projected the same dimension. The router weights also multiplied with the features for proper flow of gradients. lighter color the weight matrix indicate sliced matrix construct the nestedness.
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.
4.1 Mixture of Nested Experts (MoNE)
Tokenization: In this paper, as our primary focus is images and videos, the model input is in #xx*3xT where T 1 for images and T > 1 for videos. After tokenization, the input to the transformer is X DxN where N is the number of tokens, and D their model dimension. For images, have N H/pn- W//w>, and for video, N T/pt H/pp- W/Pw. where H, W, T are the input height, width and duration respectively. Ph. Pw and Pt are the patch sizes along these respective dimensions. We the VIT [20] and ViViT [2] architectures to tokenize images and videos respectively, obtaining a list of tokens X (x
The Mixture of Nested Experts (MoNE) framework is 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¹ = (B),..,, B'g) denote the nested blocks at a certain layer with increasing parameter sizes, (-) being the full model block. A router network decides the appropriate nested block to use for every token. Henee 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,, router produces a probability distribution over the E nested experts, r; softmax (W,;
MoNE Block:
x, + b,) where W, and b, denote the router weights and bias respectively.
These router predictions sent to assignment algorithm, which assigns every token single appropriate nested expert. Based on the assignments, we update the features for the 7th token in the tth layer follows-
x 2 (art, + +1) (B))BBFFN BFFN. (2.)
2 = x² 135A.((x))
(I)

---

where the 5th nested expert is chosen by the Expert Preferred Router (EPR(.)] algorithm for the t token as per Eq. 2:
J EPR (M))))
(2)
Note that the multiplication of the router predictions with the model output in Eq. 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 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 is used for downstream applications. For classification tasks, we apply global average pooling on all the token features and apply linear classifier layer to predict the categories.
Features and Loss:
4.2 Token to Nested Expert Assignments
Within the MoNE framework, the routing strategy is erucial for achieving an optimal balance between performance and computational efficiency. Traditionally there are two primary routing strategies token 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 heing routed to one 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 computation. Potential conflicts due to token selection by multiple experts are resolved by prioritizing based on model size.
Formally, we consider given distribution of nested models applied to the tokens, represented as ; C. s.t., Σ, G 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 discussed in Sec. 4.3. Given router probabilities r; for N tokens across E experts, we employ an Expert Preferred Routing algorithm (Algorithm 1). This is greedy assignment approach that gives higher preference 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 k3- [cg,N] of the remaining tokens to jdh nested model. Any remaining tokens, arising from integer packing constraints, are assigned to the smallest model. Algorithm presents the proposed Expert Preferred Routing (EPR) algorithm.
Algorithm 1 Expert Preferred Routing (EPR)
Require: r R ExN (router predictions), c (capacity distribution, s.t., c² 1),
Ensure: Me F1)....F. (nested model index)
1: M + ln
Default assignments to the smallest model
2: for to do
3:
kj + +cylc--No
4: I T=Too-k-Index (raddex(((jj......
Returns value and indices of Top-K
5:
M[[] -j
6:
r[: -0
Null out assigned ones
7:
end for
8: return M
4.3
Capacity Distribution Across Experts
The Expert Preferred Routing (EPR) as described Section 4.2 needs the individual expert's capacity bounds to be specified. To get this, we define metric called the effective capacity e, Σ e.d./DD where d; R/2 is the model dimension of the 5th nested model. Given certain inference FLOP requirement, we can translate that to an equivalent effective capacity e.. Since every token gets processed through exactly one nested expert, this along with the given budget imposes two constraints on the unknown capacity distribution However, since the individual expert capacities vary log- linearly, multiple distributions can lead to the same e for E > 2 and it is non-trivial to choose
5

---

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:
maximize
Σ
Ci 8-11
8
C
logg
(3)
subject to
Σ c==1
Σ
C 2E-1
=ce
0sasl
Vie (L..., b}
1-1 E,8>1, 8>0
given 0<ec <,
In practice, (8,5) to (10.2) and use Sequential Least SQuares Programming (SLSQP) optimizer to solve Eq. for the capacity distribution c, which then used by EPR (Algorithm 1) to get token to expert mappings. We empirically verify these choices in Section 6.
4.4 Videos
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,, layers, culminating in a single global token index. Subsequently, a temporal transformer processes these global tokens across L, 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- trie expert assignments.
5 Results
In this section, we empirically evaluate MoNE on multiple datasets spanning images and videos for different model sizes, assess its adaptability to stringent FL.OP constraints, and depict real-time latency gains achieved by MoNE during inference.
Implementation details: We empirically evaluate MoNE on image and video classification. For image classification, 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 Big Vision [7] for image classification and Scenie [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 this section, we place a single router at the first transformer layer. and propagate the router decisions 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 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 four nested models that we consider this work ( : 2 D). Vat is equivalent to MoNE, with deterministic router to pass all tokens to the same nested
D

---

42
PLANT VI
MIGHT
MAAT (CAPO)
25
RLOPS
FLOPM
FLOP:
(a) S/16
(b) B/16
(c) L/66
Figure 3: Image classification: Performance comparison of MoNE with baselines on ImageNet-21k for different model sizes. MoNE performs significantly better than MatViT and Mixture-of-Deee -De-thhh (MoD) and even benefits from isoFLOPs training (see fig a).
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 also token routing algorithm, but proposed for language tasks. MoD takes the extreme decision of either processing or skipping for every token in layer. MoNE, on the other hand, makes fuzzy decisions to choose intermediate-Sen sized models, instead of skipping, which helps to retain significant information at the expense 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 traditional MoEs like Sparse VMoE [39], and inference-ting tetrencee-imeeeee optimization techniques like Token Merging (ToMe) [8]. We present an extended discussion and further results 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-22 1k [18] classification using VIT. We experiment with S, and I. models to showcase efficacy of MoNE across model sizes. ImageNet-21k can have multiple labels for image, we report the commonly used precision 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 2x reduction in FLOPs.
Following the literature on language models [37, 27], we experimented with isoFLOPs training, which involves training for the same number 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 FL.OPs models.
we first train a baseline model with nested Something-ve 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 temporal clips [2]. Figure illustrates the results of the MoNE framework, significantly outperforming the individual nested models. MoNE offers 2 reduction in FLOPs compared the ViVIT baseline, without any accuracy drop (On SSv2, the
Since video models rely on heavy pre-training [2]. structure on the benchmark datasets Kinetics-400 [3]] and
Videos:
74
73
69
68
MatVVVT
MalVMIT
VMTT
VMT
66
MoNE (Ours)
MoNE (Ours)
100
250
350 210 270
110
(a) Kinetics-400
(b) Smmething-Something.we
Figure Video classification: MoNE vs. baselines on video datasets. Finetuning with the isoFLOPs training regime leads to matching baseline with > 2x FLOP improvement.
7

---

ArcuPay 8
22
e_m0.33
Taln c-00
52
Tain c006
50
Tran Assptive
12
14
100
120
140
1.00
160
200
PLOPS
FLOPs
(a) ImageNe121k
(b) SoanettinngSomettiigg---o
Figure 5: Capacity adaptation during inference: Performance changes when model trained at certain capacity (denoted as *) is evaluated at other capacities. The "Train Adaptive" plot for SSv2 denotes single model evaluated at different inference-time budgets.
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, the infer- ence time budget is often dynamic, changing based on user needs. Ideally, a model should adjust with little to retraining. To evaluate this ability, test how MoNE, trained at specific effective capacity (e.) 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 the capacity budget between train and eval. The performance degradation is steeper while adapting model trained 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.15, 0.25,.... ..... 0.95) at each training step. The results SS-v2 (Figure 5b) demonstrate our framework's s strong capability to adapt to any inference-time budget using a single model. is interesting to note that the training FLOPs of this adaptively trained model are equal to those of baseline model (isoFLOPs training). The model adapts extremely well even to capacities that are significantly different 0.3.,0.3,....... from those sampled during training.
Real Time Latency Gains: addition to the theoretical FLOP gains, Table presents the real-time latency/throughput gains of MoNE-based ViVIT model as compared to its baseline variant. The absolute wall clock times and throughput are compared on single V100 GPU, achieving nearly two-fold improvement in both FLOPss as well as runtime, whilst maintaining accuracy.
Table 1: Real Time Latency and Throughput gains for MoNE on single V100 GPU
Method
FLOPs (G)
Throughput (clips/sec)
Latency (ms)
Top-1 Accuracy
VIVIT-FE-B/16
376
15.8
129.2
64.4
MoNE (e.- 0.3)
162
30.7
65.5
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. implementation and while simple high-level efficient implementation of our framework yields gains of this scale, we believe that further improvements can be obtained by optimizing low-level GPU kernel implementation for MoNE
is important to note that inference gains depend heavily on
In addition, worth noting that the proposed Expert Preferred Routing (EPR) in Algorithm 1 loops only over the number of experts, which is typically small number and fixed to 4 in our framework. the computation any further,
While the nature of the EPR algorithm does not allow parallelization the time taken by the algorithm is negligibly small as compared to the total time taken by the model. For comparison, on V100 GPU, the EPR algorithm adds just 0.5 ms to the forward propagation time of VTT-B/16 model (190 ms), accounting for 0.3% of the total computation time.

---

NoME VVVT
MONE
120
1.00
12
00
23
20
GFLOPs
GFLOPs
(a) FLOPSs vs Latency
(b) FLOPPs vs Throughput
WVTT
64.6
64.6
84.5
Accuraty
Accuracy
64.4
64.4
64.3
64.3
64.3
WVIT
64.1
(Duns)
120
Throughput
(c) Latency vS Accuracy
(d) Throughput vs Accuracy
Figure 6: Wallclock realization of MoNE's computational savings with varying effective capacities, depicted on the Something-Something v2 dataset.
Router Analysis
In this section, 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 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 input, this also diminishes the compute gains, 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 the curve.
Number of Routers: 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.
47
56.4
85.0
262
02.
54.5
AC300
0.5
Lyar2
234
Up3r2
Liar4
NotG (005))
100
(a) Router Position
(b) Number Routers
(c) Comparing with random router
Figure 7: Router Analysis: Effect of router placement and learning on Homething-Somethie v2.

---

(a) Images from ImageNet-21k
(h) Video frames from SomethingSomethingy2
Figure 8: Tokens routed to the full model: Highlighted regions are the tokens sent the full model, while rest of the tokens are sent to the smaller nested models. (a) shows examples on images and (b) shows example on a video multiple temporal indices. As we can see, the necessary and important tokens are sent to the full model.
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 (cz), computed in Section 4.3. Results in Figure 7c suggests that with lower effective capac- ities, the random router performance degrades while the learned router still manages to understand relevant patterns from the input, thus upholding performance.
The above claim is further backed by visualizing the token impor- tance during inference low effective capacity (cz). We highlight the tokens selected by the largest expert, i.e., the full model on 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.
Visualizing Important Tokens:
Capacity Allotment: Given fixed input capacity e.. we demonstrate the superior performance of our heuristic-ber buaiet 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 adapt 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 routed through the smallest expert, exhibit inferior performance.
Table 2: SSv2 Performance of different capacity distribution methods
Static budget
Dynamic budget Greedy
Distribution
Proportionae
Uniform
MoD Greedy [37]
MoNE
MoNE
Effective Capacity (e,)
0.27
0.47
0.3
0.4
Accuracy
64.3
64.6
63.9
64.2
64.2
64,6
7 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.
Extending this to rego-segressive decoding in LLMs is non 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 role democratization of AI, allowing broader access to trained models without the need for large resources.
Limitations:
10

---

Acknowledgements
We grateful Debapriya Tula, Jeevesh Juneja, Matthew Wallingford, Pradeep Shenoy and Ali Farhadi for helpful discussions and feedback.
References
[1] U. Allingham, F. Wenzel, Z. F. Mariet, Mustafa, Puigcerver, N. Houlsby, G. Jerfel,
V. Fortuin, B. Lakshminarayanan, J. Snoek, et al. Sparse moes meet efficient ensembles. arXiv preprint 10.00330 2021.
[2] A. Amab, M. Dehghami, G. Heigold, C. Sun, M. Lučić, and C. Schmid. Vivit: video vision transformer. In ICCV, pages 6836-6846, 6846, 2021.
[3] A. Banino, J. Balaguer, and C. Blundell. Pondernet: Learning to ponder. arXiv preprint arXiv:2107.05407, 2021.
[4] I. Beltagy, M. E. Peters, and A. Cohan. Longformer: The long-document transformer. In arXiv preprint axXiv:200.005150 2020.
[5] Y. Bengio. Deep learning of representations: Looking forward. In International Conference on Statistical Language and Speech Processing, 2013.
[6] G. Bertasius, H. Wang, and L. Torresani. Is space-time attention all you need for video understanding? In JCML, volume 2, page 4, 2021.
[7] L. Beyer, X. Zhai, and A. Kolesnikov. Big vision. https: wtt.:u...coo//oooaaeeeeeo wtt.ttbb..cm/gggleeeeeeeerree big_vist vision, 2022.
[8] D. Bolya, C.-Y. Fu, X. Dai, P. Zhang, C. Feichtenhofer, and Hoffman. Token merging: Your vit but faster. In ICLR, 2022.
[9] J. Bradbury, R. Frostig, P. Hawkins, M. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, VanderPlas, Wanderman-Minn Miilee, and Q. Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http::/tptt /ttp://gghtttt com/ ggogle/jax..
[10] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam,
G. Sastry, A. Askell, et al. Language models are few-shot learners. In NeurIPS, 2020.
[11] R. Cai, Muralidharan, G. Heinrich, H. Yin, Z. Wang, J. Kautz, and Molchanov. Flextron: Many-in-one flexible large language model. ICML, 2024.
[12] K. Choromanski, V. Likhosherstov, D. Dohan, X. Song, A. Gane, T. Sarlos, Hawkins, J. Davis,
A. Mohiuddin, I.. Kaiser, et al. Rethinking attention with performers. In ICLR, 2021.
[13] Z. Dai, Z. Yang. Y. Yang, J. Carbonell, Q. V. Le, and R. Salakhutdinov. Transformer-xl: Attentive language models beyond fixed-length context. In ACL, 2019.
[14] Darcet, M. Oquab, Mairal, and Bojanowski. Vision transformers need registers. arXiv preprint arXiv:2309. 16588, 2023.
[15] M. Dehghani, S. Gouws, O. Vinyals, J. Uszkoreit, and E. Kaiser. Universal transformers. In ICLR, 2019.
[16] M. Dehghani, A. Gritsenko, A. Arnab, M. Minderer, and Y. Tay. Scenie: A jax library for computer vision research and beyond. In Proceedings of the IEEE/CVF Conference on Computer Vision and Patlern Recognition (CVPR), pages 21393-21398, 2022.
[17] M. Dehghani, J. Djolonga, B. Mustafa, P. Padlewski, J. Heek, J. Gilmer, A. P. Steiner, M. Caron, R. Geithos, L. Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters. In ICML, 2023.
[18] J. Deng, W. Dong, R. Socher, Li,J, K. Li, and Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 JEEE Conference on Computer Vision and Pattern Recognition, pages 248-255, 2009, doi: 10. 10.110//CVPR222995526663

---

[19] Devvrit, S. Kudugunta, A. Kusupati, T. Dettmers. K. Chen, I. Dhillon, Y. Tsvetkov, H. Hajishirzi, S. Kakade, A. Farhadi, P. Jain. et al. Matformer: Nested transformer for elastic inference. arXiv preprint arXiv:2310.07707, 2023.
[20] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, al. An image is worth 16x16 words: Transformers for image recognition at scale. 2021.
[21] M. Elbayad, J. Gu, E. Grave, and M. Auli. Depth-adaptive transformer. In arXiv preprint 1910 10073, 2019.
[22] W. Fedus, Zoph, and N. Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-39) 2022.
[23] Y. Gong, Y.-A Chung, and J. Glass. Ast: Audio spectrogram transformer. In arXiv preprint arXiv:2104.1 NDV:2104.01775, 2021.
[24] R. Goyal, S. Ebrahimi Kahou, V. Michalski, Materzynska, S. Westphal, H. Kim, V. Haenel, I. Fruend, P. Yianilos, M. Mueller- et al. The" something something" video database for learning and evaluating visual common sense. Proceedings of the IEEE international conference on computer vision, pages 5842-5850, 2017.
[25] A. Graves. Adaptive computation time for recurrent neural networks. In arXiv preprint 1603.:6003.08900 2016.
[26] D. Hendrycks and K. Gimpel. Gaussian error linear units (gelus). In arXiv preprint arXiv: 15, 2016.
[27] J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. d. L. Casas,
L. A. Hendricks, J. Welbl, A. Clark, et al. Training compute-optimal large language models. arXiv preprini arXiv:2203. 15556, 2022.
[28] G. Huang, D. Chen, T. Li, F. Wu, L. Van Der Maaten, and K. Q. Weinberger. Multi-scale dense networks for resource efficient image classification. In ICLR, 2018.
[29] A. Jaegle, F. Gimeno, A. Brock, O. Vinyals, A. Zisserman, and J. Carreira. Perceiver: General perception with iterative attention. In ICML, 2021.
[30] A. Q. Jiang, A. Sablayrolles, A. Roux, A. Mensch, B. Savary, C. Bamford, D. S. Chaplot, D. d. I. Casas, E. B. Hanna, Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401.044 000088, 2024.
[31] E. Kim, C. Ahn, and S. Oh. Nestednet: Learning nested sparse structures in deep neural networks. In Proceedings of the IEEE Conference on Computer Vision and Patlern Recognition, pages 8669-8678, 2.18.
[32] Z. Liu, Y. Lin, Y. Cao, H. Hu, Y. Wei, Z. Zhang, S. Lin, and B. Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In ICCV, 2021.
[33] Z.. Liu, J. Ning, Y. Cao, Y. Wei, 7. Zhang, S. Lin, and H. Hu. Video swin transformer. In CVPR, pages 3202-3211, 202..
[34] Z. Liu, J. Wang, T. Dao, T. Zhou, B. Yuan, Z. Song, A. Shrivastava, C. Zhang, Y. Tian, C. Re, al. Deja vu: Contextual sparsity for efficient Ilms at inference time. International Conference Machine Learning, pages 22137-22176. PMLR, 2023.
[35] Y. Lou, F. Xue, Z. Zheng, and Y. You. Sparse-mlp: A fully-mip architecture with conditional computation. arXiv preprint arXiv:2109.000088 1:12, 2021.
[36] C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. JMLR, 2020.
[37] D. Raposo, S. Ritter, B. Richards, T. Lillicrap, P. C. Humphreys, and A. Santoro. Mixture-of- depths: Dynamically allocating compute in transformer-basee language models. arXiv preprint arXiv:2404.0222 02255228, 2024.
12

---

[38] M. Reid, N. Savinov, D. Teplyashin, D. Lepikhin. T. Lillicrap, J.-b. Alayrac, R. Soricut, A. Lazaridou, O. Firat, J. Schrittwieser, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. In arXiv preprini arXiv:2403.055:00 2024.
[39] C. Riquelme, J. Puigcerver, B. Mustafa, M. Neumann, R. Jenatton, A. Susano Pinto, D. Keysers, and N. Houlsby. Scaling vision with sparse mixture of experts. NeurIPS, 34:8583-8559 2021.
[40] M. S. Ryoo, A. Piergiovanni, A. Arnab, M. Dehghani, and A. Angelova. Tokenlearner: What learned tokens do for images and videos? In NeurIPS, 2021.
[41] T. Schuster, A. Fisch, J. Gupta, M. Dehghani, D. Bahri, V. Tran, Tay, and D. Metzler. Confident adaptive language modeling. In NeurIPS, 2022.
[42] N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In ICLR, 2017.
[43] A. Steiner, A. Kolesnikov, X. Zhai, R. Wightman, Uszkoreit, and L. Beyer. How to train your vit? data, augmentation, and regularization in vision transformers. arXiv preprint arXiv:2106.10220, 2021.
[44] Y. Tay, M. Dehghani, D. Bahri, and D. Metzler. Efficient transformers: A survey. ACM Computing Surveys, 2022.
[45] A. Vaswani, N. Shazeer, N. Parmar, Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and
I. Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[46] A. Veit and Belongie. Convolutional networks with adaptive inference graphs. In ECCV, 2018.
[47] C. Wan, H. Hoffmann, S. Lu, and M. Maire. Orthogonalized sgd and nested architectures for anytime neural networks. In International Conference on Machine Learning, pages 9807-9817. PMLR, 2020.
[48] S. Wang, B. Z. Li, M. Khabsa, H. Fang, and H. Ma. Linformer: Self-attention with linear complexity. In arXiv preprint 2xXiv:2006.04766 2020.
[49] Y. Wang, K. Li, X. Li, J. Yu, Y. He, G. Chen, B. Pei, R. Zheng, J. Xu, Z. Wang, et al. Internvide02: Scaling video foundation models for multimodal video understanding. In arXiv preprint arXiv:2403. 15377, 2024.
[50] F. Xue, Z. Shi, F. Wei, Y. Lou, Y. Liu, and Y. You. Go wider instead of deeper. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 8779-8787, 2022.
[51] Xue, V. Likhosherstov, A. Arnab, N. Houlsby, M. Dehghani, and Y. You. Adaptive computa- tion with elastic input sequence. In ICML., 2023.
[52] H. Yin, A. Vahdat, J. M. Alvarez, A. Mallya, J. Kautz, and P. Molchanov. A-vit: Adaptive tokens for efficient vision transformer. In CVPR, 2022.
[53] J. Yu, L. Yang, N. Xu, J. Yang, and T. Huang. Slimmable neural networks. arXiv preprint arXiv:1812.0892 089288 2018.
13

---

A Appendix
A.1 MatFormer Structure on Model Dimension
Following MatFormer convention, we define E VIT blocks B, such that B: Bi+1 for all i [.] meaning that the parameters of B; are contained those of B++1- With d, denoting the hidden dimension corresponding to nested model B; such that d1 d2 <...dg dg D, the block operation for a nesting B, on an input token set X {x,) for x; RR is given by:
(X) B(x, di) 8 (Z, di), Z= +X,
di = (di,d)....,dd ,i)
(4)
N times
W
Concatenate
W Wy--o
Scaled Not-Product Attention
W
W
w
W FF4
X;
X;
Figure 9: MatFormer Structure on Model Dimension
The modified Self-A Attentienn BSA and Feed-Forward BIFN subroutines are shown below.
SA (x. =LN (0 (X
W (Xoa Wi) Ed WSA
(5)
dem
FTN (x, LN od WFF,) Bd WFF.
(6)
where dd and E respectively denote the sliced in and out projection operators, such that:
(X oa W), (x;) Widdi
(X W), = (Wiaz))
In the general Mixture of Nested Experts (MoNE) setting discussed in Section 4.1, the overall block computation for the set of tokens X requires knowledge of the expert assignments for each token beforehand. Given these assignments m N such that m, E (1,2...., E), the tth token processed by the 7th expert can be represented as:
the computation for
(x,x,) [B(X.
d (dm.)
(8)
In Eq. 8, the block update for token x; dependent on the complete input set and their respective expert assignments m, but omit these in the definition B, for notational convenience. Additionally,
14

---

Table 3: MoNE Comparison on ImageNet- 1K with other Adaptive Baselines
Method
FLOPs (G)
Throughput (clips/sec)
Top-1 Accuracy
ViT [20]
1.3
3410
71.3
PonderNet [3. 21]
1.0
66.2
DepthAdapt [21]
1.1
69.4
ACT [25]
71.0
A-VIT [52]
0.8
71.0
MoNE (Ours)
0.8
4333
71.4
this definition directly extends to the sub-routines BSA and BEFN as presented in Eq. 1. Here, the weight matrices of SA are Wa. WK, Wv, WSA. RR (D( and the weight matrices of FFN are WFF,, WFF. E DDddff, ignoring bias terms for simplicity. W[sk] denotes the first rows of W. Here, n, denote the number of heads in the attention mechanism, dyf denotes the feed forward dimension, and a denotes a non-lincarity, typically set to GeLU [26].
A.2 Comparing MoNE with Other Adaptive Baselines
Firstly, we establish that MoNE, complements rather than competes, with the traditional Mixture of Experts framework. Traditional MoEs like Sparse VMoE [39] route inputs in each layer to one out of k independent experts (typically the FFN block), each having the same parameter footprint, thus increasing the parameter space k-fold for the expert blocks. On the other hand, independent MoNE do not increase the parameter space, and thus MoNE blocks can potentially be used as experts in the MoE framework.
MoNE acts as an in-place replacement for a dense model like ViT, hence all our comparisons maintain the same parameter space. VMoE frameworks show cross-scale results the expense of increased parameter space (e.g., equivalent performance of VMoE-L/16 to ViT-H/14 in Table 2 in [39], and similar cross-scale comparisons in Figs. 4 to 8 in [1].. 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 certain tasks, this 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-sizee vanilla model. Additionally, as shown in Table of [19], joint optimization of shared experts leads to better performance than having independent experts of the same size.
In Table we compare MoNE with other baselines, particularly with adap- tive computation of dense models. We perform this experiment on ImageNet- 1K with T sized model. ACT [25]. PonderNet [3], Depth/dapt [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].
Drop in Top-1 Acc
-2
ToMe+VIT
We also highlight that MoNE can be uti- lized as a baseline for further inference- time 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- IK. For this experiment, we train model with full capacity till the third layer and
ToMe MoNE
0.6
0.7
0.8
0.9
1.0
1.1
1.2
GFLOPs
Figure 10: Latency gains by using ToMe on top of MoNE, compared against ToMe vanilla VIT. We use ViT- Ti/16 model trained on ImageNet-1K for this comparison.
15

---

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 distinct sets of nested tokens, indicating that ToMe is complementary MoNE.
A.3 Understanding the Router Behaviour
Number of Routers: 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, 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 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 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 higher number of routers might lead to better results.
To further analyse the decisions made by the MoNE router, we study the visual inputs from the ImageNet-2IK 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.c. harder to understand inputs are sent larger experts). Therefore, instead of using the greedy EPR Algorithm 1, we take an argmar over the router logits to make decisions. The results presented in Fig. 11 depict two sets of images, the t 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.
Task Difficulty:
88
120
Lowest compute
Highest compute
Figure 11: A few images, which have been predicted by the router require lowest and highest computation, from over 2000 randomly selected images in the ImageNet21k dataset.
Learned 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 lower capacities, from 0.1% at 0.6 to 1.3% at e = 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 few tokens can utilize the heavy experts. Interestingly, ViTs inherently shuffle information [14], potentially even the "Random" router setting as well, acting as intrinsic information router. We note that model trained with learned router when evaluated with random router, performs significantly worse (~ 6% drop Top1 Accuracy on Ti/16 trained on ImageNet- IK).
16

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Yes, verify our claims experimentally in Sec. 5.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Please check See 7 for more details.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: All results presented in this work are empirical.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
We mention all the codebase used in this work, along and use their default for the respective networks. We also present our work in great detail, which we believe should be enough to reproduce results.
Justification:
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification:
The datasets and codebase top of which we build algorithm are open- sourced, we will open-source the code for this paper upon acceptance.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: See "Implementation Details" for more details.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [NA]
Justification: The experiments are extremely compute intensive barring us from running each of them multiple times. However, the datasets we work on standardized and the gains are stable across multiple runs.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
17

---

Justification: We use maximum of 64 TPU v3 chips per training experiment.
9. Code Of Ethies
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: :
Answer: [Yes]
Justification: Yes. We do not have anything in conflict with the NeurIPS code of ethics.
Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: See Section 7 for a short overview on how MoNE can help democratize large image and video models in practice.
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: We do not work generative media.
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: [NA]
Justification: We use popular public research datasets and models which are under permissive license for research and fair use. We shall open source our code and models.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: No datasets are introduced in this work.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing research with human subjects. 15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
18


Reviews:
[{"review_id": "drvdAw960n", "reviewer": "Reviewer_E3Y8", "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.\n\nThe results show that MoNE achieves a favorable accuracy-efficiency trade-off curve for image and video classification tasks compared to the baselines.", "questions": "Please review the weaknesses section.", "limitations": "The authors mention one of the limitation of their work regarding extension to sequence modeling tasks.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written and easy to follow\n- The extension of the MatFormer architecture into an MoE architecture with dynamic routing seems quite intuitive. The Matformer architecture is inherently composed of nested subnetworks and enabling dynamic routing in this architecture makes sense.\n- Experiments on Image and Video Classification tasks shows MoNE outperforms the baselines.", "weaknesses": "- 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.\n\n- 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.\n\n1. Meng, Lingchen, et al. \"Adavit: Adaptive vision transformers for efficient image recognition.\" CVPR 2022.\n2. Li, Yi, et al. \"Svitt: Temporal learning of sparse video-text transformers.\" CVPR 2023.\n3. Yin, Hongxu, et al. \"A-vit: Adaptive tokens for efficient vision transformer.\" CVPR 2022.\n\n- 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.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the insightful review. We are glad to hear that the reviewer found the paper to be well-written, the MoNE framework to be intuitive, and compelling experimental results. Below we answer some of the questions raised by the reviewer. \n\n**Dynamic Routing vs MatFormer and the need for MoE comparisons**\n\nWe agree with the reviewer that our dynamic framework allows for more informed token routing than the static strategy in mix’n’match in MatFormer [16]. We would like to point out that MatFormer’s mix’n’match happens on a per-layer basis, each layer processing all tokens uniformly using a particular nested model. Our method, in contrast, takes decisions per-token. MatFormer does not explicitly explore input-dependent dynamic decision-making for tokens. It is non-trivial to come up with a learning framework which would allow inference time mix’n’match where tokens are processed through different models.\n\nRegarding the concern about comparison with other MoE methods like Sparse VMoE, we discuss the comparison of MoNE with these MoE frameworks in the **MoE Comparison** section of the Global Author Rebuttal. In addition, MoNE can be extended to have multiple disjoint experts as in Sparse Vision-MoE, offering compute efficiency in such models. This can be considered as a future work. We will update the manuscript with this clarification.\n\n**Clarifying the Advantages of Nested Experts and Comparison with Dynamic Token Processing Approaches**\n\nThe key advantage of having a nested set of experts over non-overlapping experts is that the computational gains can be achieved without increasing the parameter space. Non-overlapping experts for a given parameter space would lead to a limitation of representation power of each expert. However, as in MoNE, overlap between experts allows the largest expert to enjoy the full parameter space. Additionally, as shown in Table 5 of the MatFormer [16] paper, joint optimization of shared experts leads to better performance than having independent experts of the same size.\n\nWe additionally compare our method with some of the methods mentioned by the reviewer, and show results in the **Baseline Comparisons** section of the Global Author Rebuttal. The comparisons show that MoNE performs better than other adaptive processing algorithms. \n\n**Understanding Expert Specialization: Contrasting MoNE and Conventional Sparse MoE Architectures**\n\nThe primary motivation behind the nestedness in MoNE is input adaptivity thus offering compute efficiency while keeping the same parameter space. MoNE still remains fully compatible with traditional MoE approaches (Mixture of MoNE). In that setup, intuitively speaking, not all concepts may need the same amount of compute and using MoNE in the MoE setup would offer compute efficient MoE models.\n\nWhile one of the expected outcomes of MoE architectures is specialization of experts, this is not always the case in literature. Quoting a few lines from Section 5 of Mixtral of Experts [b], a MoE method that has non-overlapping experts - “Surprisingly, we do not observe obvious patterns in the assignment of experts based on the topic. For instance, at all layers, the distribution of expert assignment is very similar for ArXiv papers (written in Latex), for biology (PubMed Abstracts), and for Philosophy (PhilPapers) documents.“ Moreover, in Sparse VMoE [34] the authors observe very weak correlation of router decisions to categories.\n\n[b] A Jiang et al. Mixtral of Experts. arXiv:2401.04088, 2024.\n\n**Understanding the Impact of the Number of Routers in MoNE**\n\nThis is a very important point raised by the reviewer and we clarify in the **Number of Routers** section of the Global Author Rebuttal.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xQ88h3Xs0M", "reviewer": "Reviewer_kK61", "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.", "questions": "See weaknesses.", "limitations": "Limitations are discussed in the paper.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Overall, the paper is well-motivated and the idea of combining nested structers with MoE is very interesting; The problem of information redundency do exist in image or video classification tasks.\n\n2. Empirically, the MoNE models attain very competitive results with less FLOPs and parameters, which could support their theoretical analysis in vision information redundency.\n\n3. The paper is well-written and easy to understand. The method is simple and easy to follow.", "weaknesses": "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?\n\n2. 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?\n\n3. 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive feedback. We are glad to hear that the reviewer found the core idea of the work – nested structures to exploit information redundancy – to be well motivated and interesting, the experimental results to be promising and the paper to be well-written. Below we address some of the questions raised by the reviewer.\n\n**Actual inference speed comparisons (Throughput/Latency)**\n\nWe agree with the reviewer about the importance of realizing the FLOP gains in terms of latency/throughput. We provide real inference speedups in the **Latency/Throughput** section of the Global Author Rebuttal. We will include these results to the main paper.\n\n**EPR algorithm implementation**\n\nThe EPR algorithm contains a loop only over the number of experts, which is quite low and fixed to 4 in our framework. While the nature of the EPR algorithm does not allow to parallelise 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 GPU), for a ViT B/16 that takes 190ms for forward propagation, the EPR algorithm takes 0.5ms, less than **0.3%** of the total computation time. We will mention this in the revised paper.\n\n**Applying MoNE to dense prediction tasks**\n\nThank you for this suggestion. As dense prediction tasks typically operate on higher resolution images, we believe MoNE can offer further computational gains as the number of input tokens increases. Dense prediction tasks also provide stronger supervision at each pixel, which may help our model learn redundancy in the data. Leveraging MoNE for tackling redundancy in denser tasks would be an exciting direction for future work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VuL7uW7lPy", "reviewer": "Reviewer_YGXi", "summary": "This paper tried to use Matryoshka mechanism to assign tokens to different experts.", "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?\n2. 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?\n3. In figure 6b, do you mean that there is one router in the whole model? This is suprising. \n4. Seems like random router is only a bit worse than the learned one...", "limitations": "1. Important Tokens: is there a way to measure it?\n2. The method description part is not clear to me.\n3. 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.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. Seems like the proposed approach can learn some effective components in images, shown in visualizations .\n2. The empirical performance is good compared to mavit.", "weaknesses": "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 .\n2. Why is imagenet accuracy so low compared to a normal ViT? Such as Figure 3.\n3. Why is there no experiments on imagenet1k?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's positive feedback and are pleased they found our method produces informative visualizations and better results than MatViT.  We'll address their questions below.\n\n**Comparison with Single-Scale FF and Fine-tuning with Same Compute**\n\nWe appreciate the reviewer's suggestion regarding fine-tuning our framework from a MatViT model with the same compute. We exactly follow this procedure for videos, as detailed in Section 5, Lines 246-254. We utilize isoFLOPs training [24, 32] across all methods in Figure 4, meaning equal training FLOPs. While the x-axis in Figures 3, 4 depict the inference FLOPs consumed by the individual expert models (MatViT and MatViViT) and MoNE models with varying capacities, they are trained in isoFLOPs manner. We'll clarify this further by updating the figures and their captions.\n\nFor images, we show that MoNE can be trained from scratch for the same number of epochs as ViT (thus taking much lower training FLOPs), and still perform favorably to ViT and MatViT (Figure 3). In Figure 3a, we show that MoNE’s performance can be further improved by using the same training FLOPs as ViT (MoNE-isoFLOPs). Note that MatViT performs 4 forward passes per training step, while MoNE performs just one. Regarding comparison with a single scale, Figures 3 and 4 depict the performance of individual MatFormer models. We will highlight these points in the revision.\n\n**Low ImageNet Accuracy**\n\nThe numbers in Figure 3 correspond to validation accuracy on ImageNet-21k, and they are on-par with those reported in Fig. 4 of AugReg paper [38], with ~2x reduction in compute, for all three S, B and L models.\n\n**Experiments on ImageNet1k**\n\nWe primarily experiment on ImageNet-21k to showcase our model's effectiveness in large data regimes. These have been elaborately benchmarked earlier in the AugReg [38] paper. But, as ImageNet1k is a widely used benchmark for smaller scale ViT models, we compare MoNE with other adaptive ViT models in the **Baseline Comparisons** section of the Global Author Rebuttal.\n\n**Router Placement**\n\nOur method substantially differs from traditional MoE frameworks. Unlike MoE, where parameter count increases with routers in more layers, our method maintains a constant parameter count regardless of the number of routers. This methodological difference explains why MoE router settings don't directly apply to our framework.\n\nWe experiment with routers placement and present results in Figures 6a,b. Experiments suggest the best setup is a single router at the first transformer layer. While having multiple routers (expert choice) is the norm in the MoE literature, the MoNE setup is different from MoE. We further extend this discussion in the **Number of Routers** section of the Global Author Rebuttal.\n\n**Router Stability and its effectiveness in identifying to task difficulty**\n\nThe router is jointly trained with the model to optimize only the classification loss, and we find the training to be stable, as suggested by the high performance of the framework and relevant visualizations in the paper. Figure 1 and 7 show that MoNE can learn to assign important tokens to the higher sized models, and the redundant ones to the lower sized nested models.\n\nThe question raised about task difficulty is a really interesting one. While we fix a capacity and route tokens according to importance, as would be necessary for real-time applications, we do see the model is inherently able to learn the difficulty level of a specific data point. We discuss and visualize these results in the **Task Difficulty** section of the Global Author Rebuttal. These results validate that the router is able to associate compute requirements with task difficulty.\n\n**Random vs Learned Router**\n\nFigure 6c 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 [c], 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 Acc on Ti/16 trained on ImageNet1k). We will update the paper with these discussions.\n\n[c] Darcet et al., Vision transformers need registers, ICLR 2024\n\n**Measuring Token Importance**\n\nThe router leads to intuitively sensible decisions on a per-token basis, as depicted in Figures 1 and 7, on both images and videos. The router decisions correlate well with the important tokens in the image/video. In Figure 7a, we see that the relevant image regions are processed by the largest expert. In Figure 7b, we see that the tokens sent to the largest expert correlate well with the motion in the video. This is a qualitative measure of token importance, and the router decision logits can thus be taken as a quantitative measure as well, which the EPR algorithm does while assigning tokens to experts. \n\n**Comparison with Token Merging**\n\nSince TokenMerging algorithms can work on top of any ViT-like models, it is complementary to our method and can potentially be applied to a pre-trained MoNE model to further reduce computation. We empirically validate this claim in **Baselines Comparisons** in the Global Author Rebuttal and compare against other adaptive algorithms as well. Note that a naive implementation already performs well and this can be further improved by considering the MoNE architecture into account, as discussed. We will add this result to the revised manuscript.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FBgRT9QMHy", "reviewer": "Reviewer_LDxY", "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.", "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.\n\n* 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.\n\n* 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.\n\n* 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).\n\n* 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?\n\n* 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)\n\n* for each of the different FLOPs levels in the performance curves, what are the expert assignment ratios?\n\n* 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.\n\n* 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\n\n* 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\n\n* 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.\n\n\n\nSmaller comments:\n\n\n* 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)\n\n* 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\n\n* 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?\n\n* 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": "addressed in the final discussion section", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "This is a well-written paper that describes an in interesting and effective idea.  The approach is evaluated convincingly on three datasets for image and video classification tasks, and in enough settings to profile its performance characteristics and anecdotal visualizations of its behavior.", "weaknesses": "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.\n\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive feedback and glad to see that the reviewer found our work to be well-written, interesting and effective with comprehensive eval.\n\n**Projections and compute save. Projection from D/m to D for (QK)V.**\n\nIn a Transformer model, there are total 8 primary linear projection operations - \n(1) 3 to obtain $Q=W_qx$, $K=W_kx$, $V=W_vx$, \n(2) 2 to obtain $y=(QK^T)V$, \n(3) 1 in applying $z = W_oy$, \n(4) 2 in MLP: $W_2\\sigma(W_1x)$\n\nMoNE saves compute in all of these ops, except (2). In transformers, the weight matrices are of size $\\mathbb{R}^{d_a\\times d_b}$, where $d_a, d_b$ represent the feature dim (model dim, MLP dim, etc). In MoNE, these are either $\\mathbb{R}^{d_a\\times \\frac{d_b}{m}}$ or $\\mathbb{R}^{\\frac{d_a}{m}\\times d_b}$, depending on the op. $m$ denotes the reduction in the feature dim for the nested expert. Hence, in all projections, nested experts have FLOP gains by a factor of $m(>1)$. Denoting the model dimension as $D$, the features are extracted from $D$ to $\\frac{D}{m}$. For eg, for the nested model with dim $\\frac{D}{m}$, in (1) the input $x$ is $\\frac{D}{m}$ which is projected to $D$. For that, we slice a weight matrix of dim $\\mathbb{R}^{D\\times \\frac{D}{m}}$ from the full matrix of $\\mathbb{R}^{D\\times D}$, and multiply it with $x$, as explained in Lines 93-94. Hence we save by a factor of $m$, by choosing the nested matrix instead of the full matrix, for which the input would be $D$. Similar explanations for (3) and (4). These operations are discussed in Appendix A.1.\n\n**Layer Normalization**\n\nAs in Line 98, after linear projections, the features are padded to D, to add with the residual features, and LayerNorm on full D dim. In the next layer, a part of the feature $\\frac{D}{m}$ is sent and weight slicing operations are applied to project, as discussed above.\n\n**Leveraging redundancy**\n\nBy redundant, we mean that the info of a certain token set is redundant/less needed, given a few of important/needed ones, as measured by final accuracy for a given compute. Our router does learn to prioritize features like edges for the largest expert, confirming the reviewer's hypothesis. Such features can be learned by low level convolutional filters. Tasks which need to compress redundancy at a higher level might need a router deeper into the network. Additionally, ViTs are known [c] to shuffle info among tokens, acting as an intrinsic router, thus adapting itself given the initial router decisions.\n\n[c] Darcet et al., Vision transformers need registers, ICLR 2024\n\n**Eq. 3 relation between dim and capacity**\n\nYes, the last line of the reviewer for this point is indeed correct. In the MLP dimension, the full model would do two projections: $D \\rightarrow 4D$ and $4D \\rightarrow D$ amounting to 2 * 4D * D compute. The nested model would do $\\frac{D}{m} \\rightarrow 4D$ and $4D \\rightarrow \\frac{D}{m}$, hence amounting to 2 * 4D * $\\frac{D}{m}$ compute, a $m$ fold reduction.\n\n**Eq. 3 Rationale**\n\nThe weights $\\delta^i$ are exponential to account for the exponentially distributed accuracy differences (Fig 3, 4) between consecutive MatViT models. While this is empirical, this first term is as a proxy to accuracy obtained by a particular expert. The larger experts are thus favored more. However, in without the entropy term, this objective would lead to a greedy assignment of either the largest or the smallest expert. This will not leverage the framework’s full potential. The entropy term balances this, so that each expert is provided a significant capacity. We empirically reason some of our choices in Table 1 of paper, showing that this heuristic leads to competitive results across capacities.\n\n**FLOPs to Expert Capacity**\n\nThe expert assignment ratio, $c_i,i\\in[1,E]$ for expert $i$ is obtained by optimizing Eqn 3 for a given capacity $e_c$. Given FLOP needs, $e_c$ can be obtained by solving: GIVEN_FLOPS = FLOP_b + e_c * FLOP_a where FLOP_b corresponds to the model operations that are not reduced by MoNE (patch embedding, attention value ($(QK)^TV$) computation), and FLOP_a correspond to the linear projections.\n\n**Single expert accuracy**\n\nThe individual MatFormer expert accuracy are in Fig 3 and 4 as MatViT/MatViViT.\n\n**Zero-capacity expert**\n\nThis is close to skipping entire layers for some tokens in Mixture of Depth (MoD). We do compare with MoD in Fig 3, and find MoNE to perform better.\n\n**Runtime/latency**\n\nWe agree that latency is heavily implementation dependent. Our high level Jax implementation resulted in latency gains scaling linearly to FLOP gains. We present this in the **Latency/Throughput** part of Global Author Rebuttal. The largest computations that bottleneck throughput can be resolved by efficient implementation. The ops are tensor products of different sizes, which can be decomposed into smaller sub-tensor products and efficiently parallelized (Fig 3 [14]), avoiding bottlenecks. We believe that further latency improvements can be obtained by low-level implementations.\n\n**Sec 4.4 videos**\n\nOur exploration for MoNE in videos is a first step based on the popular factorized-encoder ViViT [2] model, for which the gains in the spatial dimension are more readily realized, given that temporal information here is integrated late in the architecture so a router has 196*16 choices spatially vs. 16 temporally (over the 16 frame-wise CLS tokens). We will expand and clarify this discussion further; integrations with other space-time video architectures with more interleaved operations is an exciting direction for future work.\n\n**Smaller comments:**\n\nWe thank the reviewer for noting the small comments. The reviewer's understanding is correct. We will clarify them in the paper. Yes, remaining tokens due to integer packing are not dropped but assigned to the smallest expert.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3ECKtuNenQ", "reviewer": "Reviewer_R1sX", "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.\nThe proposed method is then evaluated on image and video classification.", "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) ?\n* 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": "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.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The routing strategy is learned in the first transformer layer and propagated to all subsequent layers. This can be better from a hardware perspective to know early on which channels need to be turned on/off.\n- Having experts with different compute costs/\"capacity\" is well motivated and reasonable to improve the accuracy/efficiency trade-off.\n- Good ablation experiments on the location of the router", "weaknesses": "* **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.\n\n* **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:\n\t  * 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.\n\t  * 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.\n\tIn addition, some of the metrics reported in the paper could be improved:\n\t  * As far as I know, ImageNet-21k evaluation is not very standard as there is no official train-validation split.\n\t  * 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.\n\n* **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\n\n\n\n#### references\n  * (1) Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time, Liu et al, 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful comments. We appreciate the recognition of MoNE to have hardware-efficient routing design, the use of compute-aware experts to optimize the accuracy/efficiency trade-off, and the thorough ablation study demonstrating the impact of router placement. Below we answer some of the questions that the reviewer raised.\n\n**Fixed number of experts**\n\nThis is a very astute point. Since the MatViT (MatFormer [16]) experiments suggest that the nested dimensionalities interpolate smoothly, it is fair to assume that MoNE is in-fact not training just a fixed number of experts (set to 4), but everything between the smallest and the largest experts. MoNE uses a fixed number of nested blocks to develop a tractable framework for capacity allocation and routing. As the reviewer mentioned, there is nothing stopping from extending this to all the interpolated dimensionalities (D) of nested experts barring the added complexity of precise capacity allocation and load balancing, which is generally non-trivial. The design choice to route between 4 models also relies on potential better serving capabilities given the individual block sizes. In short, yes, we agree that we can do MoNE with all the intermediate nested experts (not just the 4) but prototyped with the current design choice to show the benefits to begin with.\n\nWe thank the reviewer for pointing to the dynamic sparsity work, we shall include them in the revised version of the manuscript. Compared to dynamic sparsity, MoNE and MatViT may offer a structural advantage, potentially leading to improved system efficiency by organizing sparsity in a more organized manner. As evident in the **Throughput/Latency** section in the Global Author Rebuttal, even with very high level implementation, the FLOP gains directly translate to throughput and latency gains.\n\n**Strengthening experimental evaluation with dynamic routing baselines and latency**\n\nMoNE, 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. In contrast, MoNE reduces computation while keeping the same parameter space. Given the constant parameter space, it allows for complementary use of other MoE methods together. We discuss this point further in the **MoE Comparison** section in the Global Author Rebuttal.\n\nWe compare against MoD because the method is similar in spirit to ours, conditionally computing the inputs based on complexity. We extensively experimented with MoD and found that the settings mentioned in MoD translate best to these vision datasets as well. \n\nThe reviewer correctly points out that for classification, dynamic routing and token pruning methods are valid benchmarks. We therefore compare against some of these methods, and report results in the **Baseline Comparisons** section in the Global Author Rebuttal. This shows that MoNE can offer superior performance compared to other dynamic routing methods. \n\nWe 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]. Additionally, we show results on ImageNet-1K as well compared to other efficient architectures for smaller scale ViT models in the **Baseline Comparisons** section in Global Author Rebuttal.\n\nThe reviewer correctly points out that FLOPs gains may not always translate to latency/throughput gains. Therefore, we present latency and throughput results as compared to baseline models in the **Throughput/Latency** section of the Global Author Rebuttal. We also present the variation of latency vs FLOPs, at different model capacities and observe that latency gains scale close to linear to FLOP gains. Also note that the”pad” mentioned in Figure 2a do not actually result in any additional computation, and can be avoided in implementation.\n\n**Significance of dynamic budget allocation**\n\nIt is true that the uniform allocation performs as well as our heuristic method, but at a higher capacity requirement. More importantly, a uniform allocation $c_i=\\frac{1}{4}, \\forall i$ leads to a model with a fixed capacity based on the individual compute requirements of the experts, i.e., $e_c=\\frac{1}{4}(\\frac{1}{8}+\\frac{1}{4}+\\frac{1}{2}+1)=0.47$. In contrast, the proposed heuristic\nallows for flexible choice of model capacity ($e_c$) as per compute requirements. As discussed in the paper, given a user specified $e_c$, many solutions exist for $c_i$, and it is non-trivial to choose one over the other. The heuristics add constraints to reach a solution which offers promising results. The high performance of MoNE across capacities is depicted in Figures 3 and 4. \n\n**Clarification on input for computing r_i in Eqn (1)**\n\nThe router logits $r_i$ are computed on the block input $x_i$, because unlike traditional MoE architectures, MoNE leverages the nested structure not just in the FFN layer, but also in the self-attention, to further increase compute gains. The exact changes to these operations are briefly described in the Appendix A.1 (Eqn. 5 and 6).\n\n**Padding to dimension D**\n\nYes, padding is only to make sure two different dimensional vectors can be added. 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.\n\n**Limitations**\n\nWe have addressed the limitations mentioned by the reviewer in this discussion and in the Global Author Rebuttal. We will include a more elaborate discussion of our method and areas for future work in our final revision.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
**Latency/Throughput**

We present the latency/throughput gains of MoNE compared to baselines here, in addition to the FLOP gains mentioned in the paper. In the table below, we show absolute wall clock times and throughput for MoNE compared to a baseline ViViT model, on a single V100 GPU, achieving ~2x improvement in both FLOPs as well as runtime, whilst maintaining accuracy. We additionally show the variation of latency and throughput with FLOPs for varying model capacities $e_c$ of MoNE in Fig 1, 2 (attached pdf). The plots show that latency and throughput gains scales close to linearly with FLOP gains. Inference gains depend heavily on implementation. A simple high-level efficient implementation of our framework yields gains of this scale, and we believe that further improvements can be obtained by optimizing a low-level GPU kernel implementation for MoNE.

|Method|FLOPs (G)|Throughput (clips/sec)|Latency (ms)|Top-1 Acc|
|-|-|-|-|-|
| ViViT-FE-B/16| 376| 15.8| 129.2| 64.4|
| MoNE $e_c=0.3$|**162**| **30.7**| **65.5**|**64.6**|

**MoE Comparison**

Our framework MoNE, unlike traditional MoEs, does not increase the parameter space. Traditional MoEs, like Sparse VMoE, route inputs in each layer to one out of k independent experts (typically the FFN block), each having the same parameter footprint, thus increasing the parameter space k-fold for the expert blocks. On the other hand, independent MoNE blocks can potentially be used as experts in the MoE framework.  Therefore, our work _complements_ VMoE and other similar works. MoNE in this paper acts as an in-place replacement for a dense model (ViT), and hence all our comparisons maintain the same parameter space. VMoE frameworks show cross-scale results at the expense of increased parameter space (e.g., equivalent performance of VMoE-L/16 to ViT-H/14 in Table 2 in [34], and similar cross-scale comparisons in Figs. 4 to 8 in [1]). MoNE, in contrast, matches baseline performance with limited inference compute while working with the same parameter space.

**Baseline Comparisons**

We compare MoNE with more baselines, particularly with adaptive computation of dense models, as shown in the table below. We performed this experiment on ImageNet1k with a Ti/16 sized model. ACT, PonderNet, DepthAdapt, A-ViT are works with similar motivation of input adaptivity as MoNE, and MoNE shows superior performance. Latency gains on bigger models e.g., ViT-B in our paper) are even higher, as also observed in literature [7].

| Method|GFLOPs|Throughput (img/sec)|Top-1 Acc (ImageNet1k) |
|-|-|-|-|
| ViT|1.3|3410|71.3|
| PonderNet [a, 18]| 1.0| - | 66.2|
| Depth Adapt [18] | 1.1| - | 69.4|
| ACT [22]| 1.0| - | 71.0 |
| A-ViT [47]| 0.8| - | 71.0 |
| MoNE (Ours)| **0.8**| **4333**| **71.4**|

Additionally, in Fig. 3 (attached PDF), we apply Token Merging (ToMe) [7] on top of the MoNE style ViT-Ti/16 model trained on ImageNet1k. We trained a model with full capacity till layer 3 and a router with $e_c=0.5$ after that. We applied ToMe 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. Our preliminary results demonstrate that this implementation improves performance compared to ToMe on ViT, and this can be made better by extending this approach to all MoNE layers, applying it to distinct sets of nested tokens. This indicates that ToMe is complementary to MoNE. We will add these results to the revised manuscript.

[a] Banino et. al, PonderNet: Learning to Ponder, arXiv:2107.05407, 2021.

**Number of Routers**

Note that the number of routers does not have the same implications in MoNE as compared to traditional MoE. 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, in Sparse VMoE [34] (Table 2, 8), 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 due to two reasons: (1) it brings in additional optimization challenges (as prevalent in the MoE literature as well [34]), (2) when we reassign a token from a smaller to bigger nested expert, its information content may be bounded by the representation power of the smaller expert, hence not improving performance. The opposite may happen 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. We will update the revision with this discussion.

**Task Difficulty based on router decisions**

We visualize the inputs that require the least and most compute per the router decisions, in a setting without capacity constraints in Fig. 4 (attached PDF). This is to understand if the router decisions correlate with task difficulty, i.e., send harder samples to larger experts. Specifically, instead of performing token to nested expert assignment using the EPR algorithm in the paper, we use the router decisions as is by just an argmax. Using this, we get an estimate of the computation that the router wants to assign to an image. The two sets of images indicate the top-3 images that demand the lowest and the highest compute, according to the router decisions. We observe that the images demanding less compute are visually simple, while the ones demanding highest compute are relatively complex.

Meta-review: 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.
Decision: Accept (poster)