Tight Clusters Make Specialized Experts
Authors: Stefan Nielsen
Rachel Teo
Laziz Abdullaev
Tan Minh Nguyen
Venue: ICLR 2025 Poster (2025)
Keywords: Mixture of Experts
robustness
clustering
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose a novel Mixture-of-Experts routing method that computes token-expert assignments in a transformed space that promotes separation of latent clusters in the data and more easily identifies the best-matched expert for each token.

Abstract: Sparse Mixture-of-Experts (MoE) architectures have emerged as a promising approach to decoupling model capacity from computational cost. At the core of the MoE model is the router, which learns the underlying clustering structure of the input distribution in order to send input tokens to appropriate experts. However, latent clusters may be unidentifiable in high dimension, which causes slow convergence, susceptibility to data contamination, and overall degraded representations as the router is unable to perform appropriate token-expert matching. We examine the router through the lens of clustering optimization and derive optimal feature weights that maximally identify the latent clusters. We use these weights to compute the token-expert routing assignments in an adaptively transformed space that promotes well-separated clusters, which helps identify the best-matched expert for each token. In particular, for each expert cluster, we compute a set of weights that scales features according to whether that expert clusters tightly along that feature. We term this novel router the Adaptive Clustering (AC) router. Our AC router enables the MoE model to obtain three connected benefits: 1) faster convergence, 2) better robustness to data corruption, and 3) overall performance improvement, as experts are specialized in semantically distinct regions of the input space. We empirically demonstrate the advantages of our AC router over baseline routing methods when applied on a variety of MoE backbones for language modeling and image recognition tasks in both clean and corrupted settings.

Published as a conference paper at ICLR 2025
TIGHT CLUSTERS MAKE SPECIALIZED EXPERTS
Stefan K. Nielsen" FPT Software AI Center stefannvkp fpt com com
Rachel S.Y. Teo" Department of Mathematics National University of Singapore rachel. tsyuu. nus edu
Laziz U. Abdullaev
Tan M. Nguyen
Department of Mathematics National University of Singapore
Department of Mathematics National University of Singapore tanmn nus. edu. sg
laziz. abdullaevee nus edu
ABSTRACT
Sparse Mixture-of of of-eperrss- (MoE) architectures have emerged promising ap- proach to decoupling model capacity from computational cost. At the core of the MoE model is the router, which learns the underlying clustering structure of the input distribution in order to send input tokens to appropriate experts. However, latent clusters may be unidentifiable high dimension, which causes slow conver- gence, susceptibility to data contamination, and overall degraded representations as the router unable to perform appropriate token- expert matching. We examine the router through the lens of clustering optimization and derive optimal feature weights that maximally identify the latent clusters. We use these weights to com- pute the token-expert routing assignments in an adaptively transformed space that promotes well-separated clusters, which helps identify the best-matched expert for each token. In particular, for each expert cluster, compute set of weights that scales features according to whether that expert clusters tightly along that feature. We term this novel router the Adaptive Clustering (AC) router. Our AC router enables the MoE model to obtain three connected benefits: 1) faster convergence, 2) better robustness data corruption, and 3) overall performance improvement, as experts are specialized in semantically distinct regions of the input space. We empirically demonstrate the advantages of our AC router over baseline routing methods when applied on variety of MoE backbones for language modeling and image recognition tasks both clean and corrupted settings. The code publicly available at https: / //g. thub. com/ stefvk/ ACMoE.
INTRODUCTION
Scaling up model capacity continues to deliver substantial performance gains across a wide range of tasks, with particularly impressive results in visual representation learning and language modeling (Alexey, 2020; Bao et al., 2021; Radford et al., 2019; Raffel et al., 2020; Nguyen et al., 2023). How- ever, larger models incur growing computational costs, prompting increasing research into Sparse Mixture-of-Experts models (MoE), which offers promising avenue to balancing model scale with efficiency by activating only sub-modules, termed experts, of the network during training and in- ference (Shazeer et al., 2017; Fedus et al., 2022; Lepikhin et al., 2020; Nguyen et al., 2025). This approach has been shown to achieve better performance than dense models with nearly constant computational overhead on tasks from speech recognition, image recognition, machine translation, and language modeling (Riquelme et al., 2021; Kumatani al., 2021; Lepikhin et al., 2020; Teo & Nguyen, 2025a).
At the core of the MoE layer is the learned router which assigns inputs to the relevant experts. The router must learn segment the input space appropriately such that inputs and experts are well matched, enabling the experts to be trained on semantically similar data. This expert specialization
Co-first authors. Please correspond to: stefannvkp fpt.com and tanmn nds.edu.sg

---

Published as a conference paper at ICLR 2025
Topur ACMoF.
Txput
ACMaF (Ouns)
Figure 1: ACMoE discovers semantically distinct regions. We show 14x14 image reconstructions where patches are colored by assigned experts. Top row: Swin assigns large chunks of foreground and background to one expert (red), while ACMoF accurately discovers the bird and relevant foreground. Bottom row: When the background and foreground are hard to distinguish, Swin's 's router fails to register the stingray (left) or shark (right) and allocates one expert for virtually the entire image. ACMOE, however, discovers the semantically dis- tinct regions, using one expert (green) specialize the foreground and different experts for the background.
allows MoE models to produce better representations than their dense counterparts while activating only fraction of the total parameters. Recently, various methods have been proposed to find optimal expert-toke tookn matches, including linear programs (Lewis et al., 2021), cosine similarity- (Chi et al., 2022), soft assignments via convex combinations of inputs (Puigcerver et al., 2023), and both top-k experts per token (Shazeer et al., 2017) and top-k tokens per expert (Zhou et al., 2022b). We note that the above approaches fundamentally rely on dot-products between inputs and experts learn the corresponding assignment, which might be suboptimal in cases where the semantic regions are not easily discoverable in the nigh-dimensional feature space. Typically, we expect that the true underlying clusters present in the data will cluster on different, potentially disjoint, subsets of features, and may not be discoverable when using the full feature set. This phenomenon can lead to slow convergence as the experts are unable to specialize on semantically similar regions of the data, poor robustness as data contamination can spuriously assign inputs to unsuitable experts, and degraded overall downstream performance due to suboptimal input-cxpert matching.
rules
Contribution. In this work, propose the Adaptive Clustering (AC) router and corresponding Adaptive Clustering Mfxture-of-Experts (ACMoE), novel MoE method in which the router com- putes token-expert assignments in transformed space that maximally identifies latent clusters in the data and more easily discovers the best-matched expert for each token. More specifically, we adaptively learn for each input which features best determine its cluster assignment and scale its features accordingly such that features that promote tight expert clusters are upweighted, and fea- tures that produce dispersed expert clusters are downweighted. This transformation accentuates the relevant characteristics of each input according to the specialization of the experts, thereby allowing the router to more easily discover the optimal input-expert allocation. Computing the routing assign- ments following this scheme produces three benefits: 1) faster convergence as experts are able to specialize more quickly by being allocated semantically similar inputs, 2) better robustness as latent clusters are better separated, thereby minimizing the risk that data corruption erroneously assigns tokens to unsuitable experts, and 3) betler overull representations and downstream performance due to improved expert specialization. In order to discover the key features per token and their corre- sponding weights, we present a feature-weighted clustering optimization perspective on the MoE framework and demonstrate how the clustering solution obtains the required feature weights. We show how these weights can be integrated into the routing mechanism such that routing takes place in cluster-adaptive transformed space. We theoretically prove that our proposed routing mechanism learns the latent clustering structure of the data faster than standard routing mechanisms and that our mechanism is more robust to data contamination. Furthermore, our proposed method involves no learnable parameters and can be computed highly efficiently. In summary, our contributions are fhree-fold:
1. We develop the novel Adaptive Clustering router, routing method in MoE architectures that computes token-expert assignments in a transformed space that promotes separation of latent clusters in the data and more easily identifies the best-matched expert for each token.
2

---

Published as a conference paper at ICLR 2025
Text PML
GLaM
38.27
4.0
Switch Transformer
03.02
Validation Loss
GLaM-ACMoE (Ours)
36.26
3.5
Switch ACMol (Ours)
86.01
3.0
2.5
2.0
1.5
1.0
0.5
100
120
30
Epochs
Epochs
Figure 2: Fast Convergence of ACMoE. Left: Convergence speed on WikiText-105 pretraining using the Generalist Language Model (Du et al., 2022) backbone. Right: Convergence speed on Banking-77 finetuning using the Switch Transformer (Fedus al., 2022) backbone. Across both backbones and tasks, we observe substantially faster convergence. We display final test perplexity (PPL.) and accuracy (Acc.), showing better overall performance well.
We propose a feature-weighte clustering optimization perspective on token-expert assign- ment and derive the optimal feature weights for adaptively transforming the input data for routing.
3. We derive theoretical framework demonstrating how MoE robustness and convergence de- pend on the shape of latent clusters and the clustering geometry of the input space.
We empirically demonstrate that 1) the Adaptive Clustering router outperforms baseline routing methods in MoE architectures in large-scale tasks such as WikiText-100 language modeling and downstream finetuning, and ImageNet- k object classification in both clean and contaminated set- tings, 2) the Adaptive Clustering router exhibits faster convergence than baseline methods, and 3) the Adaptive Clustering router attains these performance improvements for free that is, with no learnable parameters and negligible computational overhead.
Preliminaries. We consider Transformer (Vaswani, 2017) based MoE architectures and follow the approach of previous work where the MoE layer is inserted after the self-attention layer within the Transformer, replacing the traditional feed-forward network (Fedus et al., 2022; Du et al., 2022; Liu al., 2021). Let æ be an input token with hidden representation the N learnable expert embeddings for model hidden dimension d. The MoE layer selecting the top experts is described by the following equations:
RR and e1, 62, ....... eN Rd be
K = tooopp, (sk) (h' ex)
SMoB (h) h. Σ (FFF))
keR.
where JFFN is the th expert forward network, 8s hTe is the similarity score between token representation and the Wthh expert eş and 9(-) gating function often chosen as softmax, g(sx) exp(*x)/ Ejex exp(s;). We refer to Eqn. as the router, which learns the top k best matched experts per token, and Eqn. 2 as the overall standard MoF layer.
ej
Organization. We structure this paper as follows: In Section 2, we present clustering optimization problem and show that its solution adaptively scales the feature space according to which dimensions promote tight clustering. In Section 3, we present how the solution to clustering optimization problem can built into our proposed AC router and we provide the full technical formulation of AC routing and Adaptive Clustering Mixture-of of - (ACMoE). We then present theoretical propositions on faster convergence and robustness. We empirically validate the advantages of AC- MoE in Section 4 and discuss related work in Section 5. We end with concluding remarks and future work in Section 6. Proofs, technical details, and further experiments are provided in the Appendix.
2 A CLUSTERING OPTIMIZATION PERSPECTIVE
We begin by examining the MoE router through the lens of feature-weighted clustering (Witten & Tibshirani, 2010; Friedman & Meulman, 2004; Brusco & Cradit, 2001; Gnanadesikan et al.,
3

---

Published as a conference paper at ICLR 2025
1995). We explicitly model the router's task as learning token assignment that groups together similar tokens. We consider the role learnable feature weights in solving a clustering optimization problem to optimally reveal latent clusters and present analytical solution for the optimal weights for any given routing assignment. We finally discuss how this solution improves the MoR router before providing the full formulation of our AC router and ACMoE in the next section.
2.1 CLUSTERING OPTIMIZATION
Let h; [la..... h be the 5th hidden representation and Dij denote the distance between h; and hj. Given distance metric Pijo between hig and hig over the 4th dimension, the distance between h; and h can be defined as Dis(w) = Ege[d] WyPija for weights w = [wi,..., wa] with Ege[d] Wy = and Wa O for all q [d]. The weights determine the global importance of the qth feature the overall distance among representations.
Cluster analysis aims to divide the input set of N objects into groups, where objects within the same group are more similar to each other than to those other groups. This formalized using a classifier r(i) = ki, assigning the 7t object to a group k. Then the optimal classifier r" minimizes criterion Q(r) that evaluates clustering quality:
arg min Q(1)
Σ kE[E]
Σ E r(i)-kr(j)-k
Σ Di((w) Di(ww.
(3)
N
We expect that different groupings will cluster different subsets of features. In particular, wish to model the scenario that groupings exist different latent subspaces with varying dependence on possibly disjoint subsets of features. We therefore replace the global feature weight ws in Eqn. 3 with cluster-dependent feature weights, { (w))es for E groups, which allows us to capture the differing feature dependencies of each cluster. Then, we can adapt the optimization problem with these cluster-dependent feature weights as follows:
(r", (wi)Ra) arg min E Σ Σ ('w())), (.(w [E[ NO
(())
such that Wigk = 1, VA [E], [c[d]
E
where D (wx) Ei= WgkPiij (wx) denotes the weighted distance between and combined with some regularization J and regularization strength X.
To avoid toint-mass solutions in which we assign all weight the single best-clustering feature, set the regularizer to the bullback-Leiblee divergence between the feature weights w and the uniform distribution = (1/d,... 1/d) æ denoted by J(w)) D (u wx). The regularization parameter À reflects our preference to maintain more or less features in the solution set.
2.2
MoE AS CLUSTERING OPTIMIZATION
Within the MoE framework with learnable routing, the router performs the role of the classifier T + [E]. which is learned via gradient descent optimize the final output loss'. Therefore, modify Eqn. 4 by fixing r and focusing just on optimizing the criterion with respect to cluster-wise feature weights wx. Under this interpretation, the router learns via backpropagation optimally allocate representations to experts, with representations adaptively transformed to maximally reveal the clustering structure of the input data. Eqn. 4 then becomes
{ 1 = arg min Σ (Wa) kc[
Σ Σ (ws),
N2
-())kr())-k
such that
Σ Wab = & [E].
(5)
ge[d]]
The following theorem presents the optimal weights per feature g and cluster k:
A top-k router straightforwardly be the classifier in Eqn. 4 Rd [E]]
4

---

Published as a conference paper at ICLR 2025
N Er(s)=k Er( Er(j)=kPij be measure of dis- persion on the t feature for the representations assigned to cluster Then, for given router function : I [E], the corresponding optimal weights (e that minimize the feature- weighted clustering optimization problem in Eqn. 5 are given by
Theorem 1 (Optimal feature weights). Let Syk := N
A/d
Wykk
(6)
Sqk + Olky
for (q, ()) [d] [F], where {ax] are constants that for any X satisfy
d
Σ [c[d] 8yk
O X
The existence of ag satisfying Eqn. 7 and the proof of Theorem 1 is provided in Appendix A.1. The optimal weights for a cluster k given in Eqn. 6 take an intuitive form in that they are inversely proportional the measure of dispersion in cluster & along each dimension, Wh [i .......... sus Hence, the optimal cluster-wise feature weights scale features according their contribution to forming tight clusters. Specifically, the solution weights upweight a feature if cluster & clusters tightly (has small dispersion Syk) along the feature g and downweights a feature p if cluster 1: clusters loosely (has large dispersion Spk) along feature p.
This method enables the MoE router to perform better expert matching. The cluster wise feature weights we capture the features on which the hthh expert is specialized, as large weights indicate those features are highly important to the identification of that expert cluster and small weights indicate those features are unimportant to identification of that expert cluster. Then, we can use wy to scale the tokens to accentuate their features according to the specialization of the experts, thereby allowing the router to best identify the most suitable expert for each token. Note that this solution is local in that we learn the optimal weights adaptively per cluster, obtaining wa for all l e [E]. and we compute unique scaling of the feature space adaptively per cluster as well. Integrating these cluster-dependent weights which scale the feature space according to the identification of each expert into the MoE router obtains our AC routing method and corresponding ACMoE. We detail the AC router and ACMoE fully in the next section.
3 A TIGHT CLUSTER IS A SPECIALIZED EXPERT
In this section, we demonstrate how we implement the solution weights from the clustering op- timization problem in Eqn. into the MoE routing mechanism, thereby obtaining the Adaptive Clustering router. We then provide the full technical formulation of our proposed routing method and corresponding ACMoE model. We also present theoretical results on how computing the routing assignments according to our framework promotes faster convergence and robustness.
3.1 FULL TECHNICAL FORMULATION
We integrate the weights from Eqn. into the Adaptive Clustering router transformation in Defini- tion 1 which, for cluster k scales the dimensions of the feature space according the 1eth expert's specialization on those features. Formally this is:
Definition 1 (Adaptive Clustering Router Transformation M₃). Let Cf. (hi,... n be the representations assigned to expert k at layer f. Let syk E F. be a measure of a spread in the qtdh dimension for cluster k, such as mean absolute deviation syy No Evect [hi, -hyl- Then, the cluster-dependent router transformation for expert k at layer é is given by diagonal matrix M := diag(1/si 1/s
We use the transformation Mj in Definition I to adaptively scale the feature space in which we perform token-expert matching. This obtains our Adaptive Clustering router and corresponding ACMoE layer, described in the following definition.
Definition (Adaptive Clustering Router and MoE Layer). Let h Idd be the hidden representation of input, ei...., ey y Rdd be expert embeddings at layer f. Let h. h-- Cf-1 have been assigned
5

---

Implementation details. Given ACMoE requires the expert assignment from the previous layer to compute the routing assignment (Egn. 8). ACMoE is only implementable after the first layer. Furthermore, we scale the measures of dispersion in M = ( 1/8 to have mean 1. Published as a conference paper at ICLR 2025
expert k' in the previous layer. Let M E &Radd be the Adaptive Clustering transformation
(Definition 1) for imput h at layer 2-1. 1. Let g(-) be the softmax function. Then the following equations describe the Adaptive Clustering router (Eqn. 8) and overall ACMoE layer (Eqn. 9):
== ((x) topk; (h M ( e
(8)
FACMOE (n') Σ g(h M f S (n kee
(9)
We see from Eqns. 8 and 9 that the standard MoE layer is recovered by setting the AC router transformation to the identity matrix, M2- Ia for all lo E [E]. Within our framework then, standard routing schemes implicitly assume all experts k E [E] depend equally on all dimensions.
Remark 1.
Remark 2. The Adaptive Clustering router computes a dot dot-product between h and experts eş with the dimensions scaled by the weights in M& and so is proportional to Mahalanobis distance. Under this interpretation, we soft project the tokens and expert embeddings onto the axes of the feature space that best identify the expert cluster k".
This is remove the effect of different clusters features having different absolute magnitudes. Our method is concerned with identifying the key sets of features that contribute more or less to identification of the expert clusters, and so we wish compute our scaling in relative sense.
3.2 ADAPTIVE CLUSTERING PROMOTES ROBUSTNESS AND FAST CONVERGENCE
We now present theoretical propositions on the improved robustness and convergence speed of our method. The robustness of our method follows from better separation of expert-clusters. This pro- duces more stable assignment in which the probability of erroncously sending token to unsuitable nearby experts decays cxponentially with increased inter-cluster distance. Faster convergence fol- lows from our AC routing method improving the conditioning the Hessian of the loss with respect to the expert embeddings, enabling faster and more stable convergence of the router.
Promoting robustness. We begin with Lemma stating that our AC transformation (Definition 1) increases the separation between clusters in the transformed space, followed by Lemma 2, which provides an explicit expression for the probability of incorrect expert assignment. To give the prob- ability bound an exact form, we assume the cluster structure can be modeled as a Gaussian mixture model (GMM). We note that GMMs are a highly expressive and general framework, so this assump- tion does not place sigmificant restrictions on our robustness analysis. We further assume that though clusters may overlap, they are well-separated along the features for which they cluster tightly².
Lemma 1 (Adaptive Clustering Router Transformation Increases Cluster Separation). Let the data be generated from a Gaussian mixture model with components, 9e N (pe: E.) c E [F]. Without {a, b) where token representation h 9o. belongs to cluster a. Let M. diag(1/81ar- sia...., 1/saa) be the router transformation constructed from the feature-wise dispersions, Sqa. of cluster 9a for each feature q [d] as given by Definition 1. Then the distance between cluster means in the M. transsoorrd space, defined as || -Ma Mall 1 M,," :=
loss of generality, consider two expert clusters
(pu-pa.) Ma, (1u--po) is larger than in the original Euclidean space: (ux-pall Mr. 2 || p
The proof is provided in Appendix A.2. In Lemma 2, we derive the probability of mis aasigmetnnmennt as a function of inter-cluster distance, showing how separation mitigates the effect of perturbations.
Lemma 2 (Incorrect Assignment Probability). Let h ~ NN (E Σi> be representation belong- ing to cluster k* Let h' h + e. be contaminated by some O-mean noise e.- ~ E.). Let le be the nearest, incorrect cluster to k*. Let the inter-cluster mean distance between k. and k be given by |8μ1| == ||p -. Let the routing assignment be given by r R2d + [E] and denote the cumulative density of standard normal distribution by d.. Then the probability of incorrect assignment is:
1
Pr(r(h') +
-1-) ((1()>==)=11--)) 2/6 (En+ E.jōy
(10)
Intuitivelyy this assumption captures the natural property that the semantic regions of the input space distinct along the dimensions that best identify them.

---

(22x))-1e for large and Published as a conference paper at ICLR 2025
Remark 3. It is worth noting that since 1 - (x)
Voja (Ex+ Egjōp O(1 we find that the probability of incorrect cluster assignment as given by Eqn. 10, Pr(r(h') = is exponentially decreasing function in ||ôμي)).
=006bu))))
The proof is provided in Appendix A.2. Combining Lemmas 1 and 2. we directly obtain that the probability of erroneous assignment using the AC router is exponentially smaller than under a stan- dard routing scheme. This is formalized in Proposition I, given by:
Proposition 1 (Robustness of ACMoE), Consider an expert assignment setting for the representa- fion h ~ g- (1g2e ; Ege) ) as in Lemma 2 with two routers given by r R - [E] and rAC, :R" + [E] for standard (Eqn. 2) and AC routers (Definition 2), respectively. Then the probabilities of incorrect assignments of routers T and nCC satisfy Pr (ACCh') k*) S PPr((hh') + k*).
Promoting faster convergence. For an expert embedding ej: E 12 and associated cluster C2., our AC router in Definition adaptively spheres Cs by stretching the feature space with weights inversely proportional to the coordinate-wise dispersion in Ct. This reduces the conditioning number of the Hessian of the loss with respect to the expert Ek, improving the loss landscape and enabling faster and more stable convergence of the router. This notion is formalized in Proposition 2:
Proposition 2 (Faster convergence of ACMoE). Let L MoE, O - I2, and LACMoE: - R, be the network loss functions defined on the whole parameter set O when employing the standard (Eqn. 2) and AC roulers (Definition 2), respectively. Let k(A)= Amax/ Amin Amin denote the conditioning number of matrix A with largest and smallest eigenvalues Amax and Amin respectively. Let the Hessian of an hh expert be given by V Ve. Then for each i [E] the following holds with high probability
K. CAMMEE) S
<VV2 (MBB)
(11)
Remark 4. Faster convergence of ACMoE can also be argued from the perspective of learning Gaussian mixture models with Expectation Maximization (Dempster et al., 1977). The classic result of Ma et al. (2000) shows the convergence rate to the true parameters depends on the overlap between component Gaussians. Our AC method adaptively transforms the input space with by M& (Definition 1), which decreases component overlap by increasing inter-cluster distances.
The proof is provided in Appendix A.3. We find this result empirically supported as shown by the
rapid convergence in Fig. 2.
4 EXPERIMENTAL RESULTS
In this section, we empirically justify the advantage of ACMoE over baseline MoE models. We evaluate our method on large-scall tasks including Wikitext- 103 (Merity et al., 2016) language modeling and ImageNet (Deng et al., 2009) object classification. We implement our AC router into Switch Transformer (Fedus et al., 2022). Generalist Language Model (GLaM) (Du et al., 2022), and Swin Transformer (Liu et al., 2021) backbones and compare our router against the standard Sparse Mixture-of cftuure-of-Experte Experts (SMoE) router (Shazeer et al., 2017) and the XMoE router (Chi et al., 2022). We show that i) ACMoE obtains substantive improvements over baseline models across both language and vision tasks; ii) ACMoE offers robust improvements on contaminated and out-of-distriiution samples; and iii) ACMoE attains these gains without introducing any learnable parameters and with negligible additional computational overhead. Results are averaged over runs with different seeds.
4.1 LANGUAGE MODELING
Experimental Setup. We adopt the experimental setup of Pham et al. (2024). We compare AC- MoE with Switch Transformer and GLaM baselines with 16 total experts in small (70M parameters) and medium (220M parameters) configurations with top-2 expert routing. We present pretraining test perplexity (PPL) results for Wikitext-103 and test perceacpr- (BPC) for character-level EnWik-8. We report top-1 accuracy for finetuning classification tasks on the 2-class Stanford Sen- timent Treebank-2 (SST2) (Socher et al., 2013), 5-class Stanford Sentiment Treebank-5 (SST5) (Socher et al., 2013), and 77-class Banking-77 (B77) (Casanueva et al., 2020). Full experimental details are provided in Appendix C.
7

---

Published as a conference paper at ICLR 2025
Table 1: Test Perplexity (PPL) and bytes-per-character (BPC) pretraining and top-1 test accuracy on Stanford Sentiment Treebank 2, 5 (SST2, SST5), and Banking-77 (B77) finetuning classification.
Model
Test BPC/ PPL (1)
SST2 (t)
SST5 (t) B77 (1)
En-ik-8 Pretrain
Switch Trunsformer (Fedus et al., 2022)
1.153
63.27 64,45
32.21 33.79
53.48
Switch-ACMoE (Ours)
1.137
54.26
WikiText- 103 Pretrain
Switch Transformer (Fedus et al., 2022)
35.48
76.27
39.13
83.82
Switch-ACMoE (Ours)
34.42
77.32
40.04
86.01
GLaM (Du et al., 2022) GLaM-ACMeE (Ours)
38.27
69.97
33.69
80.89
36.26
71.90
34.24
82.33
Table 2: Perplexity (PPL) WikiText- 103 contaminated by Text Attack.
Model
Clean Test PPL (L)
Contaminated Test PPL ())
Switch Transformer (Fedus et al., 2022)
35.48
48. 12
Switch-ACMoE (Ours)
34.42
47.61
GLaM (Du et al., 2022)
38.27
50.84
GLaM-ACMoE (Ours)
36.26
47.91
Pretraining and Finetuning. Table 3 shows ACMoE attains top test PPL. on WikiText- 103 language modeling in Switch and GLaM backbones at small and medium configurations under baseline SMoE and XMoE routers. The improvement the GLaM-medium architecture is a particularly substantive 4.8% over the next best baseline. Table shows ACMoE pretrained models both WikiText-105 and EnKik-8 surpass the performance of baselines in finetuning tasks, with strong, consistent im- provements of approximately 3%, showing ACMoE's s strong performance carries over to finetuning.
Robust Language Modeling. Table 2 show test PPL Wiki contaminated by Text At- tack, where words are randomly swapped with generic token AAA'. We follow the setup of Han et al. (2024); Teo & Nguyen (2025b); Ab- dullaev & Nguyen (2025) and assess models by training them on clean data before attacking the test data using an attack rate of 2.5%. ACMoE outperforms baseline Switch and GlaM with par- ticularly robust performance in the GLaM aM back- bone, surpassing GLaM by 5.8%.
Table 3: WikiText- 103 test PPL of ACMoE and baseline GLaM and Switch.
Router
Test PPT. (J)
Switch Transformer (Fedus et al., 2022)
SMoE-small (Shazeer et al., 2017) XMoE-small (Chi et al., 2022)
87.94
87.21
ACMoE-small (Ours)
85.07
SMoE Eneoiimd (Shazeer et al., 2017)
35.48
XMoE-medium (Chi et al., 2022)
35.88
MubleMooE-mddiue (Dai et al., 2022)
35.33
ACMoE-medium (Ours)
34.42
4.2 IMAGE CLASSIHICATION
GLaM (Du al., 2022)
Experimental Setup. We adopt the experimen- tal setup of Liu al. (2021) for pretraining and evaluation on ImageNet. We evaluate ACMoE against the Swin Transformer baseline with 16 total experts in both top-1 and top-2 expert rout- ing settings. The Swin backbone has 280M pa- rameters. We additionally conduct experiments on ImageNet under white box adversarial attacks
SMoE-small (Shazeer et al., 2017) XMoE small (Chi et al., 2022)
58.27
54,80
ACMoE-small (Ours)
54.55
(Shazeer al., 2017)
38.27
XMoE-mediun (Chi et al., 2022)
38,10
StableMoE mediau (Dai al., 2022)
38.04
ACMoE-medium (Ours)
36.26
fast gradient sign method (FGSM) (Goodfellow et al., 2014) and projected gradient descent (PGD) (Madry et al., 2017), and black box attack simultancous perturbation stochastic approximation (SPSA) (Uesato et al., 2018). We also present results out-of-diitribution (OOD)(Hendrycks al., 2021a;b). In all robust image classification tasks, image classification using ImageNet- A/O/R we adopt the conventional setup of pretraining on ImageNet and evaluating the trained models the contaminated/OOD datasets (Han el al., 2024; Zhou et al., 2022a; Puigcerver et al., 2022; Nguyen al., 2024; Nielsen al., 2025). Full experimental details are provided in Appendix C.
8

---

Published as a conference paper at ICLR 2025
Table 4: Test Accuracy on ImageNet corrupted PGD, FGSM, and SPSA.
Clean Data
PGD
FGSM
SPSA
Model
Top
Top 5
Top
Top 5
Top
Top 5
Top
Top 5
Swin-Top / (Liu et al., 2021)
75.22
92.51
39.69
74.59
52.84
83,86
59.92
82.63
Swin-ACMoE-Top (Ours)
75.39
92.56
40.66
73.46
53.43
82.80
59.97
82.47
Swin-Top 2 (Liu et al., 2021)
76.10 76.31
92.99
40.85
75.51
54.70
85.22
60.57
82.75
Swin-ACMoE-Top (Ours)
93.14
43.74
78.55
55.78
85.80
63.47
86.05
Table 5: Test Accuracy on Image Classification in Imagenet-
Im-A
Im-R
Im-O
Model
Top-1 Ace. (t)
Top-1 Acc. (t)
AUPR (t)
Swin Transformer-Top / (Liu et al., 2021)
6.83
30.60
17.89
Swin-ACMoEE-Top 1 (Ours)
7.13
30.85
18.45
Swin Tronsformer-Top (Liu et al., 2021)
9.38
32.07
18.51
Swin-ACMoF:TTop (Ours)
9.42
32.35
19.55
Image Classification under Adversarial Attack. Table 4 shows performance on ImageNet clas- sification against FGSM, PGD, and SPSA. Compared with Swin baseline, ACMoE-Top 2 attains noteworthy 7% and 5% improvements against PGD and SPSA in top-1 accuracy respectively.
Out-of-distribution Image Classification. Table shows ACMoE improves over the baseline Swin Transformer image classification on hard OOD and real-world adversarially filtered images. Eval- uation on ImageNet-A/OIF shows consistent improvements over the baseline in top-1 and top-2 expert choice, with particularly strong improvements in ImageNet-O under top-2 routing with performance gain in area under precision recall (AUPR) of almost 6%.
4.3 EMPIRICAL ANALYSIS
Load Balancing. We analyze in Table 6 the effect of ACMoE on expert load balancing. Load bal- ance is calculated as the percentage of tokens assigned to each expert. The load balance score is then taken as the standard deviation over these percentages. A standard deviation of O, where all ex- perts activated in exactly equal proportions, is therefore perfect load balance. We compute this statistic per MoE layer and present the overall load balance averaged over all layers. ACMoE attains better overall load balancing compared to Switch and Swin transformers. Against all backbones, ACMoF achieves smaller spread in the load balances over layers, shown by smaller standard devi- ation. Visually we see how better expert specialization can aid load balance Fig. 1, where better identification of the semantic regions of the input space leads to more experts being activated.
Efficiency Analysis. Computing the cluster-wise feature weights { [we) )=( requires no learnable parameters and is obtained by computing the mean absolute deviation for each set of tokens assigned to the h expert. This can be computed using just two computations of the mean one for the mean per cluster and one for the mean of the absolute deviations per cluster done in parallel over all clusters. This is of order ((nd) = (n) ) for n tokens, hence the upper-bound time complexity of the MoF layer is unaffected. Table 7 provides empirical efficiency analysis in terms of compute speed, memory allocation, and parameters, which shows changes in speed and memory are within a margin of approximately 1 5% less, implying there is significant efficiency loss.
5 RELATED WORK
Routing Methods. Recent studies have proposed token-expert assignment algorithms based on reinforcement learning (Bengio et al., 2015), deterministic hashing (Roller et al., 2021), optimal transport (Liu et al., 2022), linear programs (Lewis et al., 2021). cosine similarity (Chi et al., 2022), soft token mixing (Puigcerver et al., 2023), greedy top-k experts per token (Shazeer et al., 2017) and greedy top-k tokens per expert (Zhou et al., 2022b). Existing work has predominantly consid- ered dot-products between inputs and experts as a suitable metric for similarity (Lewis et al., 2021; Puigcerver et al., 2023; Shazeer et al., 2017; Zhou al., 2022b; Chi al., 2022). This work contin- ues with dot-prott oot-product based learnable routing but computes the routing assignments in an adaptively transformed space to maximally identify the latent expert clusters.

---

Published as a conference paper at ICLR 2025
Table 6: Load Balance Analysis of ACMOE and Baseline MoE Models
Model
Layer-Averagee Load Balance ())
Switch Transformer (Fedus et al., 2022)
5,577 +: 4.131
Switch-ACMoE (Ours)
5.317 +: 2.622
GM (Du et al., 2022)
2.901 + 1.434
Gl Gl (Ours)
2.938 + 1.221
Swin Transformer (Liu et al., 2021)
2.134 + 1.110
Swin-ACMoB (Ours)
2.127 # 0.968
Table 7: Efficiency Comparison between ACMoE and baseline MoE Models
Model
Compute Speed (ms/it)
Max Memory (K)
#Params (M)
GLaM (Du et al., 2022)
422.62
25.69
220
GLaM-ACMoE (Ours)
425.15
25.72
220
Switch Transformer (Fedus et al., 2022)
391.93
34.64
216
Switch-ACMoE (Ours)
393.29
34.68
216
Swin Transformer (Liu ct al., 2021)
403.36
22.00
280
Swin-ACMoF (Ours)
408.56
22.19
280
MoE and Cluster Analysis. The MoF. framework traces its roots back to Gaussian mixture models where the input space is assumed divisible into separate regions with expert specializing in each region (Jacobs et al., 1991). Recent studies show that the router can recover the clustering structure of the input space and each expert specializes in a specific cluster (Dikkala et al., 2023; Chen al., 2022). Our work leverages the clustering perspective on MoE to consider adaptive transformations of the input space to more easily distinguish latent clusters. We learn these transformations via weathre-weighted cluster analysis, which has been studied in the clustering literature (Brusco & Cradit, 2001; Witten Tibshirani, 2010; Gnanadesikan et al., 1995; Van Buuren Heiser, 1989; Friedman & Meulman, 2004). Friedman & Meulman (2004) consider cluster-dependent feature weights to augment iterative clustering algorithms. Our approach similarly uses cluster-dependent feature weights but uses a different optimization problem to derive optimal weights.
Robust MoE. The robustness of MoE architectures is a newly emerging research area. Puigcerver al. (2022) provide the first study in this direction from the perspective of model capacity and the Lipschitz constant, finding conditions under which MoE models are provably more robust than their dense counterparts. Zhang al. (2023) examine the effect of adversarial training and propose al- ternating optimization adversarial defence. Teo & Nguyen (2024) integrates heavy-ball momentum in SMoE to improve the model's stability and robustness. Our work differs from these approaches by examining the robustness of MoE models purely through the lens of the latent clustering structure of the input space. To the best of our knowledge, this a novel lens on robustness in MoE models.
6 CONCLUSION AND FUTURE WORK
In this paper, we present the Adaptive Clustering (AC) router and ACMoE layer, a novel MoE rout- ing method that computes token-expert assignments transformed space that maximally identifies latent clusters in the data and more easily discovers the best-matched expert for each token. We adap- tively learn for each input which features are relevant to determining latent cluster assignment and scale its features accordingly such that features that promote tight clustering are upweighted and features that produce dispersed clusters are downweighted. This transformation accentuates the relevant characteristics of each input according to the specialization of the experts, thereby allowing the router to more easily discover the optimal input-expert allocation. Our AC routing method en- ables faster convergence by improving the Hessian conditioning of the router and better robustness by increasing the separation of latent clusters the transformed space. This approach makes no assumptions on the downstream task, requires no learnable parameters, and can be applied within any MoE architecture to boost performance on clean and contaminated data. limitation of our method is that the AC router requires estimates of each token's 's cluster assignment. We obtain these by using the expert assignments in previous layers, which means we require the embedding size to remain the same between adjacent MoE layers. For ongoing work, we are investigating improved methods for estimating the latent cluster memberships without reliance on previous layers and with provable consistency guarantees.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS AND DISCLOSURE OF FUNDING
This research project is supported by the National Research Foundation Singapore under the Al This research project is supported by the Ministry of Education, Singapore, under the Academic Research Fund Tier (FY2023) (A-8002040-00-00, A-8002039-00-000 This research / project is also supported by the NUS Presidential Young Professorship Award (A (A-0009807701-00),
Singapore Programme (AISG Award No: ALSGG-TC-2002-012SSGLII
Thanks to our anonymous reviewers, who provided valuable feedback which improved the paper substantially. Thanks also to Loi Xuan Ly for lending his eye for design.
Reproducibility Statement. Source code for our experiments are provided in the supplementary material. We provide the full details of our experimental setup including datasets, model specifi- cation, train regime, and evaluation protocol for all experiments Appendix C. All datasets are publicly available.
Ethics Statement. Our work considers fundamental architectures, and in particular their robustness and convergence properties. Given this, we foresee no issues regarding fairness, privacy, security, or any other harmful socictal or ethical implications in general.
REFERENCES
Laziz Abdullaev and Tan Minh Nguyen. Transformer meets twicing: Harnessing unattended residual information. In The Thirteenth International Conference Learning Representations, 2025.
URL htthttss:///ppneeee /penrevieww net forum? id=1 6KG5aN1eS.
Dosovitskiy Alexey. An image is worth 16x 16 words: Transformers for image recognition at scale. arXiv preprint arXiv: 2010.11929, 2020.
Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei. Beit: Bert pre-training of image transformers. arXiv preprint arXiv:2106.08254 2021.
Emmanuel Bengio, Pierre-Lue Bacon, Joelle Pineau, and Doina Precup. Conditional computation in neural networks for faster models. arXiv preprint 1 1,00069 2015.
Michael Brusco and Dennis Cradit. variable-selection heuristic for k-means clustering. Psy- chometrika, 66:249 270, 2001.
Iñigo Casanueva, Tadas Temčinas, Daniela Gerz, Matthew Henderson, and Ivan Vulić. Efficient intent detection with dual sentence encoders. arXiv preprint arXiv:2003.04807 2020.
Zixiang Chen, Yihe Deng, Yue Wu, Quanquan Gu, and Yuanzhi Li. Towards understanding the mixture-of-experts layer deep learning. Advances in neural information processing systems, 35:23049 23062, 2022.
Zewen Chi, Li Dong, Shaohan Huang, Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj, Xia Song, Xian-Ling Mao, et al. On the representation collapse of sparse mixture of experts. Advances in Neural Information Processing Systems, 35:34600-3444 3460030--11,.-0011100001111 2022.
Damai Dai, Li Dong, Shuming Ma, Bo Zheng, Zhifang Sui, Baobao Chang, and Furu Wei. Sta- bleMoB: Stable routing strategy for mixture of experts. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 7085-7095, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 10.18653/v 1/2022.acl-2221e1 long.oo,,,,,,, URL https: http:///aaaaaatttaaaa/ee
prg//02222 a.1-222..ccoon acl-long 489.
Arthur P Dempster, Nan M Laird, and Donald B Rubin. Maximum likelihood from incomplete data via the em algorithm. Journal of the royal statistical society: series (methodological), 39(1): 1-22, 1977.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scall scale hi- erarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255 leee, 2009.

---

Published as a conference paper at ICLR 2025
Nishanth Dikkala, Nikhil Ghosh, Raghu Meka, Rina Panigrahy, Nikhil Vyas, and Xin Wang. On the benefits of learning to route in mixture-ou of experts models. In Proceedings of the 2023 Conference 2023.
on Empirical Methods in Natural Language Processing, pp.
Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with offeert-of-experts. In International Conference on Machine Learning, pp. 5547- 5569. PMLR, 2022.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1 2022.
Jerome H Friedman and Jacqueline J Meulman. Clustering objects on subsets of attributes (with discussion). Journal of the Royal Statistical Society Series B: Statistical Methodology, 66(4): 815-849, 2004.
Ram Gnanadesikan, Jon R Kettenring, and Shiao Li Tsao. Weighting and selection of variables for cluster analysis. Journal of classification, 12:113-136, 1995.
lan J Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. arXiv preprint arXiv:1412.6572. 2014.
Yongxin Guo, Zhenglin Cheng, Xiaoying Tang, Zhaopeng Tu, and Tao Lin. Dynamic mix- ture of experts: An auto-tuning approach for efficient transformer models. arXiv preprint arXiv:2405.
14297, 142997777.4 2024.
Xing Han, Tongzheng Ren, Tan Nguyen, Khai Nguyen, Joydeep Ghosh, and Nhat Ho. Design- ing robust transformers using robust kernel density estimation. Advances in Neural Information Processing Systems, 36, 2024.
Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, et al. The many faces of robustness: A criti- cal analysis of out-of-distribttion generalization. Proceedings of the IEEE/CVF international conference on computer vision, pp. 8340-8349, 2021a.
Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song. Natural adversarial examples. In Proceedings of the IEEE/CVF conference on computer vision and pattern recogni- fion, pp. 15262-15271, 2021b.
Robert A Jacobs, Michael I Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 1991.
Kenichi Kumatani, Robert Gmyr, Felipe Cruz Salinas, Linquan Liu, Wei Zuo, Devang Patel, Eric Sun, and Yu Shi. Building great multi-lingual teacher with sparsely-gated mixture of experts for speech recognition. arxiv preprint arXiv:2 1rXiv:2/12.05820 2021.
D Lepikhin, H Lee, Y Xu, D Chen, O Firat, Y Huang. M Krikun, N Shazeer, and Z Gshard. Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006. 16668, 2020.
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models. In International Conference on Machine Learning, pp. 6265-6274. PMLR, 2021.
Tianlin Liu, Joan Puigcerver, and Mathieu Blondel. Sparsity-constrained optimal transport. arXiv preprint arXiv:2209, 2e0e222999.4444 15466, 2022.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings the IEEE/CVF international conference on computer vision, pp. 10012-10022, 2021.
Jinwen Ma, Lei Xu, and Michael Jordan. Asymptotic convergence rate of the em algorithm for gaussian mixtures. Neural Computation, 12(12):2811-290 2000.
12

---

Published as a conference paper at ICLR 2025
Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. Towards deep learning models resistant to adversarial attacks. stat, 1050(9), 2017.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXiv: (eeo::609.77850 07843, 2016.
Tam Minh Nguyen, César Uribe, Tan Minh Nguyen, and Richard Baraniuk. Pidformer: Trans- former meets control theory. In Forty-first International Conference on Machine Learning, 2024.
Tan Minh Nguyen. Tam Minh Nguyen, Nhat Ho, Andrea L. Bertozzi, Richard Baraniuk, and Stanley Osher. A primal-dual framework for transformers and neural networks. In The Eleventh Interna- fional Conference Learning Representations, 2023. URL https: wpwp/eeeeeeeeeeeeeeeeeeeeeeeeeee nett/
Viet Dung Nguyen, Minh Nguyen Hoang. Luc Nguyen, Rachel Teo. Tan Minh Nguyen, and
Linh Duy Tran. CAMEx: Curvature-aware merging of experts. In The Thirteenth International Conference on Learning Representations, 2025, URL https:/ /wwpnreveewwwweee//o/ee/ee/oeeee net/ forum?
id=nT2u0M0nEB.
Stefan Nielsen, Laziz Abdullaev, Rachel SY Teo, and Tan Nguyen. Elliptical attention. Advances in Neural Information Processing Systems, 37:109748-109777 2025.
Quang Pham, Giang Do, Huy Nguyen, TrungTin Nguyen, Chenghao Liu, Mina Sartipi, Binh T Nguyen, Savitha Ramasamy, Xiaoli Li, Steven Hoi, et al. Competesmoo-effeetin training of sparse mixture of experts via competition. arXiv preprini arXiv:2402.002526, 2024.
Joan Puigcerver, Rodolphe Jenatton, Carlos Riquelme, Pranjal Awasthi, and Srinadh Bhojanapalli.
On the adversarial robustness of mixture of experts. Advances in Neural Information Processing Systems, 35:9660-9671, 2022.
Joan Puigcerver, Carlos Riquelme, Basil Mustafa, and Neil Houlsby. From sparse to soft mixtures of experts. arXiv preprint arXiv:2308.000 00951, 2023.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140): 2020.
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. Advances in Neural Information Processing Systems, 34:8583 8595, 2021.
Stephen Roller, Sainbayar Sukhbaatar, Jason Weston, et al. Hash layers for large sparse models. Advances in Neural Information Processing Systems, 14:1555-17756, 2021.
N Shazeer, A Mirhoseini, K Maziarz, A Davis, Le, G Hinton, and J Dean. The sparsely-gated mixture-of-experts layer. Outrageously large neural networks, 2017.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over sentiment treebank. In Proceedings of the 2013 conference on empirical methods in natural language pro- cessing, pp. 1631-1642, 2013. 2013.
Rachel Teo and Tan Minh Nguyen. MomentumSMoe: Integrating momentum into sparse mixture of experts. In The eighth Annual Conference on Neural Information Processing Systems, 2024. URL https: https::/openrevieeeee wwwp//veeeeeeeeee/eeeeeee net/ wttet//oommnnnddy2eeeeeeeeee
Rachel Teo and Tan Minh Nguyen. MoLEx: Mixture of layer experts for fine-tuning with sparse upcycling. In The Thirteenth International Conference Learning Representations, 2025a. URL
n forum wewwwiiiieeeii Www.riWiiii.......oo
13

---

Published as a conference paper at ICLR 2025
Rachel SY Teo and Tan Nguyen. Unveiling the hidden structure of self-attention via kernel principal component analysis. Advances in Neural Information Processing Systems, 37:101333-101444 2025b.
Jonathan Uesato, Brendan O'donoghue, Pushmeet Kohli, and Aaron Oord. Adversarial risk and the dangers of evaluating against weak attacks. In International conference on machine learning, pp. 5025-5034. PMLR, 2018.
Stef Van Buuren and Willem Heiser. Clustering objects into groups under optimal scaling of variables. Psychometrika, 54:699-706, 1989.
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Daniela M Witten and Robert Tibshirani. A framework for feature selection in clustering. Journal of the American Statistical Association, 105(490):713-726, 2010.
Yihua Zhang, Ruisi Cai, Tianlong Chen, Guanhua Zhang, Huan Zhang, Pin- Yu Chen, Shiyu Chang, Zhangyang Wang, and Sijia Liu. Robust mixture-of-expert training for convolutional neural net- works. In Proceedings of the IEEF/CVF International Conference Computer Vision, pp. 90- 101, 2023.
Daquan Zhou, Zhiding Yu, Enze Xie, Chaowei Xiao, Animashree Anandkumar, Jiashi Feng, and Jose M Alvarez. Understanding the robustness in vision transformers. In International Conference on Machine Learning, pp. 27378-27394. PMLR, 2022a.
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, Quoc V Le, James Laudon, et al. Mfxture-of-expert. with expert choice routing. Advances in Neural Information Processing Systems, 35:7103-7114, 2022b.
14

---

Published as a conference paper at ICLR 2025
Supplement to "Tight Clusters Make Specialized Experts"
Table of Contents
Technical Proofs
15
A.1 Proof of Theorem
15
A.2 Proof of Proposition
17
A.2.1 Proof of Lemma
17
A.2.2 Proof of Lemma 2
17
A.3 Proof of Proposition 2
18
Implementation Procedure and Computational Efficiency
20
Experimental Details and Additional Experiments
20
C.1
Language Modeling
20
C.1.1
Datasets
C.1.2 Model, Optimizer, & Train Specification
20
C.2 Image Classification
................................
C.2.1 Datasets and Attacks
21
C.2.2 Model, Optimizer, & Train Specification
................................
22
C.3 Adversarial Attack At Higher Perturbation Budget
22
C.4 Cluster Visualization
22
C.5
Ablation Studies
23
C.5.1 Measures of Dispersion
................................ 23
C.5.2 Layer Placement
................................
23
C.5.3
Random Ablation
................................
24
C.6
Cluster Weight Mixing
25
C.7 Adaptive Clustering Integration into Soft Mixture of Experts
25
C.8 Image Classification in Swin Transformer Base Configuration
26
C.9 Router Stability
26
C.10 Dynamic Routing
26
D Broader Impact
27
A TECHNICAL PROOFS
A. PROOF OF THBOREM ]
To begin with, we present the following lemma to show the existence of constants ag for & [E] that satisfy Eqn. 7:
Lemma 3. For any X > 0, Eqn. 7 has exactly d real solutions with respect to ox-
15

---

Published as a conference paper at ICLR 2025
Proof of Lemma 3. Without loss of generality, assume that $1 82k ... 8dk- Denote
(a))= == Σ ge[d] 8qk
(12)
X
Then, the existence of solutions to Eqn. 7 is equivalent to the condition ş(ar) = 0. Note that yp(a) a strictly decreasing function its connected continuity domains since
1 0
((a)==
Σ q=[d]
(13)
(8gh a)²
for all a IR ( (-818, sak) Further, we observe that
lim p(a) -00
lim (a) = +00
(14)
ak
for all q [d]. and
d
lim ((a) a
> <
(15)
-U.------ Now consider the domain of continuity of (a)), namely (-oo, -81k (- -821 u(-)) o0)
Due to the monotonicity and limits 14 & 15, there exists a unique solution in each of the intervals except for (-o0,-S) -1kk) where the function is always strictly negative, thus, yielding d roots in total.
Now we follow up with the main proof of this section.
Proof of Theorem 1. First, let Za {i r(i)
k} for convenience. Now let restate the clustering
optimization problem (4) here once again:
A log WakPiio
min Q(c, (wk (=aE))
N2
d
dw
E
4.jeZy qE[d]
such that
Σ Wigh 1, ge[d]
Vk [E],
(16)
where we have immediately used the fact that
DKL (u || ww) Σ log 1/d
1
(17)
(=[d]]
d
Wigk
Also, note that
1
WykPija
+ X log d dwigh
Σ q=[[]
WynPija > log( (og(dw))))
(d(d]
Σ [e[d]
WakPija
d log work Alog d.
(18)
We can ignore the term Alog since it does not depend on the optimization variable. Method of Lagrange multipliers turns this constrained optimization problem into the following unconstrained counterpart:
min L(c, (wx) (wx)
Σ [E[E]
Σ
Σ
WakPijo
log Wigh
Σ ke(
Σ [c[4]
wek
i.jeI ge[d]
where a [ai ar] is the vector of Lagrange multipliers. Note that the last optimization problem can be separated into the following L independent optimization subproblems:
min (c, Wk, a) 02.0
1
Σ K.jeI gE[d]
Σ
Wykkisg
A
log Wyk +ak
Σ [9[d]
Wigh - -1)
N2
16

---

Published as a conference paper at ICLR 2025
for & [E]. Since the objective function is positive combination of convex functions, the opti- mization problem is also convex. By setting the derivatives of Lx with respect to both optimization variables to O, we obtain the following system of equations:
0L* = Sgk
X
+ + =0, 0,
Owgk
d Wigh
0L*
Eve(ao Wyk
dack
for all & [E], where Sgb is the data dispersion measure defined in the theorem statement, The first
equation yields
Wyg
(19)
d Syk Ork
where og is found from Ege[a] Wigk
1 which fact gives
d Sqke akk X
Σ
(20)
ge[d]
for all k [B] as desired.
A.2 PROOF OF PROPOSITION 1
Since Proposition composition of Lemma I and Lemma 2, we proceed by providing their proofs.
A.2.1
PROOF OF LEMMA
Proof of Lemma 1. Notice that we can expand incquality (1) as
Σ
2 Σ 8p4
#[d
s[[d]
where we let 8ÿ pi 14a- Since Ma entries are mean-scaled, we can rewrite them as
dm
(21)
Eyean
for some initial dispersion estimates {m } j4dd]- Without loss generality, assume that [d'] is the relatively much smaller than those in the
of dimension indices for which the dispersions rest of the dimensions in the sense that mn >> mí for any i E [d'] and j ε [d] [d']. Then, there exists positive a < 1/2 such that Ex((g) m, d- and Eis[d]s[r) m, < a. By the assumption that clusters are best-separated along the features for which they cluster tightly, this means that the
weight matrix M, maximizes the contribution of largest d' terms in Eve[a] m2ō2ς corresponding individual feature-wise distances dimensions where the feature dispersions are the smallest instead of giving uniform weights to all dimensions, which leads to inequality (1).
A.2.2 PROOF OF LEMMA
Proof of Lemma 2. Since we use the L2 distance between the token h and μe as similarity metric, Assume that the token h' is assign cluster Jk* the token h' iff hh' -4- < | pu
noisy observation of an underlying true token h which actually originates from cluster 9a. Then, the token h' can be decomposed as h' the decision variable D(h') h - # 1|2 1''-4 which turss the clustering condition to D(h') s for the cluster ga* Let us analyze the decision variable D as a random variable where randomness may come from the underlying sampling strategy and noise. Note that
c for random noise c ~ (O, E.). Now define
12
D(h') h 1h+c 6 M4. 12 h e 1
1 Me -Ih mell 2 2( p.a) T
DD(h) 25p c,
(22)
17

---

Published as a conference paper at ICLR 2025
where off 14 pa. Due to the assumption that h is drawn from the distribution 9k it can be rewritten as = pa + V with v ~ N NN,,, ) Then for the first term in Equ. 22, we have
(hh) h M 1|2 | M
"μ (2h- p pa)
.µ (2v -ô4)
20p -- 1ôlel??.
(23)
Substituting this back into Eqn. 22, we get
D(h') 20
_ (v +6) -
(24)
This shows that D(h') ~ . (-16p112, 46 (Ek. + E.)6 Since D(h') follows a normal distri-
bution with the derived parameters, the probability that is assigned to cluster Ok- is given by
16u|1
cr(rorrect cluster) Pr (D(h) s = +
(25)
2 6p (E Σ,)6μ
where & denotes the CDF of normal distribution as usual. Since 4 is an increasing function, the probability that the noisy token is assigned to the correct cluster is proportional the distance between the cluster centroids and inverse proportional to the covariance matrices of the cluster and the additive noise. On the other hand, for the incorrect clustering probability, we have
16
Pr(incorrect cluster) = 1 2/06² (Ex- + E.jōH
(26)
as claimed.
A.3 PROOF OF PROPOSITION 2
Proof of Proposition 2. Let the router be given by 9 and let the softmax function be given by Go TR - R. parameterized by expert embeddings (e]]) The network loss depends on expert embeddings only through the router function g. We shall explore the exclusive contribution of each expert embedding in minimizing LACMOB, In order to do this, we look at the network loss scalar function hth expert embedding vector while treating all other network parameters as fixed. Then, we can write CACMEE Rd I such that CACMo ACMMo (go (e;)). For simplicity, shall omit the subscript 0. The gradient that comes from brok-propn is then given by
Ve, ACCME = =VV (VVCCACMan Ve.g,
(27)
where Ve,g Rdd denotes the Jacobian matrix of y since for ya == (ga(e;)) k. we can write
aLACMoE 0g:.
L ACMOE
.91....,9H 9a) Σ
(28)
de,
Ögn
dee+
Note that for gr softmax (h" M Mex), we have
0gk
m,h.gs 9.) myhhbbik
(29)
des
Then, the element of the Hessian matrix of the network loss at index (s,t) [d]] x [d]]
be written
as
( (CACMOB)
CACME
Σ
ALACMOE
Ögk
desedes
dea
Ogk
Des
82 CACMOE Ög;
Ogk
ACACMo
8 gb
Σ Σ
Ogk
Dest
deis
Ogk
de
2 LACMOE
aLACMoE
-m,hymyh
Σ Σ
Σ
bis
bks
b/:s
Ogvdgg
Agk.
=m,h.m,hy Bi,
(30)
18

---

Published as a conference paper at ICLR 2025
where B, is some constant that depends only on index i. Due to Eqn. 30, the Hessian takes the following matrix form
H = B,(M))
(31)
Taking expectation from both sides, obtain
En-(H.E) Bi [(hh))M (2M
(32)
where assume h is centered. Now recall that M = diag(mm...... ma) where for each m4 1/vEx holds. Assume that the covariance matrix E is symmetric positive definite. Then, it is
diagonalizable as Σ UAUT with = (Ji,g(A1...... ,.,))..).). a diagonal matrix with eigenvalues of Σ. With the transformation M, we get
MEM = MUAU M 1-UMAMU
(33)
A
U
UT.
(34)
m Aa
Since the eigenvalues capture the variances along the principal components of the covariance matrix, m². reciprocal of measure of dimension-wind wwise dispersion, is reasonably correlated with 1/X,, as demonstrated by Lemma 4, implying X3 S A = my 2 my with high probability. Therefore, obtain that
Amnax (MEM)
m Amnax Σ)
K0 (MEM))
S <k(E), Σ
(35)
Amnin (MEM)
m
max Amin
which implies the claim.
Lemma 4 (Correlation between dimension-wen wise varainces and covariance eigenvalues). Let (b; hed be the set of normalized basis vectors of RS.. Consider a symmetric positive definite covariance matrix E and its unit eigenvectors {Vs)ie[d]- Assume that the eigenvector v; is a reasonably small perturbation of the basis vector b, such that v 2 for all [d] and a small constant e > O. Then, for all i [d], we have
| - Eal SE E max [A, Ayy
(36)
j=t
where (Ae is the set of ordered eigenvalues of E corresponding eigenvectors he[dj
Proof of Lemma 4. Note that each diagonal element of the SPD covariance matrix Σ can be written
as
En = bb = 6..
Σ
by A
Laaa'k
Σ {o)
(37)
je[c]
jo[d]
Then, the difference on the left hand side Eqn. 36 can be bounded as
1A: - Endl
A Σ (,())b,))
1o (1- (v:e;))) - Σ>((oT6,)
]0[d]
Jei
A (v b.) - {o)(b)11)
(38)
1 X₂) (0)
AA: A₂ E bo)
max
Jel
max
AA: 11 (v,bb))))
(39)
j+i
<max AA Ay1:
where we used the fact that
Σ )(E
n
Ev)b,)²
Σ
(1)b.))
(o)b.)o = b =
to obtain Eqn. 38 and Eqn. 39 since the eigenvectors of E are orthonormal.
19

---

Published as a conference paper at ICLR 2025
B IMPLEMENTATION PROCEDURE AND COMPUTATIONAL EFFICIENCY
Training and Inference. Given the AC routing scheme requires requires the expert assignment per token from the previous layer, we can only implement AC routing from the second layer on. We incorporate AC routing into both training and inference stages. This is because, firstly, AC routing both clean and contaminated data, and so even in the presence
designed to offer improvements of completely clean train and test data, it is advantageous to incorporate the AC method into both stages. Secondly, it is commonplace to encounter data contamination only the test stage and indeed highly possible to encounter it in train as well. Therefore, in the interest of robustness as well, AC routing is incorporated into both stages.
Computational Efficiency. Computing the required (wx) [a[j for number of experts E requires no learnable parameters and is obtained simply by computing the mean absolute deviation for each set of tokens assigned to the hth expert. This be computed using just computations of the mean once for the mean per cluster and once again for the mean of the absolute deviations per cluster done in parallel over all clusters using torch.index reducx)educe() and is of the order ((2nd) (nn) for n tokens. Hence the upper-bound time complexity of the MoE layer is un- affected. We provide in Table additional efficiency analysis terms of throughput, max GPU memory allocated, and parameters which shows no significant efficiency loss compared to baseline MoE architectures.
C EXPERIMENTAL DETAILS AND ADDITIONAL EXPERIMENTS
C.I LANGUAGE MODELING
C.1.1 DATASETS
103. The WikiText-103' dataset contains around 268K words and its training set con- sists of about 28K articles with 103M tokens. This corresponds to text blocks of about 3600 words. The validation set and test sets consist of 60 articles with 218K and 246K tokens respectively.
EnWik-8. The EnWik-8 dataset is byte-level dataset of 100 million bytes derived from Wikipedia that, in addition to English text, also includes markup, special characters, and text in other languages. EnWik-8 contains 90M characters for training, 5M for validation, and 5M for testing.
Stanford Sentiment Treebank-2. The Stanford Sentiment Treebank-2 (SST2) (Socher et al., 2013) is a 2 class corpus with fully labeled parse trees for analysis of the compositional effects of sentiment in language. The dataset consists of 11,855 single sentences extracted from movie re- views. It was parsed with the Stanford parser and includes 215,154 unique phrases from the parse trees, each annotated by 3 human judges.
Stanford Sentiment Treebank-5. Stanford Sentiment Treebank-5 (SST5) (Socher al.. 2013) is 5 class dataset used for sentiment analysis. It consists of 11,855 single sentences extracted from movie reviews. It includes 215, 154 unique phrases from parse trees, each annotated by human judges. Phrases are classified as negative, somewhat negative, neutral, somewhat positive, or positive.
Banking-77. Banking-77 (B77) (Casanueva et al., 2020) is highly fine-grained 77 class classifi- cation dataset comprising 13083 eustomer service queries labelled with 77 intents.
C.1.2 MODEL, OPTIMIZER, & TRAIN SPECIFICATION
Models. We use backbones the Switch Transformer (Fedus al., 2022) and Generalist Lan- guage Model (Du et al., 2022). Table 8 contains the specification over self-attention (SA) layers, feed-forward network (FFN) layers, ofxture-of-Expertt (MoE) layers, attention span (Att. Span),
20

---

Published as a conference paper at ICLR 2025
embedding size and parameter count for both backbones small and medium configurations for each pretraining task. All backbones use 16 experts with top-2 expert routing.
Table 8: Language Modeling Backbone Specifications
Model
SA Layers
FFN Layers
MoE Layers
Att. Span
Embed Size
Params
Wikifext- 103 Pretrain
Switch-small
3
3
256
128
70M
Switch-medium
6
6
1024
352
216M
GlaM-small
6
3
3
2048
144
79M
Gl.M-mmdium
12
6
6
2048
352
220M
En Wik-8 Pretrain
Switch
8
8
2048
352
36M
Optimizer. All experiments use Adam with a base learning rate of 0.0007. Small configurations use 3000 iterations of learning rate warmup while medium configurations use 4000 iterations.
Pretrain Specification. For WikiText- 103 pretraining, small Switch backbones are trained for 40 epochs with batch size of 96 and medium Switch backbones are trained for 80 epochs with a batch size of 48. Small GLaM backbones are trained for 60 epochs with a batch size of 48 and medium GLaM backbones are trained for 120 epochs with a batch size of 48. We use 0.01 auxiliary load balancing loss.
For E pretraining, both Switch and GLaM backbones are trained for 80 epochs with batch size 48. We 0.01 auxiliary load balancing loss.
Finetune Specification. For SST2 and SSTS finetuning, we finetune for epochs using Adam and a base learning rate of 0.001 without warmup and batch size of 16. For B77 we finetune for 50 epochs using Adam and a base elarning rate of 0.00001 without warmup and a batch size of 16.
Compute Resources. All models are trained, evaluated, and finetuned on four NVIDIA A100 SXM4 40GB GPUs.
C.2 IMAGE CLASSIFICATION
C.2.1 DATASETS AND ATTACKS
ImageNet-1K. We use the full ImageNet dataset that contains 1.28M training images and 50K validation images. The model learns to predict the class of the input image among 1000 categories. We report the top-1 and top-5 accuracy on all experiments.
ImageNet-A/O/R. ImageNet-A (Hendrycks et al., 2021b) contains real-world adversarially fil- tered images that fool current ImageNet classifiers. A 200-class subset IK's 1000 classes is selected so that errors among these 200 classes would be considered egregious, which cover most broad categories spanned by ImageNet-IK.
of the original ImageNet-
ImageNet-O (Hendrycks et al., 2021b) contains adversarially filtered examples for ImageNet out-of- distribution detectors. The dataset contains samples from ImageNet-22K but not from ImageNet where samples that are wrongly classified as an ImageNet-IK class with high confidence by ResNet-50 are selected.
Imagenet-R (Hendrycks et al., 2021a) contains various artistic renditions of object classes from the original ImageNet dataset, which is discouraged by the original ImageNet. ImageNet-R contains 30,000 image renditions for 200 ImageNet classes, where a subset of the ImageNet-1 1K classes is chosen.
21

---

Published as a conference paper at ICLR 2025
Adversarial Attacks. We produce corrupted ImageNet samples using white box attacks fast gradient sign method (FGSM) (Goodfellow et al., 2014) and projected gradient descent (PGD) (Madry al., 2017), and black box simultaneous perturbation stochastic approximation (SPSA) (Uesato et al., 2018). FGSM and PGD use a perturbation budget of 1/255 while SPSA uses a per- turbation budget 1. All attacks perturb under Lx norm. PGD and uses 20 steps with step size of 0.15 and SPSA uses 20 iterations.
C.2.2 MODEL, OPTIMIZER, TRAIN SPECIFICATION
Models. Our results are based off of the Swin Transformer (Liu et al., 2021) architecture. This backbone uses 4 base layers of depth 2, 2, 18, and 2. The first two base layers each contain 2 self- attention layers and 2 feed-forward layers. The third base layer contains 18 self-attention layers with alternating feed-forward and MoE layers. The final base layer contains 2 self-attention layers with one feed-forward and one MoE layer. The embedding dimension is 96 and the heads per base layer are 6, 12, and 24. We use 16 total experts and present results for both top-1 and top-2 expert routing. The total parameter count is 280M.
Optimizer. We use A with a base learning rate of 1.25e-4, minimum learning rate of 1.25e- 7, 0.1 weight decay and cosine scheduling.
Train Specification. We train for 60 epochs with a batch size of 128 and 0.1 auxiliary balancing loss.
Compute Resources. All models are trained and evaluated on four NVIDIA A100 SXM4 40GB GPUs.
C.3
ADVERSARIAL. ATTACK AT HIGHER PERTURBATION BUDGET
80
ACMoE (Ours)
ACMoE (Ours)
Swin Transformer
Swin Transformer
Test Test Accuracy
60-
02
25-
40-
1.-
10
20
1.0
1.5
30
3.0
35
4.5
4.5
1.5 2.0
2.5
3.5
4.5
5.0
Perturbation Budget
Perturbation Budget
Figure 3: ACMoE and Swin Transformer under PGD attack at increasing perturbation budgets. ACMoE widens its performance gain over Swin at increasingly severe attacks in both top-1 test accuracy (Ieft) and top-5 test accuracy (right), starting approximately 7% improvement 1/255 and ending at just over 10% 5/255.
Figure shows that for PGD perturbation budgets 1/255 through to 5/255, ACMoE widens its al- ready substantive robust performance gain over Swin, with top-1 1 and top-5 test accuracy improve- ments increasing from 7% to approximately 10%.
C.4 CLUSTER VISUALIZATION
We pass random ImageNet batches through Swin and ACMoE and plot the representations along with their assigned experts, using t-sne to represent the high dimensional data in 2 dimensions. The result is shown in Fig. 4. where we see Swin learns overlapping and indistinguishable expert clusters. ACMOE, on the other hand, performs better in learning the clusters, producing much clearer and detter-distinguishee clusters.
22

---

Published as a conference paper at ICLR 2025
-13
13
00
65
Swin Transformer
ACMoE (Ours)
Cluster Visualization on ImageNet. Each token is represented as point and colored by its assigned expert. Left: Swin identifies one cluster clearly (yellow/gold) but otherwise fails to distinguish remaining clusters Right: ACMoE learns better-defined expert clusters.
Figure 4:
Table 9:
Ablation Measure of Spread in Switch Transformer (Fedus et al., 2022)
Table 10: Ablation Layer Placement in
Switch Transformer (Fedus et al., 2022)
Measure of Spread
Test PPL (4)
Layer Placement
Test PPL (L)
Variance
34.87
Back Half
34.95
MAD
34.42
Alternating
34.80
Skip Full
34.42
34.88
C.5 ABLATION STUDIES
C.5.1 MEASURES OF DISPERSION
We present in Tables 9 and 11 results for Switch-ACMoE and Swin-ACMoE when changing the measure of dispersion used in AC routing transformation (Definition 1) from mean absolute deviation (MAD) to variance. We see mean absolute deviation outperforms variance as a measure of spread. This is an intuitive finding given that squared distances, as used in variance computa- tions, highly sensitive to outliers. Using mean absolute deviation as alternative measure of spread reduces this issue and produces a more robust estimate of dispersion. We note that MAD addi-
not the only robust measure of spread. We conjecture that taking interquartile range tionally robust measure spread may produce good results in both clean and contaminated data. We, however, leave this interesting direction future research interquartile range poses im- plementation challenges as it requires designing concurrent linear scans over the expert clusters. MAD, by contrast, requires just two computations of the mean which is easily parallelizable using torch. index ().
C.5.2 LAYER PLACEMENT
We consider the effect of layer placement in the Switch-medium configuration and in the Swin Transformer (see Sections C.1.2 and C.2.2 for the full model specifications). In particular, Switch is a layer model and Swin is a 24 layer model. With regard to Swin, we focus on the deepest block of depth 18 implement our ACMoE layers. This is due the change embedding size between base layers, meaning we are restricted to this base layer of depth 18. Note further that Swin only uses MoE layers an alternating pattern with feew-forward networks between each MoE layer. For example, for Switch, a full ACMoE specification would mean placing ACMoE on layers 2,3,4,5,6. For Swin, a full specification means placing ACMoE on layers 10,,,,10,,2,11,,6, To examine the effect of layer placement we consider the following models:
23

---

Published as a conference paper at ICLR 2025
Table 11: Ablation on Measure of Spread in
Table 12: Ablation on Layer Placement in Swin Transformer
Swin Transformer
Test Acc. Top 5
Test Acc.
Measure of Spread
Layer Placement
Top
Top
Top 5
Swin-Top/ (Liu et al., 2021)
Swin-Topl (Liu et al., 2021)
Variance
75.06 75.39
92.49
Back Half
75.16
92.46
MAD
92.56
Skip 2
75.34
92.42
Skip Full
92.45 75.35
Swin-Top2 (Liu et al., 2021)
75.39
92.56
Variance MAD
76. 11 76.31
93.08
Swin- Top2 (Liu et al., 2021)
93.14
Back Half
76.16
93.02
Skip 2
76.10
92.93
Skip
76.29
92.98
Full
76.31
93.14
Alternating: For Switch this means we place ACMoE on layers 2,4,6. For Swin this means we place ACMoE on layers 1,8.12,16.
Back Half: For Switch this means we place ACMoE on just the last 3 layers of the network. For Swin this means we place ACMoE on just the last 5 layers of the network.
Skip 2: For Swin this means we palce ACMoE on layers 8,10,12,14.16,18
Skip I: For Switch this means we place ACMoE on layers 3,4,5,6. For Swin this means we place ACMoE on layers 60,11.,12.14.1.111
* Full: We place ACMoE on every possible layer.
We present Table results for Switch and Swin ACMoE models when changing the positions of the ACMoE layers throughout the network. The results agree with our expectation that, generally speaking, more ACMoE layers improve performance, but a in some circumstances a threshold is met at the point where ACMoE layers are used too early in the network such that the model has not been able to learn reasonably good approximations of the cluster membership of the tokens yet.
We find that in the Switch backbone, performance improves the more ACMoE layers add, which agrees with our expectation that more ACMoE layers improve performance. However, we find that top performance attained when allowing two standard MoE layers to go before the first ACMOE, as opposed to the minimum of I standard MoE layer. We conjecture this is because we need to give the model few layers before the first ACMoE order learn decent representations such that we have good enough estimated cluster assignments for use in the ACMoE layer. Encouragingly, we find just one additional standard MoE. layer is sufficient for the benefits of ACMoE to be obtained.
We find in Table 12 that with Swin, best performance is obtained using ACMoE on every possible layer, again agreeing with our expectation that more ACMoE layers improve performance. With Swin, however, we do not face any drop in performance from placing ACMoE too early in the network, and indeed we see Full attaining top performance. We conjecture that Swin does not encounter this issue since Swin uses four layers of feed forward networks before the first MoE layer, and so by the first MoE layer the representations are of reasonably good quality to produce good estimates of the cluster membership.
C.5.3 RANDOM ABLATION
We show the efficacy of the adaptive clustering transformation M (Definition 1) in our AC router at capturing meaningful feature-wise information by ablating it against an alternate x diagonal matrix made up of normal random variables with mean I and standard deviation 0.5 (where we clip any negative values to prevent negative weights). We present in Tables 13 and 14 results for lan- guage modeling (using Switch) and image classification (using Swin), which show fairly substantial drops in performance in both backbones. This offers evidence to the claim that our AC routing transformation meaningfully weighting features to improve routing, and that performance gains
24

---

Published as a conference paper at ICLR 2025
of our proposed method do not flow from kind of implicit regularization of introducing noise into the router.
Table 13: Random Ablation in Switch (Fedus
Table 14: Random Ablation in Swin (Liu et al., 2021)
et al., 2022)
Model
Test PPL
Model
Top Acc.
Top 5 Acc.
Switch- (Fedus et al., 2022)
38.17
Swin-Random Swin-ACMoE
74.22
91.87
Switch-ACMoE
34.42
76.31
93.14
C.6 CLUSTER WEIGHT MIXING
The AC routing scheme estimates the cluster membership of each token based on its highest affinity cluster assigned in the previous layer. We could also further leverage the top-k structure of the MoE models by mixing the cluster-wise feature weights with weights corresponding to the affinities the top-k routing. For example, if h has affinity scores and 1-a to clusters & and k' respectively, then we could also obtain the required AC routing transformation for as M = o., + (1 Mg-. This approach therefore factors in the confidence with which we believe h belongs to cluster k or K, and can be used for integrating ACMoE into higher expert granularity backbones (i.e higher top-k settings). Tables 15 and 16 show results for computing Mg- by mixing the top-affinity cluster weights (Mix 2) in S Switch and GLaM with top-2 routing, versus our presented results which compute M&. just based off of the highest affinity cluster (Mix 1). We see that GLaM-ACMof benefits substantially from cluster weight mixing whereas Switch-ACMoE prefers just using its top affinity cluster weights. For consistency across models, we present in our main body the Mix I results, as Gl GLaM-ACMoF. already performs extremely strongly using Mix 1 and so we prefer to opt for the added performance gain the Switch backbone.
Table 15: Results on Cluster Weight Mixing in Switch (Fedus et al., 2022)
Table 16: Results on Cluster Weight Mixing in GLaM (Du et al., 2022)
Clusters Mixed
Test PPI. ())
Clusters Mixed
Test PPI. (L)
Mix Mix
34.66
Mix 2
35.29
34.42
Mix 1
36.26
C.7 ADAPTIVE CLUSTERING INTEGRATION INTO SOFT MIXTURE OF EXPERTS
We present here results for integrating ACMoE into SoftMoE (Puigeerver al., 2023). To use ACMoE in the SoftMoe setting, which can be be understood as a top-E routing setting where all experts are active for every token, we compute M₁- using cluster weight mixing (Section C.6) over the top-8 highest affinity clusters. We present the performance of Soft-ACMoE on clean data, adversarially attacked data, and ImageNet-A/O/R in the following Tables 17 and 18.
Table 17: Test Accuracy on ImageNet corrupted PGD, FGSM, and SPSA using SoftMoE
(Puigcerver et al., 2023) backbone
Clean Data Top
PGD
FGSM
SPSA
Model
Top 5
Top
Top 5
Top 1
Top 5
Top 1
Top 5
SoftMoE (Puigcerver et al., 2023)
72.86
90.92
45.29
78.91
56.95
85.60
66.59
88.70
Soft-ACMoE (Ours)
73,21
91.23
48.25
80.49
59.01
86.69
70.63
93.22
We see in Tables 17 and 18 the efficacy of ACMoE in the SoftMoE backbone, offering evidence of the adaptability of our framework into further MoE setups. In particular, the SoftMoE framework models a setting in which expert clusters are highly overlapping, as each token is soft assigned to all experts. Therefore, the performance gains shown in clean and contaminated data of SofE-ACMoE demonstrates that our AC router is well-suited to modeling such clustering structure.
25

---

Published as a conference paper at ICLR 2025
Table 18: Test Accuracy on Image Classification in Imagenct-A/O/R using SoftMoE (Puigcerver et al., 2023) backbone
Im-A Top-1 Ace. (t)
Im-R
Im-O
Model
Top-1 Acc. (t)
AUPR (1)
SoftMoE (Puigeerver al., 2023)
6.69
31.63
17.97
Soft-ACMoF (Ours)
6.93
32.18
18.35
C.8 IMAGE CLASSIFICATION IN SWIN TRANSFORMER BASE CONFIGURATION
We further evaluate the performance ACMoE when scaling up model size in Table 19, We integrate ACMoE into the Base configuration of Swin (0.5B parameters) and evaluate on clean ImageNet-IK as well as under adversarial atacks.
Table 19: Test Accuracy on ImageNet corrupted PGD, FGSM, and SPSA using Swin Base (Liu
al., 2021) backbone
Clean Data
PGD
FGSM
SPSA
Model
Top 1
Top
Top
Top
Top 1
Top
Top 1
Top
Swin-Base (Liu et al., 2021)
79.06 79.25
94.37 94.42
44.61 46.28
79.20
59.91
87.72
68.94 70.18
89.00
Swin-ACMoE Masse (Ours)
80.24
61.78
87.55
89.33
C.9 ROUTER STABILITY
We present in Fig. 5 the routing stability of ACMOE, SMoE, XMoE, and StableMoE in the Switch backbone evaluated on WikiText-103 Routing instability computes over adjacent layers the propor- tion of tokens that are assigned different experts across the two layers. Specifically, for tokens (hi,..., J. we compute at layer & the matrix S RRax such that S if the 7th and jth tokens are assigned the same expert in layer & and otherwise. The router instability at layer can then be calculated as of = me mean ([S 1).. This metric therefore captures the degree to which tokens that are assigned to the same experts remain together through the model. A high of indicates the router doesn't i maintain consistent expert assignments, tokens that it considers semantically similar at one layer it considers different at the next.
SMoE
StableMoE
XMcE
ACMoE
Router Instability
Layer
Figure 5: Router Instability of ACMOE, SMoE, XMoE, and StableMoE, ACMoE maintains consis- tent routing, while baseline routers more frequently change the expert assignments of tokens.
In Fig. 5, we see that baseline routers reach high levels of instability, where the case of SMoE and StableMoF, at the last layer over 60% of tokens are assigned to different expert. ACMoE, by contrast, maintains a more consistent, stable assignment through the model, with no more than 20% of tokens changing expert assignment across any layer.
C.10 DYNAMIC ROUTING
We further test the compatibility our Adaptive Clustering routing scheme dynamic top-p rout- ing. In this setting, rather than routing each token its top-k highest affinity experts in each MoE
26

---

Published as a conference paper at ICLR 2025
layer, route each token to all experts that have affinity over certain threshold p. This setting permits activating more or less experts for different tokens at different layers throughout the model, therefore dynamically assigning experts to tokens. We integrate our AC routing directly into this set- ting using the same setup as in Section 3, where the AC routing transformation computed based on the estimated cluster membership of each token using the top affinity assignment of the previous layer. We present the results for Switch transformer on WikiText- 103 language modeling in the following Table 20.
Table 20: Results on Top-p Dynamic Routing Switch Backbone (Fedus et al., 2022)
Model
Test PPL (4)
Fixed top-k routing (Shazeer al., 2017)
Switch-medium (Fedus et al., 2022)
35.48
ACMoE-medium (Ours)
34.42
Dynamic top-p routing (Guo et al., 2024)
Switch-Fixed p
35,20
Switch-ACMoEF-iixll p (Ours)
34.14
Seitch-Learrable p
34.29
Baaechh-CME--Laaaabaee (Ours)
33.49
For fixed p. we set p 0.05. For learnable p. we initialize the parameter 0.05. We select this ini- tialization it reproduces approximately similar performance in the Switch backbone under default top-2 routing, thereby aiding direct comparison between fixed top-k and dynamic top-p routing. We in the dynamic routing setting. ACMoE maintains the same consistent improvement Switch baseline of roughly full PPL. These results suggest ACMoE is well-suited to the dynamic routing setting.
the
D BROADER IMPACT
Our research offers benefits to Mixture-of-Exper (MoE) architectures in both clean and contami- nated settings. In particular, our work offers socially beneficial outcomes with regard to defense against adversarial attack, which we hope can be used protect important Al systems from mali- cious actors. Furthermore, as large language models, many of which are built on MoE backbones, continue profligate and be used in important societal settings, to data contamination can aid this promising technology to continue to grow and improve in realistic settings of noisy training and evaluation data. Our research also shows substantially faster conver- gence than comparative baselines. We believe this faster convergence can deliver significant social benefit in terms reducing the energy requirements large model training, thereby helping to ease the growing environmental burden of Al training runs. We recognize there will always be risk of misuse with AI systems, however we hope that our work can be used to enhance and protect socially beneficial AI while also decreasing the environmental impact of this technology. We furthermore hope that our research can spur others on to continue building on robust and efficient AI for social good.
hope our improved robustness
27


Reviews:
[{"review_id": "TacEwoibZf", "reviewer": "Reviewer_UqYP", "summary": "This paper proposes an Adaptive Clustering (AC) router for Sparse Mixture-of-Experts (MoE) architectures to improve routing efficiency. It focuses on optimizing the matching of tokens to experts by adaptively transforming the input feature space. This method promotes faster convergence, increased robustness against data contamination, and improved overall performance. The authors demonstrate the AC router's effectiveness in large-scale language and image tasks, outperforming baseline routers in robustness and efficiency without added learnable parameters.", "questions": "1.\tHow does the AC router handle scenarios with overlapping or dynamically evolving clusters where expert specialization may be less clear?\n\n2.\tWhat mechanisms could be integrated to enhance adaptability in routing without relying on previous layer assignments?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tThe proposed router uses a unique, feature-weighted approach for adaptive clustering, enabling more efficient specialization in MoE models. It achieves improvements without extra learnable parameters or added computation cost, making it ideal for scalable applications.\n\n2.\tThe paper provides robust theoretical proof supporting the method's faster convergence and increased robustness\n\n3.\tThe experiments cover various large-scale tasks, such as WikiText-103 and ImageNet, and include backbones like Switch Transformer, GLaM, and Swin Transformers, evaluating performance under both clean and corrupted data conditions.", "weaknesses": "1.\tThe approach seems to tailor to settings with well-structured clusters, which may limit performance on datasets or tasks where latent clusters are less distinct or not aligned with chosen features.\n\n2.\tThe AC router relies on prior layers for initial expert assignments, potentially constraining flexibility and requiring consistent embedding sizes across layers.\n\n3.\tAlthough tested on large-scale datasets, further application to more dynamic datasets could strengthen claims of robustness and adaptability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "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": "49SQNxbpYG", "reviewer": "Reviewer_LDye", "summary": "This paper proposes a novel routing mechanism called Adaptive Clustering (AC) router for Mixture-of-Experts (MoE) architectures. The key idea is to compute token-expert assignments in an adaptively transformed space that better reveals latent clusters in the data. The transformation is derived from a feature-weighted clustering optimization perspective, where features that promote tight clustering for each expert are upweighted. The authors demonstrate both theoretical and empirical advantages of their method, showing improved convergence speed, robustness to data contamination, and overall performance across language modeling and image classification tasks.", "questions": "1. The choice of mean absolute deviation (MAD) as the measure of dispersion is justified, but the paper does not deeply explore the impact of other potential measures, such as variance or interquartile range, on performance. Have you considered or tested other measures of dispersion besides MAD? How sensitive is the method to the choice of dispersion measure, and could alternative measures potentially improve performance or robustness?\n\n2. How does the AC router handle situations where the assignments from previous layers are noisy or not well-defined? Is there a way to initialize or adjust the method to be effective in early layers or when prior assignments are unreliable?\n\n3. Since the method requires consistent embedding dimensions between layers, how can it be adapted for architectures where the embedding size changes, such as in some convolutional neural networks or transformers with variable dimensions?\n\n4. The theoretical analysis relies on Gaussian mixture models. How does the AC router perform when the data clusters have non-Gaussian distributions or are not well-separated?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The paper introduces a novel perspective by viewing the MoE routing mechanism through the lens of feature-weighted clustering optimization. This theoretical framing allows for the derivation of optimal feature weights that improve cluster separability.\n\n2. The authors provide rigorous theoretical analysis that support their claims on the improved robustness and convergence speed of their method.\n\n3. Extensive experiments on large-scale datasets for both language modeling and image classification demonstrate the effectiveness of the AC router.\n\n4. The proposed method requires no additional learnable parameters and introduces negligible computational overhead. And it can be integrated into any existing MoE architecture", "weaknesses": "1. The AC router relies on expert assignments from the previous layer to compute the adaptive transformation. This dependence may limit its applicability in scenarios where embedding sizes change between layers or in networks without consistent layer structures.\n\n2. The effectiveness of the AC router is closely tied to the quality of expert assignments in earlier layers. Since the method relies on routing decisions from prior layers, it may be less effective if the early layers have not yet learned meaningful or well-separated cluster structures. In practice, early layers in deep networks often learn basic features that may not be semantically distinct enough to form tight clusters. As a result, the AC router may struggle to make optimal routing decisions in later layers if the cluster assignments in the initial layers are poor.\n\n3. The method assumes that the input data naturally clusters in the feature space. While this assumption holds in many cases (e.g., language modeling and image classification), it might not generalize to tasks where the input data is more uniformly distributed or lacks clear clustering patterns.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "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": "K6YnjTaexr", "reviewer": "Reviewer_D6YM", "summary": "the paper studies policies for selection of experts in sparse MoE models, and casts the problem as optimal clustering\n\nthe main contributions are 1) to derive closed form expression for the cluster weights. the formulation elegantly extends the classic topk MoE formulation using a diagonal matrix using inverse of cluster dispersion measure,  2) to theoretically bound the incorrect assignment probability and 3) prove faster convergence (though is not clear under which assumptions the prove holds\n\nthe paper presents experimental results on  text and image classification, mostly presenting results over baseline models", "questions": "no further questions", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- the theoretical framework is solid, with clear explanations and remarks on the main body, and abundant material (only skimmed through) in the appendix concerning proofs  of the framework \n\n- the paper is clearly written and motivated. \n\n- the experimental section tackles two complementary domains and, whereas has less critical mass than the theoretical part, it is well executed\n\n- extensive material in the appendix complement and completes the paper", "weaknesses": "# necessity of the framework\n\nwhile the theoretical framework provides and the analysis in Tab 7 shows the method to only provide a rather modest additional complexity,  and the theoretical proofs shows it to be (in a sense) sufficient to solve the routing task, an open question is whether it is also necessary for the task.\n\notherwise stated, in the case of top-1/top-2 expert selection, an open question is whether it would be possible to significantly simplify the operations to be carried on (same way you do not need a full sort if you only need the maximum value of a vector) to attain similar performance ?\n\nsimilarly, a more systematic comparison with related competitor work would better allow to appreciate the practical impact.\n\n# structural difference of the solution \n\nauthors do a fair job in visually shown differences in the top experts (Fig 1), but that remains at an anecdotal level.  The load balancing analysis in sec 4.3 provides some complementary view, but is too aggregated (as it looks at measures of variation, but the baseline of perfect load balancing is clearly not a target as it likely would not be attainable).\n\nas such, assessing and quantifying the /differences/ at individual expert level, or input level would provide a factual quantitative assessment of the number of times ACMoE expert selection differs from competing selections. (ie. instead of maps in Fig, providing a distribution of the times where top-k selection differed).  \n\nAs in many cases the end-to-end performance improvement is slim (few percentage points) this intermediate assessment could allow to perceive the amount of top-k choices that were changed wrt the baselines models.\n\n\n# deviation measures\n\nit is unclear to what extent the proofs (and the practical performance) of the framework depends on subtle hyperparameter choices that are fixed and  not considered in the paper: e.g., measure of dispersion used in sec 3 is the mean absolute deviation, which is easier to compute but less robust than median absolute deviation (admittedly, this would go in the opposite direction as for simplification; but is unclear to what extent practical details as this one can have a performance impact in practice).\n\nto have no impact on the selection, the top-k rank would need to be maintained across a family of measures of dispersion -- the impact of which is far from being clear after reading the paper. If the proof holds for families with given properties, the paper could be reinforced by making that clear.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "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": "Q8wPtpRNMY", "reviewer": "Reviewer_2Ypo", "summary": "The paper presents an Adaptive Clustering (AC) router for MoE models. The Ac router offers faster convergence, better robustness, and overall performance improvement.", "questions": "None", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The solution is theoretical based along with extensive evaluation.", "weaknesses": "- The paper targets MoE experts only, some of the notations and terms seem to be used specifically in the MoE domain.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "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": {}}]

Author rebuttal:

Decision: Accept (Poster)