Representation Surgery in Model Merging with Probabilistic Modeling
Authors: Qi Wei
Shuo He
Enneng Yang
Tingcong Liu
Haobo Wang
Lei Feng
Bo An
Venue: ICML 2025 poster (2025)
Keywords: Model Merging
Representation Learning
Primary area: general_machine_learning->transfer_multitask_and_metalearning

Abstract: Model merging aims to achieve multitask performance by merging multiple expert models without the need to access the raw training data.
Recent research identified the \textit{representation bias} of model merging, characterized by a discrepancy in the representation distribution between the merged and individual models, hindering the performance of model merging methods. To mitigate the representation bias, a task-specific MLP, Surgery, was built to model the bias that is subsequently decreased on the merged representation. However, this strategy is still suboptimal due to the limited modeling capability within the deterministic manner. 
To address this issue, we present ProbSurgery, a probabilistic module specifically designed to accurately model the representation bias.
This module generates an embedding distribution for each sample and outputs the representation bias through a sampling process.
ProbSurgery offers superior representational capacity by naturally handling the uncertainty resulting from parameter interference of merging multiple models. Besides, we provide a theoretical analysis to reveal the advance of the probabilistic manner and propose an extension of ProSurgery for adapting to the task-sharing setting. Extensive experiments verify the effectiveness of ProbSurgery for representation surgery while maintaining generalization capabilities in real-world scenarios, including out-of-distribution and domain shift challenges.

Representation Surgery in Model Merging with Probabilistic Modeling
Qi Wei' Shuo He' Enneng Yang' Tingcong Liu
Haobo Wang? Lei Feng Bo An
Abstract
solely on specific task, MTL leverages the shared knowl- edge and representations among tasks, enhancing the overall performance (He & Lawrence, 2011). However, constrained by rigid training paradigm, i.c., "data collection first and then jointly training", MTL raises several concerns: 1) the high cost of data collection as well as the risk of privacy leakage, and 2) limited flexibility, as introducing a new task necessitates retraining the MTL model on combination of both old and new (tasks') datasets.
Model merging aims to achieve multitask perfor- mance by merging multiple expert models without the need to access the raw training data. Recent re- search identified the representation bias of model merging, characterized by discrepancy in the representation distribution between the merged and individual models, hindering the performance of model merging methods. To mitigate this bias, task-specific MLP, Surgery, was built to model the bias that is subsequently decreased on the merged representation. However, this strategy still suboptimal due to the limited modeling capa- bility within the deterministic manner. To address this issue, we present ProbSurgery, probabilistic module specifically designed to accurately model the representation bias. This module generates an embedding distribution for each sample and out- puts the representation bias through a sampling process. ProbSurgery offers superior representa- tional capacity by naturally handling the uncer- tainty resulting from parameter interference of merging multiple models. Besides, we provide theoretical analysis to reveal the advance of the probabilistic manner and propose extension of ProSurgery for adapting to the task-sharing setting. Extensive experiments verify the effec- tiveness of ProbSurgery while maintaining gener- alization capabilities in real-world scenarios. The code is now available this url.
Recently, model merging has attracted increasing attention (Matena & Raffel, 2022; Ilharco et al., 2023; Yadav et al., 2023; Yu et al., 2024; Yang et al., 2024c;a), as it offers a new approach to fulfilling MTL without requiring training data. This approach is encouraged by the growth of open- source communities like Huggingface and recognizes that published trained (or fine-tuned) models can be treated as resources, similar to datasets. Therefore, a model can obtain multi-task knowledge by adapting the parameters of multiple existing expert models that share a consistent network structure instead of training from scratch on a task- combined dataset. the era of foundation models, the advantages of data independence and training free make model merging more practical.
However, models merged through existing methods (Matena & Raffel, 2022; Ilharco et al., 2023; Yadav et al., 2023; Yang et al., 2024c) often experience performance degradation, making them hardly comparable to traditional MTL meth- ods (Sun et al., 2020b; Royer al., 2024), let alone that of individual models. recent study (Yang al., 2024a) analyzed this issue from representation perspective and discovered discrepancy in the representation distribution between merged and individual models, termed representa- fion bias, leading to degenerated performance. To mitigate this bias, the authors introduced a task-specific, learnable module after the encoder, termed Surgery, to refine the bi- ased representation. With the Surgery module, the merged model attains representational capabilities close to those of individually trained models.
1. Introduction
Multi-task learning (MTL) (Caruana. 1997) machine learning approach that simultaneously learns multiple re- lated tasks to improve model generalization and efficiency. Compared with traditional single-task learning that focuses
'Nanyang Technological University. Singapore 'Zhejiang Uni- versity, China 'Southeast University, China "Skywork AI, Singa- pore. Correspondence to: Lei Feng <fenglei scu.edu.cn>.
Despite notable improvements achieved, we argue that there are still two concerns about the Surgery module:
'ttpps: https:///hgggiggggcccee co. For example, HuggingFace hosts collection of over 1.7 million pretrained models.
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Representation Surgery in Model Merging with Probabilistic Modeling
task arithmetic
Surgery
w/ Ours
task arithmetic
w/ Surgery
w/ Ours
(a) GTSRB
(b) RESISC45
Figure 1. Comparisons of varying generated representations, given Task Arithmetic as the baseline and VIT- ITT-B/22 the backbone. We see that the feature rectified by Surgery (the red points) is not totally aligned with the optimal (the blue points).
(Insufficient performance) Has the representation discrep- ancy been fully addressed? We visualize the represen- tations learned with Surgery across different methods in Figure 1. The results indicate that the rectified representa- tions via Surgery still remain considerably distant from the optimal ones, which cannot be correctly tackled by the subsequent classification head.
We propose a probabilistic representation surgery module, named ProbSurgery, as post-calibration to correct biased representations resulting from model merging, which can consistently enhance the performance of existing model merging methods.
We provide theoretical analysis for ProbSurgery, which shows that modeling representation bias as a distribution yields lower classification error.
(Limited applicability) Can multiple task-specific Surgery modules be substituted with a single task-shared one?
The Surgery module is task-specific which relies on prior knowledge about the task ID during inference. Ideally, we aim construct single but powerful module, akin to a meta-model, capable of handling biased representations across all tasks. It provides fewer parameters and better applicability in real-world applications.
We propose two strategies for extending ProbSurgery to the one-to-all setting, i.e., learning ProbSurgery mod- ule for all merged tasks, which offers practical values.
Advancing performance and extensive analyses demonstrate the effectiveness and applicability of ProbSurgery.
To address these concerns, we deeply dig into modeling the representation bias by the Surgery mechanism and discover its probabilistic understanding in this paper. Specifically, rather than the deterministic manner employed in Surgery, which is constricted by limited representation capability, we reformulate the representation bias as distribution whose parameters are determined by MLPs in an amortized fashion. We refer to this module as ProbSurgery in our work. Then, the representation bias can be obtained by sampling from the generated distribution. In contrast to Surgery, ProbSurgery presents two advantages, including 1) better representation capability: Due to significant interference among the param- eters of multiple models when merging, the generated repre- sentation exhibits uncertainty. Probabilistic representation learning can effectively model the uncertainty. 2) Superior applicability: One single ProbSurgery can be directly uti- lized for tackling all merging tasks (named one-to-all setting in this paper), whose performance even outperforms that of Surgery trained by the task-specific setting. This proba- bilistic reformulation not only improves the flexibility and robustness of the model but also broadens its applicability to diverse and challenging real-world scenarios. Eventually, we further provide a theoretical analysis for ProbSurgery based on PAC-Bayes framework (McAllester, 2003), un- covering that modeling representation bias as a distribution yields lower classification error.
2. Related Work
2.1. Model Merging for Multi-Task Learning
Referring to Yang et al. (2024a), we divide existing model merging methods into three types based on the stage they focus on: before, during, and after merging.
1) The methods that operate before merging mainly focus on providing valuable prior knowledge for subsequent model merging. For example, Ortiz-Jimenez et al. (2024) add a constraint when fine-tuning the pretrained model in different tasks, i.c. updating the model's s parameters in independent Tangent spaces (Jacot et al., 2018). Similarly, one recent work, DARE (Yu et al., 2024), proposes to force the offsets of randomly selected parameters (above 90%) to remain unchanged. Thus, different tasks can benefit from indi- vidual parameter subspace and reduce performance losses.
2) The majority of works in model merging pay attention to mitigate the issue of interference and conflicts with dif- ferent models, which belong a family of during merging. Previous studies such as Weighted Averaging (Li al., 2023), Ties-Merging (Yadav et al., 2023), and AdaMerging (Yang et al., 2024c) typically pre-set or learn task-specific weights to adjust the contributions of different tasks to the final merged model. In the past year, a lot of methods (Lu et al., 2024; Tang et al., 2024; Huang et al., 2024) based on Mixture-of-experts or Router technique (Masoudnia &
Our contributions are summarized as threefold:
2

---

Representation Surgery in Model Merging with Probabilistic Modeling
Ebrahimpour, 2014) were proposed to split shared and task- specific knowledge and then dynamically select them for tackling downstream tasks after merging.
90-
performance
-0330
90-
Distance
"Ist Acc (%)
0.26
80
75-
0.24
75
3) The methods that operate after merging (Yang et al., 2024a;b) are proposed to rectify the biased representations generated after merging. The most representative method, Surgery (Yang et al., 2024a), constructs a lightweight, task- specific module to capture the bias between the optimal and interfered representations. A key merit of these methods lies in their seamless integration with weight-space merging methods, thereby enhancing the overall performance.
70
0.22
70
merge wi olou merge w w/Sueery
65
0.20
65
only merne
60-
Averoge
Antiomattc
Manjink
Layae webo
Adansrging
Average
Antometicc
plersing Lugraredion
Merging Methods
Merging Methods
Figure 2. VIT-B/32 and 8 vision tasks, Left: Comparison of av- erage performance and LI representation distance regarding four methods on all vision tasks. Right: Performance improves after rectifying representations using Surgery or ours. Upper perfor- mance denotes the average performance of various expert models on corresponding merged tasks.
In this paper, we design a post-calibration approach derived from the Surgery framework in a probabilistic manner, offer- ing enhanced generalization and representation capabilities. ProbSurgery can be integrated into existing merging meth- ods and further boost their performance.
0 base and O,, i.e., To 02 0 Obase
Task vector can be regarded as the task-specific knowledge and has been applied for various scenarios, such as knowl- edge transfer and adaptation, modeling task relationships, and task fusion in model merging (Ihaarco et al., 2023).
2.2. Probabilistic Representations
Learning representations in a stochastic embedding space first emerged in word embeddings (Vilnis & McCallum, 2014) and has since been widely applied to various natu- ral language processing tasks (Li et al., 2018; Neelakantan et al., 2015). Meanwhile, the probabilistic representations are also adapted into various vision tasks like face recog- nition (Chang et al., 2020; Shi & Jain, 2019), human pose few-shot
Based on these notations, we give the formal introduction to four model merging methods. 1 Weighted Average. which directly averages the weight of models on T tasks, is repre- sented by Ounit = 1 Ei= 0,. @ Task Arithmetic (Ilhacco et al., 2023), which merges the total T task vectors into the base models' parameters, denoted by 0unit Obase + A E Tr. where À is a haperrparameter. @ Ties-Merging (Yadav et al., 2023), which splits the merging processes into three steps: TRIM, ELECT SIGN, and MERGE, where the first two steps aim to eliminate the symbol conftict problem in task vectors, denoted by <(-) for clarity. The merged weight can be written as 0unif = Obase + > ((T)). 4 AdaMerging (Yang et al., 2024c), which learns set of task- wise or layer- wise coefficients for Task Arithmetic or -ies-Merging. denoted by Ounif = 0 Obase + Σt= λeT: and Quiif- = (0) base Et= NoT/)= (L denotes the total num-
estimation (Sun et al., 2020a), prototype- learning (Scott et al., 2019) and video understanding (Park et al., 2022). Existing methods demonstrate that probabilis- tic representations contribute to uncertainty modeling and improving the model's generalization (Li et al., 2018).
3. Preliminaries
Problem Definitions. Consider a set of models denoted as (fe...... }. fine-tuned from a base model fer.. on T individual tasks. In the context of model merging, our objective is to integrate these parameters {= into a unified one Qunif, where the resulting model, fe..rr can achieve consistently strong generalization across all tasks.
ber of layers/ / blocks in the model f), respectively.
3.2. Representation Bias in Model Merging
Formally, we aim to obtain a merged model fe..s that can fulfill the minimal loss value on the test dataset (D It=1 of all tasks, which can be formulated by this objective min 1 Σ EE (a {( (z), y), where & denotes the loss function, such as the cross-entropy loss.
Here, we first provide a formal definition of representation bias (Yang et al., 2024a) and then explain how it impacts the performance of merged models.
Definition 3.2. (Representation hias). Consider T ex- pert models and the merged model, the representation bias on t-th task is defined as & - where z denotes the extracted feature representation and (b()) is the arbitrary distance measure function.
'.. EiDi"l (2, Ounit xx."
3.1. Representative Model Merging Methods
In this section, we first define some notations within model merging and introduce four representative merging methods.
Observation 3.3. Consider T test datasets and individ- ual expert models {Die, fe. TT 1. the overall performance gap between the merged model fe..u and these expert mod- els obeys G(fomar: (fo,)T = Mccc///Acoffoo fo., D1O)
Definition 3.1. (Task vector). Assume a model fo. obtained by men-ttning the base model fer... on t-th task. The task vector Tt is defined as the element-wise difference between
3

---

Representation Surgery in Model Merging with Probabilistic Modeling
Acc(Semmer D EE--ETEE
Task
Task B
Task
Task B
Surgery
ProbSurgery
This observation be empirically validated (see Figure 2), which indicates that model merging methods that mini- mize the gap between expert models and the merged model are likely to achieve superior performance. Consequently, representation surgery emerges as a promising approach for further performance enhancement for the merged model.
(Marged) Encodet
Block L
cargery
BlockL
Block 2
Block 2
Block 1
Block
Dnaen
Townn
Input
Input
Imput
Input
4. Methodology
(a)
[b)
(c) Module comparison
Figure 3. Comparisons of different MTL frameworks. (a) Typ- ical model merging methods, which utilize the parameter fusion technique fulfill one unified model that can handle multiple tasks simultancously. (b) Post-merging representation surgery framework. The (last-layer) representation surgery module can be integrated with current model merging methods and boost their performance. (c) We compare our proposal, ProbSurgery. with Surgery (Yang et al., 2024a), which provides probabilistic repre- sentation leamning strategy for modeling the representation bias.
We divided our proposed methodology into two core parts:
A probabilistic representation surgery module, named ProbSurgery, which enables more accurate modeling for representation bias and thus achieves better performance. A theoretical analysis for ProbSurgery and two strategies for extending ProbSurgery to the one-to-all setting, i.e., learning a single ProbSurgery module for all merged tasks.
ProbSurgery not only outperforms its counterpart, Surgery, as a post-calibration technique for model merging, the pri- mary goal of our work, but also showcases markedly supe- rior performance in tackling out-of-distribution and domain- shift challenges. Figure shows the overall framework.
where μ and are learned by ProSurgery, lightweight, three-layer, fully connected MI.. From this representation distribution p(ElQ), we sample an embedding & to represent the latent representation bias of the original one &. where a reparamcterization trick (Kingma et al., 2015) is used as:
4.1. ProbSurgery
= a with e ~ ~~(00 1),
(3)
Previously, a post-merging representation surgery technique was proposed by Yang et al. (2024a), which aims to align the merged feature 20uut with the individual model-generated feature 2 for the task f via a surgery module denoted by (-), where the network g is modeled as three-layer fully connected MLP and w denotes the learnable parameters (see Figure 3 (b)). The training objective of Surgery is:
Correcting the biased embedding 2 S by decreasing the probabilistic bias Ç offers the main advantage: the prob- abilistic approach inherently incorporates uncertainty to effectively model the biased embedding. This significantly enhances the capacity to resolve ambiguities arising from the merging of multiple models in parameter space.
D
Overall, the learning objective in our framework is
arg min Lalign = arg min (wi,....rr> (Wi.....TT } 1=1 2=1
26 (2
Ounit - Er,,, )
arg min falign + X (wi....orr
where Eiee Iw. (2
(1)
|D
arg min
Σ Σ
2 (=20uait Punit - E,,,, 2. + LkLI (4)
Totally, there are trainable modules (gwi. 9ws) } that are independently trained on T' task validation sets. Due to limited representation capability, one task task-specific Surgery module can hardly be extended to other tasks.
(wI)...(T)) t=1
Ê., = MVa., (2..1"r) and L'L = (L(((EII) 1)).
Note that À is a hyperparameter set to 1x 10-4, (b-) denotes the reparameterization trick, Vw, denotes ProSurgery module with learnable parameters wi, and LKL denotes the Bellback-Leiblee (KL) divergence.
To solve this issue, we propose simple method that adapts the Surgery module in a probabilistic manner (see Figure 3). Firstly, we consider the same setting with Surgery (Yang et al., 2024a), i.e., a one-to-one setting that learns T task- specific modules for T individual tasks. Specifically, for specific task?, we formulate probability distribution p(EIQ) to represent the latent space Q of the representation bias &,
4.2. Theoretical Evidence
We theoretically demonstrate that, under a PAC-Bayes framework (McAllester, 2003), modeling representation bias as a distribution yields lower classification error.
which is denoted by a Gaussian distribution:
p(EIQ) ~ (H,
Let h(-) denotes the fixed classifier (after the encoder). To correct the encoder's output, we introduce a bias vector
FFor clarity, we omit the symbol of the task index
4

---

Representation Surgery in Model Merging with Probabilistic Modeling
4.3. Extending ProbSurgery to One-to-All Setting
Table Comparisons of different methods from four perspectives. Note that larger number of "O" denotes better performance. Combining ProbSurgery with DFA/PDA achieves trade-off be- tween parameter efficiency and performance.
In this section, we move beyond the "one-for-one" constraint in Surgery and propose two strategies to extend single ProbSurgery module for all merged tasks while preserving competitive performance.
Methods
No Task
Parameter Efficient
Uncertainty Modeling
Performance Robustness
Surgery ProbSurgery w/ DEA w/ PDA
x
I: Direct feature alignment (DFA). Owing to the strong performance of the probabilistic model, we empirically un- cover that a single ProbSurgery module alone can handle all biased representation from multiple tasks. Thus, the optimized objective in Eq. (4) is reformulated as
x
/
000
DO
/
00
0OO
5
V
00
00O
& E IR' in subtractive form: the input to the classifier is 2 - S. We compare 1) the deterministic bias, a single vector Edet. and 2) the probabilistic bias, a distribution Q: from which we sample for each input. Our goal is to show that our probabilistic approach can yield a strictly smaller classification error bound than the deterministic one.
arg min Lalign + XRrez {w}
=arg min Σ Σ
Bunit - Eis 2.:1) + LKL,
(8)
where the original optimized target [wi,.... ) is replaced by w. Experimental results show that existing model merg- ing approaches can achieve better performance improve- ments even when only a single ProbSurgery module used, compared to scenarios with multiple Surgery modules.
Firstly, we formally give the expected 0-1 classification error for our probabilistic method, as follows:
Rprob = (( E Q(()) [1(h(= - () + w)], (5)
I:: Proxy distribution alignment (PDA). To further boost the performance of the probabilistic model in uncertainty prediction, we change the mode of "feature-to-feature alignment to the "distribution-to-ccauuue alignment by mini- mizing scoring rules (Gneiting & Raftery, 2007), function S.) that evaluates how well a predicted distribution P over random variable X' aligns with the actually observed real- izations of X.
where 1(-) is an indicator function and D. denotes the whole sample set of all merged tasks with a size of N.
Theorem 4.1. (PAC-Bayes Theorem, adapted from (McAllester, 2003)) Let Qo(E) be a prior distribution over bias corrections and Q(E(z) be a data- dependent posterior. With the probability - 5 over the set DAV we have
N
Rerob
1
Σ
Eg_o [1 (h(z E.) # v.)
To be specific, we change the alignment target from the optimal representation zt to the bias 2 - 2 Then, by scoring rules, we can calculate the score between the generated bias distribution Q(4, o) and the actual observed bias. Eventually, the optimization objective
N
im1
Qo) + (24
N
The bound highlights that the probabilistic method achieves a balance between fitting the training data (low empirical risk, the first term) and maintaining simplicity (low KI-- divergence, the second term). By controlling the trade- off between these two terms, the bound ensures that the model generalizes well to unseen data, even when the output representations are biased.
arg min S(Qw, pbias) == arg min lE pbino (Qw, z biss w
(9)
In practical implementation, we have only one realized ob- servation for the bias, i.e., 2 biss 20mit - -20, Then, approximate S(.) to two terms and write the objective
T
IDM
Theorem 4.2. (Proofs: Appendix A) Based on the theorem above, the probabilistic method achieves a strictly tighter generalization error bound than the deterministic baseline:
arg min
Σ 1=1
Σ
( bias Zi,
(10)
where K E [0, 1] and is empirically set as 0.5, and r denotes the sampling number that is set as 5. Note that the first term is function of the sampled bias 2d and the observation bias z bias The second term is function of two embeddings, z' and z² drawn from the distribution Qu. More analyses for the sampling number r can be found in Appendix C.
(LI(Q)Q). N
Rprrab < RRiet + +Ol
(7)
This theorem demonstrates that the probabilistic method (our ProbSurgery) not only achieves tighter error bound but also explicitly reduces the generalization error compared to the deterministic baseline (i.e. Surgery).
A detailed comparison of varying methods is shown in Table
1. Experimental results showcase that the performance of
5

---

Representation Surgery in Model Merging with Probabilistic Modeling
Table Multi-task performance when merging ViT-B/32 models eight vision tasks. The green value denotes the improvement of ProbSureer compared with the counterpart method, Surgery. Results with more backbones are shown in Appendix C.S.
Method
SUN397
Cars
RESISC45
EuroSAT
SVHN
GTSRB
MNIST
DTD
Avg
Pretrained
59.7
45.5
31.4
32.6
48.5
43.8
Individual
75.3
77.7
99.7
97.5
98.7
99.7
79.4
90.5
Traditional MTL
73.9
74.4
93.9
98.2
95.8
98.9
99.5
77.9
88.9
Fisher Merging (Matena & Raffel, 2022)
68.6
69.2
70.7
66.4
72.9
51.1
87.9
59.9
68.3
RegMean (Jin et al., 2023)
65.3
63.5
75.6
78.6
78.1
67.4
93.7
52.0
71.8
Conerete (Tang et al., 2023)
62.5
61.1
76.0
95.7
91.0
81.9
98.5
51.9
77.3
Concrete AM (Tang et al., 2023)
70.0
96.0
91.6
96.7
98.7
63.8
84.0
AdaMerging (Yang et al., 2024c)
58.0
53.2
68.8
85.7
81.1
84.4
92.4
44.8
71.1
Weight Averaging
65.3
63,4
71.4
71.7
64.2
52.8
87.5
50.1
65.8
w/ Surgery (Yang al., 2024a)
67.6
64.6
85.8
96.8
76.9
82.9
97.8
67.3
80.0
w/ ProbSurgery (Ours)
70.7
70.1
94.0
99.6
83,4
98.7
99.3
78.1
86.7 (6.7)
Task Arithmetic (Mharco et al., 2023)
55.2
54.9
66.7
78.9
80.2
69.7
97.3
50.4
69.1
w/ Surgery (Yang et ul., 2024a)
63.8
59.9
83.3
97.9
87.0
87.0
98.6
69.4
80.9
w/ ProbSurgery (Ours)
67.0
67.0
94.1
99.8
91.2
98.8
99.4
79.0
87.0
Ties Ties-Merae Merging (Yadav 2023)
65.0
64.4
74.8
77.4
81.2
69.3
96.5
54.5
72.9
w/ Surgery (Yang 2024a)
69.8
66.1
87.3
97.5
86.7
87.6
98.5
71.6
83.1
ProbSurgery (Ours)
71.5
70.6
94.4
99.7
90.6
98.9
99.4
78.9
88.9
LW AdaMerging (Yang et al., 2024c)
64.5
68.1
79.2
93.8
87.0
91.9
97.5
59.1
80.1
w/ Surgery (Yang al., 2024a)
71.2
72.0
92.3
99.0
92.2
97.9
99.0
76.1
87.5
w/ ProbSuurery (Ours)
71.7
73.1
94.8
99.7
93.6
98.8
99.5
79.3
88.8(1.3)
SOTA: EMR MERGING (Huang al., 2024)
75.2
72.8
93.5
99.5
96.9
98.1
99.6
74.4
88.7
Table Multi-tas -ask performance when merging BERT models five NLP tasks. The green value denotes the improvement of ProbSurgery compared with the counterpart method Surgery (Yang et al., 2024a).
Method
AG News
Yelp Sentiment
Amazon Sentiment
Yahoo Q&A
DBPedia Wikipedia
Avg.
Traditional MTL
59.1
55.6
71.3
98.5
75.0
KD4MTL (Li & Bilen, 2020)
91.6
59.2
57.0
71.2
98.7
75.5
Weight Averaging
79.2
49.8
45.0
50.3
55.1
55.8
Surgery
90.3
58.0
54.2
70.8
98.4
74.3
w/ ProbSurgery (Ours)
91.4
62.1
55.8
72.4
98.1
76.7)
Task Arithmetic (ITharco et al., 2023)
82.9
55.8
48.4
53.1
81.5
64.3
Surgery
58.4
55.4
70.3
98.0
74.4
w/ ProbSurgery (Ours)
90.7
60.4
56.4
72.3
98.7
75.7 (1.3)
only one single ProbSurgery module with DFA (or PDA) can remarkably outperform that of multiple Surgery modules.
2023). Concrete AM (Tang et al., 2023), AdaMerging (Yang et al., 2024c), Surgery (Yang al., 2024a), and EMR- MERGING (Huang et al., 2024).
5. Experiments
Architecture. (1) Following Surgery (Yang et al., 2024a), we adopt CLIP (Radford et al., 2021) as the backbone model for vision tasks, where various ViT family architectures are utilized for the visual encoder. For NLP tasks, assign Bert the backbone. (2) The structure of ProbSurgery is lightweight, three-layer connected MLP network, formally denoted by (h,, h22 hi where hy represents the embed- ding size output by the encoder in CLIP or NLP models and h2 is the width of the hidden layer, hyperparameter (set as 128 for all settings) in this paper.
5.1. Experimental Setup
All experiments are implemented by Pytorch library and conducted on a single NVIDIA RTX A6000, We tested three times with different random seeds and reported the average performance.
We categorize the compared methods into two parts: 1) basic (non-merged) methods: Pretrained, Individual, and Traditional MTL, 2) model merging methods: Weight Aver- aging, Fisher Merging (Matena & Raffel, 2022), RegMean (Jin al., 2023), Task Arithmetic (Ilharco al., 2023), Ties-Merging (Yaday et al., 2023), Concrete TA (Tang al.,
Other details are divided into three parts, including Datasets. See Appendix B.1, 2) Baseline. See Appendix B.2, and 3) Training implementation. See Appendix B.3.
6

---

Representation Surgery in Model Merging with Probabilistic Modeling
Table Comparison results of ProbSurgery and Surgery, X denotes "learning eight modules for eight tasks independentl" and / denotes "learning one module for all merged tasks".
Methods
One Module
SUN397
Cars
RESISC45
EuroSAT
SVHN
GTSRB
MNIST
DTD
Avg.
Task Arithmetic
55.2
$4.9
66.7
78.9
80.2
69.7
97.3
50.4
69.1
w/ Surgery
63.8
59,9
83.3
97.9
87.0
87.0
98.6
69.4
80.9
w/ ProbSurgery
67.0
67.0
94.1
99.8
91.2
98.8
99.4
79.0
87.0
w/ Surgery
60.3
56.0
70.1
92.1
82.7
76.8
97.7
55.0
73.8
w/ ProbSurgery
DFA PDA
62.2
58.4
84.8
98.9
88.7
94.2
98.8
68.4
81.8
w/ ProbSurgery
63.1
58.7
85.7
99.5
89.0
95.0
99.2
68.9
82.4
Weighted averaging
Task Arithmetic
Ties, mergingrggggg
LW Adamerging
merge only
L1 distancce
merge w/ Surgery
merge w/ Ours
0.2
0.0
SUN397
Card
RESISC45
EURSAAT
SVW
CMSRR
HIIST
DTD
sum0a7
Car
RESOCUS
EuroSAT
SVNN
GTSRB
NNIST
DTO
suN297
Care
RESICC5
EUrOSAT
SVIll
GTTSe
MWST
DTO
SUN137
Car
RESISCAS
EUTESAT
SViW
GYSRa
NNVIST
DTD
Figure 4. Performance comparisons of our ProbSurgery with Surgery (Yang et al., 2024a) regarding calibrating the representation bias.
5.2. Main Results
ProbSurgery achieves an average accuracy of 76.0% versus KD4MTL's 75.5%. These results underscore the effective- ness and broad applicability of our method in delivering state-of-the-ar performance across diverse NLP tasks.
Vision tasks. In Table 2, report the comparison results of ProbSurgery with previous state-of-the-ar art methods on eight vision tasks where the backbone is ViT-B/32. We can observe that 1) current model merging methods are signifi- cantly worse than traditional multi-task learning methods, even for the best one, LW adamerging. It because the issue of parameter interference largely degrades the task knowledge provided by different models. 2) ProbSurgery demonstrates consistent improvements across various model merging strategies, including all four base methods, achiev- ing notably higher average accuracy on a diverse set of vision tasks. Compared to standard Surgery, ProbSurgery delivers stronger and more stable gains without sacrificing performance on individual tasks. For example, applying ProbSurgery to Weight Averaging boosts the average curacy from 80.0% to 86.7%, an increase of 6.7%. These consistent gains across different model merging strategies highlight the robust performance benefits that ProbSurgery brings to multi-task learning.
5.3. Performance in One-to-All Setting
Ideally, we hope to resort to one module to rectify repre- sentations from all merged tasks, which greatly reduces the training cost and improves the practicality. In Table 4, the results demonstrate the superiority of our proposed ProbSurgery approach. Notably, using a single ProbSurgery module significantly outperforms the performance achieved by using eight independent Surgery modules (81.8% or 82.4% vs. 80.9%), as reflected in the average accuracy scores. Additionally, both strategies, DFA and PDA, consis- tently enhance performance across multiple tasks, further validating the effectiveness of our probabilistic framework. This highlights the efficiency and robustness of our method in handling diverse tasks with unified module, largely promoting the real-world application of ProbSurgery.
NLP tasks. Following (Yang al., 2024a), we adopt Bert as the backbone for five NLP tasks and then conduct model merging. We report the results in Table 3. From this table, our post-calibration method (ProbSurgery) consistently out- performs Surgery on all five NLP tasks, yielding up to a 1.7% improvement (from 74.3% to 76.0%) under Weight Averaging and a 1.3% improvement (from 74.4% to 75.7%) under Task Arithmetic. Notably, this is the first time a post- hoc calibration approach has surpassed dedicated multi- task learning algorithm (KD4MTL (Li & Bilen, 2020)), as
5.4. More Analysis
Performance on modeling the representation bias. In this work, we aim to minimize the discrepancy in the rep- resentation distribution between the merged and individual models. To demonstrate this, we plot L1 distance after post-calibrating four methods with our ProbSurgery.
The results are presented in Figure 4. Overall, our method (represented by the red bars) consistently achieves lower LI distance across various tasks and base methods, indi-
7

---

Representation Surgery in Model Merging with Probabilistic Modeling
Table Comparison results of ProbSurgery and Surgery in the setting of domain shift and out- distribbinn
Settings
Domain shift
Dut-of-distribution
Weighted Averaging
Task Arithmetic
Weighted Averaging ACC (t)
Task Arithmetic
Methods
In-domain
Other-domain
In-domain
Other-domain
AUROC (t)
ACC (1)
AUROC (1)
Only merge
84,6 89.1 (4.5)
71.4 72.1 (0.7)
86.2 90.2 (4.0)
73.1 74.6 (1.5)
82.7 90.6 (7.9)
58.5 69.5 (11.0)
84.2
62.1 72.6 (10.5)
w/ Surgery
92.7 (8.5)
ProbSurgery
91.4 (6.8)
76.8 (5.4)
93.1 (6.9)
77.2 (4.1)
93.7 (11.0)
76.1 (17.6)
95.5 (11.3)
80.2 (18.1)
scouregy accuracy (10)
Weighted Averaging
IDistar 000 Lass (Disancc)) loss
9.25
Weighted Averaging
Marge w Sungery Merge w/ Ours
89
88.22
00:47
88,52
88.52
80-
Marpe Surpury
Harpe Ours
0.20-
Test Accuracy
87-
88
87.27
07.14
86.69
87.74
86.87
87,53
06.76
87.65
76
0.15
Ties merging
0.25-
Ties merging
06.04
56.43
86.
86.25
85-
Merge w/ Ours
Merge w/ Surpery
45 #5.63
UW Adameeringg
85
84.
Ties-merging
Harge Suppery
Artthmetic
84
84.10
Weighted Averaging
100
500
2000
150d
2000
2000
4000 dlao
scaed
2048
Steps
h2, the width of hidden layer
Figure 5. Better efficiency is achieved by proposed Prob-
Figure 6. parametre selection for h2, which denotes the width the hidden layer in ProbSurgery module.
Surgery compared to Surgery (Yang et al., 2024a).
cating superior performance in mitigating representation bias. Compared with "merge only" (green) and "merge w/ surgery" (blue), our approach not only reduces L1 distance more significantly in most tasks but also demonstrates con- sistent effectiveness regardless of different merging strate- gies. This highlights the stability and general applicability of our method, which can better align the merged represen- tations with their expected distributions, ultimately offering improved generalization in multi-task settings. More results on various distance metrics are shown in Appendix D.
We merged four models trained CIFAR-1O. 10, EuroSAT, SVHN, and GTSRB. Then, we introduce 500 samples from SUN397 in the test phase as the OOD data. More experi- mental details can be found in Appendix B.5. In Table 5 right, ProbSurgery outperforms Surgery on OOD tasks by achieving higher ACC and AUROC in both methods. For instance, it boosts ACC from 90.6% to 93.7% and AU- ROC from 69.5% to 76. 16.1% %% under Weighted Averaging.
Overall, we can consider that the variance in our proba- bilistic manner introduces extra cooss information, contributing to boosting the generalization performance.
Generalization. To evaluate the generalization of Prob- Surgery, we conduct experiments on two settings, including:
Efficiency. Probabilistic models often achieve faster con- vergence than deterministic models because they effectively model and capture uncertainty, which reduces the risk of overfitting and leads to smoother optimization paths.
Domain shift. We test the performance in a domain- shifted setting, which would reflect the generalization of generated representation bias. Specifically, we merge four individual models, independently trained on SUN397, Cars, Homeeffice (Real-World) and DTD with traditional model merging methods. After post-calibrating, we com- pute the average test accuracy on the other three domain datasets Home (Art, Clipart, and Product). The results are shown in Table 5 left, reflecting that the im- provement the in-domain dataset brought by Surgery can hardly transfer to other-domain datasets. In detail, we can see that with weighted averaging, the improvement of test accuracy in Real-world domain achieves 4.5% af- ter training with Surgery, while only an average 0.7% improvement obtained in the other three domains. In contrast, we have a performance increase of 5.4%.
In Figure 5, we plot the test accuracy of four methods in- tegrated with Surgery (Yang et al., 2024a) and ours. From these curves, is evident that our method ProbSurgery (blue lines) converges much faster than Surgery (red lines) on all merging methods. With only 300 steps, ProbSurgery has already surpassed Surgery's 's best performance across all four base methods. This improvement underscores the effi- ciency of our technique, demonstrating that even with fewer training steps, it can achieve superior test accuracy.
5.5. Ablation Study
Hyper-parameters. In ProbSurgery, there exists one es- sential hyper-parameter, i.e., h2, the hidden layer's width. Results shown in Figure 6 demonstrate that the choice of h2 would impact test accuracy. The best performance among
Out-of-distributtin. We compare ProbSurgery with Surgery on the OOD Detection task, where the test phase will occur with samples from other datasets/distribuutions..
8

---

Representation Surgery in Model Merging with Probabilistic Modeling
all base methods is observed when the value of h2 is be- Tween 128 and 512. At l22 128, the hidden layer is wide and the complexity of learnable parameters is enough to capture complex patterns, leading to significant accu- racy boost. Increasing 22 to 512 further stabilizes perfor- mance, but beyond this point (e.g.. h2 = 2048), accuracy slightly decreases due to overfitting caused by excessive ca- pacity. Thus, selecting l2 within the interval of [128, 512] contributes to both model generalization and task-specific learning.
Chang, J., Lan, Z., Cheng, C., and Wei, Y. Data uncertainty
learning in face recognition. In CVPR, pp- 5710-5719, 2020.
Cheng, G., Han, J., and Lu, X. Remote sensing image scene classification: Benchmark and state of the art. Proceed- ings of the IEFE, 15(10):1865-1883 2017.
Cimpoi, M., Maji, S., Kokkinos, I., Mohamed, S., and Vedaldi, A. Describing textures in the wild. In CVPR, pp. 3606-3613 2014.
6. Conclusion
Del Corso, G. M., Gulli, A., and Romani, F. Ranking a stream of news. In WWW, pp. 97-106, 2005.
In this paper, we follow one research Surgery (Yang et al., 2024a) and study how to mitigate the representation bias inherent in the final merged model, which aims fill the performance gap between the merged model and multiple individual models. To this end, we propose to model the bias in probabilistic manner and design module named Prob- Surgery. This module outputs more robust representation bias via sampling from the generated distribution that well models the uncertainty within model merging. Meanwhile, we propose two strategies that can extend ProbSurgery to the one-to-all post-calibration setting, which brings more practical values. Extensive experiments demonstrate the effectiveness of ProbSurgery.
Gneiting, T. and Raftery, A. E. Strietly proper scoring rules, prediction, and estimation. Journal tfe American statistical Association, 102(477):359-37, 2007.
He, J. and Lawrence, R. A graphbased framework for multi- task multi-view learning. In ICML, pp. 25-32, 2011.
Helber, P., Bischke, B.,, Dengel, A., and Borth, D. Eurosat: A novel dataset and deep learning benchmark for land use and land cover classification. IEEE Journal of Se- lected Topics in Applied Earth Observations and Remote Sensing, 12(7):2217-2222 2019.
Huang, C.,, Ye, P., Chen, T., He, T., Yue, X., and Ouyang, W. Emr-Merging: Tuning-free high-performance model merging, In NeurlPS, 2024.
Acknowledgment
This research is supported by the Ministry of Educa- tion, Singapore, under its MOE AcRF Tier 3 Award MOF-MOET32022-006 and MOE AcRF Tier Award T2EP20223-0022. Any opinions, findings and conclusions or recommendations expressed in this material are those of the author(() and do not reflect the views of the Ministry of Education, Singapore.
Ilharco, G., Ribeiro, M. T., Wortsman, M., Schmidt, L., Hajishirzi, H., and Farhadi, A. Editing models with task arithmetic. In ICLR, 2023.
Jacot, A., Gabriel, F,, and Hongler, C. Neural tangent kernel: Convergence and generalization in neural networks. In NeurIPS, 2018.
Impact Statement
Jin, X., Ren, X., Preotiuc-Pietro. D., and Cheng, P. Data- less knowledge fusion by merging weights of language models. ICLR, 2023.
This paper aims to correct biased representations resulting from model merging, promising approach to multi-task learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
Kingma, D. P., Salimans, T., and Welling, M. Variational dropout and the local reparameterizztion trick. NeurIPS, 28, 2015.
References
Krause, J., Stark, M., Deng, J., and Fei-Fei, L. 3d object representations for fine- one-grained categorization. In ICCV workshops, pp. 554-561, 2013.
Auer, S., Bizer, C., Kobilarov, G., Lehmann, J., Cyganiak, R., and Ives, Z. Dbpedia: nucleus for a web of open data. In International Semantic Web Conference, pp. 722- 735. Springer, 2007.
LeCun, Y. The Mnist database of handwritten digits. http://yann. lecun, com/exdb/nmist// 1998.
Li, T., Huang, Z., Tao, Q., Wu, Y., and Huang, X. Train- able weight averaging: Efficient training by optimizing historical solutions. In ICLR, 2023.
Caruana, R. Multitask learning. Machine learning, 28: 41-75, 1997.

---

Representation Surgery in Model Merging with Probabilistic Modeling
Li, W.-H. and Bilen, H. Knowledge distillation for multi- task learning. In ECCV Workshop, pp. 163-176, 2020. Li, X., Vilnis, I., Zhang, D., Boratko, M., and McCallum, A. Smoothing the geometry of probabilistic box embeddings. In ICLR, 2018.
Tang, A., Shen, L.,, Luo, Y., Ding, L.,, Hu, H., Du, B., and Tao, D. Concrete subspace learning based interference elimination for multi-task model fusion. arXiv preprint 0rXiv:2312.0113,, 2023.
Tang, A., Shen, L.,, Luo, Y., Yin, N., Zhang, L., and Tao, D. Merging multi-task models via weight-ensembling mixture of experts. In ICML, 2024.
Lu, Z., Fan, C., Wei, W., Qu, X., Chen, D., and Cheng, Y. Twin-Merging: Dynamic integration of modular expertise in model merging. In NeurlPS, 2024.
Van der Maaten, and Hinton, Visualizing data using t-sne. Journal of Machine Learning Research, 9(11), 2008.
Masoudnia, S. and Ebrahimpour, R. Mixture of experts: literature survey. Artificial Intelligence Review, 42: 275-20 293, 2014.
Vilnis, L. and McCallum, A. Word representations via gaussian embedding. arXiv preprint arXiv: arXiv:1442.6600 1412.6623, 2014.
Matena, M. and Raffel, C. A. Merging models with wiighte averaging. In NeurlPS, 2022.
Xiao, J., Ehinger, K. A., Hays, J., Torralba, A., and Oliva, A. Sun database: Exploring a large collection of scene categories. IJV,, 119:3-22, 2016.
McAllester, D. A. Pac-Bat Bae-Bayeni stochastic model selection. Machine Learning, 51(1):5-21, 2003.
Neelakantan, A., Shankar, J., Passos, A., and McCallum, A. Efficient non-parametric estimation of multiple embed- dings per word in vector space. In EMNLP, 2015.
Yadav, P., Tam, D., Choshen, L., Raffel, C., and Bansal, M. Resolving interference when merging models. In NeurIPS, 2023.
Ortiz-Jimenez, G., Favero, A., and Frossard, Task arith- metic in the tangent space: Improved editing of pre- trained models. In NeurlPS, 2024.
Yang. E., Shen, L., Wang. Z., Guo, G., Chen, X., Wang, X., and Tao, D. Representation surgery for multi-task model merging, In ICML, 2024a.
Park, J., Lee, J., Kim, I.-J,, and Sohn, K. Probabilistic representations for video contrastive learning. In CVPR, 2022.
Yang, E., Shen, L., Wang, Z., Guo, G., Wang, X., Cao, X., Zhang. J., and Tao, D. Surgeryv2: Bridging the gap be- tween model merging and multi-task learning with deep representation surgery. arXiv preprint arXiv:2410. 14389, 2024b.
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al. Learning transferable visual models from natural language supervision. In ICML, pp. 8748-8766 8763, 2021.
Yang. E., Wang, Z., Shen, L., Liu, S., Guo, G., Wang, X., and Tao, D. AdaMerging: Adaptive model merging for multi-task learning. In ICLR, 2024c.
Royer, A., Blankevoort, T., and Ehteshami Bejnordi, B. Scalarization for multi-task and multi-domain learning at scale. In NeurIPS, 2024.
Yu, L., Yu, B., Yu, H. Huang, F., and Li, Y. Language mod- els are super mario: Absorbing abilities from homologous models free lunch. In ICML, 2024.
Scott, T. R., Ridgeway, K., and Mozer, M. C. Stochastic prototype embeddings. In ICML Workshop, 2019.
Yuval, N. Reading digits in natural images with unsuper- vised feature learning. In NeurIPS Workshop. 2011.
Shi, Y. and Jain, A. K. Probabilistic face embeddings. In ICCV, pp. 5902-6911, 2019.
Zhang, X., Zhao, J., and LeCun, Y. Character-level con- volutional networks for text classification. NeurIPS, 2015.
Stallkamp, J., Schlipsing, M., Salmen, J., and Igel, C. The german traffic sign recognition benchmark: A multi-class classification competition. In IJCNN, 2011.
Sun, J., Zhao, J., Chen, L.-C.,, Schroff, F., Adam, H., and Liu, T. View-invariant probabilistic embedding for human pose. In ECCV, 2020a.
Sun, X., Panda, R., Feris, R., and Saenko, K. Adashare: Learning what to share for efficient deep multi-task leamn ing. In NewrIPS, 2020b.
10

---

Representation Surgery in Model Merging with Probabilistic Modeling
Organization of the Supplemental Material
The main contents of this appendix are as follows:
Appendix A: We provide the theoretical proofs for the given therome in the paper.
Appendix B: We describe the datasets, training procedures, and architecture of PorbSurgery in detail.
Appendix C: We show some experimental results and analyses that are deleted due to the page limit of the main text.
Appendix D: We visualize the representation generated by Surgery and ProbSurgery, respectively, via t-SNE technique (Van der Maaten Hinton, 2008).
A. Theoretical Proofs
In this section, we first provide analysis for the limitation of the deterministic method and then give detailed proof steps for Theorem 4.2.
A.1. Analysis for the Deterministic Method
The deterministic method post-representation calibration manner, like surgery, is equivalent to Dirac delta posterior:
(Elx( = 8(5 - Edet), where 8(E Edet) = +00, 5 = Edet
(1I)
0, Edet
Substituting into the PAC-Bayes bound (McAllester, 2003), the KL-divergence term KL(Quet Qo) diverges unless Qo is also delta distribution, which violates the PAC-Bayes assumptions. Thus, deterministic representation correction cannot leverage PAC-Bayes bounds and instead relies on classical VC-dimension bounds:
d Raet S Raot + +0( N
(12)
where d denotes the complexity the three-layer MLP in Surgery, and Rotets Rast denote the expected and empirical risk, respectively.
A.2. Proofs for Theorem 4.2
From Theorem 4.1, for any prior Qu and posterior Q, with the probability 1 over the training set N we have
N B [1((2z [) # w:)l +
KL QQ) + N
Rprob I
(13)
N
N
Complexity Term
Empirical Risk Agyes
For the first term, i.e., the empirical risk term for ProbSurgery, it satisfies:
1 N
Ex-q [1(h(2 {) * ve)l <
1(h(2 Ede) * yr),
(14)
where Eaer [E-QlK] is the deterministic correction (equivalent to Surgery). This follows from Jensen's inequality applied the convex 0-1 loss. Thus, we can deduce that Rprob Ract-
For the deterministic method (Surgery), classical VC-dimension bounds give:
d (V
Raes < Rist + O
(15)
N
where d the VC-dimension (or parameter count) the deterministic network.

---

Representation Surgery in Model Merging with Probabilistic Modeling
Eventually, by substituting the empirical risk bound in Eq. (14) into the PAC-Bayes bound in Eq. (13) and then combining with Eq. (15), we have
KL(Q1Q0)
Rprob S Rdet + +V
(16)
Here, we omit the item of In N 8 since its value is very small given large scale of the training set (i.e., N is large).
B. Experimental Setup
B.1. Datasets
Following prior studies model merging, such Task Arithmetic (Ilharco al., 2023), Ties-Merging (Yadav et al., 2023), AdaMerging (Yang et al., 2024c), and Surgery (Yang al., 2024a), merge models trained on the following eight vision datasets and five NLP datasets:
SUN397 (Xiao et al., 2016); A benchmark for Scene Understanding (SUN) featuring 108,753 images across 397 classes. Each class contains at least 100 images.
Cars (Krause et al., 2013): The Stanford Cars dataset comprises 16,185 images spanning 196 classes, approximately split 1:1 into training and test sets.
RESISC45 (Cheng et al., 2017): A publicly available benchmark for remote sensing image scene classification. includes 45 scene classes, each with 700 images 2256 resolution), totaling 31,500 images.
EuroSAT (Helber et al., 2019): A dataset of Sentinel-2-bassee satellite images focusing on land-use classification. It consists of 27,000 labeled and geo-referenced images divided into 10 classes.
SVHN (Yuval, 2011): A dataset of real-world house number images in 10 classes, visually resembling MNIST (LeCun, 1998) but containing over 600,000 color images.
GTSRB (Stallkamp et al., 2011): The German Traffic Sign Recognition Benchmark, with over 50,000 images in 43 classes. Images vary in lighting conditions and backgrounds.
MNIST (LeCun, 1998): One of the most renowned datasets in machine learning: 70,000 (60k training 10k testing) grayscale images of handwritten digits in 10 classes, each sized 28 x2.28.
DTD (Cimpoi et al., 2014): The Describable Textures Dataset (DTD) comprises 5,640 labeled texture images in 47 classes, with each class containing around 120 images (ranging from 300x300 640x640 in resolution).
AGNews (Del Corso et al., 2005): A news classification dataset with roughly 120k training samples and 7.6k test samples, covering 4 classes (World, Sports, Business, Science/Technone Tecchoo/Tccho.ooy)..
Yelp (Zhang al., 2015): A sentiment analysis dataset of user-submitted reviews (positive negative), containing around 560k training samples and 38k test samples.
Amazon (Zhang et al., 2015): Another sentiment analysis dataset (positive vs. negative) from Amazon product reviews, including about 3.6M training samples and 400k test samples.
DBPedia (Zhang et al., 2015): Focused on categorizing Wikipedia article excerpts into 14 classes (e.g., Company, Artist, Educattioalllsiituiin))) with approximately 560k training samples and 70k test samples.
Yahoo (Auer al., 2007): question-and -nd--answddrrrrrrrraa classification dataset of around 1.4 million training samples and 60k test samples, covering 10 classes (e.g., Society & Culture, Science & Mathematics, Health).
B.2. Baselines
We categorize the compared methods into two parts: 1) basic (non-merged) methods: Pretrained, Individual, and Traditional MTI., 2) model merging methods: includes as follows
Weight Averaging: Simply averages the corresponding weights from multiple trained models to a merged model.
Fisher Merging (Matena & Raffel, 2022): Merges models by weighting parameters according to Fisher information, prioritizing directions erucial for each task.
12

---

Representation Surgery in Model Merging with Probabilistic Modeling
Individual models
0
Target represenlation
Alignment
0²
@
a pm
S, 20
Z gt
merging
Mcan
Representation
The bias
Samples pooling
1
sample
Rectified representation
gm
h1
h2
Merged model
h1
Variance
ProbSurgery
Figure 7. Flowchat the post-representation calibration processes with our proposal, ProbSurgery.
Table Comparison of the number of learnable parameters within Surgery and ProbSurgery. Our proposal ProhSurgery can he applied one-to-all settings, i.e., learning one ProbSurgery module for all merged tasks. Thus, the number of extra parameters will not increase with the number of merged tasks increase.
Networks
VIT-B/16
VIT-B/32
VIT-L/14
Number of the backbone's 's parameters
907, 589, 640
894, 032
2, 740, 196,992 392
Number extra learnable parameters
Surgery Module w/ merging N tasks
131, 072
196, 608 N x 196, 608
N 131, 072
ProbSurgery Module w/ merging N tasks
131, 712
197, 504 197, 504
131, 712
RegMean (Jin et al., 2023): Adds a regularization term when averaging models, penalizing large deviations from each individual task model.
Task Arithmetic (Ilharco et al., 2023): Views models as points in parameter space and combines them via vector arithmetic operations (e.g., addition, subtraction).
Ties-Merging (Yadav et al., 2023): Selectively ties and reuses shared parameters across models to preserve beneficial components during merging.
Concrete TA& & AM (Tang et al., 2023): Provides a continuous relaxation of task arithmetic, learning an optimal linear mixture of tpec-fie parameters.
AdaMerging (Yang et al., 2024c): Adaptively reweights each model's parameters based on their contribution to the merged solution.
B.3. Implementation Details
To (post)train and learn the parameters of proposal, ProbSurergy, follow the convention in Surgery (Yang et al., 2024a). To be specific, we do not require any information from the test labels rather than adopting self-training manner. For training, we adopt Adam as the optimizer with learning rate of x 10 for all training iterations. We totally train the ProbSurgery module for 5, 000 iteration with a batch size of 16. During the training phase, we utilize the stochastic sampling strategy, i.c., the reparameterization trick, to obtain the generated representation bias. For inference and testing, we consider the mean value y output by ProbSurgery as the representation bias.
13

---

Representation Surgery in Model Merging with Probabilistic Modeling
Table 7. Multi-task performance on the ViT-B/32 model when different loss functions are used in the representation surgery module. The green value denotes the improvement of ProbSurgery compared with the counterpart method Surgery (Yang et al., 2024a).
Methods
Loss Func.
SUN397
Cars
RESISC45
EuroSAT
SVHN
GTSRB
MNIST
DTD
Avg.
Task Arithmetic (Ilharco et al., 2023)
55.2
54.9
66.7
78.9
80.2
69.7
97.3
50.4
69.1
w/ Surgery
(L1 Loss)
63.8
59.9
83.3
97.9
87.0
87.0
98.6
69.4
80.9
w/ ProbSurgery Loss)
66.3
65.8
94.3
99.7
90.5
98.6
99.4
78.0
86.6 (5.7)
w/ Surgery
(MSELoss)
64.3
59.8
84.0
97.8
87.6
88.7
98.8
69.8
81.4
w/ ProbSurgery
(MSELoss)
66.5
65.9
94.1
99.7
91.1
98.9
99.5
78.5
86.7 (5.3)
w/ Surgery
(SmoothLlLoss
64.1
59.7
84.1
97.9
88.1
89.7
98.8
70.6
81.6
w/ ProbSurgery
(SooothLlLoss)
66.0
63.4
93.9
99.6
90.8
98.5
99.4
78.2
86.2 (4.6)
B.4. Architecture of ProbSurgery
Our proposal module, ProbSurgery, belongs the lightweight network, which is parameterized by a three-layer fully connected MLP and exhibited in Figure 7. We can observe that this module is simple and contains fewer learnable parameters. whose architecture is represented by {h1, h2,h1) h1). Note that h1 denotes the width of feature embedding output by the encoder in CLIP or NLP models. For example, hn is set as 512 for ViT-B/16, 32 and 768 for ViT-L/14. Furthermore, h2 is the width of the hidden layer, a hyperparameter in this paper. We set h2 = 128 for all experiments.
We report the parameter comparison between Surgery and our ProbSurgery in Table 6. Firstly, even though our method belongs a probabilistic manner, does not significantly introduce more parameters compared with Surgery. Then, when the number of merged tasks increases, the advantage of the learnable parameter would be more obvious since ours can be directly extended into one-to-all settings.
B.5. Implementation of OOD Experiments
In the OOD experiment, we introduce two metrics, i.c., test accuracy and AUROC, to evaluate the performance of Surgery and ProbSurgery. The computation details for these two metries are as follows. When we finish the model merging and the post-calibration process with (Prob)Surgerr,
test the performance of different model merging methods on OOD detection. We use the test sets of four tasks and small OOD set of 500 samples randomly sampled from the test set of SUN397, denoted by Do, D1, D2, D3, and D.
To measure test accuracy. we follow the standard model mooeiing approach and maintain separate task-specific classifiers for the merged tasks. Formally, the test accuracy on the test set of the 1-th task D, can be expressed as ACC, Ei,.)en, (()=v) where I denotes the predicted label.
[|+||++|DDI|
To report AUROC, we use a commonly adopted unsupervised technique that does not rely on the classifier. Specifically, for each task, we store its class prototypes during training (when no OOD samples are present). In the test phase, for task t, we compute the representation of every sample in the test set D, + D. and measure its distance to the stored class prototypes. The greater the distance, the more likely the sample is to be outlier. Eventually, the AUROC value be calculated on the distance and the 0-1 label (where 0 and denote the In-distribution and Out-of-distribution samples, respectively).
C. More Experimental Analyses
C.1. Impacts of Loss Functions
In this paper, to measure the discrepancy the representation distribution be- tween the merged and individual models, we utilize LI distance (torch. nn. LiLoss ()) as the optimization loss function, which keeps the same as Surgery (Yang al., 2024a). To verify that this discrepancy is agnostic with different distance metrics, we test totally four distance metrics and set them as the loss function for comparison, including L1 loss, MSE loss and Smooth L1 loss.
The experimental results in Table demonstrate the superior performance of our proposed method, ProbSurgery, compared to the baseline Surgery across various loss functions. ProbSurgery achieves consistent improvements in average performance, with gains ranging from +4.6 to +5.7, highlighting its adaptability to different loss functions. Notably, ProbSurgery excels in datasets like RESISC45 and DTD, where significantly reduces task interference and enhances representation alignment,
14

---

Representation Surgery in Model Merging with Probabilistic Modeling
Weighted averaging
Task Arithmetic
Ties merging
LW Adamerging
88
89.0
Test Test Accuracy (%)
Test Accuracy (%)
Test Test Accuracy (%)
Test Accuracy (%)
86
84
merging Sugery
merging w/ Sugery
86
merging Sugery
88.5
merging w/ Sugery
merging PDA
84
merging PDA
merging PDA
merging w/ PDA
82
88.0
82
84
80
87.5
5
10
20
10
20
10
20
5
10
20
Sampling number in PDA
Sampling number in PDA
Sampling number in PDA
Sampling number in PDA
Figure Performance impact of sampling number regarding computing scoring values our Proxy Distribution Alignment (PDA).
achieving gains as high as +11.0 and +8.6, respectively. These improvements demonstrate its robustness handling diverse tasks, from large-scale datasets like SUN397 to texture-basse datasets like DTD. Moreover, its compatibílity with different loss functions and ability to generalize across tasks establish ProbSurgery an effective approach for multi-task learning.
C.2. Impact of Sampling Number in Proxy Distribution Alignment (PDA)
In real implementation, to better estimate the discrepancy between the predicted distribution Qw with our practical representation bias observation abass the scoring rules should leverage the sampling strategy (i.e., the reparameterization trick) to compute the distance between two vectors. Therefore, larger number for sampling can better approximate the true distribution and reduce estimation uncertainty. It is worthy-noted that increasing the sampling number would not raise the learnable parameters and training costs since it does not involve extra gradient computation.
In Figure 8, we report the influence of sampling number regarding computing scoring values in our proposed Proxy Distri- bution Alignment (PDA). The experimental results demonstrate that increasing the sampling number in PDA consistently improves test accuracy across different merging strategies, with the performance gains becoming stable the sampling number reaches 10 or more. This highlights the effectiveness of leveraging larger sampling sizes better approximate the true distribution and reduce estimation uncertainty. Notably, the use of PDA consistently outperforms the baseline (merging without PDA) in all methods, with LW Adamerging achieving the highest overall accuracy. Furthermore, the approach introduces minimal computational overhead, as no additional learnable parameters or gradient computations are required, making PDA a practical and efficient method for improving model performance.
C.3. Impact of the Size of the Validation (Test) Set
In Surgery and ProbSurgery, we rely on unlabeled test/validation data to build self-supervision signals and update the parameters of the (Prob)Surgery module. Here, we made an experiment to verify how the size of the validation set impacts ProbSurgery's s performance. In the following table, we can see that using a larger proportion of the validation set helps the ProbSurgery module better capture and mitigate inherent biases, leading to higher overall accuracy. Notably, even when only 10% of the unlabeled validation data is used, method still exceeds Surgery's performance of 80.9%, demonstrating its robustness under limited data conditions.
Table 8. Comparison of performance across datasets under different validation set ratios. The backbone is VIT-B/32.
Method
Ratio
SUN397
Cars
RESISC45
EuroSAT
SVHN
GTSRB
MNIST
DTD
Avg
Task Arithmetic
55.2
54,9
66.7
78.9
80.2
69.7
97.3
50.4
69.1
w/ Ours
10%
63.5
64.5
89.6
94.9
88.4
89.4
98.4
74.5
82.9
w/ Ours
50%
65.9
66.5
92.3
97.8
90.5
95.0
98.8
77.2
85.5
w/ Ours
100%
67.0
67.0
94.1
99.8
91.2
98.8
99.4
79.0
87.0
w/ Surgery
100%
63.8
59,9
83.3
97.9
87.0
87.0
98.6
69.4
80.9
15

---

Representation Surgery in Model Merging with Probabilistic Modeling
C.4. Performance on Merging Tens of Tasks
To verify the performance of our ProbSurgery in more challenging settings, we additionally add three tasks (ImageNet100 CIFARI00, and real-world Homeoffice) original eight tasks. The results in Table 9 show that ProbSurgery consistently outperforms both the Task Arithmetic and Surgery baselines across all datasets. Notably, the performance gains especially pronounced on the newly introduced, more complex benchmarks: ProbSurgery achieves 74.8% on ImageNet100 (vs. 11.6% for Task Arithmetic and 56.8% for Surgery), These results indicate even when merging a large number of models, our method can effectively correct the representation bias and achieve superior performance.
ImageNett00
HomeOffice
SUN397
RESISC45
EuroSAT
GTSRB
MNIST
CIFAR100
SVHN
Cara
DTD
Method
Avg
Task Arithmetic
31.8
39.3
49.6
60.3
74.4
60.4
95.8
38.4
11.6
42.6
70.3
52.3
w/ Surgery
50.5
52.7
76.8
96.3
83.8
79.9
98.6
65.8
56.8
55.6
80.3
72.5
w/ ProhSurgery
53.7
55.7
86.3
97.9
85.5
96.0
98.8
72,9
74.8
63.4
86.2
79.2
Table Performance comparison of Surgery and ProbSurgery on merging 11 tasks.
C.5. Performance on More Backbones
For vision tasks, we further conduct comparison experiments on a different vision backbone, i.c. ViT-L-14, to verify the effectiveness of our proposal, ProbSurgery.
ViT-L/14. Table demonstrates the performance various methods for merging ViT-L/14 models across eight tasks. Our proposed method, ProbSurgery, consistently outperforms the baseline Surgery and other existing methods in terms of average performance (Avg.). Notably, ProbSurgery achieves significant improvements, with performance gains ranging from +1.5% to +5.1% % over Surgery. This highlights its superior ability to balance task-specific features and enhance overall stability. Moreover, ProbSurgery adapts seamlessly to different merging strategies, such Weight Averaging, Task Arithmetic, and LW AdaMerging, consistently yielding better results. These advantages make ProbSurgery a highly effective and generalizable solution for multi-task learning and large model merging tasks.
VIT-B/16. Table highlights the superior performance of our proposed method, ProbSurgery, merging ViT-B/16 models across eight tasks. Compared the baseline Surgery, ProbSurgery consistently achieves higher average scores, with improvements ranging from +2.4% +6.9%, demonstrating its ability to balance task-specifie features effectively. It excels across diverse datasets, including challenging tasks like DTD and SVHN, and outperforms other methods, such as Weight Averaging, Task Arithmetic, and LW AdaMerging, across all strategies. ProbSurgery's 's robust and generalizable approach ensures better task balance, enhanced generalization, and improved stability, making a highly effective solution for multi-task learning.
16

---

Representation Surgery in Model Merging with Probabilistic Modeling
Table 10. Multi-task performance when merging ViT-1/114 models eight tasks. green value denotes the improvement of ProbSurgery
compared with the counterpart method Surgery (Yang et al., 2024a).
Method
SUN397
Cars
RESISC45
EuroSAT
SVHN
GTSRB
MNIST
DTD
Avg.
Pretrained
66.8
77.7
71.0
59.9
58.4
50.5
76.3
55.3
64.5
Individual
82.3
92.4
97.4
100
98.1
99.2
99.7
84.1
94.2
Traditional MTT.
80.8
90.6
96.3
96.3
97.6
99.1
99.6
84.4
93.5
Fisher Merging (Matena & Raffel, 2022)
69.2
88.6
87.5
93.5
80.6
74.8
93.3
70.0
82.2
RegMean (Jin ct al., 2023)
73.3
81.8
86.1
97.0
88.0
84.2
98.5
83.7
Concrete TA (Tang ct al., 2023)
74.6
86.2
89.0
96.7
93.6
93.4
99.1
66.9
87.4
Concrete AM (Tang al., 2023)
77.8
91.2
92.1
97.0
94.4
97.9
99.0
79.5
91.1
Weight Averaging
72.1
81.6
82.6
91.9
78.2
70.7
97.1
79.6
w/ Surgery (Yang et al., 2024a)
73.7
83.9
92.0
98.4
82.4
86.3
98.7
71.9
85.9
w/ ProbSurgery (Ours)
75.7
87.1
96.5
99.6
88.9
99.0.
99.5
81.6
91.0 (5.1)
Task Arithmetic (Ilharco al., 2023)
73.9
82.1
86.6
94.1
87.9
86.7
98.9
65.6
84.5
Surgery (Yang el al., 2024a)
75.7
84.4
93.1
98.8
91.3
93.4
99.1
76.1
ProbSurgery (Ours)
77.0
87.5
96.2
99.7
94.2
99.1
99.4
81.8
92.0 (4.0
Ties- Merging (Yadav et al., 2023)
76.5
85.0
89.3
95.7
90.3
83.3
99.0
68.8
86.0
w/ Surgery (Yang et al., 2024a)
76.5
85.9
93.7
99.2
89.7
92.0
99.1
78.1
89.3
ProbSurgery (Ours)
77.7
88.4
96.7
99.7
93.2
99.0
99.5
82.4
92.6 (3.3)
LW AdaMerging (Yang al., 2024c)
79.0
90.3
90.8
96.2
93.4
98.0
99.0
79.9
90.8
w/ Surgery (Yang et al., 2024a)
80.3
90.8
94.3
98.2
94.1
98.7
99.2
82.5
92.3
w/ ProbSurgery (Ours)
80.7
91.4
96.2
99.6
95.9
99.2
99.5
84.1
93.8 (1.5)
SOTA: EER-MERGING (Huang et al., 2024)
83.2
90.7
96.8
99.7
97.9
99.1
99.7
82.7
93.7
Table Multi-task performance when merging VIT-B/16 models on eight tasks. The green value denotes improvement ProbSurgery
compared with the counterpart method Surgery (Yang et al., 2024a).
Method
SUN397
Cars
RESISC4S
EuroSAT
SVHN
GTSRB
MNIST
DTD
Avg.
Pretrained
63.8
64.6
65.7
54.5
52.0
43.3
51.7
45.1
55.0
Individual
81.8
86.8
96.9
99.7
97.8
99.1
99.7
82.0
92.9
Fisher Merging (Matena & Raffel, 2022)
68.5
69.9
75.2
80.4
73.2
61.2
94.5
50.7
71.7
RegMean (Jin et al., 2023)
69.1
71.6
77.6
88.8
83.7
70.2
96.9
54.6
76.6
Weight Averaging
67.7
70.0
75.3
79.5
74.9
60.1
94.4
43.8
70.7
w/ Surgery (Yang et al., 2024a)
70.3
72.4
88.8
97.6
82.0
83.1
98.1
68.5
82.6
w/ ProbSurgery (Ours)
74.0
79.4
95.7
99.7
87.8
98.8
99.4
81.2
89.5 (6.9)
Task Arithmetic (Ilharco al., 2023)
61.1
65.9
74.0
76.2
88.0
73.9
98.4
53.0
73.8
w/ Surgery (Yang et 2024a)
68.3
72.3
88.7
97.7
91.0
89.5
98.9
72.9
84.9
w/ ProhSurgery (Ours)
71.9
79.3
95.8
99.7
93.5
99.0
99.4
81.5
90.0 (5.1
Ties Ties-Merging (Yadav al., 2023)
69.1
72.5
80.5
84.0
85.0
71.5
98.1
$4.9
77.0
w/ Surgery (Yang et al., 2024a)
73.0
76.2
90.7
98.1
89.7
86.7
98.7
75.2
86.0
w/ ProbSurgery (Ours)
75.1
80.9
95.8
99.7
92.2
99.0
99.5
82.0
90.5 (4.5)
LW AdaMerging (Yang al., 2024c)
70.2
81.6
94.8
91.6
95.8
98.5
84.9
w/ Surgery (Yang et al., 2024a)
73.6
81.5
90.4
98.5
93.2
97.4
98.9
77.0
88.8
w/ ProbSurgery (Ours)
75.3
83.7
95.5
99.8
94.9
99.0
99.4
82.3
91.2 (2.4)
D. Visualization Results
Representation visualization. In Figures 9, we visualize the generated representation via t-SNE technique, which clearly illustrates the advantages of our proposed method, ProbSurgery, generating high-quality representations as
17

---

Representation Surgery in Model Merging with Probabilistic Modeling
post-calibration technique for model merging. Its advantages can be divided into three parts, including
Better Alignment evident as the red points produced by ProbSurgery closely align with the blue points representing individual task-specific models, preserving critical features and ensuring the merged representations stay faithful the original tasks.
Reduced Overlap is prominently shown across datasets such as DTD and SVHN, where ProbSurgery minimizes the interference between tasks by creating well-separated and distinct clusters, in contrast the significant overlap seen in both four baseline methods like Weighted Averaging.
Task-Specific Robustness is highlighted by ProbSurgery's ' ability to consistently produce clear and tightly clustered representations across diverse tasks, including both simpler datasets like MNIST and more complex ones like EuroSAT, demonstrating its adaptability and effectiveness in handling a wide range of scenarios.
These strengths establish ProbSurgery superior method for generating interpretable, task ask-preserving, and reliable representations.
task arithmetic
w/ Surgery
w/ Ours
task arithmetic
w/ Surgery
w/ Ours
(a) Cars
(b) DTD
task arithmetic
w/ Surgery
Ours
task arithmetic
w/ Surgery
Ours
(c) EuroSAT
(d) GTSRB
task arithmetic
w/ Surgery
w/ Ours
task arithmetic
w/ Surgery
w/ Ours
MNIST
RESISC45
task arithmetic
w/ Surgery
w/ Ours
task arithmetic
w/ Surgery
w/ Ours
(g) SUN397
(h) SVHI
Figure 9. Comparisons of varying generated representations, where the baseline is Task Arithmetic and the backhone is VIT-B/32.
18


Reviews:
[{"review_id": "OsYO6Gye90", "reviewer": "Reviewer_XTHh", "summary": "This paper proposes ProbSurgery, a probabilistic approach to post-merging representation correction in model merging. The authors address the representation bias that occurs when merging multiple models for multitask learning. Unlike prior deterministic approaches (e.g., Surgery), ProbSurgery models the bias as a probabilistic distribution, which better handles uncertainty in parameter merging. The paper provides theoretical insights using a PAC-Bayes framework and proposes a one-to-all extension of ProbSurgery, allowing for efficient multi-task calibration. Experimental results across vision and NLP tasks show that ProbSurgery consistently outperforms existing merging strategies, achieving better generalization and robustness, particularly in OOD and domain shift settings.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your insightful comments and give the response as follows.\n\n**Q1 Computational Complexity**: In practical implementation, we only sample once to generate the representation bias in the ProbSurgery module, which does not incur any additional cost compared to the deterministic method. Our results demonstrate that a single sample is sufficient to achieve state-of-the-art performance—outperforming the previous Surgery method. In the next version of our manuscript, we will provide further insights and discussions on how the number of samples impacts the algorithm's performance.\n \nBesides, we agree with your viewpoint that sampling multiple times can probably improve the estimation accuracy of the representation bias, thereby enhancing overall model performance. Specially, we sample $\\rho$ times and compute their average values during model training. However, even if we perform multiple samplings, it does not incur extra computational cost since this averaging operation only involves first-order derivatives during backpropagation. In the following table, we try different sampling numbers and compare their performance. We can observe that increasing the sampling number does not lead to a significant improvement in overall performance. Therefore, to ensure better scalability, we have chosen not to introduce this hyperparameter explicitly, and instead, we set it to 1 by default.\n\n|Sampling numbers|1|3|5|10|\n| - | - | - | - | - |\n|Average acc on eight tasks|86.74|86.87|86.90|86.91|\n\n**Q2 Sensitivity to Hyperparameters**: Based on our empirical experience, we have set the trade-off coefficient for the KL term to $1 \\times 10^{-3}$. To further explore the algorithm's sensitivity to this hyperparameter, we conducted experiments with Weight Average on ViT-B/32 and show the result in the following table, which shows that the algorithm's average accuracy remains nearly constant across various hyperparameter values, demonstrating that our method is insensitive to this parameter.\n\n|the trade-off coefficient $\\lambda$|$1\\times10^{-4}$| $1 \\times 10^{-3}$ | $1 \\times 10^{-2}$ | $1 \\times 10^{-1}$ |\n|-|-|-|-|-|\n|Average acc on eight tasks|86.17|86.70|86.40|86.02|\n\n**Q3 Gaussian Assumption for representation**: In this work, we treat the representation bias as a Gaussian mainly for tractability and flexibility. The normal distribution provides a compact parameterization ($\\mu$ and $\\sigma$) that is easy to optimize, while also serves as a high-entropy “baseline” choice—minimizing assumptions about the true shape of the bias. Although the real distribution may deviate from exact Gaussianity, empirical tests show that this assumption reliably captures the dominant uncertainty in merged representations.", "claims_and_evidence": "Yes, all claims are supported by both theoretical and experimental evidence. 1) The main claim that ProbSurgery improves representation correction and generalization is well-supported by both theoretical and experimental evidence. 2) The paper uses a PAC-Bayes theoretical framework to show that modeling bias as a distribution leads to a lower classification error.", "methods_and_evaluation_criteria": "Yes, the utilized criteria are reasonable. 1) The problem of representation bias in model merging is well-motivated, and the paper systematically compares ProbSurgery to existing merging methods (e.g., Weighted Averaging, Task Arithmetic, AdaMerging). 2) The use of L1 distance and test accuracy as evaluation metrics is reasonable and aligns with prior work (Figures 2, 4).", "theoretical_claims": "Yes, I checked all theoretical claims and didn't find errors. The PAC-Bayes analysis is correctly applied and suggests that a probabilistic bias correction method results in a lower classification error bound.", "experimental_designs_or_analyses": "Yes, the experimental designs are sound\n1) The experiments use strong baselines, including Weighted Averaging, Task Arithmetic, and AdaMerging, ensuring fair comparisons.\n2) The inclusion of OOD and domain shift experiments strengthens the claim that ProbSurgery improves generalization.", "relation_to_broader_scientific_literature": "Yes, the contribution and difference compared with previous works have been widely discussed. The paper builds on previous model merging works (Matena & Raffel, 2022; Ilharco et al., 2023) and extends the Surgery framework (Yang et al., 2024a). It also relates ProbSurgery to uncertainty modeling in probabilistic embeddings (Vilnis & McCallum, 2014).", "essential_references_not_discussed": "The cited works cover model merging and uncertainty modeling, making the context of the study clear.", "other_strengths_and_weaknesses": "Strengths:\n1) Novelty: Probabilistic modeling of representation bias is a significant advancement over deterministic approaches.\n2) Theoretical Rigor: Strong PAC-Bayes justification for ProbSurgery's effectiveness.\n3) Empirical Validation: Extensive experiments across vision and NLP tasks show clear improvements.\n4) Real-world Relevance: The one-to-all setting makes ProbSurgery practical for multitask learning.\n\nWeaknesses:\n1) Computational Complexity: The additional probabilistic sampling step may introduce latency, but this is not discussed.\n2) Sensitivity to Hyperparameters: The impact of hyperparameters like KL regularization $\\lambda$ is not fully analyzed.\n3) Gaussian Assumption: In this paper, the authors regularize the corrected representation as a normal Gaussian distribution. However, they do not provide sufficient motivation for this choice. The assumption that representation bias follows a Gaussian distribution requires further justification and discussion.", "other_comments_or_suggestions": "This paper can be improved from the following three parts:\n1) Discuss computational cost of ProbSurgery vs. deterministic methods.\n2) Provide ablation studies on KL regularization λ and sampling variance.\n3) Verify Gaussian assumptions for features.", "questions_for_authors": "See Weaknesses", "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": "ASk55sc4uX", "reviewer": "Reviewer_teo3", "summary": "Surgery is a method to improve the merging performance by reducing the representation bias of model merging. This paper argues that there are two main issues of Surgery. First, the representation discrepancy is not fully addressed. Second, Surgery requires multiple task-specific modules rather than a unified one. This paper proposes ProbSurgery to mitigate these. They adapt the latent variable model (e.g., VAE) to enhance Surgery and further extend it to the one-to-all setting. Empirical and theoretical results are shown to illustrate their method's benefit.\n\n# After rebuttal:\n\nMy concerns are addressed and I will raise my score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your insightful comments and give the response as follows.\n\n**Q1 The observation**: Here, $G(\\cdot)$ quantifies the difference in test accuracy, which reflects the average discrepancy between the feature representations produced by these two models. For a more formal representation, we extend the function $G(\\cdot)$ to:\n$E_{t\\sim[T]} \\Big[\\text{Acc}(f_{\\theta_t}, D_t) - \\text{Acc}(f_{\\theta_{unif}}, D_t)\\Big]\n\\propto E_{t\\sim[T]} \\xi_t,$ where $\\text{Acc}(f, D)$ denotes the test accuracy of model $f$ on the test set $D$. Detailed explanations will be added to our manuscript in subsequent revisions.\n\n**Q2 Upper performance in Fig2**: The upper performance in Figure 2 denotes the average performance of various expert models on corresponding merged tasks.\n\n**Q3 More baselines**: Due to page restrictions, we keep the same setting with Surgery that makes experiments on four methods based on task vectors. (Prob)Surgery is completely orthogonal to existing model merging methods and can be incorporated into all of them to mitigate the representation bias. In the table below (*more detailed results can be found in the response to Reviewer HxL4, Q4-1*), it shows that integrating (Prob)Surgery into existing baselines significantly improves performance. This confirms that (Prob)Surgery effectively mitigates representation bias in model merging and can be seamlessly incorporated into any merging approach.\n| Method  | Avg |\n| -  | - |\n| Fisher Merging | 68.3 |\n| w/ Surgery  | 82.6 |\n| w/ ProbSurgery | 87.4 |\n| RegMean | 71.8 |\n| w/ Surgery | 82.9 |\n| w/ ProbSurgery | 88.2 |\n\n**Q4 Format of Eq. (4)**: We follow the objective in Eq. (1) to maintain the core motivation of Surgery, namely to remove the representation bias introduced by the merged model. Concretely, $z_{i,t}^{\\mathrm{unif}} - \\xi_{i,t}$ represents the corrected (post‐calibrated) representation rather than the raw merged representation $z_{i,t}^{\\mathrm{unif}}$. Extending from a deterministic correction term $\\xi_{i,t}$ in Eq. (1) to a probabilistic treatment in Eq. (4) lets us more accurately capture the uncertainty arising from parameter interference during merging.\n\n**Q5 Statement in Lines 217-219**: The “limited representational capacity” refers to the deterministic nature of the Surgery module. While deterministic modeling can be simpler and faster to train, it struggles to capture the inherent uncertainty of multiple merged tasks. As Figure 1 shows, even after applying a deterministic correction, the merged features remain partially misaligned with each individual model’s distribution, implying that a single deterministic module does not generalize well beyond its initially learned task. Besides, “other tasks” denote the remaining tasks in the merged set, which the single deterministic Surgery module fails to accommodate. \n\n**Q6 Theorem 4.1**: In McAllester (2003), the statement closest to the paper’s Theorem 4.1 is Theorem 1, often referred to as the “PAC-Bayesian Theorem”. It provides a general PAC-Bayes bound on the expected loss of a posterior hypothesis in terms of its empirical loss plus a KL‐divergence term from the prior—precisely the structure adapted in ours. \n\n**Q7 Expectation in Eq. (9)**: Eq. (9) measures how a predicted distribution $Q_{\\omega}$ fits the observed (true) distribution $P^{\\mathrm{bias}}$. We take the expectation of that score with respect to the “true” distribution $P^{\\mathrm{bias}}$: $S\\bigl(Q_{\\omega}, P^{\\mathrm{bias}}\\bigr) =\\int S\\bigl(Q_{\\omega}, z^{\\mathrm{bias}}\\bigr)\\mathrm{d}P^{\\mathrm{bias}}(z^{\\mathrm{bias}}). $ Finally, it is written as $\\mathop{\\arg \\min}\\limits_{\\omega} \\mathcal{S}(Q_\\omega, P^{\\rm bias}) := \\mathop{\\arg \\min}\\limits_{\\omega} E_{z^{\\rm bias} \\sim P^{\\rm bias}}[Q_\\omega, z^{\\rm bias}]$, as shown in Eq. (9). Due to only one observation $z^{\\rm bias}$ (a fixed bias) in practice, this equation can omit the expectation form and be reformulated as $\\mathop{\\arg \\min}\\limits_{\\omega} \\mathcal{S}(Q_\\omega, P^{\\rm bias}) := \\mathop{\\arg \\min}\\limits_{\\omega}  \\mathcal{S}(Q_\\omega, z^{\\rm bias})$. Thus, in Eq. (10), we only have one variable $z^{\\rm bias}$ and estimate the distribution $Q_\\omega$ via sampling.\n\n**Q8 Validation set**: First, the post-training step in (Prob)Surgery relies on an extra validation set that does not require any labels. Acquiring such a small-scale set is feasible in the real world since no label annotations are required. Meanwhile, to ensure a fair comparison, we employ test-time adaptation (TTA, adopted in many model merging methods like Adamerging) to post-train a merged model with Task Arithmetic. The results in the table below clearly demonstrate that our method significantly outperforms the TTA approach.\n| Ratio of val set | 10% | 50% | 100% |\n| - | - | - | - |\n| Task Arithmetic: 69.1 |\n|w/ TTA | 72.8 | 74.1 | 75.4 |\n|w/ Ours | 84.8 | 86.2 | 87.0 |\n\n**Q9 Typos**: Thanks for pointing out these typos. We will correct these in the next manuscript.", "claims_and_evidence": "Some illustration figures and motivations are unclear to me. Please see the questions.", "methods_and_evaluation_criteria": "It may lack some important baselines. The motivation and design of the method are not clear to me. Please see the questions.", "theoretical_claims": "I wonder about the adaptation of Theorem.4.1 from literature. See the questions for details.", "experimental_designs_or_analyses": "I have no questions.", "relation_to_broader_scientific_literature": "Model merging is an emergent technique but suffers from the performance gap between merged and task-specific models. This paper's method is built on a previously proposed method to address two main issues. I believe the proposed method can reduce both the memory cost and inference time of the merged model, as well as improve the merging performance compared to the previous one. Moreover, as a post-merging method, this paper is compatible with other merging methods, which enhances the impact of this paper.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "The novelty and originality of this paper are good as they address two main issues of the previous method by applying VAE and Surgery. However, I think the clarity could be improved. Also, some important baselines are missing in their experiments, which may limit the significance of their paper. Please see the questions for details.", "other_comments_or_suggestions": "Typos:\n- Title of Sec.3.1: represetative - >  representative\n- Def.3.2: $f_{\\theta_T}$ - >  $\\{f_{\\theta_t}\\}_{t=1}^T$", "questions_for_authors": "1. How to compute $G$ in Observation.3.3.? I think a detailed formulation would be better than a description “overall performance gap”.\n2. What is the upper performance in Fig.2(a)?\n3. Why is the ablation of (Prob)Surgery not studied with Fisher and RegMean merging in Tab.2?\n4. In Eq.4, why do we align the shift of merged features $z_{i,t}^{\\theta_{unif}}-\\xi_{i,t}$ rather than $z_{i,t}^{\\theta_{unif}}$ with $z^{\\theta_t}$? I understand that Eq.4 is an extension of Eq.1, but I am not clear why we follow the objective in Eq.1.\n5. I do not understand the following sentence in lines 217- 219: \"Due to limited ... to other tasks\". Is the limited representation capability the same as what is illustrated in Fig.1? What do \"other tasks\" mean? \n6. I am curious about which theorem in (McAllester, 2003) the author adapted in their proposed Theorem.4.1.\n7. What is $E[Q_{\\omega}, z^{bias}]$ in Eq.9?\n8. As the proposed method needs to retrain the (Prob)Surgery modules after merging, I think it would be fair to compare it with the performance of the merged model after training on all validation sets.\n\n---\n\nMy concerns are addressed and I will raise my score.", "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": "nkwm3MTOdY", "reviewer": "Reviewer_ACu4", "summary": "This paper provides a probabilistic interpretation of the previous work, Surgery, which addresses the issue of representation bias. Additionally, two strategies are proposed to reduce overall training costs, supported by a theoretical analysis that highlights the advantages of the probabilistic approach. Extensive experiments demonstrate the superiority of the proposed method.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your insightful comments and give the response as follows.\n\n**Q1 ProbsurgeryV2, integrating Probsurgery into each block**: In this paper, ProbSurgery is integrated into the last layer. However, it also can be integrated into each block of the ViT-based model to surgery the layer (or block)-level representation bias, which is more refined and contributes to better performance. By adopting the design in Surgery V2, we conduct experiments on Task Arthmetric with ViT-B/32 and show the result in the following table. When our method is applied after every block, it can more effectively reduce representation bias and achieve better performance.\n\n| Method | SUN397 | Cars | RESISC45 | EuroSAT | SVHN | GTSRB | MNIST | DTD | Avg |\n| - | - | - | - | - | - | - | - | - | - |\n| Task Arithmetic   | 55.2 | 54.9 | 66.7 | 78.9 | 80.2 | 69.7 | 97.3 | 50.4 | 69.1 |\n| w/ Surgery        | 63.8 | 59.9 | 83.3 | 97.9 | 87.0 | 87.0 | 98.6 | 69.4 | 80.9 |\n| w/ ProbSurgery    | 67.0 | 67.0 | 94.1 | 99.8 | 91.2 | 98.8 | 99.4 | 79.0 | 87.0 |\n| w/ Surgery v2     | 73.8 | 67.9 | 94.5 | 99.6 | 96.8 | 98.8 | 99.5 | 78.0 | 88.6 |\n| w/ ProbSurgery v2 | 74.1 | 68.2 | 94.8 | 99.8 | 97.1 | 99.1 | 99.8 | 78.3 | 88.9 |\n\n\n**Q2 Integrating with other methods**: Due to page restriction, we keep the same setting with Surgery that only conducts experiments on four baselines based on task vectors. Our method ProbSurgery is completely orthogonal to existing model merging approaches and can be incorporated into any model merging method to solve its representation bias problem. In the following table, we integrate ProbSurgery with other three methods that are not based on task vectors, including Fisher Merging, Regmean, and EMR\\-Merging. By integrating our method, each baseline sees a substantial boost in performance, highlighting the effectiveness of our approach in mitigating representation bias.\n\n| Method | SUN397 | Cars | RESISC45 | EuroSAT | SVHN | GTSRB | MNIST | DTD | Avg |\n| - | - | - | - | - | - | - | - | - | - |\n| Fisher Merging | 68.6 | 69.2 | 70.7 | 66.4 | 72.9 | 51.1 | 87.9 | 59.9 | 68.3 |\n| w/ ProbSurgery | 71.7 | 73.1 | 91.3 | 99.7 | 91.2 | 96.1 | 99.5 | 76.8 | 87.4 |\n| RegMean        | 65.3 | 63.5 | 75.6 | 78.6 | 78.1 | 67.4 | 93.7 | 52.0 | 71.8 |\n| w/ ProbSurgery | 71.5 | 72.7 | 94.1 | 99.0 | 93.0 | 97.7 | 99.3 | 78.3 | 88.2 |\n| EMR\\-Merging   | 75.2 | 72.8 | 93.5 | 99.5 | 96.9 | 98.1 | 99.6 | 74.4 | 88.7 |\n| w/ ProbSurgery | 75.3 | 73.3 | 94.1 | 99.7 | 97.2 | 98.4 | 99.7 | 75.1 | 89.1 |\n\n**Q3 The size of the validation set**\nThanks for this constructive suggestion. We made an experiment to verify how the size of the validation set impacts ProbSurgery's performance. In the following table, we can see that using a larger proportion of the validation set helps the ProbSurgery module better capture and mitigate inherent biases, leading to higher overall accuracy. Notably, even when only 10% of the unlabeled validation data is used, our method still exceeds Surgery’s performance of 80.9%, demonstrating its robustness under limited data conditions. \n\n\n| Method | Ratio of Val set | SUN397 | Cars | RESISC45 | EuroSAT | SVHN | GTSRB | MNIST | DTD | Avg |\n| - | - | - | - | - | - | - | - | - | - | - |\n| Task Arithmetic | -  | 55.2 | 54.9 | 66.7 | 78.9 | 80.2 | 69.7 | 97.3 | 50.4 | 69.1 |\n| w/ Ours | 10% | 63.5 | 64.5 | 89.6 | 94.9 | 88.4 | 89.4 | 98.4 | 74.5 | 82.9 |\n| w/ Ours | 50% | 65.9 | 66.5 | 92.3 | 97.8 | 90.5 | 95.0 | 98.8 | 77.2 | 85.5 |\n| w/ Ours | 100%| 67.0 | 67.0 | 94.1 | 99.8 | 91.2 | 98.8 | 99.4 | 79.0 | 87.0 | \n\n\n**Q4 Merging tens of models**: Thanks for this constructive suggestion. To verify the performance in more challenging settings, we additionally add three tasks (ImageNet100, CIFAR100, and real-world Homeoffice) to original eight tasks. The results in the following table show that even when merging a large number of models, our method can effectively correct the representation bias and achieve superior performance.\n\n\n| Method | SUN397 | Cars | RESISC45 | EuroSAT | SVHN | GTSRB | MNIST | DTD | ImageNet100 | CIFAR100 | HomeOffice | Avg |\n| - | - | - | - | - | - | - | - | - | - | - | -|-|\n| Task Arithmetic  |31.8 | 39.3 | 49.6 | 60.3 | 74.4 | 60.4 | 95.8 | 38.4 | 11.6 | 42.6 | 70.3 | 52.3 |\n| w/ Surgery | 50.5 | 52.7 | 76.8 | 96.3 | 83.8 | 79.9 | 98.6 | 65.8 | 56.8 | 55.6 | 80.3 | 72.5 |\n| w/ ProbSurgery  | 53.7 | 55.7 | 86.3 | 97.9 | 85.5 | 96.0 | 98.8 | 72.9 | 74.8 | 63.4 | 86.2 | 79.2 |", "claims_and_evidence": "Yes, the paper addresses a valuable problem, and the motivation for the proposed method is clearly articulated. The algorithm is supported by the theoretical analysis.", "methods_and_evaluation_criteria": "Yes, the paper introduces a probabilistic approach to model representation bias, leading to more robust estimations. The datasets used in the experiments are widely applicable. In addition to conventional metrics, the authors incorporate visualization techniques to illustrate biased representations.", "theoretical_claims": "Yes, I checked all contents for the theoretical claims and didn't find significant errors.", "experimental_designs_or_analyses": "The authors provide rational and rigorous experimental designs of main experiments and ablation studies. The analyses are comprehensive and insightful", "relation_to_broader_scientific_literature": "This paper proposes a novel approach to bridge the gap between multiple individual expert models and the merged model through a representation rectification operation. Experimental results demonstrate that the proposed method, ProbSurgery, achieves performance comparable to that of various expert models. Consequently, I believe this approach has the potential to enhance the applicability of model merging as a multi-task learning strategy.", "essential_references_not_discussed": "This paper provides a comprehensive introduction to the task background and relevant preliminaries. There are no existing papers that are closely related to this work that require further discussion.", "other_strengths_and_weaknesses": "Strengths:\n1) The writing in this paper is exceptionally clear, with the author providing detailed background information and a well-articulated motivation for the study.\n\n2) The proposed method appears highly reproducible. The probabilistic approach is both simple and effective, successfully modeling the uncertainty that implicitly arises from parameter interference when merging multiple models.\n\n3) The theoretical analysis demonstrates that the classification error of the probabilistic approach is lower than that of a deterministic approach, enhancing the novelty of the proposed method.\n\nWeaknesses and questions:\n1) There is an enhanced version of Surgery, named SurgeryV2 [1], which incorporates the surgery module into each block of the encoder. I am curious to know whether the proposed ProbSurgery also implements this operation.\n\n2) The proposed ProbSurgery shows potential for improved performance when integrated with existing model merging methods. I noticed that the selected baselines belong to weight-based approach. I have some concerns regarding the applicability of ProbSurgery to other types of model merging methods. For instance, could it be integrated with route-based methods like EMR-Merging [2]?\n\n3) Impact of the scale of validation set. The proposed framework relies on unlabeled test/validation data to build self-supervision signals and update the parameters of the ProbSurgery module. It would be beneficial for the author to discuss the impact of varying amounts of available test data on model performance.\n\n4) Recent methods [2][3] have begun to evaluate their performance under extreme conditions, such as model merging with tens of models. I suggest the author conduct experiments under this setting to validate the generalization and robustness of ProbSurgery.\n\n[1] SurgeryV2: Bridging the Gap Between Model Merging and Multi-Task Learning with Deep Representation Surgery. ArXiv 2024\n[2] Emr-Merging: Tuning-free high-performance model merging. NeurIPS 2024\n[3] Localizing Task Information for Improved Model Merging and Compression. ICML 2024", "other_comments_or_suggestions": "None.", "questions_for_authors": "See Weaknesses.", "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": "AdAcNsBYHE", "reviewer": "Reviewer_HxL4", "summary": "This paper builds upon an earlier work, representation survey. It there are T tasks and 1 model that is trained on each task, representation surgery paper defined representation bias as the sum of the distances between the merged model and the model trained on the current task for all the tasks. To alleviate this bias, they train an MLP module for each task that ensures that the representations of the merged model are closer to the model trained on the given task. This paper uses a gaussian distribution to model the latent space of the representation bias and uses an MLP module for each task to learn the mean and the standard deviation of the distribution, where the objective is to minimize the representation bias over all the tasks. They also formulate direct feature alignment (DFA) and proxy distribution alignment (PDA) methods to learn one module for all the tasks. They show that the  PDA consistently outperforms the surgery based merging on all the model merging tasks. They also demonstrate that the model merged using method generalizes on out of distribution and on domain drift when used in conjunction with weight averaging.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your insightful comments and give the response as follows.\n\n**Q1 Generalization on ImageNet**: Thanks for your valuable suggestion. We try the generalization performance of the merged model on unseen tasks like Classification task (ImageNet1k). We follow the setting in our paper, i.e., merging eight expert models, and then verfiy its general capability on the test set of ImageNet1k. The backbone is ViT-B/32 and the baseline is Task Arithmetic. The results in the table below show that the baseline Task Arithmetic method achieves 27.6% top-1 test accuracy, while the Surgery and ProbSurgery methods improve it to 48.5% and 52.7%, respectively. Thus, we believe our proposal represents a significant advancement in enhancing the generalization performance of existing model merging techniques. \n\n| Method | Test Acc |\n| - | - | \n|Task Arithmetic   | 27.6  |\n| w/ Surgery         | 48.5  |\n| w/ ProbSurgery | 52.7  |\n\n\n**Q2 ProbsurgeryV2, tntegrating Probsurgery into each block**: Thanks for this valuable suggestion. Our proposed ProbSurgery also can integrated into each block of the ViT-based model to surgery the layer (or block)-level representation bias, which is more refined and contributes to better performance. By adopting the design in Surgery V2, we conduct experiments on Task Arthmetric with ViT-B/32 and show the result in the following table. \n\n| Method | SUN397 | Cars | RESISC45 | EuroSAT | SVHN | GTSRB | MNIST | DTD | Avg |\n| - | - | - | - | - | - | - | - | - | - |\n| Task Arithmetic   | 55.2 | 54.9 | 66.7 | 78.9 | 80.2 | 69.7 | 97.3 | 50.4 | 69.1 |\n| w/ Surgery        | 63.8 | 59.9 | 83.3 | 97.9 | 87.0 | 87.0 | 98.6 | 69.4 | 80.9 |\n| w/ ProbSurgery    | 67.0 | 67.0 | 94.1 | 99.8 | 91.2 | 98.8 | 99.4 | 79.0 | 87.0 |\n| w/ Surgery v2     | 73.8 | 67.9 | 94.5 | 99.6 | 96.8 | 98.8 | 99.5 | 78.0 | 88.6 |\n| w/ ProbSurgery v2 | 74.1 | 68.2 | 94.8 | 99.8 | 97.1 | 99.1 | 99.8 | 78.3 | 88.9 |\n\n\n**Q3 Typos**: Thanks for pointing out these typos. We will correct these typos in the next version of our manuscript.\n\n**Q4-1 Other merging methods**. Our method is applicable to all model merging approaches, not just those based on task vectors. The table below reports the performance of three non-task-vector methods—Fisher Merging, RegMean, Task Arithmetic + DARE, and EMR-Merging — using ViT-B/32 as the backbone. We can see that applying ProbSurgery consistently improves average accuracy. In conclusion, we believe that ProbSurgery's effectiveness has been fully validated.\n\n| Method | SUN397 | Cars | RESISC45 | EuroSAT | SVHN | GTSRB | MNIST | DTD | Avg |\n| - | - | - | - | - | - | - | - | - | - |\n| Fisher Merging | 68.6 | 69.2 | 70.7 | 66.4 | 72.9 | 51.1 | 87.9 | 59.9 | 68.3 |\n| w/ ProbSurgery | 71.7 | 73.1 | 91.3 | 99.7 | 91.2 | 96.1 | 99.5 | 76.8 | 87.4 |\n| RegMean        | 65.3 | 63.5 | 75.6 | 78.6 | 78.1 | 67.4 | 93.7 | 52.0 | 71.8 |\n| w/ ProbSurgery | 71.5 | 72.7 | 94.1 | 99.0 | 93.0 | 97.7 | 99.3 | 78.3 | 88.2 |\n| Task Arithmetic + DARE|  51.7 | 51.4 | 63.2 | 75.4 | 76.7 | 66.2 | 94.9 | 46.9 | 65.8 |\n| w/ ProbSurgery |  67.1 | 65.1 | 93.2 | 98.8 | 90.3 | 97.8 | 98.4 | 78.1 | 86.1 |\n| EMR\\-Merging   | 75.2 | 72.8 | 93.5 | 99.5 | 96.9 | 98.1 | 99.6 | 74.4 | 88.7 |\n| w/ ProbSurgery | 75.3 | 73.3 | 94.1 | 99.7 | 97.2 | 98.4 | 99.7 | 75.1 | 89.1 |\n\n**Q4-2 Improvement in Adamerging**: The improvement observed with Adamerging is not significant, as it is a more advanced merging method that results in representation bias during merging (as shown in Fig 4, a smaller distance is achieved by Adamerging). Therefore, when integrated with ProbSurgery, the improvement is not considerable.\n\n**Q4-3 NLP tasks**: Since different NLP tasks always enjoy the same parameters in almost all layers or blocks given a model like Bert, it exhibits fewer representation bias compared to vision tasks, which are highly sensitive to parameter changes. This difference leads to minimal performance degradation when merging multiple NLP tasks. Thus, integrating our proposal with NLP tasks yields a relatively smaller performance improvement than with vision tasks.\n\n\n**Q5 Merging fewer models**: We believe that any model fusion process, regardless of the number of models involved, induces parameter changes that lead to representation bias. Consequently, our method can alleviate this bias and enhance performance. Due to the word limit, we merged three models and present the results in the following table. Notably, our method also demonstrates performance gains even when merging a small number of models.\n\n| Method            | SUN397 | Cars | RESISC45 | Avg |\n| -                 | - | - | - | - | \n| Task Arithmetic   | 71.8 | 69.5 | 90.1 | 77.1 |\n| w/ ProbSurgery    | 74.1 | 73.8 | 94.4 | 80.8 |", "claims_and_evidence": "Claims:\n1. To calibrate the representations of the outputs of the models merged using various methods and they demonstrated the same using their experiments. They used 4 model merging methods - Weight Averaging, Task Arithmetic, TIES, ADAMerging. The evidence is in table 2,3,5. \n2. To use one unified module for all the tasks. The evidence is in table 4 and figure 8.\n3. Tackling out of domain distribution and domain shift challenges. The evidence is in table 5.", "methods_and_evaluation_criteria": "The datasets used by this method have been used in the past by other model merging methods such as adamerging. In addition to the vision domain, they also evaluated their method on NLP tasks.", "theoretical_claims": "Ya, they claimed that the probabilistic method results in a smaller classification error than a deterministic method. They proved it using PAC-Bayes theorem", "experimental_designs_or_analyses": "They did a lot of ablation studies such as the values of the hidden size h2, the sampling parameter, the various loss functions etc. They also showed that their method converges faster when compared to surgery and the efficacy of calibration by comparing the l1 distance between the representations yielded by the probsurgery and surgery methods when compared to the representation of the model trained on that task.", "relation_to_broader_scientific_literature": "This method corrects the representations of the model yielded by various model merging methods in a post-hoc manner to make the representation closer to that of the model trained on the current task. This work builds upon an earlier work called Representation surgery. They were able to show gains over the representation surgery method in multi-task learning, domain drift and out of distribution data settings. While the representation surgery method required one module for each task, this method can improve over it using only one module for all the tasks.", "essential_references_not_discussed": "All the essential references are discussed", "other_strengths_and_weaknesses": "Strengths:\n1. The paper is very well written. The experiments section is very elaborate and bolsters all the claims\n2. This method is parameter efficient and requires only 1 module when compared to the representation surgery method\n\nWeakness:\n1. There is a risk that the merged model might lose its general capabilities. It is important to also show that the merged model performs well on datasets such as Imagenet for classification, coco dataset for object detection etc.", "other_comments_or_suggestions": "You could also explore correcting the representation bias of all the layers rather than the last layer along the lines of surgeryv2\n\nTYPO:\nIn the caption of table4, both learning with eight modules and learning with one module are denoted by cross.", "questions_for_authors": "1. Your method is effective as the parameter interference hampers the model's performance, thus requiring this post-hoc calibration. Would this method still be able to improve methods such as DARE which try to reduce the parameter interference by reducing the number of parameters to begin with? Why is the performance improvement less for ada-merging while it is more pronounced for others? Similarly even for the NLP tasks, the gains are not as significant as the vision domain. Is it because of a smaller model?\n2. The performance of model merging methods deteriorates as the number of models keep increasing. While your method alleviates it, would it still be effective if the number of models are 2 or 3?", "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:

Meta-review: This paper introduces ProbSurgery, a probabilistic approach to post-merging representation correction in model merging. The method addresses representation bias—a key limitation in merging multiple task-specific models—by modeling bias as a Gaussian distribution rather than using deterministic corrections (as in prior work like Surgery). 

ProbSurgery makes a significant methodological advance in model merging, combining theoretical soundness (PAC-Bayes) with practical impact (scalability, OOD robustness). The rebuttal convincingly addressed critiques, and new experiments (e.g., ProbSurgeryV2, 11-task merging) strengthened the contribution.
Decision: Accept (poster)