Mixture of Nested Experts: Adaptive Processing of Visual Tokens

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.

Fixed number of experts

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

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

Strengthening experimental evaluation with dynamic routing baselines and latency

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

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

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

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

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

Significance of dynamic budget allocation

It 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 allows 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.

Clarification on input for computing r_i in Eqn (1)

The 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).

Padding to dimension D

Yes, 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.

Limitations

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

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.

Projections and compute save. Projection from D/m to D for (QK)V.

In a Transformer model, there are total 8 primary linear projection operations - (1) 3 to obtain $Q=W_qx$, $K=W_kx$, $V=W_vx$, (2) 2 to obtain $y=(QK^T)V$, (3) 1 in applying $z = W_oy$, (4) 2 in MLP: $W_2\sigma(W_1x)$

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

Layer Normalization

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

Leveraging redundancy

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

[c] Darcet et al., Vision transformers need registers, ICLR 2024

Eq. 3 relation between dim and capacity

Yes, 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.

Eq. 3 Rationale

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

FLOPs to Expert Capacity

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

Single expert accuracy

The individual MatFormer expert accuracy are in Fig 3 and 4 as MatViT/MatViViT.

Zero-capacity expert

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

Runtime/latency

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

Sec 4.4 videos

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

Smaller comments:

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

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.

Comparison with Single-Scale FF and Fine-tuning with Same Compute

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

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

Low ImageNet Accuracy

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

Experiments on ImageNet1k

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

Router Placement

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

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

Router Stability and its effectiveness in identifying to task difficulty

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

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

Random vs Learned Router

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

[c] Darcet et al., Vision transformers need registers, ICLR 2024

Measuring Token Importance

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

Comparison with Token Merging

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

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.

Actual inference speed comparisons (Throughput/Latency)

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

EPR algorithm implementation

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

Applying MoNE to dense prediction tasks

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

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.

Dynamic Routing vs MatFormer and the need for MoE comparisons

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

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

Clarifying the Advantages of Nested Experts and Comparison with Dynamic Token Processing Approaches

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

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

Understanding Expert Specialization: Contrasting MoNE and Conventional Sparse MoE Architectures

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

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

[b] A Jiang et al. Mixtral of Experts. arXiv:2401.04088, 2024.

Understanding the Impact of the Number of Routers in MoNE

This is a very important point raised by the reviewer and we clarify in the Number of Routers section of the Global Author Rebuttal.

Dear Reviewers and Area Chairs

We wholeheartedly thank you for your thoughtful reviews and the time you've dedicated to improving our work. We've carefully addressed all the points raised and submitted our rebuttal. We would greatly appreciate any further feedback or confirmation on our responses.