CAMEx: Curvature-aware Merging of Experts
Authors: Viet Dung Nguyen
Minh Nguyen Hoang
Luc Nguyen
Rachel Teo
Tan Minh Nguyen
Linh Duy Tran
Venue: ICLR 2025 Poster (2025)
Keywords: Sparse Mixture-of-Experts
efficiency
expert merging
Primary area: foundation or frontier models, including LLMs
TL;DR: We introduce CAMEx (Curvature-Aware Merging of Experts), a novel expert merging protocol that incorporates natural gradients to account for the non-Euclidean curvature of the parameter manifold

Abstract: Existing methods for merging experts during model training and fine-tuning predominantly rely on Euclidean geometry, which assumes a flat parameter space. This assumption can limit the model's generalization ability, especially during the pre-training phase, where the parameter manifold might exhibit more complex curvature. Curvature-aware merging methods typically require additional information and computational resources to approximate the Fisher Information Matrix, adding memory overhead. In this paper, we introduce CAMEx (Curvature-Aware Merging of Experts), a novel expert merging protocol that incorporates natural gradients to account for the non-Euclidean curvature of the parameter manifold. By leveraging natural gradients, CAMEx adapts more effectively to the structure of the parameter space, improving alignment between model updates and the manifold's geometry. This approach enhances both pre-training and fine-tuning, resulting in better optimization trajectories and improved generalization without the substantial memory overhead typically associated with curvature-aware methods. Our contributions are threefold: (1) CAMEx significantly outperforms traditional Euclidean-based expert merging techniques across various natural language processing tasks, leading to enhanced performance during pre-training and fine-tuning; (2) we introduce a dynamic merging architecture that optimizes resource utilization, achieving high performance while reducing computational costs, facilitating efficient scaling of large language models; and (3) we provide both theoretical and empirical evidence to demonstrate the efficiency of our proposed method. The code is publicly available at: https://github.com/kpup1710/CAMEx.

Published as a conference paper at ICLR 2025
CAMEX: CURVATURE-AWARE MERGING OF EXPERTS
Dung V. Nguyen Minh H. N 1*
Luc Q. N
Rachel S.Y. Teo
Tan M. Ng Linh Duy Tran
Faculty of Mathematics and Informatics, Hanoi University of Science and Technology Viettel AI, Viettel Group
*Department of Mathematics, National University of Singapore
{dung. Stung.nvv3221 nww322122211m 5M, minh. htb3hh31223333333503 hust edu. vn [luenql, linht (1attdd5feiiette.. com. rachel. tsyēu. u. nus edu, tanmnenus edu. sg
ABSTRACT
Existing methods for merging experts during model training and fine-tuning predominantly rely on Euclidean geometry, which assumes flat parameter space. This assumption can limit the model's generalization ability, especially during the pre-training phase, where the parameter manifold might exhibit more complex cur- vature. Curvature-aware merging methods typically require additional information and computational resources to approximate the Fisher Information Matrix, adding memory overhead. In this paper, we introduce CAMEx (Curvature-Aww Aware Merging of Experts), a novel expert merging protocol that incorporates natural gradients to account for the non-Fuclidean curvature of the parameter manifold. By leveraging natural gradients, CAMEX adapts more effectively to the structure of the parameter space, improving alignment between model updates and the manifold's geometry. This approach enhances both pre-training and fine-tuning, resulting in better optimization trajectories and improved generalization without the substantial memory overhead typically associated with curvature aware methods. Our contributions are threefold: (1) CAMEx significantly outperforms traditional Euclidean-based expert merging techniques across various natural language processing tasks, leading to enhanced performance during pre-training and fine-tuning: (2) we introduce a dynamic merging architecture that op- timizes resource utilization, achieving high performance while reducing computational costs, facilitating efficient scaling large language models; and (3) we provide both theoretical and empirical evidence to demonstrate the efficiency of our proposed method. The code is publicly available at: https: / /iihhuu.... com/ kwkp wwppp111/// O/AMEX..
INTRODUCTION
Sparse Mixture of Experts (SMoE) (Jacobs et al., 1991; Shazeer et al., 2017) is currently core component for constructing foundation and large language models (LLMs), whose parameters count can rise up to billions and trillions (Devlin al., 2019; Yang et al., 2019; Liu al., 2019; Raffel et al., 2020; Nguyen et al., 2022b; Fedus et al., 2022; Wei et al., 2022; Teo & Nguyen, 2025). Nevertheless, (Hoffmann et al., 2024; Kaplan et al.. 2020) recognized scaling law that underpins the LLM's evolution, which larger models require exponentially more computational resources and data to continue improving, and without sufficient scaling in all dimensions, performance gains may plateau. Thus, identifying and implementing efficient methodologies for the sustainable scaling of L.MSs is imperative. SMoE: addresses this challenge by sparsely activating parameters of large models, which can boost model performance with only minor losses in computational efficiency. The methodology is integrated chiefly into feedforward layers of transformers, processing tokens by selectively activating a small number of experts and hence trimming down the computing memory and FLOPS (Fedus et al., 2022; Lepikhin al., 2021).
Since its debut in (Shazeer et al., 2017), SMoE has gone through numerous explorations and advancements in routing mechanism development and expert architecture design. (Dai al., 2022) proposes two-phase
"Co-first authors Coollsst authors

---

Published as a conference paper at ICLR 2025
training strategy for stabilizing the gate function that the expert's selection of one token does not fluctuate between different inference times. (Zhou et al., 2022) changes the perspective of the router to experts with experts choice routing, ensuring a balancing load between experts. (Teo & Nguyen, 2024) incorporates heavy-ball momentum in SMoE to enhance the model's stability and robustness. (Chi et al., 2022) and (Chen et al., 2023) address the concern of representation collapse in SMoE by proposing cosine scoring and fixed random initialized router, respectively. Some other works view the routing mechanism as a reinforcement learning, clustering, or optimal transport problem (Nielsen al., 2025). In terms of expert design orientation, (Rajbhandari et al., 2022) and (Dai et al., 2024) introduce the concept of shared experts wherein each token is processed by a fixed expert and another selected through gating. achieving two experts engagement per layer without increasing the communication cost beyond that of top-1 gating. (Muqeeth et al., 2024) proposes to merge experts by taking the weighted mean of the expert's parameters with respect to router scores. This methodology is then extended in (He et al., 2023), (Zhong et al., 2024), and (Li et al., 2024) for causal language modeling pretraining and fine-tuning tasks.
Among existing rigorous research on SMoE, our work focuses the experts merging lines of research. Specifically, we systemically integrate natural gradient into task-specific merging protocol for SMoE. To the best of our knowledge, the current merging protocol applied for SMoE still deems the parameter space of the expert's 's parameters as Euclidean ones. Nevertheless, it has been shown that the space of neural network parameters brings the characteristic of the Riemannian manifold (Amari, 1998). Therefore, it is natural for us to make an effort in such direction for merging experts. Although some existing works on merging models have already leveraged the Fisher Information Matrix (Matena & Raffel, 2022: Jin al., 2023), we find that they require large computational space and complicated steps to perform well. In contrast, our merging protocol is simple and straightforward implement while still taking into account the curvature of the parameters manifold. We discover the superior performance of curvature- aware merging in our method compared the regular merging procedure applied to SMoE. Our main contributions are three-fold:
1. We introduce a novel rapid and efficient merging technique named Curvature- Aware Merging of Experts (CAMEx) for SMoE that includes information about the curvature of the expert's parameters manifold.
2. We propose new architecture based on CAMEx, which dynamicalizes the merging protocol along with parameters reduction. Our empirical experiments prove the dominant performance of this architecture on pre-training tasks.
3. We theoretically prove that our CAMEx obtains better alignment between experts and the training task domain.
We empirically demonstrate that 1) our proposed merging method can add in rapidness of convergence speed for pre-training and 2) when combined with other merging protocols, it can boost the model's performance a variety of practical tasks, including language modeling, text classification, question answering, and image classification.
2 CURVATURE-AWARY MERGING OF EXPERTS
This section aims to give an overview of model merging methods and their integration into SMoE architecture. We then introduce curvature-aware merging protocol stamping from the natural gradient. Finally, we perform a theoretical analysis to support our proposal.
2.1 BACKGROUND: EXPERT MERGING IN SPARSE MIXTURE OF EXPERTS
It is convenient to recall the concept of SMoE and few well-known experts merging methods for SMoE. From this point to the rest of our paper, let us use the notations summarized in Table 1.
Sparse Mixture of Experts. A SMoE layer processes the tokens series follows:
y: = Eies, G(t.i) E, - softmax (W, S.
SMEAR. (Muqeeth et al., 2024) introduces the ensemble of expert parameters through weighted average computation with the factors being the router scores.
2

---

Published as a conference paper at ICLR 2025
Ouiput (4 44)
Ship ghadeac
M
Mash Proeddav
Gegrnen (
Figure 1: Overview of CAMEx for a causal language modeling SMoE. The experts are merged through the router scores and the curvature-matrix M. During the merging protocol, we generate the masks for the domain-vectors, denoted as "Ks such as Ties or Dare. We follow the causal segmenting pipeline from (Zhong et al., 2024) achieve both memory efficiency and causal information constraints. Note that stop gradient operator is applied for the first segment router scores.
Table I: Notations and Definitions.
Symbol
Description
Symbol
Description
Number tokens
Number of selected experts
N
Total number of experts
E G(.))RTXX at [0.1]
Weights for the ith expert
hERTXd S.
Input tokens or hidden states Set of top-k experts for token h,
Gating function output
Rescalling factor
Task-specific merging in SMoE. Our work will follow the scheme of task-specific merging (Ilharco et al., 2023). In such a setting, we assume the existence N pre-trained models parameterized by 6, each of which was pre- pre-trained through the merged model Om as = 02 02-0m. The merging protocol will be performed by Eqn. Merg. Under the context of SMoF, each expert learns to handle particular subset of the input space specializes in a specific type of feature or pattern (Jacobs et al., 1991; Dai et al., 2024). We believe it is more suitable to reference this technique as domain-speeifie merging. We, therefore. will rename the tensors T, E, E,EE... as domain-vectts vector. Additionally, to take the router information into account, we will define the formulation for domain- speciffc merging in a SMoE layer as follows:
different task. We then define the task-vector for each pre-traine pre-trained model
É.. Em +o
8:7
where s; denotes the score of the router for the ith expert. We want to note that with a < domain-specifee merging aligns with soft merging.
2.2 BACKGROUND: OTHER MODEL MERGING METHODS
this section, we discuss other recent and aidely-adopted model merging methods outside the context of SMoE that will combine with our curvature merging method in experiments in Section 3.
TIES merging, (Yadav et al., 2023) improves upon task arithmetic by removing interference between the task vectors. Specifically, TIES zeros out entries in given task vector with low magnitude and resolves sign conflicts across different task vectors.
DARE merging, Different from TIES, DARE merging randomly zeros out the neurons like Dropout layer (Yu et al., 2024).
Fisher merging, Existing work Fisher merging suffers from computational complexity since computing and inverting the Fisher Information Matrix, especially for large neural networks, is often intractable. Even when using approximations like diagonal or block-diagonal Fisher matrices, these methods can
3

---

Published as a conference paper at ICLR 2025
still be computationally expensive and challenging apply at scale. Furthermore, the accuracy of Fisher approximations, such as diagonal or block-diagonal, can be problematic (Matena & Raffel, 2022).
2.3 GRADIENT INTERPRETATION OF MODELS MERGING
We want emphasize the alignments between the paradigm of gradient descent and model merging. For this, we denote 0 E R. [()), and 17 as the model's parameters, the empirical loss function, and the learning rate, respectively. During the training process deep learning model, the parameters are updated following the gradient descent formula:
8ut1 0n + (-V()
(GD)
In the aspect of deep models merging, we also have an update rule in a similar manner, which is
n T
(Merg)
gradient-like like updatte direction
where tt denotes the merged model's parameters, and f' denotes the parameters of the ith expert. Here, we interpret 6 as the optimal parameters of the model for a specific task or domain, and then the update rule gives direction toward optimizing for all tasks.
However, it has been pointed out by (Amari, 1998) that the parameter space structure of deep learning models has Riemannian characteristics. Therefore, more natural gradient updating scheme was proposed,
0 = 0 C
(NGD)
natural gradient
In this formula, G(O,,) N denotes the Riemannian metric tensor (Amari, 1998; Amari & Douglas, 1998), which characterizes the intrinsic curvature a particular manifold in N-dimensional space (Martens, 2020) or sometimes, the inversed Fisher Information Matrix. The same ideology was introduced for merging large language models in Fisher merging (Matena & Raffel, 2022) and Regmean (Jin et al., 2023). However, both methods suffer from the bottleneck in the computation cost of approximating the Fisher Information. Morcover, these methods are challenging to apply in sparse layers SMoE since they would introduce huge latency, FLOPS, and memory for computing and storing matrices whose number of entries proportional to a number of expert parameters.
2.4 EXPERTS MERGING WITH CURVATURE-
In this section, for the sake of conciseness, we focus on the language modeling task; similar methodology can be applied to other tasks, such as classification. We introduce an efficient way to merge experts within SMoE layers, based on the causal segmenting approach proposed by (Zhong et al., 2024). The goal of the causal segment routing strategy is to enhance the efficiency of expert merging operations while maintaining the autoregressive nature of language models. More details about this algorithm can found Appendix C.I and Algorithm 1. We then perform the following merging protocols:
E = E ta Σ
M. M. -s(=7///
(CA-Merg)
where M, E Redinoout dou denote the curvature matrix which performs matrix product with the gradient-like component. The curvature of the parameters manifold will be learned through these tensors while optimizing the empirical loss. This approach has also proven its effectiveness in meta-learning for few-shot classification (Park & Oliva, 2019). We further explore the computing efficiency of merging experts by proposing novel dynamic merging formula
N-1 MM... a N
E4+1
=
E m
+
N-1 1=1 +
(DyaamicMeryy
p411
=
1=1
4

---

Published as a conference paper at ICLR 2025
Output
Output
Output
Merging Protocol
Merging Protocol
Merging Protocol
Fi
R2
Fx
F
Fg
T
F
Fxy
Router
Router
Router
Input
Input
Input
Figure 2: Overall architecture of different SMoE layers. The figure presents the vanilla SMoE layer the left, the merging expert layer in the middle, and our proposed dynamic merging SMoE. layer on the right. Our architecture reduces the number parameters compared to the other two, while maintaining the same number of activated neurons per layer. Importantly, despite the dynamic merging mechanism, our architecture preserves the same number of experts at each layer as the other SMoE architectures, ensuring comparable model capacity, i.e., the number of activated parameters per layer.
The architecture corresponding to this recurrent representation can found Figure The architecture contains global expert that traverses through the SMoE layers by the updating rule in Eqn. Dynamic-Merg, Not only will this allow notable reduction model size and GFLOPS, but it also ensures the number of experts in each SMoE is the same in the full-expert setting, where each layer has the same number of experts. We refer to Appendix for a step-by-step walkthrough of key equations in CAMEx.
2.5 EFFICENCY
Parameter effficint approximation of curvature matrix. Storing and computing curvature matrix requires a whopping memory and time complexity of O(n²) and O(z²), respectively. This infeasible even for simple SMoE layer, as layer can contain many experts. To mitigate this problem, we follow (Martens & Grosse, 2015) and approximate the curvature matrix using the Kronecker product. It has been proven by (Hameed al., 2022) that can approximate arbitrary matrix using a finite of Kronecker products. For a curvature matrix M, e Ipd dout Xduadout, we present the rank- 1 approximation as below:
M &M.
33)
with M RRdi..d. and M EIR Xdouc, Still, this form of approximation is too large compute and store during training time, so we further decompose M and M using Kronecker product because of the efficient computation using tensor algebra. This form of approximation reduces the number of parameters added and only puts negligible memory and computational overhead to the training process at the cost of additional (n) memory complexity and ((n2.)) computational complexity. Although this might limit the representative capacity of the curvature matrix, empirically find that the performance of our method still surpasses other merging methods.
Efficient test-time inference with reparameterization We focus the where a= 1. To further optimize the computation of curvature-awan merging, we embed the curvature matrices into the domain- vectors using the following reparameterization trick:
Em +
In this case, the merging formula at test time becomes:
N-1 sq (E Em.
N-1
E.
Em
Σ
Mr MM-(seT))
Thus, during inference, we avoid storing the curvature matrices and recomputing their product with domain vectors, reducing the total FLOPs. This explains the computational efficiency seen in Section 3.
5

---

Published as a conference paper at ICLR 2025
2.6 THEORETICAL ANALYSIS
In this section, we investigate how optimizing the curvature matrix M in Eqn. CA-Merg can improve the generalization of the expert merging process for downstream tasks. We first recall the formulation for domain-specific merging
N-1
Bm
Em +
Σ 85
(5)
j=1
We begin by calculating the gradient of the downstream task loss function, denoted by 4, with respect to M, at a specific (th layer as follows:
ac JM, TE OM, OE
Oc TE Oc -(a8)TT))=as= 2. -E, Em)
ac
(6)
This is the outer product of the gradients of the task loss and the domain vectors. It is important to note the connection with the Fisher Information Matrix. For downstream task, if we define the loss function as the negative log-likelihood, such as in supervised classification task L(0) ((ea))plloogpppa then the empirical Fisher Information Matrix can be defined as
F T= (x.
(olgoo
Next, consider how the gradient of the curvature matrix contribute better performance. Given time step f, the standard gradient descent from an initial point M, with learning rate /3 yields the following update:
OC M M B* 8M = M a.38ss
Oc
(E -)) OF
We assume standard gradient descent for simplicity, but the argument extends to other advanced gradient algorithms, such as momentum (Polyak, 1964; Sutskever et al., 2013; Nguyen et al., 2020; He et al., 2020; Wang al., 2022; Nguyen et al., 2022a) and ADAM (Kingma & Ba, 2015; Loshchilov & Hutter, 2019). We then apply M, to the merging process in Eqn. CA-Merg and get
N-1
N-1
Oc
E +
(+1 M o
22
88
s+1
(5)
(8)
==1
domain specific merying with curvature aware
The detail for the derivation can be found in Appendix E. We classic domain-specifee merging formula with the guidance of the learned curvature. Furthermore, the second term contains the direction from the task loss gradient and the inner product between domain-vectors from two consecutive gradient steps. If M, = Ivj. this term can be seen as an auxiliary signal from task modeling the agreement of the merging direction between updating steps: if there are conflicts between current and the previous updating direction, then this signal will be alleviated, thus dampening the harm to the merging direction of the current step; otherwise they show strong agreement, this amplifies the impact of the updating direction toward minimizing the task loss with respect to the previous step, thus accelerate the training process while implicitly helping current merging direction with additional experience.
that the first term in Eqn. 8 is the
loss of the previous update step guiding the merging direction. The term 51,5t1
On the other hand, we can rewrite Eqn. 8 as follows:
N-1
OC
Êm Em +
Σ j=1
M
1+1
2
Σ
H+
gradient mutching
We now have the inner-product between the gradient of the task loss and the domain-vector. This can be interpreted as the matching between the update from the task loss gradient and the domain-specifie direction. We then have the updated domain-specifie direction for each expert whose weighting factors calculated by the inner-product. Therefore, we are performing a soft nearest distance voting to find the experts that agree the most with the task loss and enhance the merged experts with the corresponding domain-vector.

---

Published as a conference paper at ICLR 2025
Table 2: Performance of T5-baee variants on the fine-tuning tasks for GLUE All SMoE variants have 8 experts per layer. We follow (Devlin et al., 2019) in conducting experiments on the GLUE benchmark. Our curvature-awan methods outperform all baselines across tasks, while maintaining the same number of parameters and FLOPs as the SMoE models.
Methods
Params
TFLOPs
SST-2
MRPC
CoLAA
QQP
STSB
QNLI
RTE
MNLI
Vanilla
220M
4.65
93.34
89.70
58.06
88.76
89.06
92.34
74.36
86.36
SMoE
1.0B
4.65
94.26
90.87
56.78
88.69
89.44
92.07
70.75
86.38
Sennan-Speciffe
1.0B
4.65
93.57
90,19
58.07
88.77
89.40
92.51
72.56
86.40
Ties
1.0B
4.65
93,92
91,44
58.54
86.47
88.58
91.87
75.54
86.39
Dare
1.0B
4.65
93.80
89.46
58.33
88.72
89.13
92.29
73.64
86.20
Domain-speciff-CC.
1.0B
4.65
93.80
91.16
58.57
88.86
89.47
92.60
74.72
86.44
Dare-CA
1.0B
4.65
94.49
91.15
58,56
88.76
89.56
92.80
78.70
86,34
Ties-CA
1.0B
4.65
94.61
92.49
60.06
88.83
89.54
91.89
75.81
86,45
3 EXPERIMENTAL RESULTS
We perform evaluations on four major tasks, including language modeling, text classification, question answering, and image classification. For language modeling, we use the Wikitext-2 and Wikitext-100 103 (Merity et al., 2016) benchmarks. For text classification, we employ a subset of the GLUE (Wang al., 2019) benchmark, a collection of eight diverse tasks designed to test different aspects of language understanding. For question answering, we employ two famous benchmarks: SQuAD (Rajpurkar et al., 2016) and WikiQA (Yang et al., 2015). Finally, the ImageNet-Il Ik (Deng et al., 2009) dataset is chosen for image classification evaluation.
We choose GPT-2 (Radford et al., 2019) small and Swin- Transformer small (Liu et al., 2021) our backbones for language modeling and image classification, respectively. Regarding GLUE and question-answeriig tasks, T5 base (Raffel et al., 2020) is chosen.
Varila Demain Cp
55
50
45
Peplecit
Our experimental results confirm that the proposed merg- ing method accelerates pre- tre-trainig convergence and, when combined with other merging protocols, enhances model performance across tasks and settings. All re- sults are averaged over runs with different random seeds. Detailed information on the datasets, models, training procedures, and hyperparameters provided in Appendix B and Appendix D. For additional exper- iments on different routers and merging methods, we refer to Appendix H. and H.2.
40
35
30
25
20
30
Epochs
Figure 3: Perplexity of GPT2-small variants starting at the tenth epoch.
3.1 TRAINING AND EVALUATION DETAILS
We fix the number of epochs for all models on each task. For each text-related task, first undertake a comprehensive hyper-parametter search. This encompasses batch sizes (8, 16, 32, 64), learning rates f3om(3e-4, le-4, 3e-5, le-5), pinpoint the optimal fine-tuned models. Regarding image classification tasks, a batchsize of 96 was chosen for all models. In addition, we choose AdamW (Loshchilov Hutter, 2019) as the default optimizer and conduct all experiments on NVIDIA A100 GPUs. We compare our proposal to three merging basclines, including domain-specitf Ties, and Dare merging. There exists prior works on merging methods with the aid of the Fisher Information Matrix, such as (Matena & Raffel, 2022). which rely on access to validation set used to compute the Fisher matrix or fine-tune hyperparameters. To eliminate the need for a validation set, (Jin et al., 2023) proposes storing and transmitting inner product matrices derived from the training data for each task, which are of the same size as the original model. However, this approach becomes costly for large models, as storage and transmission demands increase linearly with model size and the number of tasks well as the number of experts. Therefore, we choose baselines that are needless of extra information and computational cost to perform comparisons. More details about theoretical comparison between CAMEx and Fisher-based
7

---

Published as a conference paper at ICLR 2025
merging methods can be found in Appendix A. We want to note that merging protocol can be easily integrated into other works such as merge then compress protocol (Li et al., 2024).
3.2 RESULTS
In the following tables, the row with our method's results is highlighted in grey. Results with the best and second best performance are written in bold and underline, respectively. In addition, methods with the postfix "-CA" denote the curvature-
version of the corresponding baseline.
In Table 2, the results demonstrate that CA-augmented models consistently outperform their non-CA counterparts. Ties-CA achieves the highest scores SST-2 (94.61), MRPC (92.49), CoLA (60.06), and MNLI (86.45), showing considerable improvements over both the vanilla and standard Ties models. Similarly, Dare-CA performs best on RTE (78.70), surpassing Dare (73.64), indicating that CA improves performance on smaller datasets and tasks with higher variability. Furthermore, Domain-speciffc-CC exceeds the non-CA version on QNLI and MNLI, demonstrating the broader applicability of curvature- aware methods. We provided significant t-test at Appendix G.
In Table 3, the Domain-speeific- namic outperform the Vanilla, SMoE, and Domain- specific baselines, with lower perplexity values. Domain-specific- achieves the lowest perplexity score of 21.50, showcasing superior performance in language modeling tasks when compared to all other methods. The Dynamic architecture follows closely with perplexity of 21.55 while also reduc- ing the parameter count by 9%, compared to the other methods. This highlights the Dynamic architecture's 's efficiency in maintaining strong performance with fewer parameters, making it ideal for resource-constren environments. Moreover, the Dynamic architecture is competitive with Domaan-speccffc-CCC and outperforms the rest in terms of convergence speed, which is shown in Figure 3.
and Dy-
Table 3: Performance GPT-2 small variants for the pre-training task on Wikitext-103.
Methods
Perplexity
Params (M)
GELOPS
Vanilla
23.03
125
292.5
292.5
Domain specific
21.64
522
292.5
Dpmain-peeiffe-CC
21.50
522
292.5
Dynamic
21.55
470
292.5
In Table 4, the Vanilla model reaches perplexity of 21.84. Despite increasing parameters, SMoE only slightly improves to 21.60. Domain-specific Ties, and Dare methods show small gains, with Ties reaching 21.45. However, curvature-aware (CA) methods outperform all others. Domain-specific- CA achieves the best perplexity at 21.06, followed by Ties-CA (21.11) and Dare-CA (21.42), each significantly improving over their non-CA counter- parts. All models beyond Vanilla share the same computational cost, indicating that CA methods enhance performance without added complexity. Domain-specific-Ce stands out, demonstrating the
Table 4: Performance GPT-2 small variants for the supervised fine-tuning task on Wikitext-2
Methods
Perplexity.
Params (M)
GFLOPS
Vanilla
21.84
125
292.5
SMoE
21.60
522
292.5
Domain-specific
21.56
522
292.5
Ties
21.45
522
292.5
Dure
21.60
522
292.5
Domain CPeiai--peciffeCAA
21.06
522
292.5
Dare-CA
292.5
Ties-CA
21.11
522
292.5
clear advantage of curvature-aware optimization.
In Table the baseline models, including Vanilla, SMoE, and non-CA versions of Ties and Dare, achieve solid results but show diminishing improve- ments model complexity increases. In con- trast, our curvature-awe aware methods significantly outperform their counterparts. For instance, on the SQuAD dataset, Dare-CA achieves the high- est Exact Match (EM) score of 81.76% and Fl score of 88.60%, surpassing all other methods. Similarly, WikiQA, Ties-CA attains the highest accuracy of 96,55% with Dare-CA closely follow- ing at 96.23%.
Table 5: Performance of T'-bbse variants on question answering tasks.
Methods
Params
TELOPs
SQUAD Em//FI
WIKIQA Accuracy
Vanilla
223M
2.86
81,01/88.1 11
96.06
1.08
81.25/88.00
Domain spociticc
1.08
2,86
80.21/87.4
95.5
Tiks
L.0
2.86
80.76/88.1 1
95.87
Dare
1.08
2,86
80.88/88.00
96-01
Domaln Speetll--peelleeCCC..
1.08
80.44/87
Ties-CA
1.08
2.86
81 52/888600
96.55
Daro-CA
1.08
2.86
81.76/88.00
06.23
In Table 6, while Vanilla and SMoE exhibit solid accuracy scores, they are surpassed by the curvature (CA) enhanced versions of the models. Notably, Ties-CA delivers the best top-1 accuracy at 83.38% and
8

---

Published as conference paper at ICLR 2025
Domain-spann -pseciificci
Ties-CA
Aetricc
0.00
030
0.00
0.00
0.00
0.76
0.00
0.70
36
0.7
1.2
75
76
03
1.5
0.5
03
alpha
alpha
alpha
STSB
MRPC
RTE
Figure 4: Impact of the a parameter on Curvature-Aware method performance across NLP tasks. We observe that the scaling factors that are within the range [0.8,1] consistently improve model's performance.
Domain- Dommain-peccffc-cA
Ties-CA
Dare-CA
Metric
08
0.80
0.80
0.00
0.00
0.00
Figure Impact of the Kronecker rank of curvature matrix on model's performance. We observe that as the rank increases the performance drops and then saturates. However, we would like to note that this curve might change depending on the downstream tasks and the merging protocol.
the highest top-5 accuracy at 96.96% slightly edging out Dare-CA, which achieves 83.38% and 96.94% respectively.
4 ABLATION
Impact of the scaling factor. The plot in Figure illustrates the impact of the parameter the performance of three curvature-aware (CA) model variants Dommann-peciffc-CA. Ties-CA, and Dare-CA across three natural language processing tasks: STSB, MRPC, and RTE. The a parameter ranges from 0.5 to 1.0. The overall trend suggests that increasing a leads to better generalization, particularly for complex tasks such as RTE, where sentence-level entailment and similarity benefit from stronger curvature-aware representations. Moreover, across all tasks, the model reaches its peak performance when a is inside the range [0.8.1] 8.1]. This observation aligns with that indicated by (Yadav et al., 2023). For more comprehensive analysis on the impact of a and number of experts, we direct the readers to Appendix H.4.1 and H.4.2.
Improved performance with higher Kronecker rank. Across three tasks (STSB, MRPC, and RTE), the evaluation metrics tend to improve the rank increases from 1 to 8. This indicates that higher-ranked models generally perform better, suggesting a positive correlation between rank and task performance. Notably, the Domain-specifie. CA model consistently achieves high performance across all tasks, especially in STSB, where metrics approach 0.90. Although MRPC and RTE show slightly lower metries, ranging from 0.50 0.75, there is clear improvement in performance as rank increases, particularly the lower-to-mid ranks. However, we observed a decline in performance for Ties-CA and Dare-CA as the rank increases. We hypothesize that this is due to the masking mechanism employed by these methods, which may interfere with the learning process of the curvature matrices.
Table 6:
Comparison of Accuracy for Swin- Transformer small variants on ImageNet-1k.
Methods
Parums (M)
GFLOPs
Acc@1
Acc@5
Vanilla SM&E
6.75
83.14
96.90
6.75
83.15
96.71
Domsio Domsin-specifs
157
6.75
83.15
96.91
Tes
6.75
96.33
Dare
157
6.75
83.13
96.88
Domain-speiffe-CC
6.75
Ties-CA
157
6.75
83.38
96.96
Dare-CA

---

Published as a conference paper at ICLR 2025
Robustness against noise. Table 7 demonstrates that curvature-awar models offer superior perfor- mance on corrupted ImageNet datasets compared to both Vanilla and SMoE. variants. Among the models, our best configuration (Ties-CA) stands out as the best performer, showcasing robustness comruptios acToss all datasets. These results suggest that incorporating curvature-awarene can substantially improve model robustmess in challenging conditions.
Table Comparison for Swin-Tranfformer small
variants comupted ImageNet.
Methods ImageNet-O
ImageNet-A
ImageNet-R R
Vanilla SMoE
45.88
23.68/53.10
37.34/52.34
43.34
23.72/53.15
38.02/55.17
Ours
50.69
25.45/54.24
38.37/55.42
RELATED WORK
Sparse Mixture-of-FEE Experts (SMoE), As the demand for model scaling grows increasingly widespread. there is pressing inquiry into efficient ways to optimize computing costs while minimizing the impact on model performance. To address this need, Sparse Mixture of Experts (SMoE) has emerged and undergone extensive research and exploration (Shazeer et al., 2017; Lepikhin et al., 2021; Fedus et al., 2022). Starting with (Shazeer et al., 2017), the integration of SMoE into transformer architectures followed shortly after with the works of (Lepikhin et al., 2021) and (Fedus et al., 2022). The principle of SMoE is based on a simple concept: scaling the horizontal dimension of models (i.e., the number of feedforward blocks) rather than the vertical dimension (i.e., the number of stacked layers). This allows the model to selectively activate units or parameters based on the input tokens, thereby optimizing resource usage while maintaining performance.
SMoE Efficiency Bottlenecks and Emerging Solutions. While it remains controversial whether to use Top-1 or Top-K routing. some research has highlighted the potential performance gains from increasing the number of activated experts (Shazeer et al., 2017; Chen et al., 2023). Other studies have found redundancies among experts in MoE layers (Li et al., 2024; Lu al., 2024a). Additionally, some work has proposed using low-rank experts (Wu et al., 2024b; Liu et al., 2024; Wu et al., 2024a) inspired by LoRA (Hu et al., 2022). Despite the varying research directions, these studies consistently show that traiming a robust SMoE requires substantial computational and memory resources. This has motivated researchers such as (Li et al., 2024), (He et al., 2023), and (Zhong et al., 2024) to merge experts within each MoE layer, reducing the number of experts
single one and significantly improving training and inference efficiency.
Model Merging with curvature-aware. Though numerous methods for merging models have been introduced and developed (Yadav et al., 2023; Cai et al., 2023; Ilharco et al., 2022; Matena & Raffel, 2022; Jin et al., 2022; Don- Yehiya et al., 2022; Rame et al., 2023; Tuan et al., 2024; Lu et al., 2024b), most of these works consider merging protocols in the Euclidean parameter space. However, it has been noted that the space of deep neural network models is a Riemannian one (Amari, 1998). (Matena & Raffel, 2022) and (Jin al., 2022) were the first to fuse model weights while accounting for the Fisher Information. Despite their promising results, these methods require massive computation to approximate the inversion of the Fisher matrix. Moreover, the Fisher matrix has a size proportional the dimension of the model parameters, which significantly increases memory usage. Consequently, these methods are challenging for directly integrating into SMoE layers fuse expert weights.
LIMITATION AND CONCLUSION
In this work, we introduced CAMEx, a curvature-awar approach expert merging Mixture of Experts architectures. By leveraging natural gradients to account for the parameter manifold's curvature, CAMEx enhances model performance and reduces computational costs during both pre-training and fine-tuning, outperforming traditional Fuclidean-based methods. Additionally, our dynamic merging architecture optimizes resource usage by incorporating a global expert across layers, thus minimizing model size without sacrificing accuracy. Despite the overall improvements, a minor limitation is that curvature-aware merging demonstrates reduced compatibility with Ties and Dare merging at higher Kronecker ranks. Future work could dive deeper into this phenomenon and extend CAMEx to other expert merging methods and explore survature-awar approaches in broader neural network models to further enhance our dynamic architecture. This research lays the groundwork for developing more efficient and scalable models in large-scale machine leaming.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
Dung V. Nguyen was funded by the Master, PhD Scholarship Programme of Vingroup Innovation Foundation (VINIF), code VINIF2023. ThS.024.
This research project supported by the National Research Foundation Singapore under the Al Singapore This research project is supported by the
Programme (AISG Award No:
Ministry of Education, Singapore, under the Academic Research Fund Tier (FY2023) 4-8002044000000. 1-8002039-00--00 This research / project is also supported by the NUS Presidential Young Professorship Award (A-0009807-01 and the NUS Artificial Intelligence Institute- Funding 11-8003062-000000)
Reproducibility Statement: Source codes for our experiments are provided in the supplementary materials of the paper. The details of our experimental settings and computational infrastructure are given in Section 3 and the Appendix D. All datasets that we used in the paper are published, and they are easy to find in the Internet.
Ethics Statement: Given the nature of the work, we do not foresee any negative societal and ethical impacts of our work.
REFERENCES
S. Amari and S.C. Douglas. Why natural gradient? In Proceedings of the 1998 IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP '98 (Cat. No. No.98CH36181) volume 2, pp. 1213-1216 vol.2, 1998. doi: 10.1 N0.1109/CAASSPP198..655558
Shun-ichi Amari. Natural gradient works efficiently in learning. Neural Computation, 10(2):251-276. 1998. doi: 10. 001162200099976990000177777
Ruisi Cai, Zhenyu Zhang, and Zhangyang Wang. Robust weight signatures: Gaining robustness as easy as patching weights? arXiv preprint arXiv:2302.12480,, 2023.
Daniel Cer, Mona Diab, Eneko Agirr,, lñigo Lopez-Gazpio, and Lucia Specia. SemEval-2011 task 1: Semantic textual similarity multilingual and crosslingual focused evaluation. In Proceedings of the 11th International Workshop on Semantic Evaluation (SemEval-2017), pp. 1-14, Vancouver, Canada, August 2017. Association for Computational Linguistics. doi: 10. 18656//111117-20011
Tianlong Chen, Zhenyu Zhang, AJAY KUMAR JAISWAI, Shiwei Liu, and Zhangyang Wang. Sparse moe as the new dropout: Scaling dense and self-slimmable transformers. In The Eleventh International Conference on Learning Representations, 2023. URL https: http///ppeeeeeeeeeeeeee/ee wwt.:/wwwwweeeeeeeeeee.eeee.eeee wwt.:neeetttiiineeeee forum?ird w1hwFUb_BB _8.
Zewen Chi, Li Dong, Shaohan Huang, Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj, Xia Song. and Furu Wei. On the representation collapse of sparse mixture of experts. arXiv preprint wrXiv:2204.091777 2022.
Ido Dagan, Oren Glickman, and Bemardo Magnini. The pascal recognizing textual entailment challenge. In Proceedings of the First PASCAL Challenges Workshop on Recognising Textual Entailment, 2006.
Damai Dai, Li Dong, Shuming Ma, Bo Zheng, Zhifang Sui, Baobao Chang, and Furu Wei. StableMoE: Stable routing strategy for mixture of experts. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 7085-7095, Dublin, Ireland, May 2022, Association for Computational Linguistics. doi: 10. 12.086666///02222c--oooggg
Damai Dai, Chengqi Deng, Chenggang Zhao, R. X. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y. Wu, Zhenda Xie, Y. K. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models, 2024. URL https: wrwia.iiivv...o org/ abs/ / /2401. 0 06066.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference computer vision and patter recognition, pp. 248 255. Icee, 2009.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirec- tional transformers for language understanding. ArXiv, abs/1810.04805, 2019.
11

---

Published as a conference paper at ICLR 2025
William B Dolan and Chris Brockett. Automatically constructing a corpus of sentential paraphrases. In Proceedings of the Third International Workshop on Paraphrasing (IWP200S), 2005.
Shachar Don-Yehiya, Elad Venezian, Colin Raffel, Noam Slonim, Yoav Katz, and Leshem Choshen. Cold fusion: Collaborative descent for distributed multitask finetuning. arXiv preprint arXiv:2212.01378, 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-39, 2022. URL
http:/ / jmlr. orgo wwtgg/apprs//v333 21 0998 html.
Marawan Gamal Abdel Hameed, Marzieh S. Tahaei, Ali Mosleh, and Vahid Partovi Nia. Convolutional neural network compression through generalized kronecker product decomposition. Proceedings of the AAAI Conference on Artificial Intelligence, 36(1):771-77, Jun. 2022. doi: 10.1609/aaii.v3666 1991. /ttee////9995 19958.
URL https: : /twps//js.aaa////// aaai org/ /rzg/ine index.ndxx.////// php/
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattem recognition, pp. 9729-9738, 2020.
Shwai Hc, Run-Ze Fan, Liang Ding. Li Shen, Tianyi Zhou, and Dacheng Tao. Merging experts into one: Improving computational efficiency of mixture of experts. In Proceddigss of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 14685, Singapore, December 2023.
Association for Computational Linguistics. doi: 10.18653/v1/20000 200.:///1.2222..mmmmmmmm......
Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, al. The many faces of robustness: critical analysis of out-of-distributin generalization. In 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 Computer Vision and Pattern Recognition, pp. 15262-15271, 2021b.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy. Simon Osindero, Karen Simonyan, Erich Elsen, Oriol Vinyals, Jack W. Rae, and Laurent Sifre. Training compute-optimal large language models. In Proceedings of the 36th International Conference Neural Information Processing Systems. Curran Associates Inc., 2024. ISBN 9781713871088
Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Con- ference on Learning Representations, 2022. URL https: http:///pppweeeeeeeeeeeee http:/outttt....c.c.. nZeVKeeFYf9.
Gabriel Ilharco, Mitchell Wortsman, Samir Yitzhak Gadre, Shuran Song, Hannaneh Hajishirzi, Simon Korbllihh, Ali Farhadi, and Ludwig Schmidt. Patching open-vocabulary models by interpolating weights. Advances in Neural Information Processing Systems, 35:29262-29277, 2022.
Gabriel Ilharco, Marco Tulio Ribeiro, Mitchell Wortsman, Ludwig Schmidt, Hannanch Hajishirzi, and Ali Farhadi. Editing models with task arithmetic. In The Eleventh International Conference Learning Representations, 2023. URL https: http:///ppp//eieeeee
net. forum? id= 6t OKwE8- jr
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. doi: (116///cco.199..3.....
Xisen Jin, Xiang Ren, Daniel Preotiuc-Pietro and Pengxiang Cheng- Dataless knowledge fusion by merging weights of language models. arXiv preprint /7X0v:22/..09849 2022.
Xisen Jin, Xiang Ren, Daniel Preotiuc-Pietre and Pengxiang Cheng. Dataless knowledge fusion by merging weights of language models. In The Eleventh International Conference on Learning Representations, 2023. URL https: : /ote:://opppeeee /tpp/ppeeiiieeeeeeeeeeeeeeeeeeee net forum? wawtrrummi==FFooo Cnchhuuu
12

---

Published as a conference paper at ICLR 2025
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models, 2020.
URL https:/ aarxivii org/ abss /0s//2000 200000011..2 08361.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International
Conference on Learning Representations, 2015.
Dmitry Lepikhin, HyoukJoong I Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikum, Noam Shazeer, and Zhifeng Chen. {GS hhard: Scaling giant models with conditional computa- tion and automatic sharding. In International Conference Learning Representations, 2021. URL https: / /owpnrveeeeeeeeeeeoeeeeeeeeeeeee .et.. forum? id=qrwe XXHTMYb.
Pingzhi Li, Zhenyu Zhang, Prateck Yadav, Yi-Lin Sung, Yu Cheng, Mohit Bansal, and Tianlong Chen. Merge, then compress: Demystify efficient SMoe with hints from its routing policy. In The Twelfth International Conference on Learning Representations, 2024. URL https: /wwp//pepeee...eeeeeeeeeee.e.eep neet
Forum?id=en wwc.eeeFeyyWKK..
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint uuXiv:/90) 1692, 2019.
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 of the IEEE/CVF International Conference on Computer Vision (ICCV), 2021.
Zihan Liu, Hanyi Wang. Yaoyu Kang, and Shilin Wang. Mixture of low-rank experts for transferable ai-generated image detection, 2024. URL https:
//atpp:://axiv.... /ttwi/ixivv///////////aa//aaaaaa org/ abs /2404. 04883.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, In International Conference on Learning Representations, 2019, URL https: http::///ppeeeeeeeeeeeeee wttpee///ooumndddeeeeeeeeeee
Xudong Lu, Qi Liu, Yuhui Xu, Aojun Zhou, Siyuan Huang. Bo Zhang. Junchi Yan, and Hongsheng Li. Not all experts are equal: Efficient expert pruning and skipping for mixture-of-experts large language models, 2024a. URL https: /arxivvvvva. wwwp/or//////ooo//////../o/..eee abs/ 22bs22400 22222222222 14800.
Zhenyi Lu, Chenghao Fan, Wei Wei, Xiaoye Qu, Dangyang Chen, and Yu Cheng. Twin-merging: Dynamic integration of modular expertise in model merging. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024b. URL http::/ httpa////pppeeeeeeeeeeee wtwpreeewwwwwwweefereoeeeeeeeeee fet/forum? id=8 1YII
James Martens. New insights and perspectives on the natural gradient method. Mach. Learn. Res., 21(1), jan 2020. ISSN 1532-4435.
James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In Proceedings of the 32nd International Conference International Conference Machine Learning Volume 37, ICML'15 15, 2408-2417. JMLR.org, 2015.
Michael Matena and Colin A Raffel. Merging models with fisher-weighted averaging. Advances in Neural Information Processing Systems, 35:17703-17716, 2022.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXiv:1609.07844 2016.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26. 2017, Conference Track Proceedings. OpenReview.net, 2017. URL ht http::s://///eeeeeeee/eeeee net/ forum? Bd 72udxe.
Mohammed Mugeeth, Haokum Liu, and Colin Raffel. Soft merging of experts with adaptive routing. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. URL https: / https::///ppeeeeeee /ppeneeviee. net forum? d=771 991 c54z. Featured Certification.
13

---

Published as a conference paper at ICLR 2025
Ho Huu Nghia Nguyen, Tan Nguyen, Huyen Vo, Stanley Osher, and Thieu Vo. Improving neural ordinary differential equations with nesterov's 's accelerated gradient method. Advances in Neural Information Processing Systems, 15.7712-7726. 2022a.
Tan Nguyen, Richard Baraniuk, Andrea Bertozzi, Stanley Osher, and Bao Wang. Momentummn: Integrat- ing momentum into recurrent neural networks. Advances in neural information processing systems, 33: 1924-1936, 2020.
Tan Nguyen, Minh Pham, Tam Nguyen, Khai Nguyen, Stanley Osher, and Nhat Ho. Fourierformer: Transformer meets generalized fourier integral theorem. Advances in Neural Information Processing Systems, 35:29319 2993293,,,,,,, 2022b.
Stefan Nielsen, Rachel Teo, Laziz Abdullaev, and Tan Minh Nguyen. Tight clusters make specialized experts. The Thirteenth International Conference Learning Representations, 2025. URL https: // //openreview.rec net wtweee///ooomommmomoooooe
Eunbyung Park and Junier B Oliva. Meta-curvature. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-But Buc, E. Fox, and R. Gamett (eds.), Advances in Neural Information Processing Systems, vol- ume 32. Curaan Associates, Inc., 2019. URL http://tttttpprrpceciiiigggggoo hwwp:/eeeeeeiin...eeneeeeee.ee neurips cc/ cc/paper_
Paaddaap....pe pdf.
Boris T Polyak. Some methods of speeding up the convergence of iteration methods. Ussr computational
mathematics and mathematical physics, 4(5):1117, 1964.
Quora. Quora question pairs, 2017.
https://quoradee http://uuoaaaaaa.uuuaaaaaaaaaaaa com/ First- Quora Dataset Dataset-Release Question- Pairs.
Release
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models are unsupervised multitask learners. OpenAl Blog, 18))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. arXiv preprint /9Xiv: 10683, 2020.
Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Ammar Ahmad Awan, Jeff Rasley, and Yuxiong He. Deepspeed-mce Advancing mixture-of-expee inference and training to power next-generation ai scale. arXiv preprint arXiv:2201.05596 2022.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ questions for machine comprehension of text. In Proceedings of the 2016 Conference Empirical Methods in Natural Language Processing, pp. 2383-2392, Austin, Texas, 2016. Association for Computational Linguistics. doi: 1..18863/v1/Di/ 1264. URL https: http://aaaaatttoooo.oooooo org/D16 1264.
Alexandre Rame, Kartik Ahuja, Jianyu Zhang, Matthieu Cord, Léon Bottou, and David Lopez-Paz. Model ratatouille: Recycling diverse models for out-of-distribution generalization. arXiv preprint arXiv:2212.10445, 2023.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-exppers layer. arXiv preprint arXiv:1701.06530 2017.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang. Christopher D Manning, Andrew Y Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 conference on empirical methods in natural language processing, pp.
1631-1642, 2013. 2013.
Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum deep learning. In International conference on machine learning, pp. 1139-1147. PMLR. 2013.
Rachel Teo and Tan Minh Nguyen. MomentumSMoe: Integrating momentum into sparse mixture of experts. In The eightt Annual Conference on Neural Information Processing Systems, 2024, URL https://openrent net/ffouum ddy 9esC2NJ.
14

---

Published as a conference paper at ICLR 2025
Rachel Teo and Tan Minh Nguyen. MoLEx: Mixture of layer experts for fine-tuning with sparse upcycling. In The Thirteenth International Conference on Learning Representations, 2025. URL
https://openrevier/ net/ wtwp:r/rwwweeee//ooooooee//oe//e 90hh0c.
Tran Anh Tuan, Nguyen Viet Dung, and Tran Ngoc Thang. A hyper-transformer model for controllable pareto front learning with split feasibility constraints. Neural Networks, 179:106571, 2024. ISSN 0893-6080. doi: https://oi..or//I0 http://o..oogg11.111....eeeeee 2o01
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint 18X4,08004007461 2019.
Bao Wang, Hedi Xia, Tan Nguyen, and Stanley Osher. How does momentum benefit deep neural networks architecture design? few case studies. Research in the Mathematical Sciences, 9(3):57, 2022.
Alex Warstadt, Amanpreet Singh, and Samuel Bowman. Neural network acceptability judgments. In Proceedings of the 33rd AAAI Conference on Artificial Intelligence, pp. 6732-6739, 2019.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoe Le, and Denny Zhou, Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201. 2022.
Adina Williams, Nikita Nangia, and Samuel R Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long Papers), pp. 1112-1122, 2018.
Jialin Wu, Xia Hu, Yaqing Wang. Bo Pang, and Radu Soricut. Omni-smola: Boosting generalist multimodal models with soft mixture of low-rank experts. In Proceedings of the IEFF/CVE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 14215, June 2024a.
Xun Wu, Shaohan Huang, and Furu Wei. Mixture of loRA experts. In The Twelfth International Conference on Learning Representations, 2024b. URL i id=uWWkBBCh45..
Prateck Yadav, Derek Tam, Leshem Choshen, Colin Raffel, and Mohit Bansal. Resolving interference when merging models. In NeurlPS, New Orleans, USA, 2023. Proceedings of Machine Learning Research.
Yi Yang, Wen-tau Yih, and Christopher Meek. WikiQA: A challenge dataset for open-domain question answering. In Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, pp. 2013-2018, Lisbon, Portugal, 2015. Association for Computational Linguistics. doi: 10.18653/v1/D15-. 1237. URL https: : at/ http://caaaattthooooooo org/D15-1237.
Zhilin Yang, Zihang Dai, Yiming Yang. Jaime Carbonell, Russ Salakhutdinov, and Quoe V Le. Xlnet: Generalized autoregressive pretraining for language understanding. Advances in neural information processing systems, 32, 2019.
Le Yu, Bowen Yu, Haiyang Yu, Fei Huang, and Yongbin Li. Language models are super mario: Absorbing abilities from homologous models as free lunch. In Forty-first International Conference on Machine Learning, 2024.
Zexuan Zhong, Mengzhou Xia, Dangi Chen, and Mike Lewis. Lory: Fully differentiable mmixture-ofexxpere for autoregressive language model pre-training. In First Conference on Language Modeling, 2024. URL https: wwwp://eeeeeeeeeeeeemeeeee net forum? fott:/orrmm.d-.....D........ 1d=1KE.PySnI.
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Y Zhao, Andrew M. Dai, Zhifeng Chen, Quoc V Le, and James Laudon. Mixture-of-experts with expert choice routing. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL
net forum Forummidd=jdoolHHiin
15

---

Published as a conference paper at ICLR 2025
Supplement to "CAMEx: Curvature-aware Merging of Experts"
Table of Contents
Comprision of CAMEx and Fisher-based merging methods
16
Additional Details on datasets
18
B.1
Language Modeling WikiText
18
B.2
Text classification on GLUE benchmark
18
B.3
Question-answering SQuAD and WikiQA
................................
18
B,4
Image Classification on Imagenet
18
B.5
Adversarial Examples and Out-of-distributiio datasets
19
Algorithm and implementation details
19
C.1 Causal segmenting
19
C.2 Some implementations
20
More Experiment Details
20
E
Derivation
21
F
Step-by-step walkthrough for key equations of CAMEx
21
F.1 Key equation for merging of CAMEx ................................
F2
Key equation for merging in dynamic architecture
21
F3 Curvature update
22
G Student's t-test for experiments on GLUE dataset
22
Additional Experiments
23
H.l
Integrating CAMEx into Twin-Merging
................................
23
H.2
Experiments on Token-choice vs Expert-choice routing ................................
24
H.3
Longer training for Wikitext-103 pre-training.
24
H.4
More comprehensive ablation study on byserpaneeess...................
25
H.4.1
Ablation study on
25
H.4.2 Ablation study on number of experts
25
H.5 Results on large backbone (Phi-3)
26
A COMPRISION OF CAMEX AND FISHER- -BSEDD MERGING METHODS
The pipeline comparison between CAMEX and Fisher-based merging methods is shown in Figure 6. Both approaches aim to capture the curvature of the parameter space during the merging process. (Diagonal) Fisher Merging (Matena & Raffel, 2022) applies diagonal approximation to the Fisher information matrix. this work, they estimate the diagonal the Fisher matrix as:
E [Vlogp(yyx/Vooeeepxxxxa)))
(10)
X (y-xpo(y)x)
16

---

Published as a conference paper at ICLR 2025
Merging
E
Mj (#t m))
Mj -1
Mj
0.
0C aee
(E) E
a) CAMEX
Morging
N FjE;
Bu
Storing Gradient
b) Merging
Figure 6: CAMEx merging pipeline vs Fisher-based merging pipeline. Note that Fisher merging requires the storing of the F, logpr, (w/r)) for all experts and all æ in training dataset. Furthermore, it has been pointed out that Fisher merging will have poor performance while using fewer examples estimate the Fisher.
The expectation over can be estimated via sampling from po (yllz)) computed exactly when the number of classes is small. The closed-form solution for Fisher merging (without necessarily applying the diagonal approximation) is given by:
M
N |FE
Σn
Fl
(11)
m=1
Thus, to approximate the Fisher Information Matrix for SMoE models, Fisher merging requires storing Vg, (y/æ:) for all experts and all in the training dataset. Additionally, it has been noted that Fisher merging can suffer from poor performance when fewer examples are used to estimate the Fisher matrix (Matena & Raffel, 2022).
In the case of our method (depicted Figure 6a), by denoting M, as the curvature matrix of the i-th expert, CAMEx utilizes the formula for merging experts derived from the natural gradient descent update as:
N-1
E =E +o Σ ($
(CA-Merg)
CAMEx implicitly implements the gradient-based matching between the task loss gradient and domain- vector of the corresponding expert to approximate the empirical Fisher through the dynamic of gradient descend update of M,:
OC = M -B++$*
M -M B*
(E) -,,
(12)
am
DEt
3c
where the term
EE Em) represents the outer product of gradients of the task loss and
TE
the domain vectors. This operation contributes capturing the curvature of the expert parameter space,
17

---

Published as a conference paper at ICLR 2025
ensuring curvature awareness during the merging process. This approach eliminates the need to compute the inversion of the empirical Fisher Information Matrix, thereby reducing computational overhead while maintaining sensitivity to parameter space curvature.
B ADDITIONAL DETAILS ON DATASETS
This section provides detailed information on the datasets and evaluation metrics used the experiments in Section 3.
B.I LANGUAGE MODELING ON WIKITEXT
The WikiText-103 dataset consists of Wikipedia articles designed to capture long range contextual dependencies. The training set includes approximately 28,00 38,000 articles, totaling around 103 million words. The validation and test sets have 218,000 and 246,000 words, respectively, spread across 60 articles per set, with each set comprising roughly 268,000 words. Our experiments follow the standard procedure described in (Merity et al., 2017).
-ikiText-2 a smaller version of WikiText- 103, containing million tokens and a vocabulary of 33,000 words.
B.2 TEXT CLASSIFICATION ON GLUE BENCHMARK
These tasks include MNLI (Williams et al., 2018), which assesses model's ability to determine entailment between pairs of sentences; QQP (Quora, 2017) and MRPC (Dolan & Brockett, 2005), which focus identifying sentence similarity and paraphrase detection; SST-2 (Socher et al., 2013) for sentiment analysis; CoLA (Warstadt et al., 2019) for grammaticality judgment; and QNLI (Wang et al., 2019) for question-answar classification. Additionally, STSB (Cer et al., 2017) evaluates the model's ability to measure sentence similarity, while RTE (Dagan et al., 2006) tests logical reasoning.
B.3 QUESTION-ANSWERING ON SQUAD AND WIKIQA
SQuADv1.1 (Rajpurkar et al., 2016) (Stanford Question Answering Dataset) a widely used benchmark for reading comprehension and question answering tasks. It contains over 100,000 question-answer pairs sourced from more than 500 Wikipedia articles. Fach question is paired with a paragraph from the article, where the answer is span of text extracted from the passage. The dataset consists of natural language questions that cover wide range of topics, context paragraphs from Wikipedia, and answers marked by their start and end positions within the context. The primary task is to extract the correct answer span based on the posed question. Key features of the dataset include the need for exact span extraction, the large dataset size, and its task design focused on reading comprehension. Evaluation is typically done using Exact Match (EM), which measures the percentage of predictions that exactly match the ground answers, and the score, which measures the overlap between predicted and ground-truth answers by calculating the harmonic mean of precision and recall.
WikiQA (Yang al., 2015) an open-domain question answering dataset designed for answer sentence selection tasks. It consists of natural language questions primarily extracted from search engine queries, with candidate sentences sourced from Wikipedia articles. Each candidate sentence labeled as either correct incorrect answer for the given question. The dataset contains 3,047 questions and 29,258 candidate sentences. The main challenge is selecting the correct sentence from set candidates, unlike SQuADv1.1, where the task focuses extracting text span. Key features include its real-world query origins, the sentence selection task, and the open-domain nature, which requires models to identify relevant sentences from diverse topics. WikiQA is evaluated using Accuracy.
B.4 IMAGE CLASSIFICATION ON IMAGENET
ImageNet-lk, the most widely utilized subset of the ImageNet dataset introduced by (Deng al., 2009). comprises 1. 1.28 million images for training and 50,000 images for validation, across 1,000 categories. Performance evaluation is typically based top-1 and top-5 accuracy metrics.
18

---

Published as a conference paper at ICLR 2025
Algorithm The Overall Procedures of CAMEX.
1: Initialize: A model M with SMoE layers, the total number of original experts N.
2: Let H' CR2xLXN and T' < Bxxxx denote the roder logits and the sequence tokens intermediate layer 1,
respectively.
3: for layert= do
4: K RESTTT--SSSEETTBBBKK....
Begin Causal Segmenting
H' -TT
H' ROLLandDETACI
7:
if TIES-MERGING then
> Generate mask for merging
8:
for expert i= - E. Em
9:
10:
7: f (T)(
11:
end for 7 = agmi (g (- T()
13:
for expert == I......N -
14:
1. XA7
15:
M,
16:
end for
17:
else
18:
Generate mask for DARE-MERGING
19:
end
20:
E.. +- Em +y EmdL H/ST:
Merge Experts
21:
for
B.5 ADVERSARIAL. EXAMPLES AND OUT-OF-DISTRIBUTTE DATASETS
ImageNet-A: The ImageNet-A dataset (Hendrycks et al., 2021b) contains real-world images specifically curated to fool ImageNet classifiers. focuses on 200 classes, a subset of the 1,000 classes ImageNet- Ik. Errors made within these 200 classes are considered particularly significant, they represent wide variety of categories from ImageNet-Ik.
ImageNet-O: This dataset consists of examples adversarially filtered to challenge out-of-diitriibutio (OOD) detectors on ImageNet (Hendrycks al., 2021b). It includes images from the larger ImageNet-22 22kk dataset but excludes those present in ImageNet-Ike 1k. The selected samples are those that ResNet-50 model confidently misclassifies as an ImageNet- lkk class, and the primary evaluation metric is the area under the precision-recall curve (AUPR).
ImageNet-R: ImageNet-R contains a variety of artistic renditions object classes found in the original ImageNet dataset (Hendrycks et al., 2021a). This dataset includes 30,000 artistic representations of images from 200 classes, selected from the ImageNet-1K subset. The dataset was created to challenge models with non-standard visual interpretations of the classes.
C ALGORITHM AND IMPLEMENTATION DETAILS
C.1 CAUSAL SEGMENTING
Background of Causal Segmenting: A significant advancement in SMoE design centers fully differen- tiable architectures that climinate the need for additional loss terms to stabilize training. In (Mugeeth et al., 2024), model was introduced that computes weighted average of expert fred-forwant networks (FFNs). For an input æ with corresponding routing weights, the output is defined as:
N
02 = FFN
5: Eg
where $; = Softmax (eottmaxxG(hh)):: (G(h))).
However, applying this approach to autoregressive language models is computationally costly, the merged FFN must be computed for each token in the sequence, leading to costs that scale linearly with the number of experts. An alternative based on pooling- routing via the sequence's average representation,
19

---

Published as a conference paper at ICLR 2025
follows:
8, Softmax
This, however, disrupts the autoregressive property essential for pre-training. To address this, (Zhong al., 2024) introduced causal segment routing. This technique merges FFNs in MoE layer by utilizing information from the preceding segment to process the current segment. Specifically, given a training instance X consisting of tokens (c.g., L 4096), divide the instance into N segments, each containing T (e.g.. T= 256) consecutive tokens. For the k-th segment Sk, where k > 1, we compute the average of the hidden representations from the previous segment Sk-1, denoted as h 1. By using the average hidden representation, the model can adapt to prompts of varying lengths during inference. The average hidden representation h 1 then employed to determine the routing weights, leading to a merged expert E:
hk-1
Σ T xES
h2: s; =:otmmaa Softtax((((hx-i)) E
Σ
se- E
(13)
The merged expert E is then used process all tokens in the current segment Sk, i.c., O2 (h E),Ya E S₁. This approach ensures that the model's routing decisions rely exclusively on data from preceding positions. For the first segment S1, the segment's representation is used to compute the merging weights for its FFN. To prevent information leakage, a stop-gradient operation is applied Ghi)):
1 T
ho
b #ESo
(14)
These tokens are then used to calculate the scores for the merging procedure
Sg DETAACH(G((h
(ROLLandDETACH)
$; G(
i=1,..,S-1
C.2 SOME IMPLEMENTATIONS
Implementation Kronecker product We consider the case where experts are linear layers
Calculating domain-specific vectors taus weights weight_m
output_sizo = dim_out1 * dim_out2, input_size dim dim taus = taus, view (1, -1, dim_outl, din_out2, dim
dim repoat (rank, 1, 1, 1, 1, 1)
Calculate Kronecker-product.
taus torch, torch.einsum einsum ("rbij, rbjklm->rbikin", curvel_out, taus)
taus = torch.einsun ("rbik, rbjklm=>rbjiln", curve2_out, taus)
taus = forch.einsum ("rbii, rbjkim->rbjkin", curvel_in, taus) farcheeinsum ("rbim, rojkin=>rojkli", curve2 in, taus)
taus
Summation along the Kronecker rank dimension and reshape taus taus.sum (0)
taus = taus. reshape (-1, output_size, input_size)
D MORE EXPERIMENT DETAILS
Supervised Fine-Tune Fine-Tuning Eyper-Paraametess Besides {batch size, learning rate, epoch counts) which vary for each task, we keep other of supervised fine-tuning fixed for all tasks. These are shown in Table 8.
20

---

Published as a conference paper at ICLR 2025
Table 8: Fine-tuning hyper- perameeraeeesss models Section
Hyper- Parameear
Values
Optimizer
ADAMW
Adam
1e-6
Adam B
(0.9, 0.98)
Warm-up steps
16
Weight decay
0.01
LR scheduler
LINEAR DECAY
Scaling factor or
Kronecker rank
1
E DERIVATION
This is the derivation for Eqn in Section 2.6
N-1
Em Em a
M
M++1
(15)
5=1
Em +
N-1
Σ
M -
UC $
Ty
(8/ 25/t() (z(t)))))))
(16)
in
N-1
N-1
= Em +a
+M
- a 2
$ x+1 * [j
=
dc
(17)
j=l
domain-spocifie merging with curvature-aware
F STEP-BY-STEP WALKTHROUUH FOR KEY EQUATIONS OF CAMEX
F.1 KEY EQUATION FOR MERGING OF CAMEX
E E +o Σ Mr (sto7))+
(CA-Merg)
1=1
In (CA-Merg) equation, we consider the merging of experts at 1-th layer of the model. E), denotes the "base" expert that is not included in the routing process. 7. = E₃ -E!, denotes i-th the domain-vector that adapts the "base" expert to the corresponding domain. Finally, of denotes the score of the i-th domain vector w.r.t the imput. We view the merging of experts as a optimization problem where a a sf acts as the adaptive learning rate. Therefore, it is straightforward to integrate natural gradient approach into this equation by introducing curvature matrices M,. Due to the challenging tractability of the Fisher Matrix in the intermediate layers of deep models, we proposed to learn them empirically through backpropagation as indicated by Eqn. in the main text and a simmilar approach using meta-learning (Park & Oliva, 2019).
F.2
KEY EQUATION FOR MERGING IN DYNAMIC ARCHITECTURE
N-1
E++1
El
M
N
1=1
(Dynamic-Merg)
B1+1
E1+1
M Mr (+41++++ +1
to
The Jlynamic-ergy system perform two steps which are calculating base expert for the next layer and perform (CA-Merge), respectively. For the first step, eliminate the score and take the average of curvure-aware domain vector instead to avoid information leakage. The result then takes the role as the base expert for the next layer.
21

---

Published as a conference paper at ICLR 2025
F.3 CURVATURE UPDATE
In the main text, we try to give an explaination of how our method we update the curvature matrix with the curvature information of the parameters space. To achive that, first take the derivative of equation (CA-Merge) w.r.t the curvature matrix M;:
= (asy*T,) = (Ej Em)
(18)
J
To evaluate the gradient of the task loss L w.r. MI we apply the chain-rule:
0C
OC
oE
Uc
as,*
(Ez
Em)
(19)
2Mj
DE
0Mj
EEm
G STUDENT'S T-TEST FOR EXPERIMENTS ON GLUE DATASET
We report the t-tsst results, beginning with the null hypothesis Hg: The performance between each pair of T5-Tees-CA CA vy 75-Ties, and 75 on GLUE SST-2, MRPC, CoLA, and MNLI are the same.. In this test, we choose the significant value to 0.05.
Table 9: Evaluation results on SST-2 with different random seeds.
Index
TA
Ties
Vanilla
94.44
93.77
93.31
2
94.86
94.13
93.33
94.62
93.90
93.21
4 94.60
94.12
93.46
94.54
93.70
93.41
6 94.55
93.87
93.56
94.37
94.03
93.67
Table 10: T-statistic and p-value when evaluating on SST-2.
Test
t-statistic
p-value
Ties-CA vs Vanilla
13.72
1.08e-8
Ties-CA vs Ties
7.36
8.74c-6
Table 11: Evaluation results on MRPC with different random seeds.
Index
Ties
Ties
Vanilla
92.35
91.35
89.85
92.61
91.30
89.65
3 92.55
91.40
89.74
92.40
91.55
89.49
5
92.54
91.62
89.76
92.44
91.77
89.85
7
92.33
91.43
89.62
Table 12: T-statistic and p-value when evaluating on MRPC.
Test
t-statistic
p-value
Ties-CA vs Vanilla
42.91 12.95
1.67e-14
Ties-CA vs Ties
2.06e-8
22

---

Published as a conference paper at ICLR 2025
Table 13: Evaluation results on CoLA with different random seeds.
Index
Ties_CA
Ties
Vanilla
61.01
57.95
57.74
2
59.53
58.63
57.82
60.36
58.90
58.03
4
60.13
58.92
58.23
59.41
58.31
58.51
6
59.33
57.38
57.36
60.03
58.53
58.40
Table 14: T-statistic and p-value when evaluating on CoLA.
Test
t-statistic
p-value
Ties-CA Vanilla
7.14
1.18e-5
Ties-CA vs Ties
5.16
2.00e-4
Table 15: Evaluation results on MNLI with different random seeds.
Index
Ties.Cl
Ties
Vanilla
86.52
86.25
86.22
86.45
86.32
86.31
3
86.37
86.39
86.36
86.59
86.46
86.41
5
86.32
86.53
86.50
86,54
86.38
86.34
7
86.47
86.41
86.34
Table T-statistic and p-value when evaluating on MNLI.
Test
i-statistic
p-value
Tics-Ca vs Vanilla
2.29
0.04
Tics-CA vs Ties
1.49
0.16
Based on the p-values in the tables above, we draw the following conclusions:
The T5-Ties-A variant significantly outperforms T5-Ties and T5-Vanilly on SST-2, MRPC, and CoLA.
While T5-Ties-CA does not statistically significantly outperform T5-Ties on MNLI, it still demonstrates significant improvement over T5- Vanilla.
H ADDITIONAL EXPERIMENTS
H.1 INTEGRATING CAMEX INTO
We expand our experiments include broader range of most recent merging expert methods. Specifically, we integrated our CAMEx method with the Twin-Merging approach (Lu et al., 2024b). Key distinctions between CAMEx and Twin-Merging lie in their core mechanisms:
Our method is non-Buclidean merging method, which utilizes the curvature-aware matrix, whereas Twin-Merging is model merging method, which relies on Fuclidean merging.
Our approach is specifically designed for finetuning, in contrast to Twin-Merging, which is intended for post-training.
23

---

Published as a conference paper at ICLR 2025
Finally, our dynamic mechanism performs inter-layer to form the merged expert, unlike Twin- Merging, which uses within-layer pre-calculations for merging. To integrate our method with Twin-Merging, first fine-tune the Curvature Aware model for a specific GLUE task. At test time, we apply the Twin-Merging algorithm to merge experts, referring to our approach Twin-CA. Notably, easy to implement and helps reduce memory usage during inference. We adhere to the original implementation settings, using sparsity density value of 0.2.
found Tiin-Megging be simple yet powerful technique that is
Table 17: Performance of Twin-Merging and its Curvature Aware (CA) variant on GLUE tasks.
Method
MRPC
RTE
STSB
Twin-Merging
91.97
72.20
88.56
Twin-CA (Ours)
92.30
74.73
89.55
The results in Table 17 demonstrate the effectiveness of our CAMEX approach when integrated with the Twin-Merging mechanism GLUE tasks, highlighting its strong potential for incorporation into more advanced merging techniques.
H.2 EXPERIMENTS ON TOKEN-CHOICE VS EXPERT-CHOICE ROUTING
We also demonstrate our merging approach with the following routing mechanisms. We compare the baseline performance (i.e., the ties merging expert without Curvature Aware) under different routing mechanisms with the corresponding Curvature-Awar counterparts to see how different routing functions affect CAMEx performance. It is worth noting that Expert Choice routing is not compatible with the experts merging method, as discussed by Lory (Zhong et al., 2024) in their Subsection 5.3.:
Stable MoE routing (Dai et al., 2022).
Naive routing (Shazeer et al., 2017).
X-MoE routing (Chi et al., 2022).
Note that the Curvature Aware model leverages the segment routing strategy (the causal segmenting strategy) proposed in Lory (Zhong et al., 2024), enabling direct comparison between our model and the expert choice method. The results in Table suggest that Curvature- advanced routing strategies. The CA model consistently outperforms the baseline with Ties merging across all routing mechanisms. Additionally, we observe that both Naive routing CA and X-MoE routing CA deliver robust performance across GLUE tasks, while Stable MoE routing CA emerges the most reliable choice overall.
merging benefit from more
Table 18: Performance T5-baae variants on the finetuning GLUE tasks
Method
MRPC
RTE
STSB
SST-2
Expert Choice MoE
93.10
66.78
89.19
93,80
Stable MoE routing Ties
91.92
75.48
89.48
93.37
Stable MoE routing CA
92.96
78.76
89.64
94.63
Naive routing Ties Naive routing CA
91.44
75.54
88.58
93,92
92.49
78.70
89.56
94.61
X-MoE routing Ties
91.99
75.29
88.42
93.26
X-MoE routing CA
92.79
78.20
89.26
94.38
H.3 LONGER TRAINING FOR WIKITEXT-IOS 103 PRE-TRAINING
We conduct additional experiments by training for longer iterations the Wikitext-1033 dataset. The performance gaps between methods remain stable starting around epoch 40.
As shown in Figure the trends demonstrate consistent improvements of our method over the baseline, with the gap remaining significant after prolonged training.
24

---

Published as a conference paper at ICLR 2025
400
Vanilla
Perplexity: 22.29
Domain Specific
Perplexity: 19.68
350
Domanin CC
Perplexity: 17:39
300
250
150
100
50
10
20
30
40
50
60
Epoch
Figure Performance of Vannila, Domain- speciffc,, Domain Seeaiin-specffi-CCA. under longer pre- pre-training.
H.4 MORE COMPREHENSIVE ABLATION STUDY ON HYPERPARAMETERS
H.4.1 ABLATION STUDY ON
We extend the range of for the ablation study, specifically evaluating Dare-CA and Ties-CA with a [0.1,1.6). The evaluation is conducted using 5 different seeds, and the results are averaged.
STSB
A5R
Texck
Darch
Figure Test performance of Curvature-Aware methods under varying settings C..
The results in Figure 8 lead to the following observations:
The performance of the models is suboptimal or even worse than the vanilla baseline when either too small (a E (0.1,0.4) or too large (a: > 1.1).
Dare-CA is more sensitive to the choice of a, showing sharper improvements and declines across the range.
Ties-CA exhibits more gradual changes, suggesting it more robust to variations in a. The optimal range for is [0.8,1.0).
H.4.2 ABLATION STUDY ON NUMBER OF EXPERTS
We conduct additional studies on our method using different numbers of experts in the T5 backbone.
The following conclusions be drawn from Figures 9:
Increasing the number of experts generally improves accuracy up to certain point: Accuracy improves as the number of experts increases, with the most significant gains occurring from to 8 experts.
25

---

Published as a conference paper at ICLR 2025
RTF
5135
Nanber Digerts
Figure 9: Test performance Curvature-Aware methods under varying number of experts.
After 12 experts, the accuracy either saturates slightly decreases.
We suggest using eight experts, as it provides a balanced trade-off between performance and efficiency.
H.5 RESULTS ON LARGE BACKBONE (PI11-3)
To evaluate the effectiveness of CAMEx with large backbone, experiment with Phi-3. Table 19: Performance of Phi-3-mini variants on the fine-tuning tasks for the GLUE benchmark.
Model
Params
SST-2
MRPC
CoLA
STSB
RTE
QQP
QNLI
MNLI
Phi-3
3.8B
95.76
90.12
61.36
88.7
80.51
92.38
94.84
90.39
Phi3-Ties
7.4B
96.56
92.25
62.33
89.99
87.73
94.13
95.58
91.28
Phi3-Ties-CA
7.4B
97.23
94.04
63.43
90.27
88.09
94,80
95.82
92.13
The Ties-CA and Ties variants remarkably outperform the vanilla version, creating a substantial performance gap.
Tics-CA further enhances the performance of Ties in all listed tasks.
Thus, we believe that curvature awareness has potential for improving other language models (LMs).
26


Reviews:
[{"review_id": "Sew8A7GNUp", "reviewer": "Reviewer_ukHP", "summary": "The paper presents CAMEx (Curvature-Aware Merging of Experts), a novel method for merging experts in Sparse Mixture of Experts (SMoE) architectures. CAMEx leverages natural gradients to account for the non-Euclidean curvature of the parameter space, achieving more effective model alignment and improved performance.", "questions": "1. Given CAMEx’s potential for scaling, have you tested its efficiency and performance on larger models? If such experiments were beyond the scope of this paper, could you provide any empirical or theoretical estimates of how CAMEx’s computational advantages scale with model size and expert count? This would better inform readers of its applicability to real-world, large-scale SMoE deployments.\n2. Could you elaborate on the impact of key hyperparameters, particularly the scaling factor α, on CAMEx’s performance? A sensitivity analysis or guidance on optimal parameter settings across different tasks would be highly beneficial for practitioners looking to implement CAMEx effectively.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Originality\nCAMEx introduces a novel approach by leveraging natural gradients to accommodate the non-Euclidean geometry of the parameter space, moving beyond traditional Euclidean-based merging methods. This curvature-aware design aligns model updates more naturally with the parameter manifold’s structure, enhancing both optimization and generalization. Additionally, the dynamic merging architecture optimizes resource usage without performance loss, offering a new direction for efficiently scaling SMoE architectures. This geometry-aware expert merging approach represents a significant conceptual advance with potential applications in other areas of model optimization.\n2. Quality\nThe paper is well-supported by both theoretical and empirical evidence, with robust experiments across multiple tasks demonstrating CAMEx’s effectiveness. The theoretical foundation reinforces the practical findings, adding to the reliability of the results.\n3. Clarity\nThe paper is generally well-organized, presenting its methodology and results clearly. While some technical sections could be further optimized to enhance accessibility, the core contributions are conveyed effectively.\n4. Significance\nCAMEx addresses the critical need for scalable and efficient model architectures. Its consistent performance improvements across various tasks, coupled with reduced computational demands, underscore its practical impact. The success of CAMEx may also inspire further research into non-Euclidean methods within model optimization, highlighting its broader significance.\nIn summary, CAMEx’s originality, strong experimental support, and contribution to scalable architectures make it a technically sound and impactful addition to the field.", "weaknesses": "While CAMEx introduces a novel and promising curvature-aware approach, several areas could be improved to enhance the paper's rigor and impact.\n1. Limited Baseline Comparisons\nThe paper primarily compares CAMEx with standard Euclidean-based and a few curvature-aware merging methods but lacks comparisons with a broader set of state-of-the-art SMoE techniques, especially recent adaptive or gradient-based approaches. Including a wider range of baselines would better contextualize CAMEx’s strengths and highlight its unique advantages.\nSuggested Improvement: Expand comparisons to include additional recent methods in expert merging or routing within SMoE, particularly those emphasizing efficiency or scalability. This would reinforce CAMEx’s position within the current landscape.\n2. Clarity in Theoretical Exposition\nThe theoretical basis for CAMEx, rooted in natural gradient theory and parameter manifold curvature, could benefit from added clarity. Specifically, the mathematical distinctions between CAMEx’s curvature-aware merging and Fisher-based methods may not be fully accessible to all readers, particularly regarding computational efficiency.\nSuggested Improvement: Simplify or visually illustrate the theoretical explanation, potentially with diagrams or flowcharts, to clarify how CAMEx achieves computational advantages without sacrificing performance.\n3. Limited Analysis on Hyperparameters\nThe paper provides only a brief discussion of key hyperparameters, such as the scaling factor α. Given that curvature-aware methods may be sensitive to parameter tuning, a more comprehensive analysis of these hyperparameters would offer valuable insights.\nSuggested Improvement: Conduct a sensitivity analysis on α and other critical parameters, highlighting optimal settings across tasks. This would provide practical guidance for future users of CAMEx.", "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": "G09WBzrgWn", "reviewer": "Reviewer_2ASa", "summary": "The paper proposes CAMEx, a curvature-aware approach to merging experts in Sparse Mixture of Experts (SMoE) models. The work introduces a novel merging protocol that incorporates natural gradients to account for non-Euclidean parameter manifold geometry, a dynamic merging architecture that reduces computational costs while maintaining performance, and provides theoretical analysis showing how curvature-aware updates combine traditional domain-specific merging with task alignment adjustments. The empirical evaluation demonstrates consistent improvements across multiple NLP and vision tasks.", "questions": "- How does the choice of token-to-expert routing method (token-choice vs expert-choice) affect CAMEx performance?\n- Would the performance improvements be more significant when training larger models?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The curvature-aware approach consistently outperforms traditional Euclidean-based merging methods across different tasks and architectures\n- The dynamic merging architecture achieves the same number of experts but reduces FLOPs per token, providing a practical way to improve efficiency\n- The Kronecker-based approximation for the curvature matrix is computationally practical and empirically effective based on ablation studies\n- Strong theoretical analysis in Section 2.6 shows how gradients w.r.t curvature matrix combine domain-specific merging with an adjustment term for task loss alignment\n- Thorough empirical evaluation across diverse tasks, including ablation studies examining impact of key parameters (α scaling factor, Kronecker rank)", "weaknesses": "- Technical aspects of the method need clearer explanation - the causal segmenting approach lacks sufficient background/motivation in the main paper, key equations for merging need step-by-step walkthrough, and curvature-based updates require more detailed explanation\n- Performance improvements are somewhat modest, especially on the GLUE benchmark tasks\n- Training for longer appears to reduce the performance gap between proposed methods and baselines (Figure 3)\n- Impact on larger model scales is not thoroughly explored - unclear if improvements would be more pronounced with models bigger than GPT-2\n- Missing analysis of how different token-to-expert routing strategies might affect the method's performance", "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": "wuv7te0ngr", "reviewer": "Reviewer_FRuk", "summary": "The paper proposes a curvature-aware merging method using natural gradients. The experiments show that CAMEx can outperform Euclidean-based expert merging. I have read the paper several times and I really like the ideas from the paper. However, I would like to see more experimental result analysis about the model merging.", "questions": "I may miss something. What are the training datasets for table 2?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Curvature-aware merging is an interesting idea for expert merging.", "weaknesses": "1. In the section 3.2. \"Ties-CA achieves the highest scores on SST-2 (94.61), MRPC (92.49), CoLA\n(60.06), and MNLI (86.45), showing significant improvements over both the vanilla and standard\nTies models.\" Actually, the experiments in Table 2 don't show any significantance test. Could you present the significance test(T-test) of the improvement? Otherwise, you should reclaim these conclusions. \n\n2. The experiment results are not strong enough to prove the effectiveness of CAMEx. Could you show some downstream zero-out results based on your T5 model? Please refer to https://arxiv.org/abs/2405.11157. For example, the downstream results on question answering (boolq, OpenbookQA) and reasoning (BBH), \n \n\n3. The experiments only show the T5-base model. I am a little suspective that conclusions in the paper can also be generalized to other LMs. Could you show GLUE performance based on LLama-3 or phi-3( Vanilla, Ties, and Ties-CA is fine)? \n\n\nI would like to raise my score if there is more experimental analysis.", "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)