Transformer Fusion with Optimal Transport

Thank you for your useful feedback. We are truly glad that you have appreciated our efforts for a well-presented and well-written paper. At the same time, we recognize that some explicit clarifications regarding our contribution would have helped to better convey the contribution of our paper, and we are grateful for making us aware of this. To reflect your doubts, and with the utmost goal of improving the clarity of our work, we have implemented various changes throughout the manuscript.

---

In what follows, besides addressing your specific questions, we would nevertheless like to first expand on your concerns about the contribution of our work.

Contribution of this work with respect to [1]

Key differences

1. Enabling Fusion for newer Transformer-based architectures

   [1] is the first to successfully introduce the concept of Optimal Transport for alignment and fusion of multiple models. But, broadly speaking, it is however restricted to simple architectures such as MLPs, CNNs, and instances of ResNet.

   It is not equipped in any way to align and fuse models with complex information streams and to fuse transformer-specific components such as multi-head attention layers, layer-normalization, embeddings, or the sequential nature of the data, which we systematically analyze and successfully handle.

2. New findings in regards to hard vs soft alignment

   Furthermore, we broaden the perspective on alignment introduced by [1] in the following manner. While [1] technically allows soft alignment, they discovered that for simpler architectures (MLPs, CNNs, ResNets) hard alignment outperforms soft alignment (Table S4 of [1]).

   In contrast, we find the reverse to be true when fusing Transformers (Sec. 4.3, and Sec. 5), which possibly hints at the difference in the nature of their architectures and representations, and forms an interesting question for standalone future study.

Other aspects

• Heterogeneous fusion

  As mentioned on page 2 of [1], indeed [1] is the first to showcase heterogeneous model fusion, “OTFusion accommodates the fusion of models with different widths, and in turn, different sizes”. But, we are the first to extend this concept to Transformers.

• Weighted TM combination & other minor methodological extensions:

  Lastly, we have included some other extensions for OTFusion [1], to better support transformers.

  As an example, realizing the potentially complex peculiarities of transformers when it comes to residual streams, we have extended the Averaging strategy introduced by [1], with two novel strategies, namely Weighted Scalar and Weighted Matrix, that allow for a more flexible fusion strategy depending on the properties of the residual stream of that layer (Sec. 4.2.1).

  Also, we have extended OTFusion's idea of activations-based alignment to better encompass the sequential nature of the data, introducing various novel strategies of activations filtering in the case of the ViT (Sec. 4.3).

---

Thanks for your valuable review. We are extremely glad to hear that you have appreciated our systematic investigation of the various fusion strategies. In what follows ahead, we would like to take the opportunity to address your primary concerns.

---

Accuracy of the fused model

The OT method yields comparatively lower performance when contrasted with ensemble methods.

Allow us to detail the inherent factors at play when fusing the parameters of the given two neural networks:

• Non-convexity of the neural network loss-landscape

  It is widely accepted that neural networks result in a highly non-convex loss landscape and optimization dynamics. As a result, by definition of non-convexity, it is rather unlikely that the loss $\mathcal{L}$ at, say, the midpoint between two networks (with parameters $\theta_1$ and $\theta_2$) is less than the average of the losses of the two networks. Mathematically,

  \mathcal{L}\left(\frac{\theta_1 + \theta_2}{2}\right) \nless \frac{1}{2} \mathcal{L}(\theta_1) + \frac{1}{2} \mathcal{L}(\theta_2)\.

• Empirical success of Model Fusion and conjectured Linear Mode Connectivity modulo Permutations

  Notwithstanding this inherent non-convexity, [3, 4] showed that, in practice, if the network parameters are aligned prior to fusion, like by accounting for their permutation symmetries, the loss landscape (near the basin of solutions) becomes much less non-convex. The works of [2] and [1] have provided additional evidence that when the network widths are large, the two networks (or modes) can more or less be linearly connected modulo permutation symmetries (RHS - LHS of the above equation is small).

• Recovering the exact permutation is NP-hard

  Most algorithmic approaches, either of [3] or [1] are greedy layerwise algorithms. This is because recovering the exact permutation symmetry that links the given networks becomes intractable as the search space grows in proportion to $\mathcal{O}({(m\!)}^L)$, where $m$ denotes the layer width and $L$ is the network depth. Also, see Lemma 1 in [1]. As a consequence, we see the inherent difficulty of the problem that we are tackling here.

• The added complexity of Transformers

  To the best of our knowledge, prior works have only studied fusion and linear mode connectivity mostly in the case of MLPs, CNNs, and ResNets, and not for Transformers. On the other hand, Transformers seem to further complicate the above issues in practically finding a good alignment, for instance, due to additional components such as the multi-head self-attention (resulting in the need for intra-head as well as inter-head alignment), LayerNorms, and, more often than not, larger depths. Also, other kinds of symmetries get introduced, such as scale invariance, due to the presence of softmax.

• Ensembling is implicitly using more capacity

  So far we only have outlined the difficulty of the fusion process. However, the other obvious fact is that ensembling uses $K \times$ more capacity than fusion, where $K$ is the number of networks that are ensembled. Hence, the ensembling performance can be regarded, in some sense, as the unachievable upper bound on the performance of the singular network produced through fusion.

In this light, arguably, it is rather noteworthy that like prior works:

1. We can still get highly non-trivial performance gains over vanilla fusion. In the case of Transformers, we believe this is also where soft alignment benefits since it helps relax the constraints from hard alignment (permutation matrices) to softer alignments (with transportation maps).
2. Moreover, after fine-tuning, even this gap relative to ensembling considerably diminishes. Besides, this fine-tuning procedure is usually extremely fast — requiring as few as 0.4% of the original training duration —, and since fine-tuning has to be done just once even this minor cost gets amortized over inference (while ensembling must pay $O(n)$ (where $n$ is the number of parent models) more at every inference call)

We sincerely appreciate your careful reading of our paper and giving constructive feedback. We are pleased to hear that you find our paper well structured and that you have appreciated the novelty of our work. We have seriously considered your notes and we recognize that some crucial aspects could have lacked clarity.

Owing to your insightful feedback we have implemented various modifications in the methodology section (Sec. 4) of the manuscript that, in our opinion, have improved the overall comprehensibility of our work.

---

Related Work

“[…] some papers that are slightly off: Tatro et al., 2020; Juneja et al., 2022; Kandpal et al., 2023.”

We have fixed these.

---

Residuals (Sec. 4.2.1)

“Eq. 2: What is $f$? What is its output?”

$\mathbf{f}_{residual}$ is a vector and stands for the activations coming from the residual branch.

$\mathbf{f}_{current}$ is also a vector and stands for the activations coming from the current layer l. We added a sentence for clarification in the paper and adjusted Eq. 2 to make it more clear.

“How to calculate weighted matrix?”

The calculation for the weighted matrix is very similar to the weighted scalar. The only difference is that, for weighted matrix, we compute a weighting factor for every incoming residual strand, instead of one common value. We added a sentence for clarification in the paper.

---

Multi-Head Attention (Sec. 4.2.2)

“[...] remind the formulation for Attention operation [...]”

Done

“Where do the authors remove the constraints in Section 4.2.2?”

We assume you mean the “equal transportation map” constraint here. Section 4.2.2 serves to explain how we handle the transportation map flow through the self-attention block. At the beginning of the section, we describe our analytical insight on why one should only propagate $\mathbf{T_V}$. Eq. 4 shows that if we use $\mathbf{T_Q} = \mathbf{T_K}$, these permutation matrices cancel. However, this equation is not valid anymore in the case of soft-alignment because we no longer have permutation matrices, so the $\mathbf{T_Q} = \mathbf{T_K}$ constraint is no longer needed. We added a sentence for clarification in the paper.

“It is unclear how to calculate $\mathbf{T_K}$ and $\mathbf{T_Q}$”

If we do not assume the constraint on the outgoing transportation maps $\mathbf{T_K} = \mathbf{T_Q}$ one can directly apply OTFusion to $\mathbf{W_K}$ and $\mathbf{W_Q}$ (as for any fully-connected layer). OTFusion yields $\mathbf{T_K}$ and $\mathbf{T_Q}$ as the outgoing transportation map respectively.

In the case of $\mathbf{T_K} = \mathbf{T_Q}$, we compute one common ground metric for both $\mathbf{W^K}$ and $\mathbf{W^Q}$. In the case of activation based alignment the ground metric is computed from their combined activations. For weight based alignment, the ground metric is computed from the concatenation of $\mathbf{W^Q}$ and $\mathbf{W^K}$. From this common ground metric we then compute a common transportation map that we use to align both $\mathbf{W^K}$ and $\mathbf{W^Q}$.

“What are $\mathbf{W^K_i}$, $\mathbf{W^Q_i}$, $\mathbf{W^V_i}$, and ? Does i indicate the head index?”

Yes, $i$ indicates the head index. $\mathbf{W^K_i}$, $\mathbf{W^Q_i}$, and $\mathbf{W^V_i}$ are the corresponding weight matrices. We added a clarification in the manuscript, and together with the inserted self-attention operation formulation, it should be clear now.

---

Embeddings (Sec. 4.2.3)

“What is this sentence for? “For the concatenation, we notice that the class token [...]” In addition, the class token is more important because it gathers the information from the patch.”

We assume that at the very beginning of the architecture, namely at the embedding stage of the transformer, each patch and token is of similar importance. The intuition for this assumption is that we concatenate only one class token to tens of image patches, and we can therefore assume that it is more important to propagate the transportation map of the image patches instead of the class token.

Note however that we make this assumption exclusively for the embeddings, i.e. not throughout all encoder blocks where, as you say, the information will indeed be eventually distilled into the class token.

---

We hope that in light of the modifications of the manuscript and the above clarifications, we have adequately addressed your questions and concerns. In case you might have further doubts or comments, we remain at your disposal.

Thank you for your useful feedback. We are pleased to hear that you have found our paper well written, and have appreciated the generalization capabilities of our method across different architectures.

---

"Most experiments are conducted to compare with Vanilla Fusion. More comparisons with state-of-the-art methods should be included."

In general, we share a similar sentiment as you have expressed. We would have liked to add more comparisons, however, we were forced to compare primarily with Vanilla Fusion (with and without fine-tuning; as well as parent models and ensemble), as we are not aware — to the best of our knowledge — of any other method that can handle the problem of aggregating the parameters of multiple transformer models, in particular, when they are distant in their parameter space.

More specifically, to better contextualize the state-of-the-art (SotA) in this area, we would like to reiterate here some fundamental differences and limitations of notable related work:

• OTFusion [1]
  • [1] first introduces the concept of Optimal Transport based alignment and fusion. However, its fully layerwise interpretation lacks generalization capabilities, and as such it is only applicable to simple architectures such as multi-layer perceptrons, CNNs, and instances of ResNet.
  • It is not equipped in any way to align and fuse models with complex information streams and to fuse transformer-specific components such as multi-head attention layers, layer-normalization, embeddings, or the sequential nature of the data.
• Successors of OTFusion
  • A plethora of other methods emerged from the ideas and methodology of [1] and explored various adaptations and different applications.
  • In particular, [2] focused on the linear mode connectivity perspective for MLPs, CNNs, and ResNets, while [3] focused on RNNs and LSTM architectures. These methods, too, are not applicable to transformers and a direct comparison is consequently out of reach.
• Model Soups [4] for transformers
  • Recently, [4] focused on weight averaging specifically for transformers, reaching SotA performance.
  • The inherent methodology of [4] actually relies on VF itself, namely one-to-one averaging of the parameters. However, there is a subtle but essential difference compared to our application scenario: their parent models originate from the same pre-trained model and therefore the parent models remain sufficiently close in the parameter space. This precludes the need to align them and lets them employ a simple averaging strategy while still obtaining a gain in performance.
  • For this reason, this method does not apply to our problem where we fuse transformer networks that are potentially much more distant in their parameter spaces and are diverse in nature.

Furthermore, we would like to stress that the larger point of our paper is to present the first successful model fusion technique for transformers that aggregates their parameters, with a focus on modularity, generalization capabilities, and systematic investigation of fusion of the transformers' peculiar components and characteristics (such as multi-head self-attention, layer normalization, sequentiality of the data, to name a few).

---

"Most experiments are conducted on CIFAR dataset which is relatively small."

We agree that CIFAR10 is a relatively small dataset, and we opted for this dataset for the experiments and ablations (Tab. 1, Fig. 4 and 5) that required multiple evaluations with different hyperparameters and fusion strategies, and that were prohibitive with large datasets. However, sharing your concern, and aiming to demonstrate the wide applicability and generalization capabilities of our method, we also have experiments on other larger datasets.

ViT

In particular, we have presented the results of ViT fusion also on:

• CIFAR100 (Tab. 2 for models of the same size; Tab. 4 for heterogeneous fusion)
• TinyImageNet (Tab. 15 in the Appendix),
• ImageNet-1k (Tab. 3)

where we see a consistent gain across all settings.

BERT

Furthermore, we would like to highlight that not only have we applied our method to the ViT, but also to BERT for NLP tasks, evaluated on the widely adopted GLUE benchmark (Tab. 17, in the Appendix), and the results were similar to those for the vision tasks.

---

All in all, we hope that in light of the above clarifications, we have given a proper answer to your concerns. Should you have any further questions, or comments, we will be more than happy to answer them. Thanks for taking the time.