ZipIt! Merging Models from Different Tasks without Training
Authors: George Stoica
Daniel Bolya
Jakob Brandt Bjorner
Pratik Ramesh
Taylor Hearn
Judy Hoffman
Venue: ICLR 2024 poster (2024)
Keywords: Model Merging
Mode Connectivity
Classification
Deep Learning
Primary area: representation learning for computer vision, audio, language, and other modalities
TL;DR: We merge models trained on completely different tasks, without retraining.

Abstract: Typical deep visual recognition models are capable of performing the one task they were trained on. In this paper, we tackle the extremely difficult problem of combining distinct models with different initializations, each solving a separate task, into one multi-task model without any additional training. Prior work in model merging permutes one model to the space of the other then averages them together. While this works for models trained on the same task, we find that this fails to account for the differences in models trained on disjoint tasks. Thus, we introduce "ZipIt!", a general method for merging two arbitrary models of the same architecture that incorporates two simple strategies. First, in order to account for features that aren't shared between models, we expand the model merging problem to allow for merging features within each model by defining a general "zip" operation. Second, we add support for partially zipping the models up until a specified layer, naturally creating a multi-head model. We find that these two changes combined account for 20-60% improvement over prior work, making it more feasible to merge models trained on disjoint tasks without retraining.

Published as a conference paper at ICLR 2024
ZIPIT! MERGING MODELS FROM DIFFERENT TASKS without Training
George Stoica"
Daniel Bolya'
Jakob Bjorner
Pratik Ramesh
Taylor Hearn
Judy Hoffman
Georgia Tech
{ [gstoica3, dboul daolya, wbjoooerrr, thearhe, judaen,,jddyaaatteee
ABSTRACT
Typical deep visual recognition models are capable of performing the one task they were trained on. In this paper, we tackle the extremely difficult problem of combining distinct models with different initializations, each solving separate task, into one multi-task model without any additional training. Prior work in model merging permutes one model to the space the other then averages them together. While this works for models trained on the same task, find that this fails to account for the differences in models trained on disjoint tasks. Thus, we introduce "Ziplt!", general method for merging two arbitrary models of the same architecture that incorporates two simple strategies. First, in order to account for features that aren't shared between models, we expand the model merging problem to allow for merging features within each model by defining general "zip" operation. Second, we add support for partially zipping the models up until a specified layer, naturally creating a multi-hear model. We find that these two changes combined account for 20-60% improvement over prior work, making it more feasible to merge models trained on disjoint tasks without retraining.
INTRODUCTION
Ever since AlexNet (Krizhevsky et al., 2017) popularized deep learning in computer vision, the field has thrived under the reign of massive models with an ever increasing number of parameters. Many vision problems once considered difficult or impossible are now benchmark tasks: classification with tens of thousands of classes (Deng et al., 2009; Zhou al., 2017: Gemmeke et al., 2017), fast instance segmentation (He et al., 2017; Bolya et al., 2019). realistic image generation (Karras et al., 2018; Ho et al., 2020; Rombach et al., 2022), and more.
There are an abundance of independent, carefully tuned models out there for many tasks. However, we want to expand an existing model's 's capabilities, we run into many potential issues. we try training the model on additional task, face catastrophic forgetting (Kirkpatrick et al., 2017; Li & Hoiem, 2017; De Lange et al., 2021). If we evaluate the same model on different data without adaptation, we often find it doesn't generalize
out of domain samples (Blanchard et al., 201 Muandet et al., 2013; Wang et al., 2022). We can try so called "intervention" strategies (Wang et al., 2022: De Lange et al., 2021) mitigate these effects, but these often require further training which can be expensive (Dosovitskiy et al., 2020; Zhai et al., 2022; Dehghani et al., 2023). Instead, it would be nice we could expand a model's s capacity to solve new tasks by simply "zipping" with other models trained on those tasks without additional training.
Combining multiple models into one has recently started to gain traction the vision community. Model Soups (Wortsman et al., 2022a) can add multiple models finetuned from the same pretrained initialization to improve accuracy and robustness. Git Re-Basin (Ainsworth et al., 2022) generalizes further to models trained the same data but with different initializations, though with significant accuracy drop. REPAIR (Jordan et al., 2022) improves on Git Re-Basin by adding new parameters and adjusting model batch norms where applicable. However, all of these methods only combine models trained on the same task. In this paper, we take this line of work to a logical extreme: merging differently initialized models trained on completely separate tasks (see Fig. lab). We show that this is an incredibly difficult problem for prior work and employ two simple strategies to make it feasible.
"Equal Contribution. Code: https://github. http://twwb.ccogggttaa/////////e http://toooccicct...../aa

---

Published as a conference paper at ICLR 2024
(a) Prior Work: Same Task
(b) Our Work: Different Tasks
(c) Ziplt!
Model A
Model A (ABiod
Starferd Doge
Drepe
Model B
Modal
Driss
Ploce
Deagge
Figure I: Setting and ZipIt! (a) Prior work merges differently initialized models from the same dataset with the same label sets: e.g., merging two models both trained to classify dog breeds. (b) Our setting expands this to merging models from different datasets with different label sets: e.g.. merging a model that classifies dog breeds with one that classifies bird species. (c) Ziplt! merges these models without retraining by identifying shared features.
First, we note that prior work focuses on permuting one model to the other when merging them. This creates 1-1 mapping between the two models, inherently assuming that most features across them are correlated. Since this isn't necessarily the case for models trained on different tasks, we cannot rely permutation alone. Instead, combination of correlated features within and across each model. We find that on some tasks, this alone improves accuracy by up to 20% vs. permutation-based approaches. Moreover, we prove that merging within models can yield a better result in the theoretical setting of Entezari et al. (2021).
generalize model merging support "zipping" any
Second, existing methods merge the entire network. While this might work for extremely similar models trained in the same setting, the features of models trained on disjoint tasks become less correlated over the course of the network (Kornblith et al., 2019). To solve this, we introduce partial zipping, where we only "zip" up to specified layer. Afterward, we feed the merged model's 's outputs the remaining unmerged layers of the original networks, creating multi-hees model. Depending on task difficulty. this can improve accuracy by over 15% while still keeping most layers merged.
Incorporating both of these strategies. we introduce ZipIt! (Fig. Ic), general method for "zipping" any number of models trained on different tasks into single multitask model without retraining. By deriving general graph-base algorithm for merging and unmerging (Sec. 4), we can zip models of the same architecture together, merge features within each model, and partially zip them to create multi-task model. We validate our approach by merging models trained on entirely disjoint sets of CIFAR (Krizhevsky et al., 2009) and ImageNet (Deng et al., 2009) categories, as well as merging several models trained on completely independent datasets into one, significantly outperforming prior work (Sec. 5). Finally, we ablate and analyze our method's capabilities on these scenarios (Sec. 6).
2 RELATED WORK
Model merging combines the weights of two more models into one. Our work differs from prior work in that we adapt mode connectivity techniques to target models trained on disjoint tasks (Fig. 1).
Merging Finetuned Models. If two models finetuned from the same pretrained checkpoint, they often lie in the same error basin (Neyshabur et al., 2020). Several works (Huang et al., 2017: Izmailov al., 2018; Von Oswald al., 2020; Wortsman et al., 2022b; Ilharco et al., 2022b; Don-Yehiya et al., 2023) have exploited this property to average together the weights of a model at different stages of training. Tarvainen & Valpola (2017); Cai al. (2021); Grill et al. (2020); Caron et al. (2021); Baevski et al. (2022) use an "exponential moving average" of training checkpoints as a teacher for self-supervised learning. Other works merge models initialized from the same pretrained base, but that were finetuned independently. either by simply averaging their weights (McMahan et al., 2017; Wortsman al., 2022a; Choshen et al., 2022; Ramé al., 2022). permuting one model to the other (Ashmore & Gashler, 2015; Yurochkin et al., 2019; Wang et al., 2020), combining meaningful weight regions (Ilharco et al., 2022a; Gueta et al., 2023; Yadav et al., 2023; Sung et al., 2023), or maximizing an objective (Matena & Raffel, 2021). Our setting differs, as we do not assumee the same initialization.
Merging Differently Initialized Models. Merging models with different initializations is much more challenging problem. Works in this space often rely on mode connectivity (Freeman & Bruna, 2016; Garipov et al., 2018; Draxler al., 2018; Frankle al., 2020), attempting to interpolate between models along a low loss path (e.g., Tatro et al. (2020); Singh & Jaggi (2020); Liu et al. (2022)). Most recent work follow the intuition, later formalized by Entezari et al. (2021), that
2

---

Published as a conference paper at ICLR 2024
Loss
Model B
Model
(Parmuted Al
(Permuted to A)
Ziplt!
5.0
Ziplt!
Re-Basin
5.5
Re Basin
Task A
Model A
20
Task B
Model A
Figure 2: Task Loss Landscapes for models in Tab. lb. Model A and Model B lie in low loss basins for their own tasks, but not for the other task. Thus, any interpolation between Model A and permuted Model B (e.g., Git Re-basin) lies outside the minima for both tasks and thus performs poorly. In contrast, Ziplt! improves the merge by finding a model that lies in a low loss basin for both.
models permuted to the same loss basin can be merged by averaging their weights. Most notably, Git Re-Basin (Ainsworth et al., 2022) permutes models by comparing the similarity between their weights. REPAIR (Jordan al., 2022) improves the accuracy of Git Re-Basin by instead computing the correlation between their intermediate layer feature activations, and adding several batch norms to the network. Peña et al. (2022) find permutations using global rather than local optimization, though they don't 't support skip connections. Some of these works (e.g., Singh & Jaggi (2020); Ainsworth et al. (2022)) evaluate on on a setting where each model sees varying numbers instances per class. And Peña al. (2022) evaluates a continual learning setting with disjoint categories, but their method requires training optimization. Similarly, He et al. (2018) merges models of different tasks, but requires jointly finetuning after each layer merge. As far as we are aware, we present the first general method to successfully merge models trained on disjoint tasks without additional training.
3 BACKGROUND AND MOTIVATION
Model merging stems from mode connectivity (Garipov et al., 2018), where is conjectured that models trained with SGD on the same dataset lying in the same loss basin (i.e., region or mode of low loss) can be combined into a single model that's just as performant as the original models. If these models can be combined well by linearly interpolating their weights, they are said to be linearly mode connected (LMC) (Garipov et al., 2018; Entezari et al., 2021). Our work is similar to finding LMC, but across different datasets with disjoint label sets (i.c., separate tasks in De Lange et al. (2021)).
Consider model C as a collection of layers L, E C. each which may have some parameters (c.g.. We, b, for linear layer). If LA and L are finetuned from the same checkpoint, several works (e.g.. Izmailov et al. (2018); Wortsman et al. (2022a)) find merging them is as easy as linearly interpolating their weights (i.e., they are LMC). E.g., if is linear layer, the new weight matrix W simply
W W + (1 W1?
(1)
with an interpolation constant y E j0, 1J. usually set to 1/2. However, if L and L were not finetuned from the same checkpoint, they often do not lie in the same mode (Entezari et al., 2021; Ainsworth et al., 2022) and cannot be interpolated. Indeed, Eq. 1 typically results in random accuracy.
To fix this, Entezari et al. (2021) conjecture that large enough models are likely LMC modulo permutation. This is because (1) many neural networks can be permuted internally without affecting their outputs and (2) permutation can move two models into the same basin, allowing much of the lost accuracy to be recovered. More concretely, let P. be a permutation matrix that permutes outputs of layer L to the space of L,A Then for each layer, permutation works apply
W = A + (1 - 7) P.W. P
(2)
Note that here permute the output space of W.". but also need to permute its input space to undo the permutation from the previous layer (hence the use of pT
Problems with Permutation. Eq. 2 relies on the likelihood of model B lying in the same basin as model A after permutation being high. However, this is far less likely when the models trained
3

---

Published as a conference paper at ICLR 2024
on different tasks, as each model optimizes for basins containing distinct task-specific information. In this case the optimal permutation of model B to model A still lies in a strong basin on task but doesn't 't lie in basin on task A, as shown in Figure 2. This causes the interpolated model to perform worse than either the two original models. Thus, we explore alternative merging methods.
4 ZIPIT!
In this work, we treat model merging as combining the checkpoints (i.c., collection of weights) of multiple models into single checkpoint that can perform all the tasks of its constituents. We do this by merging the layers of the models together. For instance, suppose L; E & is a linear layer with parameters W, R b, R with input features R and outputs features fi E R
= Lx)x = Wex + b;
Our goal is to take I.A E L from model A and I. E C from model B and merge them into layer L that combines their feature spaces such that information from both f and f is retained in f." We accomplish this by merging each layer of one model with the corresponding layer in the other, both merging features in one across both layers or within the same layer. This is in contrast to permutation-bast merging method, which only combine features across layers.
Why should merge within? Features of models trained on different tasks may be dissimilar, as the models solve different problems. Forcibly combining these dissimilar features can yield merges that don't 't perform well on either original task (Fig 2). Instead, those features may be more compatible with others within the same model, which would better retain performance when combined.
In fact, can prove that methods which allow merges within each model (as well across both) perform equal to or better than those which only merge across models (c.g.. permutation-reliant approaches) a limited but prevalent setting. Specifically, we obtain a tighter bound over Theorem 3.1 from Entezari et al. (2021) when redundancy exists within a model and is leveraged. Both Theorem 3.1 and our Theorem 1 (see Appendix G for formalization and proof) bound the degree to which the loss of merged two-layer model with d-input dimensions and h-intermediate dimensions increases compared to the losses of the original models. Theorem 3.1 bounds this increase to O(h~)/aa /Kaaa +)). However, if features within a model are redundant, then we reduce the bound to
(
h
22
r 0.5
Loss Increase of Merged Model S
1-2r
(4)
0
otherwise
with D E [0, 1] measuring what portion of features are redundant. This bound is 41 2 < 1
times that of Theorem 3.1 when I' < 0.5 (equal only when I = 0) and explicitly zero when 1' 2 0.5.
How do we merge features? In prior work, each of the merged features f is the result of combining one feature from f and one from f. However in our case, we can also merge features by combining two from just f. or two from just f.". To account for this, we concatenate f and f.² into single feature vector: I. E 2R. Then, like prior work (e.g. Li et al. (2015); Jordan et al. (2022)), we define feature similarity as the pairwise correlation between between neuron activations over a small set of images (without labels). However, unlike those works, we compute correlations between every activation in the full concatenated space Il Our approach thus measures the similarity of every feature A and 1 to all features in both models, rather than solely between f and s.a
Next, if two features are well correlated, we can average them without losing much information. Thus, we can construct f: by finding n4 pairs of similar features in fAl, and averaging them together. By default, we do this greedily: i.e., iteratively match the features with the highest correlation without replacement; though we explore extensions to this in Sec. 4.3 and test other methods in Tab, 4. Then can use these matches to construct f Formally, we define "merge matrix" M, E R x2n; s.t.
(5) M, averages the matched features, with each match corresponding to one output feature in f/- For instance, if uth match is between indices s, t E (1...., .,, 2n;} of f. Ilf.. then the wth row of M, would be 1/2 at columns and & and 0 elsewhere. This results in f [u] = 1 ((AA))))(())) + } ( [[].
M4 (1A))
4

---

Published as a conference paper at ICLR 2024
The Zip Operation
Find Redundant Features
(d) Propagate Mand
(e) Apply M and to Waights
Acrose madel
W"
Reatures
Rectures
WiZhki moddess
Merged telggtt
(c) Merge and Unmerge
Apoxox. e*
Approx e
C
Zipped Layer
Zipoad
[a] Two Disjoint Layers
Unmerge
Next
(M)
Prev
Zipood
Irput
Figure 3: ZipIt! merges models layer-wise by exploiting redundancy their features. (a) Output features f and f from two disjoint layers are (b) paired with other features based on the similarity of their activations. (c) We produce a merge matrix M combine the pairs into single shared output feature space, and a corresponding unmerge matrix U that undoes this operation. (d) We then propagate up the network to align the next layer's input space, and simultaneously receive the previous layer's U to align our input space. (c) We apply Eq. 8 to "zip" the layers together using the M for the output and U for the input, producing single layer (f). We then repeat (a) on the next layer.
Thus, applying M, has the effect of interpolating with ? 1/2 but is more general (e.g., allows for merging more than 2 models at once, see Sec. 4.3).
What about the next layer? After merging features in one layer, we now have the problem that the next layers, 14 L are incompatible with f Instead, we need to undo the merge operation before passing the features the next layer. Thus, define an "unmerge" matrix U, R2 22n;Xm, Xm, s.t.
U.ft R fAI!!
(6)
U; is the pseudoinverse of M; and in the case of the matching from earlier is simply 2M, Note that strict equality is unlikely here. Like in prior work, merging models a lossy operation.
We split this unmerge matrix in half along its rows into U,A, U >< that act individually to produce J and p With this, we can evaluate the next layers using the merged features:
1 ~ (UI(AAM)
1 L PU
4.1 THE "ZIP" OPERATION
We now have all the necessary pieces, and can derive a general operation to merge LA and L at an arbitrary point in the network (Fig. 3). First, compute M, and U; by matching features between f and f We then pass U, to the next layer and receive U,-1 from the previous layer. Using M; and U,-1, we "fuse" the merge and ummerge operations into the layer's 's parameters. For linear layer:
W = M, AUA, + M.W:PUP..
(8)
where M and M are M, split along its columns. bg has the same equation but without unmerging.
Note the similarity between Eq. and Eq. 2. This isn't coincidence: we only allowed merging across models and not within models, our "zip" operation would be identical to Git Re-Basin 's themute-the-innerpplate permute-then interpolate approach. Thus, Eq. 8 can be thought of as a generalization of prior work.
4.2 ZIP PROPAGATION
However, most modern neural networks are not simply collections of linear layers stacked on top of each other. In practice, we cannot combine merge and unmerge matrices into every layer of the network, as local zip (Eq. 8) expects the layer to have weight matrix-i.e. the layer has have separate input and output spaces so that we can unmerge the input space and merge the output space. Other layers (e.g., BatchNorm, ReL.U) don't have such weight matrix.
5

---

Published as a conference paper at ICLR 2024
Sendi Formurd
ReLL has no matils condinue
BN ra W matric contiinue
Conu har metria, done
(a) Stop zipping
(b) Unmerge
Result Multi-Head
Applyy
ReLu
ReLU
Hoai"
Applly
PeLLU X
ReLU
BN
Papply Merge
M
Conv
Conv
Howd!
Prestous
From
Apply
Figure 5: Partial Zip. (a) If we stop zipping early and (b) apply the latest U from the zip propagation the inputs of the first unmerged layer in each model, (c) we get a multi-head model with a head for each task.
Figure 4: Zip Propagation. We propagate M; backward until we hit a layer with weights, merg- ing merging element-wise layers (e.g., Batch- Norm) along the way.
Thus, we "propogate" M, and Ug through these layers. For instance, in Fig. 4, we show common stack of layers found in typical ConvNet. Following Jordan et al. (2022), we compute M, and U, using the activations of the network (i.e., after each ReLU). We can't fuse M, with the ReLU layer, as doesn't have any parameters. Similarly, we can merge the parameters the preceding BatchNorm layer (i.e., in the same way as bias). But it doesn't have weight matrix, we also can't fuse M; into Only once we've reached the Conv layer can we fuse M, and U, into it using Eq. (in this case, treating each kernel element as independent).
Similar care needs to be taken with skip connections, as every layer that takes input from or outputs to a skip connection shares the same feature space. However, this too can be dealt with during propagation just need to propagate M, backward and U; forward to each layer connected by the same skip connection. In general, we can define propagation rules to handle many different types of network modules (see Appendix C).
4.3 EXTENSIONS
Partial Zip. We don't always want to zip every layer of the two networks, especially if their output spaces are incompatible, or if doing would lose too much accuracy. Instead, we perform partial zip. That is, we zip most of the layers together, but leave the later ones unzipped (Fig. 5).
Implementing this operation is simple in our framework: zip normal until the specified layer i, then the remaining unzipped layers will receive U, through zip propagation. If we apply U,* to L
and U.S to L 1 the remaining unzipped layers will form "heads" that expect merged features as input. We can then ensemble the heads choose one evaluate at runtime.
Repeated Matching (o). In some cases, we'd like to merge more than two models together. To do this, we allow "repeated matches". That is, when two features are matched in our greedy algorithm, they are removed and replaced with the resulting merged feature. To ensure that one feature doesn't 't get merged endlessly, we set the correlations the new feature to be the minimum similarities weighted by (0, 1]. We find small value of typically works best.
old features'
Same-model Budget (8). To demonstrate the effectiveness of same-model merges, we introduce "budget" parameter 3 [0, 1] that denotes what percent of total merged features can come from models merging within themselves, with each model receiving an equal portion of this budget. Note that budget of 0 results in Eq. as in that case no features can be merged within models.
RESULTS
There is no standard benchmark evaluate merging approaches on models from distinct tasks, so we construct our own. We evaluate our approach in two different settings. (1) A versatile test-bed: disjoint category splits of the same dataset (i.e., same dataset and different label sets). (2) very challenging setting: completely different datasets and tasks (i.e., different datasets and label sets).
Experimental Details. For each experiment where we sample multiple disjoint splits of categories, we hold one split out for hyperparameter search and report mean and standard deviation the rest. For experiments with models trained on different datasets, we subsample the validation set into a validation and test for the same purpose. To compute correlations,
portion of the

---

Published as a conference paper at ICLR 2024
Accuracies (16)
Accurscies (16)
Method
FLOP (Ci)
Joint
Tisk A
Task
Avg
Method
FLOPS
Joim
Tisk A
Avg
Model A
0.68
48.2
27.0
45. 1ass
71.0gz
Model A
2.72
41,620
82.9
24.8
53,0
49.1-93
16.1
72.6.49
Model
41.6
25.1
82.8.13
.4.0.oz
W. Aug iyy
0.68
43.0.1
54.1
67.5
60.8.44
W. Arg thy i1
2.72
17.0
23,8
24.8.1
24.3
Git Ro-Rasin"
0.68
46.3
76.8
62.7.5
79.8
Git Re-Rasin"
40,90t
57.9
567733
57.0.0
Permuse 2ng 2
0.68
584
86.6
87.4
87.4
Permute the t
2.72
42.8
61.6
60.5
61.0zos
Zipltiara
0.68
79.1
92,9
91.2 114
92.1 118
Zipatlave
2.72
54.9
08,2 00
67.9
68.0
Ensembble
07.0
96.6244
Ensembble
73.5204
Zipli! une
0.91
83.8
05.1
94.1
94.6.aa
Ziplitioe
3.63
70.2:
g0.3
50.1
20.2
(a) CIFAR-10 (5+5), ResNet-20 (4x width).
(b) CIFAR-100 (50+50), ResNet-20 (8x width).
Table 1: CIFAR Results. Ziplt! vs. baselines combining a model trained half the classes (Task A) with one trained on the other half (Task B) without extra training. We report both joint (10/100-way) 1ay) and per-task (5/50-way) accuracy. Ziplt! significantly outperforms its baseline and closes in on the upper bound (ensemble accuracy). $ refers to Ainsworth et al. (2022).
training set for each dataset as in Li et al. (2015) (see Appendix B). For a fair comparison, we reset the batch norms for all methods (including the original models) using the training data (following the recommendation in Jordan et al. (2022)). For our method, Ziplt!um indicates that n out of the m layers in the network have been zipped (Sec. 4.3). Note, all our models have different initializations.
Evaluation. For the setting with disjoint class splits of the same dataset, we evaluate performance in two ways: joint accuracy and per task accuracy. For joint accuracy, we evaluate each model over all classes in the combined dataset. For per task accuracy, we compute the accuracy of each task individually (i.c., supposing we had task labels at runtime) and then report the average. The former is similar continual learning setting where the latter is akin to a multi-tase task setting where we know which task we'r re using at test time. For the scenario where we merge models trained on different datasets, we use the per task accuracy metric, as the label spaces are not comparable.
want augment the knowledge of the model, while
Baselines. In addition to the default Weight Matching version of Git Re-Basin (Ainsworth et al., 2022), we compare to two baselines: Weight Averaging (Eq. 1) and Permute (Eq. 2) with 1/2 using our framework (i.e., we set M, and Ug such that Eq. is equivalent). For Permute, we use linear sum assignment to find optimal permutations (following Li et al. (2015)). Note that our Permute is strong bascline we create using our framework and is more accurate than Git Re-Basin in our settings. It's also similar to REPAIR (Jordan et al., 2022), but without adding extra parameters to the model. Finally, with perfect merging, the merged model's outputs would be identical to the originals. Thus we include Ensemble as upper bound (executing and concatenating the results of both models).
5.1 CIFAR-10 AND CIFAR-100
We train 5 pairs of ResNet-20 (He et al., 2016) from scratch with different initializations on disjoint halves of the CIFAR-10 10 and CIFAR- 100 classes (Krizhevsky et al., 2009). While Zipit! supports "partial zipping" to merge models with different outputs (in this case, disjoint label sets), prior methods without retraining do not. To make fair comparison, we train these CIFAR models with a CLIP-style loss (Radford et al., 2021) using CLIP text encodings of the class names as targets. This way, both models output into the same CLIP-space regardless of the category. Note, this means the models are capable of some amount of zero-shot classification on the tasks they were not trained on.
CIFAR-10 (5+5). In Tab. 1a. we merge models trained on disjoint 5 class subsets of CIFAR-10 using ResNet-20 with a 4x width multiplier (denoted as ResNet-20x4). 4). In joint classification (i.c., 10-way), Git Re-Basin is unable to perform better than using either of the original models alone, while our Permute baseline performs slightly better. In stark contrast, our Ziplt! performs a staggering 32.9% better than Git he-Basin and 20. 7% better than our baseline. If allow the last stage of the network to remain unzipped (i.e., zip up to 13 layers), our method obtains 83.8%, which is only 3.6% behind an ensemble of model and model B (which is practically the upper bound for this setting). We also achieve similar results when merging VGG11 models in this setting (Appendix D).
CIFAR- 100 (50+50). We find similar results on disjoint 50 class splits of CIFAR-100 in Tab. 1b, this time using an 8x width multiplier instead. Like with CIFAR-10 10, Git Re-Basin fails to outperform even the unmerged models themselves in joint classification (i.e., (00-way), and this time Permute is only 1.2% ahead. Ziplt! again significantly outperforms prior work with + +14% accuracy over Git Re-Basin for all layers zipped, and substantial 29.2% if zipping 13/20 layers. At this accuracy,
7

---

Published as a conference paper at ICLR 2024
13/2t/13/20 13/20 is again only 3.3% behind the ensemble for joint accuracy and 2.6% behind for average per task accuracy, landing itself in an entirely different performance tier compared to prior work.
5.2 IMAGENET-IK (200+200)
Accuracies (8)
To test our method on the much harder setting of large-scale data, we train 5 differently initialized ResNet-50 models with cross entropy loss on disjoint 200 class subsets of ImageNet- 1k (Deng et al., 2009). To compare to prior work that doesn't t support partial zipping, we initialize the models with capacity for all 1 classes, but only train each on their subset.
Method
FL.0P% (C))
Joint
Task
Avg
Modell A
4.11
$7.2
74.3
0.5am
37.4
Model
4.11
35.3.14
0.540
70.5.3
35.511
W. Aug it
4.3L
0:3
0.6101
0.7161
0.1101
Re-Badin"
4.11
3.1
533
5.7
55111
Penmute cng
4.11
86:11
10.1:44
153
127
Zipliius
1477.1
135
Essemble
8.23
63.3
74.3
70.5
72
4.33
Zipliizes
6.39
55.8
65.9
61.1
65.0:23
Ziplliuns
7.43
60.9
70.7.3x
69.0.25
69.9.15
Table 2: ImageNet-1k (200+200) Results. Merg- ing ResNet-50 models trained from scratch on disjoint 200 category subsets (Task A and B) of ImageNet-Ik. Prior work performs poorly, but Ziplt! makes this task feasible. 1 Aimsworth al. (2022).
In Tab. 2 we show results on exhaustively merg- ing pairs from the models. To compute joint (i.e., 400-way) accuracy, we softmax over each task's classes individually (like in Ahn et al. (2021)), and take the argmax over the combined 400 class vector. On this extremely difficult task, Git Re- Re-Basin only obtains 3.1% for joint accuracy (with random accuracy being 0.25%). Both the Permute baseline and Ziplt! with all layers zipped perform better, but with each at 8.6%, are still clearly lacking. Note that we find the same-model merging budget 3 to not matter for this set of models (see Fig. 6). which suggests that there's not lot of redundant information within each model in this setting. Thus, Ziplt! chooses to merge mostly across models instead, performing similarly to the permute baseline. We find this same trend in CIFAR with smaller models (see Fig. 6). and may be an artifact of model capacity. The story changes when we increase the capacity of the merged model by partial zipping: Ziplt! 10/50 reaches close to upper bound ensemble accuracy on this extremely difficult task, while saving on FLOPs.
5.3 MULTTI-ATASST MERGING
We now take our model merging framework one step further by merging differently initialized models trained on completely separate datasets and tasks. We present two settings: merging multiple classification datasets and merging semantic segmentation with image classification.
Per:Tack Accuracies (%)
Image Classification Datasets. Merging ResNet- 50 models trained on: Stanford Dogs (Khosla et al., 201 1), Oxford Pets (Parkhi et al., 2012), CUB200 (Welinder et al., 2010), and NABirds (Van Horn et al., 2015). In Tab. 3, we show the average per task accuracy from exhaustively merg- ing each pair and the much more difficult setting of merging all four at once. We report the accuracy of our baselines by applying them up until the last layer, but we can't 't compare to prior work as they don't 't support this setting. As in all our previous experiment we merge without retraining.
Method
FLOPS (G)
SD
OP
CUB
NAB
Avg
Merging Pairs
W. W.A
4.11
12.9
18.2
13.9
0.2
11.3
Permute
4.11
46.2
47.6
35.7
Zipli
4.11
46.9
50.7
38,0
12.7
37.1
Ensemble
8.22
72.7
81.1
75.5
Ziplifsoo
6.39
62.6
71.2
62.8
53.0
62.4
Ziplthnee
7,42
66.5
75.8
68.6
66.8
68.7
Merging All
Arg the
4.12
0.8
3.0
1.2
Permute
4.12
15.7
26.1
14,0
5,3
15.3
Ziplitace
4.12
21.1
1688
Ensemble
16.4
72.7
81.2
71.0
77.2
75.5
Ziplit
55.9
45.5
Zipliare
14.1
63.5
70.8
63.7
63.1
65.3
Table 3: Multi-Dataset Results. Merging ResNet-50 models trained on completely differ- ent datasets: Stanford Dogs (SD), Oxford Pets (OP), CUB200 (CUB), and NABirds (NAB). We report average per-task accuracy over merging model pairs, and all four.
For pairs of models, Ziplt! slightly outperforms our permute baseline across all tasks and performs similarly when merging all 4 models once. How- ever, if we add capacity to the merged model through partial zipping, we perform up to 33% better on merging pairs and 50% better on merging all four models than the permute baseline. Partial zipping is a significant factor to obtain strong performance, especially with more than 2 models.
Multiple Output Modalities. In Appendix F, we combine across modalities by merging the 5e backbone of DeeplabV3 (Chen et al., 2017) segmentation model with an ImageNet- 1k classification model. The resulting combined model can perform both semantic segmentation and image classification. Even with half of layers merged, ZipIt! retains good performance on both tasks.
8

---

Published as a conference paper at ICLR 2024
(a) Ablation: Varying, /
(b) Ablation: Model Scale
CIFAR-100 (50+50)
ImageNet-1k (000-200)
CIFAR-100 (50+50)
68
74
2ipht
66
72
70
--
100xxee Aoc
Peemude
64
68
Rn Rach
Budget
64
66
Budget B
Roshatt Width Mumpplier
0
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
16x
32x
Figure 6: Varying 3 and Model Scale. Left: We find when the model has enough capacity for the task, high budget (Sec. 4.3) improves performance. Right: ZipIt! makes effective use of extra model capacity to quickly reach the ensemble on CIFAR-I00 (50+50) when we increase the width of ResNet-20 models. In contrast, our baselines only slightly benefit from the extra scale.
6 ANALYSIS
Merging within Models. A critical piece of Ziplt! compared to prior work is the ability to merge within models, not just across models. In Sec. 4.3, introduce a budget parameter /3 to limit the number of same-model merges, and here use CIFAR-I00 (50+50) and ImageNet- (2004200) to illustrate its effectiveness (Fig. 6a). On CIFAR, same-model merges are very important, with the optimal budget being above 0.8, meaning 80% of merges are allowed to be within the same model. This is not the case, however, on ImageNet, where the difficulty of the task means there likely are much fewer redundant features within each model.
Model Scale. Fig. 6b, we test the effect of model scale directly by evaluating joint accuracy on our CIFAR-100 (50+50) setting with ResNet-20 models of increasing width. Here, we explic- itly see that when the width of the models are too small for the task (c.g.. < 4x), Ziplt! and the Permute baseline perform identically (though both much better than Git Re-Basin). How- ever, when the scale increases, Ziplt! trends toward the ensemble upper bound of 75% while both the Permute baseline and Git Re-Basin plateau at around 45%. This corroborates Eq. and indicates our method uses the extra model capacity effectively, much better than prior work.
Matching Algorithm. In Tab. 4, we compare matching algorithms used to compute M, in Eq. 8. Using either the identity (weight aver- aging) or a permutation (as in prior work) un- derperforms on CIFAR-10 (5+5) joint 10-way classification. In contrast, we obtain up to 21.2% higher accuracy if we allow both permutations and merging within models. However, doing this optimally is difficult, as the standard linear sum assignment algorithm assumes bipartite matches. We could use optimal graph-based solver (e.g., Hagberg et al. (2008) instead, but doing so is prohibitively slow (11 minutes transform ResNet-20x4 model). Thus, we find matches greedily by repeatedly taking the most corre- lated pair of features without replacement. This performs almost as well, and is multiple orders of magnitude faster. we allow repeated matches (Sec. 4.3), we obtain slightly better result. Like Bolya et al. (2023), we find that matching is better for merging features than clustering (K-Means).
Algorithm Identity (Fq. 1)
AcsA/Be+R?
Acc
Time
X
45.0
ms
Permute (Bx 2)
X
58.4
28 ms
K-Means
29. 19.1
19
Zip (Eq,
Optimal Match Greedy Match Greedy, 1=
79.6
11 min
79.0
1.1 soc
79.1
1.2
Table 4: Comparing Matching Algorithms to use for M, CIFAR-10 (5+5) joint 10-way accuracy. Permuting B-+A as in prior work (Eq. 2) performs poorly. We significantly improve by merging fea- tures within each model (Eq. 8). Our greedy ap- proach nearly as accurate as the optimal algo- rithm while being two orders of magnitude faster.
7 CONCLUSION
In this paper, we tackle the extremely difficult task of merging models trained completely disjoint tasks without additional training. We find that prior work underperforms in this setting and posit that they neither fully (1) exploit model similarities nor (2) account for model dissimilarities. We introduce ZipIt!, a general framework for merging models that addresses these issues, and show it to significantly outperform prior work across several difficult settings, comprehensively analyzing each.

---

Published as a conference paper at ICLR 2024
Reproducibility Statement. To ensure reproducibility, we will release code for our algorithm, experiments, and baselines. We also include algorithm details in Section 4 and further in Appendix C, experimental details in Section 5 and Appendix B, and proof of our Theorem in Appendix G.
Acknowledgements. This work was supported in part by funding from NSF CAREER #2144194 ARL, Google, and NSF GRFP. All views and conclusions expressed in this work are those of the authors and not a reflection of these sources.
REFERENCES
Hongjoon Ahn, Jihwan Kwak, Subin Lim, Hyeonsu Bang. Hyojun Kim, and Taesup Moon. Ss-il: Separated softmax for incremental learning. In ICCV, 2021.
Samuel K Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa. Git re-basin: Merging models modulo permutation symmetries. urXiv:2209.04833 2022.
Stephen Ashmore and Michael Gashler. A method for finding similarity between multi-layer perceptrons by forward bipartite alignment. In IJCNN, 2015.
Alexei Baevski, Wei-Ning Hsu, Qiantong Xu, Arun Babu, Jiatao Gu, and Michael Auli. Data2vec: A general framework for self-supervised learning in speech, vision and language. In ICML., 2022.
Gilles Blanchard, Gyemin Lee, and Clayton Scott. Generalizing from several related classification tasks to a new unlabeled sample. In NeurlPS, 2011.
Daniel Bolya, Chong Zhou, Fanyi Xiao, and Yong Jae Lee. Yolact: Real-time instance segmentation. In ICCV, 2019.
Daniel Bolya, Cheng- Yang Fu, Xiaoliang Dai, Peizhao Zhang, Christoph Feichtenhofer, and Judy Hoffman. Token merging: Your vit but faster. ICLR, 2023.
Zhaowei Cai, Avinash Ravichandran, Subhransu Maji, Charless Fowlkes, Zhuowen Tu, and Stefano Soatto. Exponential moving average normalization for self-supervised and semi-supervised learning. CVPR, 2021.
Mathilde Caron, Hugo Touvron. Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In ICCV, 2021.
Liang-Chieh Chen, George Papandreou, Florian Schroff, and Hartwig Adam. Rethinking atrous convolution for semantic image segmentation. arXiv preprint arXiv: 055877 2017.
Leshem Choshen, Elad Venezian, Noam Slonim, and Yoav Katz. Fusing finetuned models for better pretraining. arXiv:2204. 03044,, 2022.
Matthias De Lange, Rahaf Aljundi, Marc Masana, Sarah Parisot, Xu Jia, Aleš Leonardis, Gregory Slabaugh, and Tinne Tuytelaars. continual learning survey: Defying forgetting classification tasks. TPAMI, 2021.
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters. arXiv:2302.05442, 2023.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In CVPR, 2009.
Shachar Don- Yehiya, Elad Venezian, Colin Raffel, Noam Slonim, and Leshem Choshen. ColD fusion: Collaborative descent for distributed multitask finetuning. Association for Computational Linguistics, 2023.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x 16 words: Transformers for image recognition scale. arXiv:2010. 1929, 2020.
10

---

Published as a conference paper at ICLR 2024
Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred Hamprecht. Essentially no barriers in neural network energy landscape. In JCML, 2018.
Rahim Entezari, Hanie Sedghi, Olga Saukh, and Behnam Neyshabur. The role of permutation invariance in linear mode connectivity of neural networks, arXiv:21 2021.
M. Everingham, L. Van Gool, C. K. Williams, J. Winn, and A. Zisserman. The pascal visual object classes (voc) challenge. International Journal of Computer Vision, 88(2):303-338, June 2010.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode connectivity and the lottery ticket hypothesis. ICMI,, 2020.
Daniel Freeman and Joan Bruna. Topology and geometry half-rectified network optimization. 11105516,, 2016.
Timur Garipov, Pavel Izmailov, Dmitri Podoprikhin, Dmitry P Vetrov, and Andrew G Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns. NeurlPS, 2018.
Jort F Gemmeke, Daniel PW Ellis, Dylan Freedman, Aren Jansen, Wade Lawrence, Channing Moore, Manoj Plakal, and Marvin Ritter. Audio set: An ontology and human-labeled dataset for audio events, In ICASSP, 2017.
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self- learning, NeurIPS, 2020.
Almog Gueta, Elad Venezian, Colin Raffel, Noam Slonim, Yoav Katz, and Leshem Choshen. Knowl- edge a region in weight space for fine-tuned language models. 2xXiv:2302.04863, 2023.
Aric Hagberg, Daniel A. Schult, and Pieter J. Swart. Exploring network structure, dynamics, and function using networkx. In Proceedings of the 7th Python in Science Conference, 2008.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. 2016.
Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick. Mask r-cnn. In ICCV, 2017.
Xiaoxi He, Zimu Zhou, and Lothar Thiele. Multi-task zipping via layer-wise neuron sharing. NeurIPS, 2018.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. NeurIPS, 2020.
Gao Huang, Yixuan Li, Geoff Pleiss, Zhuang Liu, John Hopcroft, and Kilian Q Weinberger. Snapshot ensembles: Train 1, get for free. arXiv: /74400090 2017.
Gabriel Ilharco, Marco Tulio Ribeiro, Mitchell Wortsman, Suchin Gururangan, Ludwig Schmidt, Hannanch Hajishirzi, and Ali Farhadi. Editing models with task arithmetic. arXiv:2212.04089, 2022a.
Gabriel Ilharco, Mitchell Wortsman, Samir Yitzhak Gadre, Shuran Song, Hannaneh Hajishirzi, Simon Kornblith, Ali Farhadi, and Ludwig Schmidt. Patching open-vocabulary models by interpolating weights. NeurIPS, 2022b.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. UAI, 2018.
Keller Jordan, Hanie Sedghi, Olga Saukh, Rahim Entezari, and Behnam Neyshabur. Repair: Renor- malizing permuted activations for interpolation repair. arXiv:22 1,08403, 2022.
Tero Karras, Samuli Laine, and Timo Aila. style-based generator architecture for generative adversarial networks. In CVPR, 2018.
Aditya Khosla, Nityananda Jayadevaprakash. Bangpeng Yao, and Fei-Fei Li. Novel dataset for fine- grained image categorization: Stanford dogs. In CVPR Workshop on Fine-Grained Visual Categorization (FGVC), 2011.
11

---

Published as a conference paper at ICLR 2024
James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho. Agnieszka catastrophic forgetting in neural networks. PNAS, 2017.
et al. Overcoming
Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. In ICML, 2019.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009,
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolu- tional neural networks. Communications of the ACM, 2017.
Yixuan Li, Jason Yosinski, Jeff Clune, Hod Lipson, and John Hopcroft. Convergent Learning: Do different neural networks learn the same representations? arXiv: /51::///1.075444 2015.
Zhizhong Li and Derek Hoiem. Learning without forgetting. TPAMI, 2017.
Chang Liu, Chenfei Lou, Runzhong Wang, Alan Yuhan Xi, Li Shen, and Junchi Yan. Deep neural network fusion via graph matching with applications to model ensemble and federated learning. In ICML., 2022.
Michael Matena and Colin Raffel. arXiv:2111.09832, 2021.
Merging models with fisher-weighted averaging.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera Arcas. Communication learning of deep networks from decentralized data. In Artificial intelli- gence and statistics. PMLR, 2017.
Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. In /CML, 2013.
Behnam Neyshabur, Hanie Sedghi, and Chiyuan Zhang. What is being transferred transfer learning? NeurlPS, 2020.
Omkar M. Parkhi, Andrea Vedaldi, Andrew Zisserman, and C. V. Jawahar. Cats and dogs. In CVPR, 2012.
Fidel A Guerrero Peña,, Heitor Rapela Medeiros, Thomas Dubail, Masih Aminbeidokhti, Eric Granger, and Marco Pedersoli. Re-basin via implicit sinkhorn differentiation. arXiv:2212.12042 2022.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In ICML, 2021.
Alexandre Ramé, Kartik Ahuja, Jianyu Zhang, Matthieu Cord, Léon Bottou, and David Lopez-Paz. Recycling diverse models for out-of-distribution generalization. arXiv:2212. 2022.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In CVPR, 2022.
Berfin Simsek, François Ged, Arthur Jacot, Francesco Spadaro, Clement Hongler, Wulfram Gerstner, and Johanni Brea. Geometry of the loss landscape overparameterized neural networks: Symme- tries and invariances. In Proceedings of the 38th International Conference on Machine Learning. PMLR, 2021.
Sidak Pal Singh and Martin Jaggi. Model fusion via optimal transport. NeurlPS. 2020.
Yi-Lin Sung, Linjie Li, Kevin Lin, Zhe Gan, Mohit Bansal, and Lijuan Wang. empirical study of multimodal model merging, 2023.
Antti Tarvainen and Harri Valpola. Mean teachers are better role models: Weight-averagee consistency targets improve cemi-supervised deep learning results. NeurIPS, 2017.
Norman Tatro, Pin- Yu Chen, Payel Das, Igor Melnyk, Prasanna Sattigeri, and Rongjie Lai. Optimizing mode connectivity via neuron alignment. NeurIPS, 2020.
12

---

Published as a conference paper at ICLR 2024
Grant Van Hor, Steve Branson, Ryan Farrell, Scott Haber, Jessie Barry, Panos Ipeirotis, Pietro Perona, and Serge Belongie. Building a bird recognition app and large scale dataset with citizen scientists: The fine print in fine-grained dataset collection. In CVPR, 2015.
Johannes Von Oswald, Seijin Kobayashi, Joao Sacramento, Alexander Meulemans, Christian Henning, and Benjamin F Grewe. Neural networks with late-phase weights. arXiv:2007.12927 2020.
Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni, Federated learning with matched averaging. arXiv:2002.06440, 2020.
Jindong Wang, Cuiling Lan, Chang Liu, Yidong Ouyang, Tao Qin, Wang Lu, Yiqiang Chen, Wenjun Zeng, and Philip Yu. Generalizing unseen domains: A survey domain generalization. TKDE, 2022.
P. Welinder. S. Branson, T. Mita, C. Wah, F. Schroff, S. Belongic, and P. Perona. Caltech-UCSI Birds 200. Technical Report CNS-TR-2010-001, California Institute of Technology, 2010.
Mitchell Wortsman, Gabriel Ilharco, Samir Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Korbblith, et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In JCML, 2022a.
Mitchell Wortsman, Gabriel Ilharco, Jong Wook Kim, Mike Li, Simon Kornblith, Rebecca Roelofs, Raphael Gontijo Lopes, Hannaneh Hajishirzi, Ali Farhadi, Hongseok Namkoong, et al. Robust fine-tuning of zero-shot models. In CVPR, 2022b.
Prateek Yadav, Derek Tam, Leshem Choshen, Colin Raffel, and Mohit Bansal. Resolving interference when merging models. arXiv:2306.0/708, 2023.
Masanori Yamada, Tomoya Yamashita, Shin's ya Yamaguchi, and Daiki Chijiwa. Revisiting permuta- tion symmetry for merging models between different datasets, 2023.
Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. Bayesian nonparametric federated learning of neural networks. In ICML, 2019.
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In CVPR, 2022.
Bolei Zhou, Hang Zhao, Xavier Puig, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Scene parsing through ade20k dataset. In CVPR, 2017.
13

---

Published as a conference paper at ICLR 2024
A PARTIAL ZIPPING
In Fig. 7 we plot the average per task accuracy by the number of layers zipped in ResNet-20x8 for CIFAR-100 (50+50) and ResNet-50 for ImageNet-1k 1k (200+200). Note that to avoid adding extra unmerge modules into the network, our stopping point while unzipping has to be the end of stage.
In Table 5, we show the average neuron correlations at each partial-zipping stage between the layers of ResNet-22 (8x) models trained on the CIFAR- 100 (50+50) task. We collect results using the same models used in Table 1b, and compute correlations as described in Section 4. Overall, we find the correlations between models consistently decreases through successive partial zipping locations. This corroborates the finding of Kornblith et al. (2019) that model layers become increasingly dissimilar with depth, as they encode more task-specific features. Coupling Table 5 with Figure 7a, we observe a direct correlation between (dissdddsssiimiittiess and performance decrease. This illustrates the importance of layer similarity between two networks and strong performance.
B DATA USAGE
In our approach, we use sample of the training set in order to compute activations and match fea- tures together. For the main paper, we used the full training set for CIFAR, 1% of the training set for ImageNet, and the number of images in the small- est training set for the Multi-Dataset classification experiment (so that could use the same number of images from each dataset). In each case, we used the same data augmentations from training.
84-
1.7x
Speed-up
Avg Task Acc 196
76
1.1x
113xx
Speed-up
81
62
Awg fasx 75
48
34
20
2x
Layers Zipped
6
Layers Zipped
22
(a) CIFAR-100 (50+50)
(b) ImageNet-Ik (200+200)
Figure 7: Varying Partial Zip. By leaving some layers unzipped (Sec. 4.3), we can recover sig- nificant amount of performance while still merg- ing most of the model.
That begs the question: how much data do we actually need, and how necessary are data aug- mentations? Here we ablate the amount of data used for our CIFAR-100 (50+50) ResNet-220 (8x width) and ImageNet (200+200) Resnet-50 (3// layers) experiments. In Fig. 8, we test how much data is actually necessary to obtain a good accu- racy on CIFAR and ImageNet with or without data augmentation.
Average Stage Correlations
Layer 7/2
Layer 1//20
Layer 1920
0.50 10.1
0.37 10.00
0.27 0.000
Table 5: CIFAR-100 (50+50) Zipping Correla- tions. We show the average correlations between two ResNet-20 (8x width) models at each partial zipping stage. Correlations consistently decrease at each successive stage, indicating that the lay- ers of the two models increasingly diverge.
We ultimately find that the amount data doesn't actually matter that much. In the main paper, we use the entire training set for CIFAR-100 with a batch size of 500 (100 batches, 50,000 images), but it seems like as little as 2 batches (100 images) produces the same result. Similarly on ImageNet, using 0.05% of the data (640 images) produces the same result as 5% (64,048 images).
With Augimenfation
68
68
Without Augmentation
Aug look Aos
avg Tegk 820
66
Without Augmentation
84
64
8
62
With Augmentstion
In fact, the main consideration is whether or not use data augmentation. For less diverse datasets like CIFAR-100 100, data augmentation seems es- sential (giving an almost 4% boost in average task accuracy), and well above the variance of results without augmentation. However, for Ima- geNet, which has much more diverse images, data augmentation actually hurts slightly on average- though the two within variance. Note that despite this result, for consistency we use data augmentation in all experiments.
8
Datches
100
60
% of Data
(c) CIFAR-100
(d) ImageNet-Ik
(50+50)
(200+200)
Figure 8: Data Usage. How much data do we need to use to compute activations? We find that only a few hundred images are needed to obtain the best performance. Data augmentation is not always useful.
14

---

Published as a conference paper at ICLR 2024
C ZIP PROPAGATION DETAILS
In the main paper we described general rules for zip propagation namely, propagate through layers until you reach a module with weight matrix. Here, we describe rules more concretely for each layer type needed to define most convnets.
Linear. Apply M, and U,. Stop propagation.
Conv. Apply M; and U; each kernel location (i.e., move the it x k kernel dimensions the batch dimension). Stop propagation.
BatchNorm. Apply Mj to all parameters (weight, bias, mean, variance). squaring it for the variance term. Continue propagation. As Jordan et al. (2022) points out, we cannot compute the correct variance without knowing the covariance between the two models (which we don't have access to). Thus, we reset batch norms after merging to evaluate the variance correctly.
LayerNorm. Apply Mj to all parameters (weight, bias). Continue propagation. Since LayerNorm computes mean and standard deviation the fly, don't need to do anything special.
ReLU, Nothing to merge. Continue propagation. Note that passing the merge matrix unchanged through the ReLU ReLUlayers an approximation, since we' re using a linear merge operation on nonlinear features. Addressing this issue could be an interesting topic for future work, as even the permute and add approach of prior work has this issue (ReLU is invariant to permutation, but certainly not adding).
Avg Max Pool. Nothing to Merge. Continue propagation.
Skip Connection. Continue propagation through every input to the skip connection (using the same M, and U, for each).
D CROSS ENTROPY ON CIFAR
In the main paper, we train our CIFAR models with a CLIP (Radford et al., 2021) loss (using CLIP embeddings class names as targets). This ensures that the output spaces of the two models are aligned, which is necessary to get good accu- racy for prior work that merge the entire model together.
Accuracies (98
Metbel
FLOPs
Tisk
Model A
10.88
37.9
74.15
1.7
36.1
38.7
W. Asg
10.88
2.7
5.0
4.9
4.9
Git Re Re-Basint
10.88
3.1
5.8
5.3
5.6
20.0
Ziplitarm
10.88
27.9
40.1
39.7
39,9
Ensemble
21.76
66.5
74.2
75.2
74.7
Zipit!.co
Ziplthne
18.14
47.0
60.6
60.5
60.6
ResNet. In Tab. 6, we show results for CIFAR- 100 (50+50) where we train with the normal one-hot targets (i.e., like we did for ImageNet), instead. Immediately, we see that accuracies of the merged models are much lower across the board, with no method able to outperform just using one of the original models when merging the entire network. In fact. Git Re- fe-Basin (Ainsworth al., 2022) does almost no better than weight averaging, which gets close to random accuracy. While Ziplt! without partial zipping also performs worse than the original models, it still greatly outperforms all prior work. And with partial zipping, Ziplt! is able to exceed the accuracy of the original models.
Table 6: CIFAR-100 (50+50) Cross Entropy. ZipIt! vs. baselines using ResNet-20 (16x width). Merging the entire model as in prior work produces bad results when using cross- entropy, hence we use CLIP in the main draft. If we use partial zipping, we can recover a lot of the lost performance. & refers to (Ainsworth et al., 2022)
Thus, in the case of using cross-entropy loss, partial zipping is extremely important. Merging the entire model as in prior work fails, since the later layers of the model are incompatible with each other due to each model having different output space. Partial zipping, on the other hand, can mitigate that issue.
15

---

Published as a conference paper at ICLR 2024
Accuracies (58)
VGG. In the main paper, use ResNets for each experiment, since they are easy to train and pro- duce strong results. However, in principle Ziplt! can work on any architecture. For completeness, in Tab. 7, we show results the CIFAR-10 (5+5) setting with V (Ix width). Note that this is a much smaller and weaker model than the ResNet- 20s we use in the main paper, so its results on CIFAR-10 10 aren't 't as strong. Furthermore, we con- ducted this experiment with cross entropy loss, so merging the entire model performs worse than the original models.
Methel
FLOPPs (C)
Joint
Tisk A
Tisk B
Avg
Model A
0.15
$4.6
89.2
21.0
55.1
Model B
0.15
44.0
25.1
88.1
55.6
W.
10.2
Permute
0.15
25.4
47.2
48.5
47.8
Zipliam
015
33.2
538
59.9
56.5
0.30
Ziplthizz
0.17
35.2
56.7
0012
58.4
Ziphtiies
0.27
44.5
66.0
65.1
65.5
Table 7: CIFAR-10 (5+5) CE with VGG. Ziplt!
vs. baselines using VGG11 (1x width) using Cross Entropy instead of CLIP loss. ZipIt! displays the same behavior here as it does for ResNet-20 with low width.
Despite this, we observe very similar trend to the ResNet-20 models that Ziplt! outperforms all baselines and that partial zipping is important for reaching the accuracy of the original models (in this case, matching not exceeding). In fact, these results continue a more general trend in that Ziplt! greatly benefits from larger model scales, making effective use of the extra capacity. In this case, the scale of the model is quite small, so there is not as much room in the weights to store the potentially disjoint features of both models.
E IMAGENET WITH 1.5X WIDTH
In the main paper, we show that ZipIt! scales very well with increased width of the model for the CIFAR-10 100 (50+50) setting. While CIFAR-100 is a challenging dataset on its own, the natural question is if that same trend occurs the much harder ImageNet- Ik (200+200) setting.
Accuracies (98)
Matbed
FLOPP (C)
Joint WiRth
Task
Tisk B
Avg
Insemble
8.22
63.3
74.3
70.5
72.4
Zipltius
4.11
8.6
Ziplitaen
4.92
33.1
41.8
42.3
42.0
Zipltizzs
6.39
55.8
65.9
64.1
65.0
Zipithos
60.9
69.9
1,5x Width
In Tab. 8, we test this by comparing ZipIt! on the original 1x width ResNet-50 in the main pa- per with a 1.5x width one. In all cases, except for the fully zipped model (likely because of the Cross-Entropp loss). Ziplt! enjoys a large jump in performance from the extra width. For 37 lay- ers, 33.1% joint accuracy becomes 49.0%. For 22 layers, 55.8% becomes 64.1%. And for 10 layers, 60.9% becomes 66.8%, now only 1% away from the ensemble. Thus, even in this much more chal- lenging setting, Ziplt! is able to make full use of the extra model capacity.
Ensemble
32.6
67.8
76.7
72.6
74.7
Zipltiuso
16.3
9.7
Zipltiics
19.5
49.0
56.2
56.7
56.4
Ziplthnss
25.5
64.1
71.6
70.4
71.0
Zipltinos
74.9
73.5
Table 8: ImageNet-Ik (200+200) Width Com-
parison. We show how ZipIt! is able to make use of the extra model width when merging models together. For instance, merging 37 layers goes from 33% joint accuracy with 1x width to 49% with 1.5x, while the ensemble only improves by 4%. These models use cross-entropy, so merging the entire network results in poor performance.
F MERGING MODELS WITH DIFFERENT OUTPUT MODALITIES
In this experiment we ZipIt! to merge two models with different initializations trained on different tasks with different output modalities: semantic segmentation and image classification. Specifically, we merge the ResNet-50 backbone of a DeepLabV3 (Chen et al., 2017) model fine- al.,
Accuracy (%)
mloU (%)
Method
ImageNet- 1k
Pascal VOC
W. Avg
0.8
3.3
Zipltlawe0
23.1
6.0
Ensemble
77.8
76.8
tuned on the Pascal VOC (Everingham 2010) dataset, with a ResNet-50 model trained on ImageNet-Ik. While the DeepLabV3 back- bone was itself pre-trained was further finetuned on Pascal VOC and does not share the initialization of our classification model. Table shows the results of combining the two ResNet-50 models with Ziplt! various partial
Ziplt!
47.7
35.0
ZipIt!2250
60.9
64.4
ZipIt!
64.9
71.7
ImageNet-Ik, it
Table 9: PASCAL VOC and ImageNet-Ik merging models with different output modali- ties using DeepLab ResNet-50 backbone and ImageNet- 1k Resnet-50 model.
16

---

Published as a conference paper at ICLR 2024
merging locations. We evaluate the performance of each merged by reporting its ImageNet- Ik accuracy, and its Pascal VOC mloU as is standard. Overall, we observe that Ziplt! is capable of merging nearly half the number of ResNet-50 layers between both models while still maintaining good performance on both tasks, all without any training.
G A TIGHTER BOUND FOR LINEAR MODE CONNECTIVITY
In this section we demonstrate that merging models by supporting feature merges both across and within each, yields tighter bound than Theorem 3.1 in (Entezari et al., 2021) in its limited setting. We first introduce necessary background from prior work, including Theorem 3.1 and a particular formalization for within-model merging borrowed from (Simsek et al., 2021). Second,
introduce
Theorem I, which produces a tighter bound on Theorem 3.1 when merging within models is allowed, and prove its validity (Section G.2 & G.3). Third, provably extend Theorem to less restrictive setting, retaining its bounds (Section G.4).
G.1 BACKGROUND
We first introduce Theorem 3.1 from (Entezari et al., 2021). Second, we formalize restricted version of within-model merging necessary for our proof using the definitions from (Simsek et al., 2021).
G.1.1 THOEREM 3.1
The Theorem. Let f(v,w)(x) oWoo) f(v f(v:ww(=)= o (W"x)) be two fully-connected networks with h hidden units where af-) is ReLU activation, R² and W R are the parameters and x E Rdd is the input. If each element of W and W is sampled uniformly from [-1/Vd, 1/Vd] and each element of v and v' is sampled uniformly from [-1/v/, 1/vh), then for any æ E Id such that || = Vd, with probability - over W, v', there exist permutation such that
af(v. w) (z) (1 - a) (v.v Ö(h .हक)
where v", w" permuted version of v', W [0, is an arbitrary interpolation constant, and the left-hand-side of the equality is the amount an interpolated model differs in output compared to the interpolation of the original models. (Entezari al., 2021) show that minimizing this quantity is analogous to minimizing the barrier (as defined by Entezari et al. (2021)) in this setting. This important because it states that achieving zero output difference is equivalent to achieving zero-barrier, which implies that two models are linearly mode connected (LMC).
Implications Theorem 3.1 states that given any two two-layer models with different random initial- izations, there exists permutation for one model such that applying the permutation makes it lincarly mode connected to the second with high probability, given that the networks are wide enough (i.e. is large enough). In other words, it states that any two randomly initialized two-layer networks are LMC modulo permutation with high likelihood. Entezari er al. (2021) use this result to conjecture that most well-trained neural networks with the same architecture and trained the same task also LMC modulo permutation with high likelihood.
Notably however, permutations only allow for merging across models. We will show how adding the ability to merge within models leads to a tighter bound than Theorem 3.1 with the same likelihood.
G.1.2 A RESTRICTED FORMALIZATION OF MERGING WITHIN MODELS
The Formalization. Let 0, = (v, W) represent a parameter-set such that f(v. w) = fe., and likewise let 6% = (v', W'), s.t. -fv. w = feg. Given 0,, let θ), denote the set of all parameter- sets with functional equivalence 0₁ This means that V0 E O and Vz E {x E RRd| ||æ||2 = Vd}. fo( fea (z). Similarly, let 0,, be the set of all parameter-sets with functional equivalence to 8% Following 0,, let 0, be an arbitrary parameter-ss set for ff which has r hidden units instead. Assume On can be reduced to some 0,, < h in function-invariant manner using the definition of
17

---

Published as a conference paper at ICLR 2024
zero-type neurons from (Simsek et al., 2021). This means that there are h of (1) rows in W that are copies one another whose corresponding v elements sum to O, and (2) some zero-elements in Thus, following Simsek et al. (2021) there exists sunction-and-loss preserving affine transformation that reduces 0, to θ,- We denote this function as Ma-> E Rrxx,, with Mn-r((N)) 0,. Note that when r h, M can simply be the identity transformation.
total combinations
By definition, 0 lies in the expansion manifold of 0, (Simsek et al., 2021). This means there is a similarly defined affine transformation - E R haxr that can expand 0. back to arbitrary 0n E O lying on the expansion manifold. One simple way is to extend O, to O₂ by filling the remaining r v elements with 0 and the h W rows with arbitrary values. Because the associated elements for each W row are zero, the values of each row don't 't impact the function output. Note that because h 2 r, - can assign 6, into arbitrary new indices in 0,. Thus, U,-> act as both a permulation and expansion transformation. Let T U M = Ur- of the reduction and expansion fraffrrtoisso that produce new networks of width from 0,. By definition, any T is permutation when M is the identity and U' is the permutation matrix. For the remainder of this section, we assume that T further contains permutation (i.c. T P U M for some permutation matrix PE R
VU--n(Mn-a (Ma-r (0n)) be the coupling
We will leverage the concept of zero-type neurons presented in (Simsek et al., 2021) to obtain a tighter bound on Theorem 3.1.
A Note on Novelty. While we borrow ideas from Simsek et al. (2021), our Theorem is a differs in theoretical application. First, Simsek et al. (2021) restrict their attention to overall connectivity across points within expansion manifolds. This is important because our Theorem and proof do not require models to lie on the same expansion manifold be linearly mode connected. Second, our models need not be reducible to the same r. That is, we allow for arbitrary reducibility between any two network parameter-sets. Our theorem also differs from Theorem 3.1 in (Entezari et al., 2021) in that we extend function-invarian bounds are achievable in the process. Furthermore, we show that uniformity assumptions may be relaxed while retaining the same bounds (Section G.4).
transformations beyond the permutation matrix, and show that tighter
G.2 A THEORETICAL RESULT
We now introduce Theorem 1, an extension of Theocrem 3.1 that yields strictly tighter bound when the transformations from Section 1.2.2 are included and r < h, and exactly as tight when r = h. We leave the proof to the next section.
Theorem 1. Let (() = (Wa), ((v))))()) = vf ('''( be two fully-connected networks with h hidden units where a(.) is ReLU activation, v R and W E Axxd are the parameters and x E IRd is the input. If each element of W and W' is sampled uniformly from [-1/Vd, 1/vd] and each element of v and is sampled uniformly from [-1/vT, 1/vhi. then for any x E Rd such that ||0|| = vd, with probability transformations T, T' such that
over W, W',v, v,v', there exist
(1- (x) af(v.w) (z) - -11 (1- a- -1-1)((111f))))))) ()),www()()))
0
2041
(r+r) h 0
(1-t)))
(10)
0
otherwise
where v. W are transformed versions of v, W from T and v. W are transformed versions of v', W' from 1" respectively. < r,r' S h are the hidden unit amounts each network can be reduced to via its respective M, .M transformation before being expanded back to width h via P o U, P' o U', where P, P' are permutation matrices.
Implications. Theorem 1 states that when redundancy exists and can be leveraged in network, one can find a transformation that yields strictly lower barrier than with permutation with any h. Moreover, it approaches zero-barrier faster with increase in h compared to permutations. Although it only explicitly holds for random initializations- like Theorem 3.1, this theoretical intuition is supported by our experimental observations. For instance it explains why algorithms like Ziplt!
18

---

Published as a conference paper at ICLR 2024
appear to converge to the ensemble exponentially faster than permutation methods in Figure 6b). The ensemble achieves zero- bero-barrier, and Ziplt! is faster to approach than permutation counterparts because it can reduce models with minimal deduction in performance.
G.3 THEOREM PROOF
We now derive our proposed Theorem 1. Theorem is very similar to Theorem 3.1- we just add the reducibility property from Section G.1.2. Thus, our derivation is nearly identical to their Appendix D proof. We fully derive the novel components of Theorem for clarity, while referring to Appendix D) in (Entezari al., 2021) for the remaining identical derivations to avoid redundancy.
Let O2 = [v. } and 0/,- - [v', W'} respectively as defined in Section G.I. Suppose each can be reduced to some 0,, 02 respectively with S h, via an appropriate Ma (hh-rr,,_-- M transformation. Further, let Uy -r be as defined in Section G.1, expanding M,,r((0,). ( M (0x) back to width by filling all r,
dimensions in V, v' with respectively and all h r, h r' dimensions in W. W with some specific values. Finally, let T, ' be transformations defined in Section G. with T Po M, P o U' M' respectively.
Let Or = (O,)), 0% = T" (0%) be the new parameter-sets obtained from T, T" respectively, and let 0 = (V. W and 0/ (v', W'. By definition, 0, th, and 0¹ E 0%,- From the definitions of T, T", W has r zero v elements and W' has - r' zero- elements. Now, let us suppose that the corresponding h - r rows W are set to copy rows in W', and similarly - r' rows in W rows are set to copy rows in W. Now, interpolating between any non-zero element and a zero-element equivalent to simply scaling the non-zero element: z a0 + (1 aaz (1 x)).. Thus, so long as h (h (h-r) + (h- -r'). we achieve perfect interpolation by placing the zero-elements of 0, and h (h- r) I h r' elements from 0) into the sero-elements of 0% This yields the zero-part of our piece wise bound. However, the proof is more involved for the second 0. We continue the proof for this case below.
elements from 0% into
case when (h r) (h- r') + (r + r²)
First, note that we only need to worry about the (r + r') - h rows in W, W' that cannot necessarily be matched with perfect interpolation as shown above. Let K, and K' be the set of these rows for each network respectively, where |K| |KK| (r + r') h. These are the only rows within the two models that must still considered. For any given 5 > 0, we consider the set Se = (-1/Vā+ - 11/v111/+3300
1/v0- - g] , a discretization of the Rd which has size (2)2) )d) For
any E St. let (W) be the set of indices of rows in K of W that are closest in Euclidean distance to than any other element in Sy:
C. (W)- {ilw; K, s arg min min. [w; s' s'lloo
C. C.(W') = (i)w, E IK', arg m ||w - 10n--8llaax 811100
where for simplicity we assume that arg min returns a single element. These are the same definitions and assumptions as in (Entezari et al., 2021).
Now for every C. C,W and C. (W)).. Whenever |C. (W) * IC, (W'), we will inevitably have unmatched in- dices because permutations only allow 1-1 mappings. Let 1, and I' denote the set of total unmatched indices from K, and IK' respectively. If (W) (C))))C((W)) 2 O. we add these extra indices that are not matched to and otherwise we add them to 7'. Because and KK' are the same size, | = I'l after adding all unmatched indices across all C,. Thus, by definition |/|= |7' S (r +r') h.
S consider a random 1-1 matching (permutation) of elements in
Pick an arbitrary E S. Since each element of IK and K' is sampled uniformly from 1-1/vd, 1/va,, for each row in the respective sets, the probability of being assigned to each s E Se is a multinomial distribution with equal probability for each s: 1/18a/. Pick an arbitrary s. (C,(W) is the sum over all indicator variables W = 1 1[ E C. (W)). }. and the expected value of this sum, EI[C,(W)) [[r + r') kl/|s.1 as there are (r + r') - h total rows. Let (r + r') - = n.. Since each W is between
I Like (Entezari et al., 2021), we choose & such that it is a factor of 1/ 1/va.
19

---

Published as a conference paper at ICLR 2024
[0, 1), we can Hoeffding's 's Inequality to bound the size of CW)) with high probability
-122
P([S, E[S,]] 2 S exxp
(Hoeffding's Inequality) (11)
n
-2t
P([C.(W) - 21) S exx
S. (W. (12)
71
(II.((W) (r r) h 21) S 2eep
-212
(13)
n
P ||C, [IC,(WW] (r + +2)
1 1 S 2exp
-212
V. vn == (r (r) h (14)
|5p]
+ (r)) - h
Let n = (r +') h. By taking a union bound over all elements of K and K', with probability 1 - 6/3, the following holds for all choices of s:
[5
11
log (12) s (W,( IC, (W') S
[5q]
$$ 2
+
log (1215{1/6) (15)
Note this derivation almost exactly follows (Entezari et al., 2021), except that we have h, yielding a tighter size-bound.
Using Eq. (15), we can obtain a bound on the cardinality differences between C, (W) and C. (W') by subtracting the minimum value of one from the maximum value of the other:
[[C((W) - [Cs (WIll S sup(IC, (W)) (nC(
(16)
[[C((W) |C. (W
n
+
n 2
log(12)Se1/8)
[Sg]
log [2)5{{/0))
(17)
[Sc]]
S 2// log(12)({|/6)
(18)
Using Eq. (18) can bound the size of /' with probability 1 6/3 follows:
Σ
IC. -C.
Σ 2
1og(12)5e//6)
(19)
SES
SESg
Σ
(
n 2/5el1
2
(og(12)({|/6)
(20)
SES
-1+ Σ
( -
2/5e1/
2 log(12)S;1/8)
(21)
SES
| = 7' = 2
1 Σ [[C.(W) - (W") S [Se
n log( Se1/8)
1 = [7|
(22)
2
SESt
Note, our Eq. (22 equivalent to Eq. (6) in (Entezari al., 2021), but in terms of n instead of h.
The remainder of our derivation exactly follows (and achieves identical bounds to) (Entezari et al., 2021) until directly after their substitution of g. To avoid writing an identical derivation, refer readers to their derivation following Eq. (6) in their Appendix D until the substitution of &. and instead pick up immediately before the substitution o:
Let e 2 0 denote the value of (fav+(1-a))) ( (I) - aff - (1)-
Following the derivation of (Entezari et al., 2021), we bound e as follows:
|1 //oo-t--aoyr.aaa------ (x) - aflv.w)()) (1 (1-a) (l-a)f(v.w?(())
(4 14
5
2 2log 210g(12/0)log((22))
+ E2d
(23)
20

---

Published as a conference paper at ICLR 2024
Setting & = c/ 4dlog (ogg( gives the following bound on h:
loo 1ogg(1/)))
h
22
1log 1ogg(1/))) V log 12 /5)
Therefore, we have:
(10g (2
log
h K
(24)
e²
4 log Sely log 12 (2 2
d
-sess = {VV
(25)
4 1log(12/8) (12()//8))
d+2
(log(12/4) + dlog(1/c))(n)
(26)
2²
h? <
(1og(12/0) log(12)/0)
(1266) + dlog((//)))
(27)
2
Using the inequality in equation (27), we have
((((rri)))) zdiz).
Thus, obtain the following piece-wise bound over the barrier:
(aa==(la)V (.) (x) a af wy (x) (1 a) a)
af
0
2 (++t"))
otherwise
G.4 UNIFORMITY IS NOT NEEDED: AN EXTENSION OF THEOREM
Although Theorem demonstrates a tighter bound compared to Theorem 3.1 is possible when merging within a model is allowed, its reliance on W v, W v' being uniform random variables unnecessary. Instead, we can assume that W, W' are sampled from an arbitrary probability distribution that bounded on [-1/Vd, 1/Vd.. Similarly, assume that v, v' is sampled from an arbitrary probability distribution that both centered and bounded on [-1/Vh, 1/vh.. Note how for both v, and W any continuous probability distribution is valid, so long as it satisfies the stated conditions. We formalize this as follows:
V'To(W'a be two fully-comnected networks with hidden units where (-)) is ReLU activation, v F and W I Axx are the parameters and x E Rd is the input. If each element of W and W is sampled from an continuous probability distribution that is bounded on |-1/Vd, 1/Vd], and each element of v and v' is sampled from continuous probability distribution that is centered and bounded on [-1/Vh, 11///,1///h for any E Rdd such that ||æ]|2 Vd, with probability transformations T, T" such that
Theorem 1.1.
Let (x) v (Wox, fv.
1/vhi, then
over W, W V, v', there exist
(aa==1-o\W.aWaaa-0 (x) af(v.w) (z) (1 (1-a)f(v.w))
2014
(147))=h)
(r+r) h>0 0
(28)
otherwise
where V. W transformed versions of W from T and v' W are transformed versions of v', W from T' respectively. < r,r" r' < h are the hidden unit amounts each network can be reduced to via its respective M, M' transformation before being expanded back to width h via U, U'.
21

---

Published as a conference paper at ICLR 2024
Proof. The proof for Theorem 1.1 is takes very similar form to Theorem with two differences. For what follows, we assume the same notation and definitions as Section G.3 up to Eq. (11), with one change: each element of K, IK' need not be assigned to each E Sg with equal probability.
Despite this change, for a given 8 E Se, we can use the Hoeffding's Inequality to bound the size of C, (W) with high probability:
-2t²
P(IS, - E|S,]| 2 t) S 2 exp
(29)
-212
P((),(WW) F(IC.(W)]]] 2 S 2exp
(30)
(31)
Despite E|IC, (W)l] no longer being equal for each s, we can take the union bound (1) over the rows of W, W and (2) for each Se to obtain,
|CC, (W) - IC, S (12
(32)
with probability 1 - 8/3. Thus, we achieve the same bound the size of I, I' as in Section G.3.
The second difference is that V is no longer sampled uniformly from 1-1/Vh,1/VA) However, this does not affect anything because we assume E[v] = 0. Noting these changes, we can follow the derivation in Section G.3 and achieve the same bounds.
G.5 SIMPLIPYING VARIABLES IN THEOREM 1
Theorem 1 can be simplified to match its introduction in Section 4. First, let 7. denote the proportions of features from h that are reduced under the M, M' transformation in each model. By definition, T, E [0,1]. 1]. Then, we can define r, and r' in terms of r, T, and h:
r= h(1- -
and r' h( -1)
This means (r + r') h h(1 - < h(1 2 min(T, r"). Let I' min(r, r'), then achieve h(1 21') in the denominator our bound, which simplifies to what is written in Section 4.
H EXPERIMENTS IN SETTINGS OF CONCURRENT WORKS
Table 10 shows the results of merging 16x width ResNet20 models trained with cross-entropy loss on the CIAR5+5 task. This setting is equivalent to that of Table 2 from the concurrent work of (Yamada et al., 2023), except for two important differences. First, Yamada et al. (2023) add the REPAIR Jordan et al. (2022) algorithm to each merging method, which significantly improves the performance of each merging algorithm by adding new parameters to the merged model.
Second, Yamada et al. (2023) include merging methods that require training in their Table 2. In contrast, we report the performance of each merg- ing method using its original capabilities (i.e.. without REPAIR), and without any training (as it is outside our setting). Thus, all results shown Table 10 are a lower-bound to what is achievable either with REPAIR, or with training. To make Ta- ble 10 close to Table (Yamada et al., 2023) as possible, we report "Joint Acc" as the average of each method's s logits for the ensemble. To the best of our knowledge, "Joint Acc" is thus the same metric used by (Yamada et al., 2023). Over- all, observe that ZipIt! melly- outper- forms the nearest baseline by over 20% in "Joint Acc", and zipping up to the classification layers
Accuracies (%)
Method
PLOPS (G)
Joint
Task A.
Tisk B
Ayg
94.4
18.4
56.4.77
Model B
10.88
40.2
22.3
93.8
58.0..
W.
Gii Re Busin 1
10.88
27.3.0
46.3
46.3:13
38.9.17
Permute
10.88
45.4
81.3
79.5 m10
80.4
Zipitiaro
86.7
07.4
Ensemble
50.0234
944401
Zipithecs
10.89
72.0
90.0
88.2
19.
zapithuss
91.5
89.6
90.5-15
Ziplthice
25.5
92.7.1
90.733
91.7
Table 10: CIFAR-10 (5+5) Cross Entropy. Zipit! vs. baselines using ResNet-20 (16x width). Merging with ZipIt! up to the last layer (ZipIt!19/20) nearly achieves the ensemble "Task Avg." performance with half the FLOPPs and vastly outperforms the nearest baseline. Partially merging, brings Zipit! even closer to the ensem- ble. refers to (Ainsworth et al., 2022)
22

---

Published as a conference paper at ICLR 2024
19/20) nearly matches the ensemble "Task Avg." accuracy without requiring any training. Interestingly, Git Re-Basin performs especially poorly in this setting, likely requiring REPAIR to achieve the performance reported in Table 2 by (Yamada et al., 2023).
23


Reviews:
[{"review_id": "IfPlRp6a5U", "reviewer": "Reviewer_3tA9", "summary": "The paper proposes a method to merge two different models sharing the same architecture, possibly with different initialization and for different tasks, without re-training. The key idea is to merge the weights whose outputs are highly correlated, both within and across models. The method also allows weights to be partially merged, adding flexibility to adjust the amount of merging depending on resource budget and accuracy requirements. Experiments show that ZipIt! outperforms the existing baselines in several settings, including merging classification models and merging across tasks.", "questions": "1. The meaning of \"ensemble\" in the experiments is unclear to me. I don't think I understood how it works.\n\n2. I would appreciate more experiments to test the stability of the method. It would also be nice to provide an analysis of when the method works and when it does not. Please see the weakness section for more details.\n\n3. It would also be interesting to include additional experiments on more compact architectures such as Mobilenet for a more comprehensive analysis of the behavior of the method. It is great that ZipIt! is already evaluated on models with larger widths (x1.5 in Table 7). I think it might be harder to merge models where the architectures are already compact because there might be less overlap and less redundancy of features.\n\n4. Can BN be removed before applying ZipIt by fusing it into Conv (using that BN is an affine transform when mean and variance are fixed)? I think then the process to reset BN after merge can be avoided.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. This paper studies a new and interesting problem: merging differently initialized models trained on different tasks into a single model without retraining. This relaxes the condition of the model merging problem, recently studied in works such as Model Soups and Git Re-Basin, which assumes the same task over models. This makes the problem more challenging while allowing for broader use cases.\n\n2. The idea of merging weights both within and across models is simple, but achieves reasonably good performance under different conditions in the experiments. It's great that experiments with merging tasks of different modalities (classification, segmentation) are also done and the method still achieves reasonable performances, showing that the method is generalizable to non-classification tasks. In addition, the ability to flexibly adjust the amount of merging is also a strong point.\n\n3. The paper is well written. It was easy to follow. It is also great that the code is already publicly available.", "weaknesses": "1. My main concern is the stability of the method. I am not sure if the method works robustly across different settings, because empirical observations are reported in the github issues that a small change (e.g. changing the learning rate and epoch numbers when training the models to merge) can cause the merged model to crash to low accuracy.\n\n2. These are not weaknesses but there are some limitations. The performance drop increase when merging multiple models. Accuracy drop looks large when merging models of different tasks (Table 8). All experiments used CNNs and transformers are not tested.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JEMoVghuwO", "reviewer": "Reviewer_pbds", "summary": "This paper proposes a novel method for merging two different models trained from different tasks, maintaining the performance on both tasks. The method resolves the problem of decreasing performance on two disjoint tasks, where prior works fail, ﻿which ﻿permutes the model to another one and the average of them. Moreover, this paper proposed partially zipping to get the trade-off between the performance and the computational complexity.", "questions": "- I’ve noticed that the “permute” baseline is stronger than “Git Re-Basin”, which is a little strange. I’m wondering how it happens and what’s the experimental setting?\n- The improvement on CIFAR is remarkable, while in the larger dataset ImageNet, the trade-off between the performance and FLOPs is debatable. Compared to the ensemble, the FLOPs don’t decrease a lot while getting lower accuracy.\n- I’m curious about the comparison between “ZipIt!” method and a single model trained on both datasets. Did you think about this?\n- The presentation for “Ensemble” results which uses light grey is a little misleading.\n- Are there the results of Git Re-Basin method for ﻿Multi-Dataset in Table 3?\n- Did you consider the time cost of computing transformation matrices in your method when comparing the FLOPs?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "This paper proposes a novel method named “ZipIt!” which can maintain the performance after merging two models trained from two disjoint tasks, which prior works fail. The partially zipping method can further give the option to do a trade-off between the performance and the FLOPs. Finally, a theoretical proof is given to ensure the existence of a transformation matrix.", "weaknesses": "The method seems not to be practical for real-world usage, while lacks insights and theoretical analysis of the model properties. The method works fine in the CIFAR dataset, but doesn’t have remarkable advantages in large datasets such as ImageNet, and disjoint tasks. Moreover, the paper lacks comparisons between its method and the model trained on both datasets directly, which limits its usage scenario.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZJgJY3PCwk", "reviewer": "Reviewer_GcCE", "summary": "This paper proposes a method for merging trained models trained on different datasets without additional training based on feature correlations. The proposed method can further adjust the inference cost and accuracy by adjusting the number of layers to be merged.", "questions": "1. Is git-rebasin activation matching considered equivalent to zipit without wighin merge? I would like to know the difference between git rebasin, permute and zipit.\n 2. Does merge by zipit keep the output strictly invariant like permutation, even when using pseudoinverse matrix?\n 3. What is the same model budget and explain how beta is introduced as an equation?\n 4. Three matching methods were proposed in git-rebasin: activation matching, weight matching, and STE. Which method is the baseline in Table 1 and 2? I doubt the claim that permutation alone makes it difficult to merge models on different datasets. According to [1], using STE, which is a permutation base, and Merge models with cifar10(5+5), it is reported that 90% accuracy is achieved, which is 10% more accurate than zipit. I think the conditions are the same as zipit!20/20 in terms of combining all layers.\n 5. Sec 5.1 says \" `If allow the last stage of the network to remain unzipped (i.e., zip up to 13 layers), our method obtains 83.8%, which is only 3.6% behind an ensemble of model A and model B (which is practically the upper bound for this setting)`\", but how can ensemble be guaranteed to be uppper bound?  In [1], a merging model is reported that can achieve higher accuracy than an ensemble. The upper limit of joint accuracy is the accuracy of the model trained on all data, rather than the accuracy of the ensemble.\n 6. Why is partially zipping more accurate than fully zipping?\n 7. How much gpu is needed for Imagenet-1k (200+200)? I am concerned about the memory and computational cost required to compute the correlation matrix.  \n[1]: https://arxiv.org/abs/2306.05641\n\n\n\n\n===========================================\n\nComments after reading the rebuttal.\n\n===========================================\n\nSince I was unable to set the readers to everyone in the reply, I am writing my comments here.\n\nThank you for your kind reply.\nI did not understand that merging models using loss like STE is outside the scope of this study.\nSince this has been resolved in regards to many of my questions, I would like to raise my score.\n\nAs other reviwers have pointed out, it would be helpful to specify as a limitation in the camera-ready version of the paper that the model must be wide enough to perform.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. This paper addresses the interesting research topic of merging models on different datasets without additional training.\n2. Large-scale experiments such as imagenet1k confirm the effectiveness of the proposed method. This paper conducts experiments in multi-dataset merging, classification and segmentation merging. In particular, this is the first effort, to my knowledge, to merge the different modalities of classification and segmentation.\n3. For reproducibility, the authors mentioned that they plan to publish the source code; the cost of replication experiments would go down significantly if the checkpoints of the trained models were made public since they involve large scale experiments such as cifar100, imagenet1k, and multidataset.", "weaknesses": "1. The paper extends feature matching between models to within the model in order to merging models across different datasets. Its contribution should be examined in more detail. In particular, there is no theoretical support as to why matching the feature would lower the loss after merging. I did not understand the connection between Theorem 1 and the fact that joint Acc can be sufficiently large for merging models on different datasets. Is the extension to \"within merge\" not only more flexible in alignment, but also suitable for merging models on different datasets?\n2. The proposed method is inaccurate for meging of imagenet1k and multi dataset unless partial zip is used, while the inference cost increases as the number of models increases with partial zip.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xpKV3d1X1L", "reviewer": "Reviewer_SEiC", "summary": "This paper tackles the problem of merging two models trained on different \"tasks\". Their best performing approach is two-fold: 1) identify redundant features within/across models and average those, and 2) only merge some layers. Combined they achieve performance which approaches the ensemble, but their improvements are most notable when the two \"tasks\" come from the same task, such as partitioning CIFAR in two parts.", "questions": "- I would be curious when the authors believe this approach is best applied, e.g., when would I find myself training from scratch on two different tasks?\n- Does this approach improve accuracy in more traditional model merging settings? It's pitched for multi-task but is it more general?\n- The \"scaling\" trend in Figure 6b is very interesting! I would be curious to hear if the authors would expect the same trend with a different method to increase FLOPs on the x-axis, such as depth scaling, or joint scaling, or training iterations.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "- The presentation, including all Figures, is very clear!\n- The idea of merging redundant features is interesting and effective. It seems to me that the authors should consider exploring this as a general way of speeding up inference and not just as a technique for merging models.", "weaknesses": "- I would appreciate more motivation for this problem. When would I find myself training from scratch on two different tasks?\n- [Subjective] I suggest the authors are much more clear/transparent about their contributions early, because currently the readers get their hopes up a lot when in reality the best performance occurs either when the \"tasks\" really come from the same task like splitting CIFAR in half. I think there's still a long way to go for training individually on two distinct datasets that are ImageNet level difficulty.\n- Figure 1a/b is not exactly clear with respect to related work, for instance [1] (Sec 3.3), and [2] (Appendix E) do consider merging models trained on different tasks, albeit from the same initialization.\n\n[1] Merging Models with Fisher-Weighted Averaging (https://arxiv.org/abs/2111.09832)\n[2] Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time (https://arxiv.org/abs/2203.05482)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)