Transformer Fusion with Optimal Transport
Authors: Moritz Imfeld
Jacopo Graldi
Marco Giordano
Thomas Hofmann
Sotiris Anagnostidis
Sidak Pal Singh
Venue: ICLR 2024 poster (2024)
Keywords: Fusion
Transformers
Model Merging
Primary area: general machine learning (i.e., none of the above)
TL;DR: Fusing Transformers with Optimal Transport

Abstract: Fusion is a technique for merging multiple independently-trained neural networks in order to combine their capabilities. Past attempts have been restricted to the case of fully-connected, convolutional, and residual networks. This paper presents a systematic approach for fusing two or more transformer-based networks exploiting Optimal Transport to (soft-)align the various architectural components. We flesh out an abstraction for layer alignment, that can generalize to arbitrary architectures -- in principle -- and we apply this to the key ingredients of Transformers such as multi-head self-attention, layer-normalization, and residual connections, and we discuss how to handle them via various ablation studies. Furthermore, our method allows the fusion of models of different sizes (heterogeneous fusion), providing a new and efficient way to compress Transformers. The proposed approach is evaluated on both image classification tasks via Vision Transformer and natural language modeling tasks using BERT. Our approach consistently outperforms vanilla fusion, and, after a surprisingly short finetuning, also outperforms the individual converged parent models.
In our analysis, we uncover intriguing insights about the significant role of soft alignment in the case of Transformers. Our results showcase the potential of fusing multiple Transformers, thus compounding their expertise, in the budding paradigm of model fusion and recombination. Code is available at https://github.com/graldij/transformer-fusion.

Published as a conference paper at ICLR 2024
TRANSFORMER FUSION WITH OPTIMAL TRANSPORT
Moritz Imfeld", Jacopo Graldi", Marco Giordano", Thomas Hofmann, Sotiris Anagnostidis, Sidak Pal Singh ETH Zurich, Switzerland {noimfeld, graldij, ngiordano/@ethh.co
ABSTRACT
Fusion a technique for merging multiple independently-traiee neural networks in order to combine their capabilities. Past attempts have been restricted to the case of fully-connected, convolutional, and residual networks. This paper presents a systematic approach for fusing two more transformer-ban Optimal Transport to (soft-)align the various architectural components, We flesh out abstraction for layer alignment, that can generalize to arbitrary architectures in principle and we apply this to the key ingredients of Transformers such as multi-head self-attention, layer-normalee noyer-oommalization, and residual connections, and we discuss how to handle them via various ablation studies. Furthermore, our method allows the fusion of models of different sizes (heterogeneous fusion), providing new and efficient way compress Transformers. The proposed approach is evaluated on both image classification tasks via Vision Transformer and natural language modeling tasks using BERT. Our approach consistently outperforms vanilla fusion, and, after a surprisingly short finetuning, also outperforms the individual converged parent models. In our analysis, uncover intriguing insights about the significant role of soft alignment in the case of Transformers. Our results showcase the potential of fusing multiple Transformers, thus compounding their expertise, in the budding paradigm of model fusion and recombination. Code is available https: //g /ggttub.... con/ graldi j former fusion.
networks exploiting
INTRODUCTION
Transformers, as introduced by Vaswani al. (2017), have profoundly impacted machine learning, establishing prevailing neural network architecture across various domains. Transformers consis- tently excel different fields, including natural language processing (Lin et al., 2022), time series forecasting (Wen et al., 2022), and computer vision (Dosovitskiy al., 2020). Their success can be attributed to their scaling properties (Kaplan et al., 2020) and efficient utilization of contempo- rary hardware architectures designed for extensive parallel computing. The unification of a single architecture across tasks facilitates immediate, far-reaching applicability of any analysis that handles general properties of the Transformer architecture.
As large Transformer foundation models (Bommasani et al., 2021) continue to grow in size and complexity, the challenges associated with training, i.e., exponential increase in parameters and compute for a fixed incremental improvement in performance (Hoffmann et al., 2022: Zhai et al., 2022; Bachmann et al., 2023), become increasingly more perilous. Consequently, achieving state-of- the-art results is often confined to researchers with access to ample GPU resources. To address these issues and strive for more efficient and sustainable performance improvements, we embark on the following more compelling and alternative inquiry:
Can we combine the capabilities of pre-trained Transformer models?
Merging multiple Transformer models into single entity while preserving their unique capabilities can yield several advantages; (a) Enhanced performance by harnessing the collective capabilities of individual models. (b) Reduced inference complexity, querying single model replaces the need to query n models in an ensemble, reducing computational (FLOPs) and storage requirements by
These authors contributed equally this work

---

Published as a conference paper at ICLR 2024
factor of n. (c) The necessity train from scratch can be readily eliminated, leveraging existing public models, already available, and numerous in quantity
A straightforward way of fusing, i..., merging. models of the same architecture, is average their weight matrices one-to-one, referred to as 'Vanilla Fusion' (VF). However, this method overlooks potential misalignments between the parameter matrices, arising due to neurons at the same positions, in different models, encoding different information (Godfrey et al., 2022). Instead, we propose to use Optimal Transport fusion (OTFusion) (Singh & Jaggi, 2020), which at its core, aligns the weight or parameter matrices before fusing them.
Thus, by virtue of such alignment, OTFusion ensures that the fused model effectively integrates the knowledge and capabilities of the individual models to be merged, rather than simply averaging the weight matrices without guaranteeing meaningful information preservation. Additionally, OTFusion accommodates the fusion of models with different widths, and in turn, different sizes, which is fundamentally not possible with VF. This is a crucial feature, as such heterogeneous models are available in plenty, to better unleash the potential of existing pre-trained models. Consequently, OTFusion has been shown to be an effective method for fusing fully connected (Singh Jaggi, 2020), convolutional (Nguyen et al., 2021) and recurrent neural networks (Akash et al., 2022) on a variety of tasks, heavily outperforming VF.
Yet, despite its wide adoption (Nguyen et al., 2021; Liu et al., 2022; Ainsworth et al., 2022), the layerwise procedure proposed by OTFusion does not fit well with contemporary architectural design, that comprises of constant residual streams, normalization layers, and attention operations. It is not equipped in any way to align and fuse models with complex information streams and to fuse transformer-speciin components. Hence, the primary aim of our work is to develop techniques that help bridge these gaps and successfully generalize fusion to Transformer-basee architectures.
Our contributions are: (a) We analyze each of the idiosyncratic architectural components in Trans- formers in thorough detail, with an ultimate aim to best fuse them across different models. Throughout our discussion, we exposit our approach based on the perspective of flow of the transportation maps", that makes for intuitive visualizations and interpretation. (b) We uncover that, surprisingly, OTFusion based on hard-alignment underperforms in this context, contrary to the case of fully colle-con or convolutional architectures; and that, soft- poft-alignment plays key role in successful one-shot fusion. (c) We showcase the efficacy our approach by extensive experimentation involving the fusion and finetuning of Vision Transformers (ViT's) across multiple datasets, including CIFAR 10, CIFAR 00, TINY IMAGENET and IMAGENET-IK, 1K, as well as BERT (Devlin et al., 2018) models for natural language tasks. We consistently outperform the original converged models across tasks and datasets, by about ~ 1.0%, while significantly reducing computational and storage costs by a factor of n.
Overall, research marks an important stride in advancing model fusion techniques, that help deliver enhanced performance and efficiency for modern Transformer based architectures.
2 RELATED WORK
Model combination and ensembling. The combination of multiple models has been a timeless idea in machine learning, from classical works on bagging and boosting (Breiman, 1996) to more contemporary approaches (Mienye & Sun, 2022; Garipov et al., 2018; Allicoeur-Martinen et al., 2023). The key idea behind these works is to boost model performance, by capitalizing on the unique strengths of each model while mitigating their individual limitations. Or, more technically, one can think of model combination as a way of reducing the variance of the predictors (Geman et al., 1992). However, the main limitation is that such methods require the execution of each (parent) model for the final prediction, with a cost that scales linearly with the number of models.
Model Fusion. Model fusion (Singh & Jaggi, 2020; Wang et al., 2020; Wortsman et al., 2022; Matena & Raffel, 2022; Ainsworth et al., 2022; Nguyen et al., 2023) has emerged as a particularly notable direction in recent years, gaining significant traction in the machine-learning community. This line of work focuses on building better model combination approaches that account for the
On huggingface there are more than 339,000 models available as of the 22nd of Septemher 2023. Thii should he reminiscent of the flow of tensors in the computation graph of neural networks, and thus allows one general strategy that can be potentially he adapted for any architecture type.
2

---

Published as a conference paper at ICLR 2024
network structure and its inherent symmetries. We elaborate on some of these works, which are more relevant to the focus of our paper, below.
Singh & Jaggi (2020) propose novel approach based on the OT theory exploiting the Wasserstein distance, where the neuron association allows fusing pre-existing models with the same depth in a one-shot fashion, thus without requiring retraining. OTFusion outperforms VF and was successfully used for model compression and fusion of CNNs, residual networks (ResNets), and multilayer perceptrons (MLPs). Since its publication, OTFusion has been extended in various ways. Nguyen et al. (2021) address the same-depth requirement of OTFusion. Liu et al. (2022) generalized the work as a
task, and taking into account the second-order similarity of model weights instead of linear alignment. Recent efforts on the topic have shown theoretical insights on fusion, extensions of previous algorithms to new network topologies, in particular, Akash et al. (2022) adapted FTFusio for recurrent networks, such as RNNs and LSTMs. Further, Stoica et al. (2023) propose an algorithm, for convolutional and residual architectures, that aims at finding redundant features within the same model and across the different models to be fused, as to keep only meaningful and unique features in the fused model.
However, the fully layerwise interpretation of OTFusion (Singh & Jaggi, 2020) is currently only applicable to simple architectures such as MLPs, CNNs, and instances ResNet. is not equipped in any way to align and fuse models with complex information streams and to fuse cunsformer-specciee components such as multi-head attention layers, layer-normalization, embeddings, or the sequential nature of the data.
Fusion with a focus on Transformers. Wortsman et al. (2022), in their approach of 'model soups', consider fusing transformer models that have common backbone network that is pre-trained on the same dataset, but that are fine-tuned, say, with different hyperparameters. Owing to this, the models remain sufficiently close in the parameter space, which precludes the need align them, and lets them employ just vanilla fusion (one-to-one averaging of the parameters) while still obtaining gain in performance. Therefore, despite apparent practical gains, the 'model soup' approach is actually a poor representative of the complexity and intricacies of the general model fusion problem.
Arguably, the more empowering capability is to fuse transformer networks that are potentially much more distant in their parameter spaces and are diverse in nature. For instance, this arises when the networks have different initializations, or see examples in different batch orderings, or when they have different sizes, and more. This specific problem is tackled in this work, which is, to the best of our knowledge, the first aiming at fusing transformer architectures by aligning their weights.
The conjecture of Linear Mode Connectivity (LMC) modulo permutations. Given the recent interest around this conjecture posed in Entezari et al. (2021) and its wider demonstrations (Ainsworth et al., 2022), would like to make few clarifications: (a) The LMC barrier approaches zero only at very high widths, even for non- non-transformer architectures, see for instance Figure 4 of Ainsworth et al. (2022), and importantly, not for any arbitrary width. Thus, for typically sized residual or convolutional neural networks, the LMC barrier in loss is not zero at all, and the corresponding barrier when measured in accuracy is even more palpable.
(b) Transformers possess a more non-convex landscape, as shown by Park & Kim (2022) in a comparison of vision transformers with residual networks, which consequently brings about higher LMC barriers. This can also be seen due to the fact that transformers contain components which further proliferate the number of symmetries, such as within- and across-head permutations as well as the translation invariance of softmax, - all of which serve to interfere the linear interpolation of parameters. Thus, the barriers in (Singh & Jaggi, 2020; Ainsworth et al., 2022) non-transformer architectures do not reveal the full nature of the underlying problem being addressed here.
3 BACKGROUND
Optimal Transport (OT), OT (Villani et al., 2009) has gained prominence in machine learning for its ability to compare probability distributions effectively, with applications in generative mod- elling (Arjovsky et al., 2017), class incremental learning (Zhou et al., 2021) and model compression (Li et al., 2021). At its heart, OT aims to find a transport map (TM) T signifying how much of a discrete source distribution should be moved towards discrete destination distribution to align the two. This alignment can be hard (T is a permutation matrix and the solution to the Earth-Mover's Distance, EMD, (Rubner et al., 2000) problem) or be relaxed yielding a soft alignment (solved
3

---

Published as a conference paper at ICLR 2024
with the Sinkhorm Knapp algorithm (Knight, 2008)). The softness of the alignment is controlled by regularization parameter Asinkhern. where lower values result in harder alignment. More details about OT can be found in the Appendix A.I.
Singh & Jaggi (2020) apply this theory to align networks a layerwise fashion, using either weights activations underlying distributions. After the alignment of one more models to an anchor model, these are then averaged. Formally, for a layer of the model, the transpose of the TM of the previous layer is pre-multiplied with the weight matrix of the current layer: ( + T( w(e,-1), The current layer can then be aligned by post-multiplinng with the TM of the current layer: W((,(-1) + ((., Ainsworth et al. (2022) propose highly similar approach which, in certain cases, effectively boils down to the same linear programming problem that uncovers (provably and practically) same alignments as OTFusion; thus we continue to base our approach on OTFusion henceforth.
OTFusion.
4 METHODOLOGY AND IMPLEMENTATION
With modular architecture like the transformer, is intuitive to develop a fusion algorithm. Therefore, we first divide the architecture into its simplest building block fully connected layers that can be fused by the prevalent OTFusion strategy. The question remains; how to effectively connect these building blocks, especially heterogeneous? How to hierarchically reconstruct fully fused transformer ensuring consistency of the single fused blocks?
hivide-and-conquer approach
As we provide solutions to such open questions, we will guide our discussion in this section with transport flow perspective, which allows for intuitive and effective concatenation of blocks of any sort, and that, therefore, in principle can be applied to every architecture. Henceforth, we will use the notation from Vaswani al. (2017) for Transformers. We display our methods in the non non-masked self-attention case, but our method can generalize to the cross-attention or causal masked attention.
4.1 TRANSPORTATION MAP FLOW GRAPH
In the typical OTFusion application, the TM of the previous layer simply passed to the next layer. However, in more complex architectures, the incoming TM of a layer can depend on multiple TMs. To formalize and visualize this flow of TMs, we present the Transportation Map Flow Graph.
To introduce the concept, we use the flow graph of residual connection (Fig. 1). Rectangles represent the neural network layers; red nodes represent any non-learnable computations or permutations inside the network; edges represent the propagation of the TMs. Layers have exactly one incoming and one outgoing edge. Computation nodes always have multiple incoming edges and one outgoing edge, where the outgoing T'M must depend the incoming TMss A major contribution of this work is to handle the various complex transportation map flows throughout the transformer architecture.
4.2 TRANSFORMER FUSION
4.2.1 RESIDUAL CONNECTIONS
In residual connections, the outputs of a current layer and a residual layer are summed up. The TMs coming from these two layers will be different, therefore the ideal TM flow strategy has to be determined. We ex- plored three heuristics calculate a weighting vector (()), where each entry 7/- scales the corresponding rows of the TMs. After obtaining (() we compute the weighted average as shown in Eq. 1. Find the
Residual Connection
cumeett
residual
Figure
TM flow graph for a residual con- nection.
results in Sec. 5.1.
( out current diag(1 - (0)) + (e)
(f))
residual diag(>(())
Averaging. For plain averaging, proposed by Singh & Jaggi (2020), set Vi, 7/4 0.5. This heuristic does not depend on activations and can therefore used even the case of weight-based alignment. However, it introduces the strict assumption that the residual and the current layer TM
4

---

Published as a conference paper at ICLR 2024
of equal importance when aligning the subsequent layer. We therefore extend Singh & Jaggi (2020)
with two novel residual policies.
Weighted Scalar. To alleviate the equal contribution constraint from the averaging method,
compute a weighting factor Vi, 7. =
"scclar (Eq. 2). We use the activations of the anchor model, over batch of samples S, because only those carry information about the importance of the current p(()
and the residual branch in the anchor model to which we try to align the other models. Fresduaal (x) are the activations from the residual branch while Eeurent
(x) are the activations from the current layer L.
Exes (t())
(x)]|1
seallar
(e)
ExES
|.(() S |.curren (x)l Exes
[|((e) Fessduut (x)]l
(2)
Weighted Matrix. As opposed the Weighted Scalar method, here, calculate weight vector ((()) where each entry 7 weighs one strand of a residual connection. The computation of each (()) similar to Eq. 2 but here do not compute the e'-Norm over the whole activation vectors, instead, we take the absolute value of the corresponding i-th values of the activation vectors.
We note that Ainsworth al. (2022) propose propagate either the identity (Tout I) the residual transportation map itself (4i, 7/ = 1). In the case of hard alignment, these methods perform worse than averaging.
4.2.2 MULTI-HEAD ATTENTION
The attention mechanism (Eq. 3) poses multiple challenges when it comes to TM flow (Fig. 2): what are the incoming TMs for WQ, W and WV? Which TM is propagated to W How to handle attention with multiple heads?
QKT
Self-Attention(x) softmax(
VV,
with KQ ((K.VV)
Vdl
The first challenge is conveniently solved by the TM flow graph. We can simply the TM from the previous layer for each W W and W.V This even holds true for multiple heads. The incoming TM of W is more complex obtain because it depends on the outgoing TMs of Wº, W and WV. However, if we constrain both TMs of W and WQ to be equal permutation matrices (i.e., hard alignment with TQ = Tx TQk), we show that the permutation matrices cancel (see Eq. 4) leaving the softmax undisturbed. Therefore, we only propagate the outgoing TM of WV Wo
For soft-alignment Eq. 4 no longer holds, in that case we investigated alleviating the constraint of equal TMs for W and Wº, Removing this constraint slightly increased one-shot accuracy.
Q = QTQR and K = KTQK and QK T
QTQK T T
K
QK
(4)
For multi-head attention fusion, there additional layer of complexity be- cause one must align the weights and the heads. On top of that, there is no guarantee that hard one-to-one align- ment between heads exists. For that reason, propose head align- ment. During cross-head alignment, W W.K and WY (where is the head index) concatenated across the output dimension form three combined weight matrices (WQ, W and WV), ). OTFuuion is then applied to each of the concatenated weight matrices. Finally, Ty is propagated wo, Find visualization of our cross-head alignment method in App.
Self-Attention
Tn
Tin
w
Tk
Tin
Tv
Tw
T'n..
AA
Ti
To
w
W
Figure 2: Self-Attention flow graph.
4.2.3 LAYER NORMALIZATION, EMBEDDINGS AND BIAS
The layer normalization is learnable neural network parameter and consequently must fused. contains only two parameters (a and (3) per input and there are no interconnections between different inputs and outputs. Therefore, no TM has to computed for this layer. The parameters are only aligned w.r.t. to the incoming TM. The incoming TM is then propagated to the subsequent layer.
5

---

Published as a conference paper at ICLR 2024
Embeddings
patch
pos
The ViT embeddings fusion approach is most effectively conveyed by its TM flow graph, as depicted in Fig. 3. For the concatenation, we notice that the class token is only small fraction of the full sequence, in other words, for the integrity of the sequence, is far more important to propagate the TM of the patch embeddings than the one for the class token. After concatenation, the positional embeddings are added. We notice that the addition is the same operation for residual connections, so we can use one of the three TM flow strategies from Sec. 4.2.1.
els
patch
Figure 3: ViT embed- dings flow graph.
The bias is only connected the output of neural network layer, so we align it using the outgoing TM of the corresponding layer.
4.3 ALIGNMENT STRATEGIES
Soft vs Hard Alignment. OTFusion technically allows soft alignment for MLPs, CNNs and ResNets, but Singh Jaggi (2020) discovered that for these simpler architectures, hard aligmment outperforms soft alignment. However, we do not want to limit the search space for optimal alignment to only permutation matrices (possibly too constraining for a complex architecture such Transformers). We, therefore, broaden the perspective algorithm and tuning the softness of the TM by optimizing over the Sinkhorn regularizer, discovering that soft alignment outperforms hard alignment for Transformers.
alignment introduced by OTFusion using the Sinkhorn
Weights vs. activations aligmment. The combined methodology introduced so far, and the novel perspective on the TM flow, allow to apply OTFusion to the single fully connected layers without further adaptations in the case of weight-based approach, while the activation-based strategy needs a bit more thought. Transformers operate on sequences of tokens as opposed simpler architectures that only operate one token at a time. In our activations-bas algorithm, we treat every token of the sequence as a possible activation.
Sequence Filtering, For ViTs, it is obvious that not every token contributes equally to the final image classification. We hypothesize that activations- alignment performs best if only the most important tokens of a sequence are considered. Therefore, we explored filtering out unimportant tokens. For datasets where images are centered, we propose window filtering, where only the by a center patches are considered as activations for activations- alignment (wi nwinddww n). Additionally, we explored using only the class token for activation-based alignment yocly
5 EXPERIMENTS AND RESULTS
We evaluate the quality of our approach with two prominent mensformer-bbaeedd architectures: the ViT (Dosovitskiy et al., 2020) and BERT (Devlin et al., 2018). Our focus is to assess the performance and robustness of our proposed fusion techniques in both image and NLP domains. These models offer direct comparison as they share the same encoder-only architecture. We conducted our experiments on multiple well-known image classification datasets: CIFAR 10, CIFAR 100, TINY IMAGENET, and IMAGENET-IE K. We used Hugging Face both for the implementation of the ViT and for retrieving the datasets. Besides the image classification tasks, we showcase our fusion strategy the BERT model for an NLP task. We train from scratch multiple BERT models on the masked language modeling (MLM) task over a subset of the Wikipedia dataset, publicly available on the Hugging Face Hub.
Model Training. First, we train individual models from scratch on each dataset until convergence. We ensure model diversity by initializing each model with different seed values and different batch randomization. This results in unique models with similar performance but located in diverse parts of the landscape, and whose suitable fusion can improve performance. These diverse models, which are rather distant in the parameter space, need non-trivial alignment strategy to be successfully fused, and therefore exhibit a dramatic drop in performance when fused with a naive approach such as VF. This approximates plethora of other scenarios (c.g, models trained on different (sub)datasets Details and training parameters of all models can be found in Appendix C.
Model Fusion. We assessed the proposed fusion strategies, and their combination thereof, on the CIFARIO 10 dataset (refer to the ablation studies in Section 5.1). We measure the performance through the so-called ohe- capability, namely the performance of the fused model, without any retraining, on the same task and metric of the parents. This capability is the first important proxy of the capacity of the fusion algorithm to align and then fuse the parent models. The optimal fusion

---

Published as a conference paper at ICLR 2024
strategy identified the CIFARIO task is then applied to the other tasks and architectures. For each task and alignment strategy (i.e. weights- based and activations-bast based) we optimize the Sinkhorn regularizer separately (see Fig. 11). The fusion step runs in just seconds
purepse CPU.
Finetuning. Besides the one-shot performance, similar to Singh & Jaggi (2020); Nguyen et al. (2021), evaluate the effect of finetuning the fused model. The resulting performance is compared against the single parent models at convergence (and thus do not benefit from finetuning), their ensembling, and the VF model that also went through a round of finetuning. Both our fused model and the VF model are optimized separately over a common set of reasonable hyperparameters.
Note. We encode the model dimension (hidden-layer dimensiion/mterrmeeeeii sion/nummer of encoders). Additionally, we report the relative computational burden (latency and FLOPs) below each result table entry.
dimen-
5.1 ONE-SHOT EXPERIMENTS
1718
We optimize the fusion strategy on CIFAR1O, searching the configurations previously introduced. In contrast to the observations of Singh & Jaggi (2020) with non-transformer architectures, we observe that a soft-alignment (Sinkhorn) strategy consistently outperforms hard-alignment (EMD). The value the Sinkhorn regularizer is chosen to maximize the one-shot accuracy (separately for activations- and weights-basee alignment). The optimal strategy for handling the residual connections has proven be the averaging policy. Activations- alignment with the 6x6 window filtering (window-6) approach performs best among other filtering strategies and weights-basee based alignment.
62
Figure 4: 2D slice of the accuracy landscapes of the anchor and one-shot OT and VF fused models.
In Tab. we present the one-shor performance for the best configuration of fusion with the weights- based alignment and the activations- based aligmment, both in the scenario with two models and with five models together. VF dramatically drops at random accuracy, while our fusion methodologies are able to preserve most of the capabilities of the individual models. In particular, we achieve the best accuracy with our soft, activations- -basdd fusion.
Fig. 4 visualizes two-dimensional slice of the accuracy landscapes of the anchor model and the two fused models, OT and VF. The visualization based on the procedure outlined in (Garipov et al., 2018). The plot shows the OT model being in the same basin the anchor one, while the VF model is separated by a barrier from such basin. This representation effectively underscores the superior performance of our algorithm in comparison to VF, emphasizing its ability to facilitate more dependable knowledge transfer.
Table 1: One-shot accuracies on CIFARIO for the individual parent models, VF, weights-based soft-aligmment fusion (Asinkhorn = 0.06), activations-based soft alignment (Asinktom = 0.08) fusion, and activations- based hard-alignment (EMD) fusion. Activations- based reported with mean and standard deviations over different random seeds. For the best-performing method, we show the
absolute increase over VF.
DATASET
INDIVIDUAL MODELS
VF
OT-WTS (OURS)
OT-ACTS (OURS)
OT-ACTS EMD (OURS)
GAIN OVER VF
CIFAR1O 192.34, 92.31]
7.59
57.23
60.87 + 0.44
24.50 + 5.66
+53.28
CIFARIO [92.34, 92.31, 92.28,
9.47
44.46
46.56
0.71
43.28 + 2.81
+37.09
92.04. 91.47]
Ablation Studies. We study the effect of the different OTFusion hyperparameter choices on the one-shot performance on the CIFAR 10 dataset for two-models fusion. We find that soft alignment (Sinkhorn) outperforms hard alignment (EMD) (see Fig. 5a). We attribute this observation to the flexibility of soft alignment which better accommodates the highly complex nature of the transformer, as multi-head self attentonnn We observe a bell-shaped curve with maximum for a non-zero regularization, thus demonstrating that the optimal alignment is neither hard nor merely soft. We can
7

---

Published as a conference paper at ICLR 2024
therefore optimize this parameter with an inexpensive sweep. Furthermore, as shown in Fig. 5b, the soft alignment for the activations-bart fusion is much more stable than hard alignment (EMD) for different seeds of data, suggesting that hard alignment is much more impacted by the activations.
DT-ACTS KARBORY DIARTS SIAHBOR
Actoracy 1%
SPARROON
One-diot
1.:0
Sinkhoor Regularizrr
Nombbr Samplle
(a)
(b)
Omn-soo Accurney 750
One-doo
Average
125
Samplee
Nomber Samples
(c)
(d)
Figure 5: (a) Sinkhorn regularizer effect on one-shol performance; (b) stability with different seeds
for activations-brn based fusion over a different number of samples; (c) performance with different cntivationsfilteriii strategies for a different number of samples; (d) different transport map policies for residual connections over different number of samples.
Fig. 5c shows the impact of various filters on the one-shot accuracy of the fusion, thereby strength- ening our hypothesis that discarding irrelevant activations helps our fusion algorithm converge to better optimum. Finally, in Fig. 5d we present the impact of the various transport map policies for residuals, as presented in Section 4.2. 1. Both weighted policies perform very similarly, slightly falling behind the best accuracy given by the averaged policy.
5.2 FINETUNED PERFORMANCE
As last stage of the experimental setup, we finetune the fused models. The performance, as well as the retraining curves, offer an important insight into the quality of the fusion algorithm. While the one-shot performance can be heavily impacted by even only single problematic layer, the capacity of the fused model to effectively, rapidly, and easily recover the performance of the parents allows for deeper insight into the quality of the fusion across the whole architecture.
Table 2: Post-finetuning accuracies on the CIFAR 100 dataset for the individual parent models, their ensemble, VF, weights- and peteratiionnbaabee soft alignment. Model dimension: (384/1536/7).
DATASET
IND. MODELS
ENS.
Fr. VF
FT. OT-WTS
FT. OT-ACTS
CIFARI00
[64.94, 64.66]
68.04
64.91 (6003)
65.80 (+0.86)
65.35 (+0.41)
x1
x2
x1
x1
x1
CIFARI00
[64.94. 64. 64.66, 64.44,
70.71
63.19 (-0.55)
65.98 (+1.04)
65.25 (+0.31)
64.38, 64.34, 64.07] x1
x6
x1
x1
x
We show the finetuning results on the widely adopted datasets CIFAR 100, and IMAGENET- 1K (results TINY IMAGENET in the Appendix). We first employ our fusion approach on the ViTs trained on the CIFAR 100 dataset. As mentioned, we separately optimize the fused model on common set of hyperparameters, in this case learning rate (LR) in {10-3, 10 the number of epochs in (10, 20, 100, 200). In Tab. 2 we observe that both our soft-alignment alignment) are capable of outperforming the
10-3 1, and
strategies (i.e. with weights- and activations-batt converged parents, with the gain that increases with the number parent models. This suggests a successful knowledge transfer of the parents into the fused model. While the obtained accuracy lacks
8

---

Published as a conference paper at ICLR 2024
behind the ensembling performance, in scenario there is no computational overhead, while the cost of the ensembling model grows linearly with the number of models.
Table Accuracies on the IMAGENET- I dataset after finetuning for the individual parent models, their ensemble, VF, and weights-based soft alignment. Model dimension: (384/1536/12).
DATASET
IND. MODELS
ENS.
FT. VF
Fr. OT-WIS
IMAGENET-IK
[75.33, 74.88]
76.56
67.83 (-7.5))
75.80 (40.47)
x1
x2
x1
x1
In Tab. 3 we present further results on the challenging and widely-adopted IMAGENET-IK dataset. The results are consistent with those found in the CIFAR 100 case, strengthening the general applicability of our methods, and its scalability to larger models and more challenging datasets. We also stress the fact that, especially with this difficult dataset, even after finetuning, VF fails to recover a comparable accuracy, converging to suboptimal performance.
In this work, we focused on the vision application of the Transformer architecture, but our method is agile to architectural changes, and we demonstrate its wide applicability to the BERT model. Although preliminary explorations of our fusion strategy on the BERT model show some differences with respect to the ViT case (more details on this in App D), the results are on par with those presented above. In particular, the fused and finetuned model, outperforms both parents and VF the widely adopted GLUE benchmark (Wang et al., 2018). The results are presented in Tab. 17 of the App. E.
We want highlight an insight into the finetuning process. In particular, we have observed that the best accuracy of our fused models is achieved extremely quickly, as much as two orders of magnitude fewer steps needed train the parents from scratch, and, as a comparison, VF requires far higher computation to reach a comparable (but worse) performance. For fur- ther exemplification refer to Fig. 12 in Appendix E.2.
Table 4: Results for heterogeneous fusion on CIFAR10N VF cannot be applied here.
ANCHOR
LARGER
ENS.
Fr. OT-WTS
63.18
64.94
67.66
64.11
(40.93)
x1
x4
X5
x1
(192/76877)
(184/156/7)
(192/768/7)
Our methodology, as opposed to VF, works out of the box with models having different widths (heteroge- neous fusion). We find a consistent absolute increase in test accuracy over the performance of the smaller anchor network, thus implying successful knowledge transfer (Tab. 4). These results showease that method is effective and efficient alternative to knowledge distillation.
64.07
64.79
67.94
64.88
(+0.81)
x1
x2.3
x3.3
x1
(384/15367)
(576/2304/7)
(384/1536/7)
DISCUSSION
The fusion methodology for transformer models proposed in this paper is easily adapted to different architectural variants and is readily applicable to models of different widths. However, heterogeneous fusion of networks of different depths is a common limitation of the predominant fusion methods (Singh & Jaggi, 2020; Ainsworth et al., 2022) which are inherently based sequential layerwise alignment. Consequently, we too inherit a similar limitation when expanding fusion to the case of Transformers. Overall, this is undoubtedly fascinating research challenge extend Transformer fusion (or, broadly speaking, fusion at large) to heterogeneous depth settings which, however, is outside the scope of the current work.
In summary, we showcased how distinct independently trained transformer networks can be com- bined through the lens of Optimal Transport. Utilizing novel graph interpretation of the transporta- tion map flow, we developed an algorithm for fusing multiple transformer networks that extends the existing fusion techniques and that specifically caters to the idiosyncrasies of the transformer archi- tecture. We also uncovered an intriguing benefit of using soft alignment when fusing Transformers, which had been under-utilized in the past. Overall, we showed that our technique can retain most of the performance of the converged parent models in one-shot, and even outperforms them after finetuning, across multiple vision and NLP tasks proving the scalability and wide applicability of our methods thereby providing a highly efficient and promising alternative ensembling. Finally, our algorithm successfully applics to the fusion of models of different sizes, too, efficiently transferring knowledge from larger to smaller Transformers, and thus offering effective alternative to distillation.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
Sidak Pal Singh would like to acknowledge the financial support from Max Planck ETH Center for Learning Systems.
REFERENCES
Samuel Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa. Git re-basin: Merging models modulo permutation symmetries. arXiv preprint trXiv:2209.048820 2022.
Aditya Kumar Akash, Sixu Li, and Nicolás García Trillos. Wasserstein barycenter-based model fusion and linear mode connectivity of neural networks. arXiv preprint arXiv:2210.06677/ 2022.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein gan, 2017.
Gregor Bachmann, Sotiris Anagnostidis, and Thomas Hofmann. Scaling mips: A tale of inductive bias. arXiv preprint vrXiv:2306.13573 2023.
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportuni- ties and risks of foundation models. arXiv preprint arXiv:2108.07258 2021.
Leo Breiman. Bagging predictors. Machine learning, 24:123-140 1996.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
Jacob Devlin, Ming- Ming-Weii Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. CoRR, abs/18 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition scale. CoRR, abs/2010. 1112200011199, 1929, 2020. URL https: wxx.v.org//. abs aww.a///aa 2tabs/20100 2010. 11929.
Rahim Entezari, Hanie Sedghi, Olga Saukh, and Behnam Neyshabur.
The role of permutation invariance linear mode connectivity of neural networks. arXiv preprint arXiv:21 10.06296, 2021.
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P Vetrov, and Andrew G Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns. Advances in neural information processing systems, 31, 2018.
Stuart Geman, Elie Bienenstock, and René Doursat. Neural networks and the bias/variance dilemma. Neural computation, 4(1):1-58, 1992.
Charles Godfrey, Davis Brown, Tegan Emerson, and Henry Kvinge. the symmetries of deep learning models and their internal representations. Advances in Neural Information Processing Systems, 15:11893-11905 2022.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendrieks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 2203....222..5533 15556, 2022.
Alexia
Emy Gervais, Kilian Fatras, Yan Zhang, and Simon Lacoste-Julien. Population parameter averaging (papa), arXiv preprint arXiv:2304.0304 2023.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361. 2020.
Philip A Knight. The sinkhorn kinkhom-knoop algorithm: convergence and applications. SIAM Journal on Matrix Analysis and Applications,
2008.
10

---

Published as a conference paper at ICLR 2024
Xiaobin Li, Lianlei Shan, and Weiqiang Wang. Fusing multitask models by recursive least squares. In ICASSP 2021 2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 1640-3644, 2021. doi: 10.1090ICASSP99732222114444
Tianyang Lin, Yuxin Wang, Xiangyang Liu, and Xipeng Qiu. A survey of transform- ers sciencedirect. https: /wwww sci encedi rect. com/ sci $2666651 02200011 6, 2022. (Accessed on 12/04/2022).
Chang Liu, Chenfei Lou, Runzhong Wang, Alan Yuhan Xi, Li Shen, and Junchi Yan. Deep neural network fusion via graph matching with applications model ensemble and federated learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning. volume 162 of Proceedings of Machine Learning Research, pp. 13857-13860 PMLR, 17-23 Jul 2022. URL https: / /twpppoceeddieee ngg mlm mwr.ppp. w.l..prress... 62/ u22k. html.
Michael Matena and Colin A Raffel. Merging models with fishhr-weighted averaging. Advances in Neural Information Processing Systems, 35:17703-17716, 2022.
Ibomoiye Domor Mienye and Yanxia Sun. A survey of ensemble learning: Concepts, algorithms, 2022. doi: 10. 00.110/AACESSS.202
applications, and prospects. IEEE Access, 10:99129-99149, 2022. 3207287.
Dang Nguyen, Khai Nguyen, Dinh Phung, Hung Bui, and Nhat Ho. Model fusion of heterogeneous neural networks via cross-layer alignment. arXiv preprint arXiv:2110.15538, 2021.
Dang Nguyen, Trang Nguyen, Khai Nguyen, Dinh Phung, Hung Bui, and Nhat Ho. cross-layer alignment for model fusion of heterogeneous neural networks. In ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 1-5. IEEE, 2023.
Namuk Park and Songkuk Kim. How do vision transformers work?, 2022.
Yossi Rubner, Carlo Tomasi, and Leonidas J Guibas. The earth mover's distance as a metric for image retrieval. International journal of computer vision, 40(2):99, 2000.
Sidak Pal Singh and Martin Jaggi. Model fusion via optimal transport. Advances in Neural Information Processing Systems, 33:22045-22055, 2020.
George Stoica, Daniel Bolya, Jakob Bjomer, Taylor Hearn, and Judy Hoffman. Zipit! merging models from different tasks without training, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Cédric Villani al. Optimal transport: old and new, volume 338, Springer, 2009.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. CoRR, abs/1800 804.07461, 2018. URL http: / /aaxiiiivv org/ abs /1804 07461.
Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. Federated learning with matched averaging. arXiv preprint arXiv:2002.06444 2020.
Qingsong Wen, Tian Zhou, Chaoli Zhang, Weiqi Chen, Ziqing Ma, Junchi Yan, and Liang Sun. [2202.07125] transformers in time series: A survey. 07125, 2022. (Accessed on 12/04/2022).
/2022.
Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International Conference on Machine Learning, pp. 23965-23998 PMLR, 2022.
11

---

Published as a conference paper at ICLR 2024
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12104-12113, 2022.
Wei Zhou, Han-Jia Ye, and De-Chuan Zhan. Co-transport for class-incremental learning. In Proceedings of the 29th ACM International Conference on Multimedia, MM '21, pp. 1645-1654, New York, NY, USA, 2021. Association for Computing Machinery, ISBN 978145038611 doi: 10.1 1445/3447408534755006 URL https: 00......0//...115333444.... 3475306
12

---

Published as a conference paper at ICLR 2024
A BACKGROUND ON OPTIMAL TRANSPORT AND OTFUSION
A.1 OPTIMAL TRANSPORT THEORY
At its core, Optimal transport (OT) provides way to compare two (or more) probability distributions
:= (a, X) Σ 8(x,) and == (b, Y) = -ド= by o(y)), where 8(.) is the Dirac-delta.
These distributions are typically supported in high-dimensional space, i.e., X, X' R and y, E y/ = 144, Vi,j. and also where, being distributions, Σ a = given distributions, in our case, may correspond to neurons or weights in a particular layer of the two networks. OT aims find transport plan (or map) that signifies how much of these weights of the source model, should be moved towards the destination model, while adhering to the geometry of the underlying "ground' space, usually available the form of "ground metric', e.g.. y) ||x y1ı in the Fuclidean case. Mathematically, one can formulate OT through an equivalent linear program:
m Σー=ı bj = 1. These
(T(,, C) min (T, F s.t., TLm a, T =b and T E P (nxm)
where appropriate mass conservation and positivity constraints are met. Here, (:: () 1 is the Frobenius inner product and I, R" denotes vector containing all ones of size n. While the above problem will find a solution at the vertex of the polytope, one can relax the search to smooth solutions by
regularizing the entropy h of the transport plan (Cuturi, 2013), i... h(T) Eis Tip log(Tiy)
OTx C) = min (T, F) - X (T( s.L., TLm = =a, T = and T I (2(mm)
Besides allowing for a soft assignment, also allows for an efficient solution via the Sinkhorn-Knapp algorithm (Knight, 2008) that results in a speed-up by an order of magnitude in the dimension di (or d2) and be parallelized on GPUs. In contrast, the unregularized problem, which is also commonly referred to as the Earth-Mover's 's Distance (EMD; Rubner et al. (2000)), scales cubically in the dimension.
A.2 OTFUSION
OTFusion (Singh & Jaggi, 2020) first aligns several models: B, C..... to anchor model A. Then, the aligned models are averaged. Alignment implemented through transportation maps, obtained by calculating the minimal transport cost between activations weights of the neurons that should be aligned, giving rise to two different approaches, namely activations- and weights- based respectively. The OTFusion process works in sequential fashion; assuming models with specific depth L, each of the models' layers, at layer f, are aligned before moving to the next layer f++ 1. First, the transpose of the transportation map of the previous layer is pre-multiplied with the weight matrix of the current layer: Will--1) + ( w!ll--), The current layer can then be aligned by post-multiplying with the transportation map of the current layer: Will--)) W(ll-))
B CROSS-HEAD ALIGNMENT VISUALISATION
Fig. 6 visualizes the cross-head alignment algorithm for tiny multi- self-attention block. The aligned weights can then averaged with the corresponding weights of the anchor model to get the weights for the OTFused model.
C EXPERIMENTAL SETUP
C.I VISION TRANSFORMER CIFARIO. CIFARI00, Tiny ImageNet AND ImageNet-Ik
Model Details We use the ViT implementation available on Hugging Face? and we train it from scratch, without using any pre-trained weights. The architectural details of the model can be seen in T.ble 5.
13

---

Published as a conference paper at ICLR 2024
W
Concat
WK
W
Align
X
W
K
X
TK
W
WK K
Separate
W
W
Figure 6: Visualization of the cross-head alignment algorithm for multi head attention block with 1= 2, dhead = 2, dmadet 4, where h is the number of heads, dheed is the head dimension and dmodel is the model dimension.
Table 5: Parameters for the ViT models.
CIFARIO/100
32x32x3
Input image size
Tiny ImageNet
64x64x3
Patch extraction
Convolutional
Patch dimension
4x4
Number of layers
7
Number of heads
12
Size of embeddings
384
Intermediate size
1536
Non Non-line linearity
GELU
14

---

Published as a conference paper at ICLR 2024
Image Augmentation We applied two different image augmentation policies the CIFAR 10/100 and Tiny ImageNer datasets to achieve satisfactory training performance. For the CIFAR datasets, the augmentations have been adapted from open-source implementation', while for Tiny ImageNer the Autoaugment class from Pytorch has been used.
Training Details Training details are reported in Table Figures 7, 8, 9 show the training curves for the CIFARIO, CIFARIOO, and Tiny ImageNer respectively.
Table 6: Training details for the VIT models trained on CIFAR and Tiny ImageNer models.
Optimizer
AdamN
Weight decay
5- 10
Learning Rate
Maximum value of 10
LR Scheduler
Cosine scheduling
Warmup
0.025% epochs warmup
CIFAR
2500
Training Epochs
Tiny ImageNet
250
CIFAR
1024
Batch size
Tiny ImageNet
256
CIFAR
2
Gradient accumulation
Tiny ImageNet
8
Random seed
0-4
Seal
Seal
Sead
Lo00
3.5
Sead
Seal
20
34-
Scad
68-
30
5301
1000
2000
2501
900
3000
2000
20ce
2001
Epoch
Epach
(b)
Figure 7: Training curves for the CIFARIO dataset over five different seeds. (a) Validation loss; (b)
validation accuracy.
C.2 VISION TRANSFORMER IMAGENET
Model Details We use the ViT class from vit-pytorch and we train it from scratch, without using any pre-trained weights. The architectural details of the model can be seen in Table
http://hhhh...ccme/ee/eaaaaaaaaa AutoAugment. htm1 http:///gg.gtabccooo//////aaaa/a
hhttps://githie
a on. transforms. www.rraasooorrrrrrrroo..o....
15

---

Published as a conference paper at ICLR 2024
4.-
Sead 0
Seal
Seat
Seed
L200
3.5
15-
19
2k
5000
1500
3000
30x1
301
3000
1000
2000
5001
Epoch
Epoch
(b)
Figure 8: Training curves for the CIFARI00 dataset over five different seeds. (a) validation loss; (b) validation accuracy.
Loo0
Seot
30-
13-
d0
10
150
900
Epoch
Epoch
(a)
(b)
Figure 9: Training curves for the Tiny ImageNet dataset five different seeds. (a) validation loss; (b) validation accuracy.
Table 7: Parameters for the ViT models.
Input image size
224x224x3
Patch extraction
Linear
Patch dimension
16x16
Number of layers
12
Number of hcads
6
Size of embeddings
384
Intermediate size
1536
Non- lincarity
GELU
16

---

Published as a conference paper at ICLR 2024
Image Augmentation We first applied RandomResizedCropll and RandomHorizconallFiip/ to the imput image form Pytorch transforms sub- sub-package Then we applied the Autoaugment class from the same Pytorch sub-package Images are then normalized with μ [0.485, 0,456, 0.406) and [0.229, 0.224, 0.225].
Training Details
Training details are reported in Table
Table Training details for the ViT models trained on Imagenet.
Optimizer
Adamw
Weight decay
1-10" 10
Learning Rate
Maximum value of 1 10-3
IR Scheduler
Cosine scheduling
Training Epochs
90
Batch size
1000
Random seed
2,4
C.3 PROFILING INFORMATION
In Tab. 9 we provide profiling information for our most used ViT configuration.
Table 9: Profiling information for our most used ViT configuration. The experiments were run on an RTX 4090. We count one fused-multiply accumulate instructions as one FLOP. Different datasets have different image resolutions, leading to different sequence lengths propagating through the transformer, which affects the computational expense of a forward pass.
MODEL
#PARAMS
DATASET
#PATCHES
FLOPs
TP
MODEL DIM.
(M)
(B)
(IMAGE/S)
VIT
12.4
CIFARIO0
65
0.8
13.2 K
(384/1536/7)
Tiny ImageNet
257
3.5
C.4 BERT
Model Details We the BERT implementation available on Hugging Face together with the pre-trained bert-base-uncased tokenizer 9, Our BERT model has the architectural details presented in Tab. 10.
Training Details We train the BERT models, from scratch, over five different seeds. Training details are shown Tab. II.
We use MLM task on subset of the Wikipedia dataset, available on Hugging Face with an MLM probability of 0.15.
The training curve of the loss, for one seed, is presented in Fig, 10.
/ http:://ssttatooooosssssssssssts
zer
simple
17

---

Published as a conference paper at ICLR 2024
Table 10: Parameters of the architecture for the BERT models.
Number encoders
Number heads
Size of embeddings
768
Intermediate size
3072
Maximum position embedding
512
Attention dropout probability
0.1
Hidden dropout probability
0.1
Non- Non-linearee lincarity
GELU
Table 11: Training details for the BERT models.
Optimizer
AdamW
Learning Rate
cosine scheduling with epochs of warmup; maximum value of 5 10-5
Training Epochs
40
Batch size
16
Random seed(s)
6-
Loss
10
16
20
255
30
35
40
Epoch
Figure 10: BERT pre-training validation loss for random seed O.
18

---

Published as a conference paper at ICLR 2024
D SINKHORN REGULARIZER ABLATIONS
The Sinkhorn algorithm, and in general the soft alignment paradigm, has been heavily underused in literature and therefore there is little information about its impact on OTFusion. As presented above, we uncover intriguing behaviors, that require reconsidering use. In the following Sections, we extend our findings related to soft alignment, in particular with the role of the regularization parameter.
D.1 ABLATION ON RESNET
To compare the findings for the transformer architecture, we also investigate effect of the Sinkhorn regularizer the ResNet architecture (Fig. Ila). In agreement with the findings of Singh & Jaggi (2020), the best result is achieved with EMD, and a small regularizer is preferred as it approaches the hard alignment. This result is thus suggesting opposite behavior when it comes to soft alignment since the transformer benefits from a soft alignment.
D.2 ABLATIONS ON CIFAR100, Tiny ImageNet, BERT MLM TASK
In Fig. 11 we present the effect of the Sinkhorn regularizer on the other considered datasets, namely CIFARIO0 (Fig. 11b) and Tiny ImageNer (Fig. 11c) for the ViT, and the MLM task on the Wikipedia subset. for BERT (Fig. 11d).
The outcomes for CIFARI00 and Tiny ImageNet are in line with the results of the CIFARIO case, namely non-zero regularizer achieves the optimal performance.
As hinted in Sec. 5.2, we have observed some differences in the regularization effect on the BERT model. This difference can be observed in Fig. 11d, where we plot the effect of the regularization parameter the validation loss. We observe that, in contrast to the observations for the VIT, the loss curve shows no inverted bell curve, suggesting that there is no finite optimal regularizer, i.e. that completely soft alignment is best suited for this model.
22
One-hht Avcuracy
Onv ahot Accuracy
15
080
015
8.15 Sinkhoru Regulatioer
0.35
6.0
0.0
Siukhoen Regularizer
(b)
2.55
One shot Accuracy 50
2.0-
Saklve
1.75
L.5-
1.55
1.00
675-
0.75
613
0.50
0.0
0.66
8.10
0.15
0.00
9.20
0.09
8/6
0.38
0.15
0.00
0.25
6:20
Sinkchom Regulariver
Sinkhom Regularizer
(c)
(d)
Figure 11: Sinkhorn regularizer effect on one-shot performance. EMD-fusion performance is shown as a reference. (a) Accuracy for ResNet on CIFARIO (higher is better); (b) accuracy for ViT on CIFARI00 (higher is better); (c) accuracy for ViT on Tiny ImageNet (higher is better); (d) loss for BERT on MLM task (lower is better).
19

---

Published as a conference paper at ICLR 2024
D.3 WHAT HAPPENS AT THE EXTREME EDGE OF SINKHORN REGULARIZATION?
As presented above, the softness of the alignment is impacted by the Sinkhorn regularizer. the regularizer close zero, the algorithm converges a permutation matrix (i.e. hard alignment); in contrast, if the regularizer very large, the algorithm converges to a unit-matrix divided by the dimension of itself.
D.3.1 SINKHIORN REGULARIZER TO ZERO
In general, have observed that the smaller the regularizer becomes, the harder the alignment gets. However, for very small Sinkhorn regularizer values the algorithm breaks down. This especially visible in Fig. 1lb and 1lc where for the smallest regularizer the one-shot accuracy falls below the one-shot accuracy of EMD. We found that normalizing the cost matrix and the activations/weights to calculate the cost matrix, pushes the breakdown closer zero and thus improving stability.
D.3.2 SINKHORN REGULARIZER TO INFINITY
We conducted experiment show that even in the case of extreme regularization (i.e. completely soft alignment) information transferred from model B to the anchor model. In this experiment, we fuse randomly initialized model (10% accuracy on CIFARIO) with model convergence (92% accuracy on CIFARIO). The one-shot accuracy for this experiment is 10%. On the other hand, if we fuse two converged models, we get one-shot accuracy of 47% for completely soft alignment. This suggests that, even in the highly regularized case, our algorithm allows knowledge transfer.
20

---

Published as a conference paper at ICLR 2024
E FURTHER RESULTS
In this section, we provide more results from our experiments. We report both one-shot and finetuned accuracies the datasets of choice.
E.1 One-shol
Tab. 12 and Tab. 13 report the one-shot accuracies for Tiny ImageNet and CIFARIO0 datasets, respectively.
Table 12: One-shot accuracies on the Tiny ImageNet dataset for the individual parent models, their ensemble, VF, weights-based soft-alignment fusion, and activations-basee soft alignment fusion. The last column shows the highest finetuned performance with mean and standard deviations over different data seeds. The figure beneath the test accuracies signifies how much computation is required by the model ensemble with respect to our fusion technique.
comparison. Activations-base is reported
DATASET
INDIVIDUAL
ENS.
VF
OT-WTS
OT-ACTS
Fr. OT-wTS
MODELS
(OURS)
(OURS)
(OURS)
Tiny ImageNet [45.30, 45.22, 44.50.
51.28
0.44
1.64
3.03 1 0.27
45.90
44.36, 43.78]
x5
x1
x1
x1
x
Table 13: One-shot accuracies the CIFARI00 dataset for the individual parent models, their ensemble, VF, weights-based soft-alignment fusion, and activations-based soft alignment fusion. The last column shows the highest finetuned performance with mean and standard deviations over different data seeds. The figure beneath the test accuracies signifies how much more computation required by the model ensemble with respect to our fusion technique.
comparison. Activations-based is reported
DATASET
INDIVIDUAL
ENS.
VF
OT-WTS
OT-ACTS
FT. OT-WTS
MODELS
(OURS)
(OURS)
(OURS)
CIFARIOO
[64.94, 64.66]
68.04
0.77
13.32
11.70 0.13
65.80
x2
x1
x1
CIFARIO0
[64.94, 64.66, 64.44,
70.71
0.98
11.16
7.45 + 0.25
65.98
64.38, 64.34, 64.07]
x1
x1
E.2 FINETUNING
After fusing the models, we finetune them. Finetuning parameters and results are reported in the subsections below.
21

---

Published as a conference paper at ICLR 2024
E.2.1 FINETUNING DETAILS- VIT
As mentioned in Sec. 5, we finetune VF and fused models separately on common set of hyperparameters. In the following paragraph the subset used over the different datasets and models:
ViT CIFARIOO: LR in {10-3 -3 10 10-4, 10-),, number of epochs in {10, 20, 100, 2003
ViT - Tiny ImageNet: LR {10-3, 10-4, 10-0).00000 number of epochs in {1, 10, 20)
Finetuning the ImageNet- Ik dataset is inherently expensive. We have thus finetuned for just 8 to 10 epochs the fused models, with an LR of 10- The boost in performance presented in Tab. is thus even noteworthy given the limited capacity exhaustively find suitable hyper-parameters for finetuning.
E.2.2 RESULTS
Vision Transformer In Tab. 14 we report the finetuning results for the fusion and ensemble of two and six models the CIFAR/00 dataset. The results show how weight-based soft alignment outperforms both weight-based hard alignment and activation-bass soft alignment. Furthermore, in Tab. 15 we present further results the Tiny ImageNet dataset.
Table 14: Accuracies the CIFARIO0 dataset after finetuning for the individual parent models, their ensemble, VF, weights- based soft alignment, weight-based hard alignment, and oetivations-base soft-aligmment. The figure beneath the test accuracies signifies how much more computation is required by the model ensemble with respect to our fusion technique.
FT.
FT.
DATASET
INDIVIDUAL MODELS
ENS.
VANILLA
OT-WTS
OT-WTS
OT-ACTS
(OURS)
EMD (OURS)
(OURS)
CIFARIO0
[64.94, 64.66]
68.04
64.91
65.80
64.72
65.35
x
x1
x1
x1
x1
CIFARIOO
[64.94, 64.66, 64.44,
70.71
63.19
65.98
65.42
65.25
64.38, 64.34, 64.07]
x1
x1
Epochs
VF 10 Epochs
0.9
VP 200 Epocks
OT--WTS (ours) 200 Epochs
OT-WTS (ours) Epocks
0.7
Brror
0.5
0.3
25
125
150
175
Epochs
Figure 12: Finetuning curves on the validation set. Cosine scheduling is used. Validation error on the CIFAR100 dataset.
BERT
The results after finetuning for the BERT model are presented in Tab. 16 and Tab 17.
22

---

Published as a conference paper at ICLR 2024
Table 15: Accuracies on the Tiny ImageNer dataset after finetuning for the individual parent models, their ensemble, VF, weights-based soft alignment, and activations-basee soft alignment. Model dimension encoded as (hidden-layer dimension//mtermccittee-eee dimension/number of encoders). The figure beneath the accuracies indicates the relative computational burden (latency and FLOPs) of the model(s).
FT.
Fr.
FT.
DATASET
IND. MODELS
DIMENSION
ENS.
VF
OT-WTS
OT-ACTS
Tiny ImageNet [45.30, 45.22, 44.50, (384/1536/7)
51.28
38.82
45.44
45.90
44.36, 43.78)
x1
X5
X1
X1
x1
Table 16: Loss values for BERT the MLM task after finetuning for the individual parent models, their ensemble, VF, and weights-based alignment fusion. Both VF and our fused model are trained with LR of 5 10-5 for only epochs. This shows the much faster speed of recovery our approach, compared to VF. The figure beneath the test accuracies signifies how much more computation is required by the model ensemble with respect to our fusion technique.
FT.
FT.
DATASET
INDIVIDUAL MODELS
ENS.
VANILLA
OT-WTS
(OURS)
MASKED WIKT [1.612, 1.761, 1.776,
1.665
2.946
2.224
1.794, 1.807J
x5
x
x1
23

---

Published as a conference paper at ICLR 2024
Table 17: Results for BERT evaluation on GLUE benchmark, after finetuning for 14 epochs. Ac- curacy is the metric for SST2, QNLI, RTE and WNLL Matthews corr. the metric for COLA. F1/Accuracy is the metric for MRPC and QQP. Pearson/Spearman corr. is the metric for STSB. Matched acc./Mismatched acc. the metric for MNLI.
TASK
PARENT
OT
VF
MRPC
0.852/78.2
0.853/77.7
0.807/72.1
STSB
0.828/0.827
0.841/0.838
0.771/0.771
QQP
0.844/88.2
0.847/88.5
0.840/88.1
MNLI
76.1/76.4
75.9/76.1
74.1/74.6
COLA
0.263
0.275
0.236
QNLI
84.1
85.1
83.0
WNLI
26.8
29.4
27.6
SST2
85.6
86.5
84.9
RTE
62.1
63.4
51.6
24


Reviews:
[{"review_id": "vLVevtPSbU", "reviewer": "Reviewer_GM8g", "summary": "This paper proposes a systematic approach for fusing two or more transformer-based networks exploiting Optimal Transport technique. The proposed method can generalize to arbitrary architectures for CNNs and Transformers. Extensive experiments involving the fusion and finetuning of Vision Transformers (ViTs) across multiple datasets demonstrate the effectiveness of the proposed method.", "questions": "See the weakness part.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper is well written.\n- The proposed method shows good generalization across different architectures.\n- The proposed method show strong performance for several benchmark.", "weaknesses": "- Most experiments are conducted to compare with Vanilla Fusion. More comparisons with state-of-the-art methods should be included.\n- Most experiments are conducted on CIFAR dataset which is relatively small.", "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": "9tniV8L56p", "reviewer": "Reviewer_5GYY", "summary": "This paper proposed a systematic approach for fusing two or more pretrained transformers by studying the flow of transportation maps in each specific component of Transformer. The authors empirically showed that when working with Transformers, hard alignment underperforms soft alignment in one-shot fusion, which is in contrast to the cases of fully connected and convolutional neural networks. Finally, they showcased the efficiency of the proposal in fusing and finetuning ViT and BERT.", "questions": "- Section 2: The model fusion literature has some papers that are slightly off: Tatro et al., 2020; Juneja et al., 2022; Kandpal et al., 2023.\n- Eq. 2: What is $f$? What is its output?\n- Section 4.2.1: How to calculate weighted matrix?\n- The authors should remind the formulation for Attention operation in either Section 3 or Section 4.2.2.\n- Section 4.2.2:\n  - Where do the authors remove the constraints in Section 4.2.2?\n  - It is unclear how to calculate $T_Q$ and $T_K$. Did the authors check the assumption $T_Q = T_K$ in the experiments?\n  - What are $W_i^Q, W_i^K$, and $W_i^V$? Does $i$ indicate the head index?\n  - Additional visualizations may help to demonstrate the method here.\n- Section 4.2.3: What is this sentence for? “For the concatenation, we notice that the class token is only a small fraction of the full sequence, in other words, for the integrity of the sequence, it is far more important to propagate the TM of the patch embeddings than the one for the class token.” In addition, the class token is more important because it gathers the information from the patch.\n\n\n**Minors**: \n- Eq. 3 should be moved up a paragraph.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- This paper is well-structured.\n- To the best of my knowledge, this is the first work that aims to fuse transformer architectures by aligning their weights.\n- The proposed method is successfully backed by theoretical results.", "weaknesses": "- The methodology part is not well-written and lacks some details.", "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": "JbEmqNhn18", "reviewer": "Reviewer_KSPk", "summary": "This paper introduces a systematic fusion technique for transformer-based networks by leveraging Optimal Transport to align architectural components. It offers a flexible approach applicable to various architectures, including key Transformer components. Heterogeneous fusion enables efficient compression, with superior performance compared to vanilla fusion and individual parent models, as demonstrated in image classification (Vision Transformer) and natural language tasks (BERT). Our analysis underscores the significance of soft alignment in the context of Transformers, highlighting the potential for combining multiple Transformers to enhance their capabilities in the emerging field of model fusion and recombination.", "questions": "Please refer to the weakness section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The authors examined various strategies (weight vs activation, hard vs soft etc) for applying optimal transport (OT) methods\n2. The authors conducted experiments employing both Vision Transformer (ViT) and BERT architectures across multiple datasets.\n3. The OT method demonstrates particular efficacy in one-shot scenarios.\n4. OT methods exhibit versatility, as they can be effectively applied to models of varying widths, presenting a viable alternative to distillation.", "weaknesses": "1. The OT method yields comparatively lower performance when contrasted with ensemble methods.\n2. The suitability of the OT method for achieving solid results on larger datasets, such as ImageNet-1K, in one-shot scenarios remains uncertain.", "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": "Vh0r2LAvbD", "reviewer": "Reviewer_hByk", "summary": "The authors propose a method for fusing multiple independently trained transformer architectures using optimal transport to align their respective architectural components. To this end, authors analyse predominant Transformer architectures based on their components, and provide OTFusion methods for each. The proposed approach allows for fusing transformers of different sizes. \nExperiments are conducted on a range of image classification datasets; CIFAR10, CIFAR100, TinyImagenet and ImageNet-1k. Authors show results for models obtained through both zero-shot (without fine tuning) and with fine tuning. In zero-shot fusion, the proposed approach outperforms Vanilla Fusion methods. With fine tuning, the proposed approach is able to beat either parent model. Authors conclude with a number of limitations of their approach and suggestions for future research.", "questions": "-What do you mean by “This diversity offers a challenging fusion problem requiring a non-trivial alignment strategy, and thus effectively recreates a plethora of other scenarios” (under 5 - Model Training). Can you explain e.g. how varying random seed equates to model training on different subsets?\n-How does your work relate to [2]? You indicate that [2] is very similar to OTFusion, but looking at zero-shot performance of your method (and your VF baseline) on CIFAR10 classification it seems performance is drastically different (~93% vs ~60%). If essentially identical, why does [2] yield zero-barrier LMC where your approach does not?\n-Could you give an intuition for soft-alignment, what resulting network is actually being constructed  in this case and why could it be beneficial compared to hard alignment approaches?\n-Do you have an intuition for why your method performs better with soft-alignment, where [1] shows better results with hard alignment?\n\n[2] Samuel K Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa. Git re-basin: Merging models\nmodulo permutation symmetries. arXiv preprint arXiv:2209.04836, 2022.\n\n---\n\nUpdate after rebuttal: I thank the authors for their thorough rebuttal. I'm pleased to say my concerns are adequately addressed. Also considering the largely positive reviews by the other reviewers, I'd like to update my recommendation to an accept.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "This paper is very well-written, authors give clear and concise descriptions of their approach and illustrate its complex aspects through figures and examples. This makes the manuscript easy to read and the ideas it expands upon easy to understand despite their complexity. The method seems to work well, drastically outperforming Vanilla Fusion (naive model averaging). The authors show a range of valuable ablations, and motivate most of their design choices well.", "weaknesses": "My main concern is to do with the clarity of the contribution of this work. The authors refer to [1] a lot in their paper, where the concept of OTFusion is introduced. It seems like a lot of the techniques used in this work were actually introduced there. Although I understand the need for reintroducing these concepts in the manuscript for contextual clarity, I think it would be good to give a clearer picture of the actual contributions made in this work and the methods proposed in previous works. From the description under 4.3 it seems [1] uses hard alignment where you find soft alignments to outperform. Are these contributions of your work? What about the TM combination approaches (Averaging/Weighted Scalar/ Weighted Matrix)? Or heterogeneous fusion?\n\nI hope the authors are able to address this in their rebuttal, in which case I see this work as an interesting and strong submission.\n\n[1] Sidak Pal Singh and Martin Jaggi. Model fusion via optimal transport. Advances in Neural Information\nProcessing Systems, 33:22045–22055, 2020.", "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)