Update Your Transformer to the Latest Release: Re-Basin of Task Vectors
Authors: Filippo Rinaldi
Giacomo Capitani
Lorenzo Bonicelli
Donato Crisostomi
Federico Bolelli
ELISA FICARRA
Emanuele Rodolà
Simone Calderara
Angelo Porrello
Venue: ICML 2025 poster (2025)
Keywords: weight interpolation
transfer learning
model rebasin
model editing
model patching
Primary area: general_machine_learning->transfer_multitask_and_metalearning
TL;DR: We propose a data-free method to transfer fine-tuned Transformer model updates to newer pre-trained versions using weight permutations, enabling fine-tuning updating without retraining.

Abstract: Foundation models serve as the backbone for numerous specialized models developed through fine-tuning. However, when the underlying pretrained model is updated or retrained (e.g., on larger and more curated datasets), the fine-tuned model becomes obsolete, losing its utility and requiring retraining. This raises the question: is it possible to transfer fine-tuning to a new release of the model? In this work, we investigate how to transfer fine-tuning to a new checkpoint without having to re-train, in a data-free manner. To do so, we draw principles from model re-basin and provide a recipe based on weight permutations to re-base the modifications made to the original base model, often called task vector. In particular, our approach tailors model re-basin for Transformer models, taking into account the challenges of residual connections and multi-head attention layers. Specifically, we propose a two-level method rooted in spectral theory, initially permuting the attention heads and subsequently adjusting parameters within select pairs of heads. Through extensive experiments on visual and textual tasks, we achieve the seamless transfer of fine-tuned knowledge to new pre-trained backbones without relying on a single training step or datapoint. Code is available at https://github.com/aimagelab/TransFusion.

Update Your Transformer to the Latest Release: Re-Basin of Task Vectors
Filippo Rinaldi"
Giacomo Capitani"
Lorenzo Bonicelli
Donato Crisostomi Federico Bolelli
Elisa Ficarra'
Emanuele Rodolà 2 Simone Calderara'
Angelo Porrello
Abstract
Foundation models serve as the backbone for nu- merous specialized models developed through fine-tuning. However, when the underlying pre- trained model is updated or retrained (e.g., on larger and more curated datasets), the fine-tuned model becomes obsolete, losing its utility and re- quiring retraining. This raises the question: is it possible to transfer fine-tuning to new release of the model? In this work, we investigate how to transfer fine-tuning out having to re-train, in data-free manner. To do so, we draw principles from model re-basin and provide recipe based weight permuta- tions to re- be-base the modifications made to the original base model, often called task vector. In particular, our approach tailors model re-basin for Transformer models, taking into account the challenges of residual connections and multi-head attention layers. Specifically, we propose two- level method rooted in spectral theory, initially permuting the attention heads and subsequently adjusting parameters within select pairs of heads. Through extensive experiments on visual and tex- tual tasks, we achieve the seamless transfer of fine-tuned knowledge to new pre-trained back- bones without relying datapoint. Code is available at https:// gi thub. com/a a wwgelaa/////a....... aaaaaa///rraausttaa
high
low
(0A))
Loss
0 B
x
ideal
O A
π (T)
0
T
0fF
0 (T)
new checkpoint with-
(I-) re basin
V f fine tuning
0^ +
Figure 1: Transporting task vector 7 from fine-tuned base model 0 = 0 + to new release 0g.
ably decreasing the amount of data and training time needed to tailor models for specific downstream tasks. For this reason, pre-trained backbones such as OpenAT's CLIP (Rad- ford et al., 2021) are being extensively utilized as base foun- dation models. As result, the corresponding fine-tuned versions play a crucial role in numerous real-world appli- cations like medical imaging (I.u et al., 2024) and satellite image analysis (Mall et al., 2024).
However, while these pre- preitrained backbones are widely adopted, their evolution poses new challenges, with tech companies and academic institutions frequently releasing updated checkpoints. Often, these updates do not modify the underlying architecture but simply consist of weights, trained on increasingly large datasets compared their pre- decessors (Ilharco et al., 2021). Moreover, the additional training data may be more curated or specifically tailored to specialized domains, boosting their zero-shot capabilities considerably.
single training step
1. Introduction
Recently, there has been notable shift among researchers and practitioners towards fine-tuning pre-trained models, rather than building them from scratch. This method lever- ages backbones developed on large-scale datasets, consider-
To take advantage of newly released checkpoints, the typ- ical approach is retrain them on the downstream task. This means fine-tuning the new checkpoint on the same data already used to adapt the original model. Besides the consid- erable costs associated with re-training the new model, this strategy is also unviable in certain scenarios. Indeed, the data for the downstream task might be no longer available due compliance with privacy or storage constraints.
Eqqual contribution 1 'AlmageLab, University of Modena and Reggio Emilia, Italy. "Sapienza, University of Rome, Italy. Corre- spondence to: Filippo Rinaldi <fiiipoo.inaldi @enimore.it>.
Proceedings of the = International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Update Your Transformer to the Latest Release: Re-Basin of Task Vectors
This raises an important question: can we re-use the fine- tuning that has already been performed on the newly re- leased model? Precisely, the overall aim of this paper is to investigate whether we can transport the previous fine- tuning, in a training-free manner. To understand the idea of transport, we consider the weights of the original base model as 0x, and their fine-tuning as 0% = 0x + T. The task vector (Ilharco et al., 2023; Drtiz-Jimenez et al., 2024) 01-- 0 represents direction from A that embodies all adjustments made during the fine-tuning process. Hence, our goal is to find a procedure TT-) that can transport the task vector T into an appealing basin of the newly released model 0g (see Fig. 1). In this approach, the procedure (()) must be designed to ensure that the modified weights Of = 0B + (T) achieve low loss on the downstream task.
and treated as a single unified projection. This way, after applying permutations, each head of the edited model may incorporate rows from different original heads an issue we refer to as head contamination. This is problematic be-
cause, without preserving the logical separation of heads during permutation, it becomes impossible to invert the permutation process and recover the original, unpermuted output of the attention block. Furthermore, existing meth- ods struggle when two addends in the computational graph rely on distinct permutation matrices, situation common in residual connections such as h = h + s(h). Differently, we avoid averaging the respective permutation matrices, thereby preserving their discrete nature.
To address these issues, we propose a structured two-level approach for effective re-basin of Transformer-bas based mod- els, called TransFusion. To avoid head contamination, TransFusion firstly seeks optimal mappings between pairs of heads (inter-head permutations); subsequently, we re- strict permutations to only the rows within these coupled heads (intra-head permutations). We mathematically prove that this two-level permutation strategy prevents head con- tamination and preserves functional equivalence between the original and permuted models. Notably, the inter-head permutations are optimized leveraging distance metric that is invariant to permutations of the rows and columns within the heads. Such a metric is founded on spectral theory (Jo- vanović & Stanié, 2012) and employs the singular values of the weight matrices, which are unaffected by orthogonal transformations like those induced by permutations.
When designing the transportation function π T()), the ideal approach should be data-free and training-free to meet the concerns above. Nevertheless, if the two base models A and 8g differ significantly (due to varying initialization, training strategies, or datasets), the knowledge acquired during fine-tuning of A may not transfer to 0g with a mere addition of the original task vector (i.e.. of = 0n + r). To bridge the gap in representation spaces and facilitate the transfer, intuitively, we have to make the two base models "compatible", such that they "speak the same language". To address this challenge, we could rebase one of the two models (for instance, 4)), such that any linear interpolation between the weights of the edited 6', and 0g yields an intermediate model that performs comparably to both 0 and 0g. This indicates that the models are now aligned and thus share a common low-loss basin. Notably, this concept of re-basin models shares similarities with the approach described in (Ainsworth al., 2023), where alignment is achieved by finding optimal permutations of the rows in the weight matrices. In this paper, we build upon this idea and explore its application in the context of fine-tuning, with a task vector being permuted and finally applied to 0B.
We
We show that transporting task vectors allows the transfer of knowledge into a new checkpoint in a data-free manner. In practice. this means we can improve the zero-shot per- formance of the new version on the downstream task. We also demonstrate that the transport retains the generalization capabilities on a support set - a crucial factor to justify updating the base model to the new release. To summarize:
While model re-basin presents an appealing framework, it currently faces several technical hindrances. To date, suc- cessful applications of model re- re-basin have been limited to Multi-Layer Perceptrons (MLPs) and Convolutional Neural Networks (CNNs) (Ainsworth et al., 2023). Unfortunately, the application of model re-basin to multi-head attention layers, despite their widespread use in Transformer based architectures, has been largely overlooked, with few very recent attempts based on Optimal Transport (see See. 5). However, these recent methods do not guarantee functional equivalence between the permuted model and the original, unpermuted model. The primary obstacle lies in managing the weights associated with multiple attention heads. As we discuss in Sec. 3, indeed, to apply standard permutation- based approaches (Ainsworth et al., 2023; Singh & Jaggi, 2020; Imfeld et al., 2024), the heads must be concatenated
Motivation. Given a fine-tuned Transformer model s(, = 0a +7) we aim to transfer the task vector into a new release f(-, 0g) ) in data-free manner. Contributions. i) We introduce a novel data-free method for aligning Transformers. ii) We propose a permutation -nnvaaaaiooninvariaatt spectral measure to manage multi-head attention layers. iā) We reveal the practi- cality of applying task vector from one base model 0 new release 0g.
2. Background
Given two models with weights 0 and fp. model re- basin (Ainsworth et al., 2023) investigates how to permute the units of one model facilitate the alignment of two mod-
2

---

Update Your Transformer to the Latest Release: Re-Basin of Task Vectors
els. The two models are then merged in the weight space, resulting in an interpolated model that achieves performance comparable to that of the two original ones.
P related to the permutation of the previous layer see Eq. (4). This results in global optimization across layers:
argmax (w/")) ,. ==(P=/1
( pTT)
+
Following the notation of (Ainsworth et al., 2023), re-basin is defined as any operation defined on the weights of two models A and Og that maps one of the two models onto the local loss region (basin) of the other one. To assess the effectiveness of re-basin, one common approach is to check the property of linear mode connectivity (Frankle et al., 2020; Entezari et al., 2022) between the permuted model and the other, reference model. Informally, this involves checking if the model weights laying on the linear path connecting A and 0g also result in a low loss value.
(W( (B) WL W PT..).
(5)
where () = -i. A Big is the inner product between real-valued matrices. As discussed in (Ainsworth al., 2023), the optimization problem described in Eq. (5) corre- sponds to the Symmetric Orthogonal Bilinear Assignment Problem (SOBLAP), which is unfortunately NP-hard. Its relaxation re-casts it as a series of Linear Assignment Prob- lems (LAPs), focusing on one permutation Pe at a time while keeping the others fixed. In formal terms:
To reach such a property, existing model re-basin techniques leverage the permutation symmetries inherent in neural net- works (Entezari et al., 2022). These symmetries allow the swapping of the units within layer without changing the functionality of the network. To show that, we consider the activation of the ε-th layer of an MLP:
argmax (wl ( P.W! ( + Pe
41 )
Notably, each LAP can be solved with efficient, polynomial- time methods like the Hungarian algorithm (Jonker & Vol- genant, 1988). The outcome is set of permutation matrices π = {P:] . which, when applied to model 0.. result in a new model 0 = T(0A). Notably, this model is function- ally equivalent and, theoretically, it resides in the low-loss basin of 0p. However, as optimizing a series of LAPs is a coarse approximation of the SOBLAP, there are no strong guarantees regarding the optimality of the solution.
26+1 = (Weze + be), 20 = I,
(1)
where Wt and by are the weight matrix and bias vector and denotes an element-wise activation function. In this case, the following relation holds for any permutation matrix P:
1 PT P2e+1 = PT Po + be),
= =P -Poo(PWeze + Pbt), where P E Sa,
with Sa denoting the set of d x d permutation matrices. Thanks to this relation, we can essentially permute the weights and biases of a layer using matrix P. Therefore, when we apply the permutation P to the parameters of layer, the resulting output undergoes the same permutation. However, to ensure that the transformed model remains funetionally equivalent to the original, the next layer must process the output in its original, unpermuted form. This can be achieved equivalently by permuting the weights of the subsequent layer using the inverse permutation pT. Ac- cordingly, we define transformed set of weights 0' as:
3. TransFusion: An Approach For Re-basin Transformer Models
Objective. Our approach, named TransFusion, is designed to transfer task-specific knowledge between transformer- based models that have undergone different pre-training. Specifically, it starts with an initial weight set 0 and a task vector т = 0% - A.. derived after fine-tuning on a downstream task. The goal is to adapt 7 to a new parameter configuration 0g. This process aims to preserve the inher- ent properties of 0p - for example, its superior zero-shot capabilities compared to A - and to integrate specialized knowledge carried out by r for the downstream task. Finally, we aim to enable model transfer in a data-free manner.
W = PWe, by -Pbe, W W+ = WT. (4)
Git Re-Basin. Ainsworth et al. (2023) exploit Eq. (2) to induce weight alignment between A and 0z. Formally, we consider the f-th feed-forward layer, with weight matrices w((4) and w-( for 0 and O respectively. Given that each row of W. and w((B) represents a distinct feature, if (W(A)) I.. ~ [W( J.... then it makes sense to associate the units and j. Therefore, we could formalize the alignment as finding the permutation matrix that maximizes the dot product between (cW and ((BB) However, to preserve functional equivalence, we have to account for the term
Weight Matching. To achieve these objectives, we start by aligning the weights of 0a with those of 0p. This is accomplished with novel data-free weight matching strat- egy tailored for Transformer architectures. The procedure is deeply discussed in Sec. 3.1. Briefly, we tackle various shortcomings of existing methods and handle two building blocks of attention- networks, namely the residual paths and the multi-head attention mechanism. To manage the latter, we introduce a novel two-step process that em- ploys a permutation invariaan spectral metric to match pairs
3

---

Update Your Transformer to the Latest Release: Re-Basin of Task Vectors
of heads within the same layer of 0A and 0g. Subsequently, permute features within the matched heads to optimize weight alignment, as detailed in Sec. 2.
0
A
SVD
SVD
Ho
E
H2
Transport. We end up with a functionally equivalent model 8'A = a)) where TT-) yields permutation of every layer in A.. Afterwards, x(.) is used to transport the task vector f A into the low-loss basin of 0g (see Sec. 3.2).
INTER
H1
Q
Σ
P
H1
Ha
3.1. Attention Alignment for Transformer Models
θB
0^
0
INTRA
A Transformer-based block consists of a multi-head atten- tion layer and an MLP block, connected through residual connections. Considering the MLP, this builds upon stan- dard linear projections, which we treat as discussed in Sec. 2. Instead, we adopt novel, tailored approach for multi-head attention layers addressing a common pitfall. Considering multiple heads, current methods view their projections as a whole linear layer, thereby joining the corresponding weight matrices before applying permutations. However, such approach does not reflect the organization of attention in dis- tinct, parallel heads. For example, this can result in artifacts, where units from separate heads in the original model are mixed together- -a- an issue we call head contamination. This compromises the structural separability of attention heads and precludes the preservation of functional equivalence, that is, the ability to permute and subsequently unpermute the weight matrices while yielding identical model outputs.
Ho
Q
H2
P
H2
Intne
Figure 2: Inter- (Step 1) and intra-head alignment (Step 2).
To achieve the permutation-ines -iniaaaiiioonnnvaiinncee property, we employ a distance based on the singular values of the sub-matrices representing the heads. Specifically, given two heads h,5 [W k,i E Rd,xdm from model Og and hit [W] from model A we compute the distance as:
dig = |E, - Eyحl,
where E, and Σ denote the singular values of h." and ligh respectively. These can be computed through the Sin- gular Value Decomposition (SVD); in formal terms, con- sidering the i-th head, the SVD decompose its weight h = U. where U, and Vi2 are orthogonal matri- ces, and Σ, is a diagonal matrix containing the singular values of h4. As demonstrated in Appendix A. the Eu- clidean distance between singular values remains invariant to permutations.
In the following, we present our proposal against head con- tamination (see Step and 2) and a practical approach to handle residual connections (Step 3). The complete method- ology is outlined in Algorithm 1.
To take into account the distance for query, key and value projections jointly, we construct a distance matrix D e # where each element Di, = df, +d; + represents an head alignment cost that is calculated as the sum of the pairwise distances across q, k and e matrices. We hence employ D to find the optimal inter-head permutation:
Step 1: Inter-Head Alignment
Consider the q(uery), k(ey), and v(alue) projection matrices WG WR. and Wn, E I dm xdm - with dm denoting the total embedding dimension of the attention module. We partition each matrix into H = #heads matrices (one for each head) of shape dę x dm, where dy = d.
H² This results in a tensor W. = sWl g, H) E TR. H xdz xdm for the query projection matrix Wg. The same operation is applied for Wt and W. to obtain W₁ and W..
"
Panter bead = aarggin PES#
D..P())
(8)
The first step involves defining a distance metric between pairs of heads, such that we can identify and execute the optimal swap between heads in A and 0g (see Fig. 2). We employ distance metric that is invariant to permutations of rows and columns within the H sub-matrices in W. g. We and W.. In this respect, one might question why invariance is crucial for comparisons between different heads. We note that the initial, natural order of units does not necessarily correspond to the optimal alignment that could be achieved. Consequently, the metric used in this initial phase must be insensitive to the specific ordering of head features, thereby ensuring an agnostic comparison of the heads.
where Dr, P[j] is the distance between the i-th head of model 0g and the Pli]-th (candidate) head of model 0 The solution Pimter can be practically determined with the Hungarian algorithm. The corresponding permutation is applied to each W., WAA, and W.A. thereby reordering the heads of A to increase alignment with those of f3-
9
Step 2: Intra-Head Alignment
After matching each pair of heads (h." h M(jy)) we aim to swap individual units between h, and hppp- To do so, as in Git Re-Ba (Ainsworth et al., 2023), we seek for permutations that maximize the inner product between the
4

---

Update Your Transformer to the Latest Release: Re-Basin of Task Vectors
Algorithm 1 Weight Matching
Algorithm 2 Attention Alignment
Require: 0^ = { [ { L
and 0B { (w(11)
(B) 1
Require: Weights (/41 P 1 E-11 W! P
/ 1=1 Ensure: A permutation πT (PI..... .P.....,P. PL-1) of f..
J t=1
A and W W(B) wi! By for multi-head atten- tion projection layer l and previous layer l l-1. 1.
1: Initialize: P1 <- ,..... PL-1 +-1
Ensure: Permutation Pt-sttn for W.(A) W(()), W((A)
2: repeat
3: for L E 1, ....., I.....,II....... -1-1 do
1: Step I: Inter-Head Alignment
4: if f is a MHA layer then
2: Create spectral distance matrix D (Eq. (7))
5: Pt + Multi-Head Attention Alignment (Alg. 2)
3: Pimerr <- Solve LAP on D (Eq. (8)).
6: else
4: Step 2: Intra-Head Alignment
7: P Solving LAP as in Eq. (6)
5: for h = 1 to H head pairs from Pinter do
8: until convergence
6: p f- Solve LAP for head pair h (Eq. (9)).
7: Paam + Pinter o (Pith) } H
> compose permutations
k=1
corresponding H sub-portion of the projection weights h. and h hpppa as follows:
Step 3:
Managing of Residual Connections
Pantra head (P((t) HHH
Each transformer block incorporates two residual connec- tions: the first bypasses the multi-head attention layer, and the second bypasses the feed-forwatd network:
intra head 1i=1
where p(i)
intra argmax (1, PhP1a)
(9)
PESa,
Zattn Wo MHA(x);
In this formula, the cost considers the dot products across query, key, and value sub-matrices. In the end, the optimiza- tion yields H permutation matrices p(i) intra E Sd.-
2; Zuttn X,
(10)
2 W
Zoat Z + Z;:
In summary, this two-step process ensures both global inter- head and local intra-head alignment while preserving the structural integrity of the multi-head attention mechanism. The comprehensive procedure for aligning multi-head atten- tion layers, combining these alignment stages, is formalized in Algorithm/ 2. Specifically, these individual alignment steps are unified into a single composed permutation, de- noted as Patn, which is applied directly to the projection matrices of the multi-head attention layer. Crucially, it can be shown that structured composed permutation pre- serves functional equivalence: despite reordering and per- muting the heads and their internal dimensions, the attention computation remains unchanged (up to a corresponding per- mutation of its outputs). We formalize this in the following theorem.
where x is the input, Wo is the weight of the attention mechanism, and W1 and W2 those of the feed-forward layer. For simplicity, we omit layer normalization as it can be regarded as a standard linear projection.
In each residual block, the input and the intermediate output are summed to produce the final output. However, we note that there are several sources of potential mismatch between the two addends: intuitively. if the two addends have un- dergone different permutations, it is reasonable to suspect a potential mismatch in their representations.
To elarify the interaction between permutations in residual blocks, con- sider Fig. 3. which represents the first residual z; = Zatta + x. When the weights of the model are permuted, the input x comes with its own per- mutation Pm- which has to be ac- counted using PiT- Moreover, the attention projection Wo adds its own permutation matrix Pw.. This leads to the following relation:
Theorem 3.1 (Equivariance of Multi-Head Attention to Structured Permutations). Let P Pinser head E Sh be per- mutation over the H attention heads, and let P Penrashead = P Jead) H be a set of independent permutations acting
Pwo
MHA
Pin
Pin x
within each head (of size dy = Ax ). Then applying the composed block permutation Patn to each of the projection matrices Wg. WR, W. E Rd xdm is functionally equivalent to permuting the output of the multi-head attention module. The resulting attention output O' satisfies: O' = =OPam aamm where O is the unpermuted output.
Figure 3: Residu- als block and per- mutations.
Zi = Pw; Zart + P.x.
(1))
When examining the permutations/sumaatoos that impact on 2,, there are two main issues: issue I) the residual branch lacks transformations that could account for the matrix P.:
A complete proof of Theorem 3.1 is provided in Ap- pendix A.2.
5

---

Update Your Transformer to the Latest Release: Re-Basin of Task Vectors
Table 1: Comparison of permutation-based methods on visual tasks, in terms of task accuracy [1] and support accuracy [1]-
EUROSAT
DTD
GTSRB
SVHN
Method
TAsk
SUPP.
TASK
SUPP.
TASK
SUPP.
TASK
SUPP.
p zero-shot
49.02
68.73
47.50
68.73
43.42
68.73
45.97
68.73
Ag + On + x [T) (Optimal Trunsport)
-7.62
-16.15
-0.15
-0.10
-5.39
-0.70
-22.00
-16.45
-14.05
-5.28
-0.53
-1.18
-2.43
-1.30
-12.30
-2.70
OB (T) (GiT Re-Basin)
+0.95
-0.48
-0.91
-0.02
+0.76
-0.05
+0.79
+0.30
TRANSFUSION (OuRS)
+4.95
-0.06
+0.21
-0.08
+1.10
-0.40
+3.64
-0.48
issue II) the projection Wo adds its own permutation Pwo of which the residual branch has no information about.
reference model 0 allows their combination at destina- tion 0B, which basically means we could still apply model merging (Wortsman al., 2022a) after ra-basin.
To maintain coherence between the two addends, they must be transformed under identical permutations. To enforce this consistency, we redefine the identity mapping made by the residual connection. We replace it with a composition, I = Pw. consisting of two permutations - one to address issue and another for issue Il as follows:
3.3. Complexity Analysis
In this subsection, we assess the computational complexity of the proposed weight matching procedure. The key insight is that the method is highly efficient compared to full re- training, and scales polynomially with model size.
2, = PWo + Im = Pwozum + PwoX, (12)
Proposition 3.2. Let L be the number of layers and dm the embedding dimension of each transformer block. The over- all computational complexity of our weight matching pro- cedure is dominated by O d d This complexity matches that of Git Re-Basin making our approach comparably efficient in terms of computational cost.
which highlights how the two addends now share the same permutation. An analogous process applies to the second residual connection yielding Zou (see Appendix A.3 for the full procedure). As a final technical note, we remark that the permutation matrix Pw, associated to the second residual block in Eq. (10) has to be considered as input permutation for the subsequent layer.
The proof is provided in Appendix A.5 and illustrates the per-layer contribution of both MLP and attention blocks.
3.2. Transporting Task Vectors from 0 to 0B
4. Experiments
By applying to model A. we would have a functionally equivalent model 6'A = ((AA) with stronger linear-mode connectivity with On compared to the original 0 How- ever, to allow knowledge transfer from the fine-tuned model off = 0a + T to 0B. we do not apply the permutations directly on f... but rather on the task vector T, as follows:
This section is structured into three main parts. Initially, we empirically assess the transportation of task vectors, involy- ing extensive experiments across both visual and natural language processing (NLP) tasks (Sec. 4.1). Subsequently, we examine the capability of our methodology to align the weights of two Transformer models while maintaining func- tional equivalence (Sec. 4.2). Finally, several ablative stud- ies show the impact of our techniques on addressing multi- head attention layers and residual connections (Sec. 4.3).
task vector: :
T T= 0 A
(13)
transport:
0 = 0g + (T),),
(14)
where a is a non-negative scaling factor (Wortsman et al., 2022b) modulating the influence of [T) on 0g.
4.1. TransFusion of Task Vectors
By leveraging the concept of transporting task vectors, we have several notable advantages, especially in a scenario with multiple models fine-tuned on distinct tasks from the same base model A In this scenario, the weight matching process between A and 0g needs to be conducted only once. Indeed, permutation set can be established and reused to transfer any number of task vectors. This ap- proach avoids the additional computational costs associated with learning separate transport functions for each transfer. Moreover, transporting multiple task vectors using the same
Visual Classification Tasks. As reference architecture, we consider the CLIP ViT-B/16 Vision Transformer (Radford et al., 2021) from Open-CLIE (Cherti et al., 2023). We refer to 0 as the original pre-training weights and @g as those used for the re-basin. We use CommonPool pre-training for O and Datacomp for 0B. both cited in (Gadre et al., 2024).
Considering the base model 9 we fine-tune the correspond- ing model on several computer vision tasks (Radford et al., 2021: Ilharco et al., 2023). We employ DTD (Cimpoi et al.,
6

---

Update Your Transformer to the Latest Release: Re-Basin of Task Vectors
Table 2: Comparison of permutation-basee methods on NLP
tasks, in terms of task accuracy [t]-
SVHN
EUROSAT
Task Gain (%)
Method
QQP
SST2
RTE
COLA
QB
55.00
50.69
51.51
40.94
8g
-8.29
+0.23
-2.53
-0.77
Task Vector
-10-
0g +1 T (OT)
-8.31
+5.39
-1.08
-1.25
Ours
8B + T (GIT Re- Re-Basin)
+3.58
+5.73
+2.17
+1.44
TRANSFUSION (OurS)
+6.50
+5.96
+3.61
+2.49
Supp. Gain (%)
0-
-5
2014), EuroSAT (Helber et al., 2019), GTSRB (Stallkamp et al., 2011), and SVHN (Netzer et al., 2011) and obtain multiple, independent fine-tuned models like of 0x + T. Afterwards, we empirically assess the transportation o to the new weights 0B. In this respect, we adopt two metries to characterize the quality of the transported model 0g + (T): the original task (special-
-10
-15
2
Alpha
Alpha
Figure 4: Zero-shot gain/drop relative to 0g of naive 0g + ar (blue) and our strategy 0B + am(()) (red) varying a.
the zero-shot performance ized knowledge), and ii) the zero-shot performance on support, unseen set to evaluate the preservation of broader capabilities. In our experiments, ImageNet-R (Hendrycks et al., 2021) serves as a support dataset.
NLP Classification Tasks. Herein, investigate different setting that involves closed-vocabulary text clas- sification specifically, set of tasks from the GLUE benchmark (Wang et al., 2019). We consider a model o = ( () composed of a pre-trained Transformer en- coder and a classification head w. We then evaluate the transport of the learned task vector Tg & Φ a new feature extractor ΦB. As access to data of the downstream task is restricted, we are unable to train new classifier for 0R: consequently, we re-use the originally fine-tuned classifier, denoted as w The goal is to evaluate whether transporting the task vector Tø aligns the repre- sentation yielded by ΦB + *(Ta) with the original, fine- tuned classifier w In our experiments, we employed two variants of the ViT-B-16 text encoder, pretrained re- (0x) and
We report the results in Tab. as drops (-) or gains (+) in accuracy relative to the zero-shot performance of 0p. As baselines, provide the results of vanilla transporta- tion (no permutations applied on 7) and those of Git Re- Basin (Ainsworth et al., 2023) and Optimal Transport (OT) (Imfeld et al., 2024), two existing methods for model re-basin. Specifically, the comparison with OT is notewor- thy since this approach is designed for Transformer models (like ours).
As can be seen, method enhances zero-shot performance on the downstream tasks and preserves generalization on the support dataset, outperforming existing permutation- based methods. Considering the results of our approach, it is particularly noteworthy that we enhance performance on the downstream task while maintaining generalization, all achieved without the use of any data.
spectively on the cwtmmmmpooo---s---e-- -atacomp-1-sib-ie (0n) (Gadre et al., 2024).
Tab. 2 presents the evaluation for the GLUE benchmark. Unexpectedly, simply applying the classification head from the original feature extractor Φ yields poor performance (see first line of Tab. ²B). On the other hand, transporting To with Git Re-Basin and Optimal Transport performs rea- sonably, with good gains on QQP and SST2. Moreover, approach leads to the highest and more consistent perfor- mance gains, highlighting the potential of our framework.
In the experiments shown in Tab. 1, we consistently set the scaling coefficient for the (permuted) task vector as = (see Sec. 3.2). To investigate sensitivity and performance changes while varying a, we kindly refer the reader to Fig. 4 (more datasets are in Appendix A.4). This illustrates the drop/gain in accuracy for +ar (blue) and our 0g (TT(F) (red). This drop/gain is measured w.r.t. the zero-shot accu- racy of fg. and varies within the range (0.01, 2.0]. The outcome is that applying the permuted (T)) to On leads to tangible improvements in the downstream task (top row), especially a x 1. Moreover, when a 2 0.5, the permuted task vector is considerably more reliable in terms of gener- alization (higher accuracy on the support set).
4.2. TransFusion Improves Alignment and Preserves Functional Equivalence
While the previous analyses focus on transferring task vec- tors, we now delve into the effectiveness of our approach in terms of weight alignment. In detail, we consider two VIT-B/16 models (Dosovitskly et al., 2021) A and B trained
7

---

Update Your Transformer to the Latest Release: Re-Basin of Task Vectors
Table 3: TransFusion results by head alignment strategy.
VIT-B/16
3.0
Vanilla Residual
EUROSAT
GTSRB
SVHN
Head Align.
TASK
SUPP.
TASK
SUPP.
TASK
SUPP.
Ours
Loss
2.0
0g zero-shot
49.02
68.73
43.42
68.73
45.97
68.73
Brute Force
+1.32
-0.21
+0.60
-0.46
+3.39
-0.40
1.0
No All-Align
+2.22
-0.47
+0.71
+0.05
+0.24
-0.08
OURS (FULL) +4.95
-0.06
+1.10
-0.40
+3.64
-0.48
0.0
0.2
0.4
0.6
0.8
1.0
3.0
- OT
Git Bassin Ours
4.3. Ablative Analysis
Loss
2.0
On the Strategy to Manage Multiple Heads. We herein explore the significance of an appropriate policy for permut- ing the projection layers within multi-head attention mecha- nisms. Specifically, we present the outcomes of transferring T with varying strategies to permute attention projection layers. As potential alternatives, we firstly consider brute force alignment, which considers all possible head pair combinations within each attention layer. Then, for each candidate pair, the intra-head alignment cost is computed by optimizing the objective in Eq. (9). The final permutation is then derived with the Hungarian algorithm, which selects pairs with the highest intra- alignment scores.
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Alpha
Figure 5: Loss values on CIFAR-10 test set during model interpolation. Top: Our permutation approach vs. vanilla interpolation and no residual variant. Bottom: Comparison with Optimal Transport and Git Re-Basin methods, which fail to preserve functional equivalence as a -> 0.
independently on CIFAR-10 (Krizhevsky et al., 2009) from scratch, which means they underwent different initializa- tions and batch orders. After training, we apply our permu- tation strategy to A and analyze the resulting alignment of o (f.) and 0B in terms of linear mode connectiviry: follow- ing (Ainsworth et al., 2023), we evaluate the loss landscape C ((1- - a (0))= + a0g), E [0,1] while interpolating between the two models.
A ffter, we compare our approach with one that pairs heads in A and B according to their natural order, thereby avoiding head contamination by design. Nevertheless, the preserva- tion of original head ordering comes at the cost of ignoring functional mismatches between attention units. We refer to this further baseline as no attention alignment.
The results of these experiments are detailed in Tab. 3. Our attention-alignment strategy achieves superior performance the downstream task (0p + T(T)) compared to alterna- tive approaches, while maintaining comparable zero-shot capabilities. The comparison with the brute force approach underscores the effectiveness of our cormutation-nvvaiine costs modeling inter-head relationships, demonstrating superior performance over a brute force alignment that opti- mizes for the best match within each candidate pair of heads. Furthermore, our results suggest that preserving the origi- nal head ordering (as in the no-alignment strategy) yields better performance than brute-force inter-head matching for two of the three experimental tasks. This underscores the menace of head contamination and the importance of structure-awars alignment methodologies over indiscrimi- nate similarity maximization.
As shown by Fig. 5 (top), applying the permutation æ yields an interpolated model that exhibits consistent lower loss compared to the vanilla approach that does not permute 9 (#A) = 4A). Moreover, the no residual approach underscores the critical role of properly handling residual connections in our method- - both during the permutation of Oa and throughout the interpolation process between models. Similarly, in Fig. (bottom) we assess the loss landscape using derived from the Optimal Transport (OT) hard alignment method (Imfeld et al., 2024) and Git Re- Basin (Ainsworth et al., 2023). While these works show a more favorable loss landscape than naive interpolation, we observe that the resulting interpolated model struggles and achieves high loss when a - O, highlighting that both OT and Git Re-Basin do not preserve functional equivalence. We conjecture that such a lack stems from potential weak- nesses in effectively permuting layers featuring residual connections and multi-head attention blocks. In virtue of the results achieved by our approach, we can claim it rep- resents the first successful data-free method to interpolate between two Transformer models in weight space, while ensuring the functional equivalence of (0.)).
Few Shot Fine-tuning, There are practical scenarios in which retaining data is infeasible. If such constraints are not present, our proposed method can be effectively combined with fine-tuning. To illustrate this, we follow (Zhang et al., 2024) and start with a small subset consisting of 10 shots per class, learning scaling coefficient per layer, denoted
8

---

Update Your Transformer to the Latest Release: Re-Basin of Task Vectors
Table 4: Fine-tuning results with permuted task vectors.
Multi-task merging via parameter averaging (Ilharco et al., 2022; 2023; Wortsman et al., 2022a; Yadav et al., 2024) not only circumvents catastrophic forgetting (French, 1999;
Method
EUROSAT
DTD
GTSRB
SVHN
McCloskey & Cohen, 1989; Porrello et al., 2025) but syn- thesizes models that retain diverse expertise, even serving as superior starting points for future adaptation (Choshen et al., 2022). The benefits of weight ensembles and interpolations extend beyond just fine-tuned models; they also apply to models that are trained from scratch. Techniques such as those proposed by (Ainsworth et al., 2023; Singh & Jaggi, 2020), leverage permutation symmetries to facilitate coher- ent interpolation between models trained in different ways. Collectively, these findings position weight-space manipula- tion as a scalable toolkit for resource-efficient model engi- neering, where arithmetic and interpolation replace brute- force retraining.
gg zero-shot
49.02
47.50
43.42
45.97
Ag + oT
+7.93
-1.44
+4.70
-15.98
8g + (T((T
+10.00
+1.21
+6.80
+10.52
Li]. The results in Tab. 4 clearly indicate a substantial improvement when fine-tuning a model that has undergone re-basin using our approach, represented as 0g + (r) In contrast, fine-tuning directly from 0g + ar, without permutation, yields inferior outcomes. This empha- sizes that re- basing and fine-tuning should not be considered mutually exclusive but complementary strategies.
as a = [ai....... Li]
5. Related Work
6. Discussion and Conclusions
Mode Connectivity occurs when paths of nearly constant loss connect different solutions within the loss landscape of neural networks (NNs) (Garipov et al., 2018; Freeman & Bruna, 2017: Garipov et al., 2018; Draxler et al., 2018). When such paths are linear, we refer to linear mode connec- tivity (LMC) (Frankle et al., 2020). Entezari et al. (2022) conjecture that solutions found by Stochastic Gradient De- scent (SGD) can be linearly connected when accounting for permutation symmetries. Motivated by this, several works first align the models into a shared optimization space by permuting their neurons, and then merge them through a simple average (Ainsworth et al., 2023; Jordan et al., 2023; Stoica et al., 2024; Peña et al., 2023; Crisostomi et al., 2024; Navon et al., 2023; Singh & Jaggi, 2020). Most relevant to our work, Imfeld et al. (2024) applies optimal transport to align activations in transformer-based networks. How- ever, unlike the latter, our method accounts for the logical division of multi-hear attention projections, preserving the functional equivalence of the permuted models.
For TransFusion to succeed, the source expert must perform strongly on the target task. This insight directly explains the diminished results on DTD in Tab. 1: while our fine-tuning achieves performance well above 95% on all other datasets, DTD reaches only 75% accuracy. We attribute this degraded task vector quality to DTD's inherently challenging char- acteristics - featuring merely 40 examples per class on average, compared to roughly 1,000 examples per class in our other datasets. Although our proposal delivers substan- tial gains without requiring access to data from the original task, we acknowledge sigmificant room for improvement. We believe most of the lost performance could be recov- ered through two complementary directions: incorporating modest amount of additional data, or employing more so- phisticated matching metrics for permutation discovery that better reflect activation distributions rather than relying on simple dot product similarity for linear layers. Such exten- sions represent promising avenues for future investigation.
Weight Interpolation and Task Arithmetic. Emerging re- search reveals that the output of NNs can be manipulated through algebraic operations in weight space (Ilharco al., 2022; Wortsman et al., 2022a). Central this paradigm are task vectors 7 (Ilharco et al., 2023), which encode task- specific knowledge and exhibit compositional properties: combining vectors through addition enables multi-task gen- eralization while their negation can suppress learned behav- iors without significantly impacting unrelated tasks.
A Acknowledgments
This work was supported by the Italian Ministerial grants PRIN 2022: "B-Fair: Bias-Free Artificial Intelli- gence Methods for Automated Visual Recognition" (CUP E53D22008010006) and by the University of Modena and Reggio Emilia and Fondazione di Modena through the "Fondo di Ateneo per la Ricerca FAR 2024" (CUP E93C240020800077, The work also received funding from DECIDER, the European Union's Horizon 2020 research and innovation programme under GA No. 965193 and "AIDA: explAinable multImodal Deep learning for person- Alized oncology" (Project Code 20228MZFAA). We ac- knowledge the CINECA award under the ISCRA initiative for providing nigh-peeformance computing resources.
Beyond arithmetic. weight interpolations further unlock un- expected capabilities: blending fine-tuned and pre-trate traitrieed weights often yields single-task performance superior to standalone fine-tuning (Frankle et al., 2020; Izmailov et al., 2018; Matena & Raffel, 2022; Ramé et al., 2023; Rame et al., 2022: Wortsman et al., 2022b), suggesting a reconciliation of specialized adaptation with generalization capabilities.

---

Update Your Transformer to the Latest Release: Re-Basin of Task Vectors
Impact Statement
ceedings of the IEEE conference on Computer Vision and Pattern Recognition, 2014.
Our approach to transferring fine-tuning between model versions without requiring re-training or additional data holds the potential to significantly lower the barriers to maintaining state-of-the of-the-art AI technologies. By enabling seamless updates, this method could empower organizations, especially those with limited resources, to keep pace with technological advancements, thereby democratizing access to sophisticated AI tools. This democratization can foster innovation across diverse sectors, potentially leveling the playing field between large entities and smaller, resource- constrained organizations.
Crisostomi, D., Fumero, M., Baieri, D., Bernard, F,, and Rodola, E. C²M²: Cycle-Consistent Multi-Model Merg- ing. Advances in Neural Information Processing Systems, 2024.
Dosovitskiy, A., Beyer, L.,, Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M.. Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. In International Conference on Learning Representations, 2021.
The idea of aligning models with minor variations in archi- tecture, such as different layer counts, is also worth explor- ing. One simple approach could involve selectively pruning layers from the model with more layers to match its coun- removiigg redundant or unimportant layers-wiih an alternative strategy might replicate the last block of the smaller network multiple times to achieve parity. We plan to investigate these strategies further in future work.
Draxler, F., Veschgini, K., Salmhofer, M., and Hamprecht, F. Essentially No Barriers in Neural Network Energy Landscape. In International Conference on Machine Learning, 2018.
Entezari. R.,, Sedghi, H.,, Saukh, O., and Neyshabur, B. The Role of Permutation Invariance in Linear Mode Connec- tivity of Neural Networks. In International Conference on Learning Representations, 2022.
However, as we facilitate easier transitions between model versions, it is crucial to ensure that these updates do not compromise the integrity of the models. Dependence on outdated or poorly validated models poses significant risks, particularly when used in critical applications. Therefore, it is imperative that as this technology is adopted, continuous efforts are made to monitor, validate, and refine these models to safeguard against biases and errors that may arise from rapid model evolution. Future research should focus on developing robust frameworks for evaluating the ethical implications, fairness, and transparency of AI models as they evolve, ensuring that advancements in AI technology are implemented responsibly and ethically.
Frankle, J., Dziugaite, G. K., Roy, D., and Carbin, M. Linear Mode Connectivity and the Lottery Ticket Hypothesis. In International Conference on Machine Learning. 2020.
Freeman, C. D. and Bruna, J. Topology and Geometry of Half- Rectifeedi Network Optimization. In International Conference on Learning Representations, 2017.
French, R. M. Catastrophic forgetting in connectionist net- works. Trends in Cognitive Sciences, 3(4), 1999.
Gadre, S. Y., Ilharco, G., Fang, A., Hayase, J., Smymis, G.,
Nguyen, T., Marten, R., Wortsman, M., Ghosh, D., Zhang, J., et al. Datacomp: In search of the next generation of multimodal datasets. Advances in Neural Information Processing Systems, 2024.
References
Ainsworth, S., Hayase, J., and Srinivasa, S. Git Re-Basin: Merging Models modulo Permutation Symmetries. In International Conference on Learning Representations, 2023.
Garipov, T., Izmailov, P., Podoprikhin, D., Vetrov, D. P.. and Wilson, A. G. Loss Surfaces, Mode Connectivity, and Fast Ensembling of DNNs. Advances in Neural Information Processing Systems, 2018
Cherti, M., Beaumont, R., Wightman, R., Wortsman, M., Ilharco, G., Gordon, C., Schuhmann, C., Schmidt, L., and Jitsev, J. Reproducible scaling laws for contrastive language-image learning. In Proceedings of the IEEE conference on Computer Vision and Pattern Recognition, 2023.
Helber, P., Bischke, B., Dengel, A., and Borth, D. AuroSAT: A Novel Dataset and Deep Learning Benchmark for Land Use and Land Cover Classification. IEEE Journal of Se- lected Topics in Applied Earth Observations and Remote Sensing, 12(7), 2019.
Hendrycks, D., Basart, S., Mu, N., Kadavath, S., Wang, F., Dorundo, E.,, Desai, R., Zhu, T., Parajuli, S., Guo, M., Song, D., Steinhardt, J., and Gilmer, J. The Many Faces of Robustness: A Critical Analysis of Out-of- -iit----iitiiutionn Generalization. IEEE International Conference on Com- puter Vision, 2021.
Choshen, L., Venezian, E., Slonim, N., and Katz, Y. Fusing finetuned models for better pretraining. arXiv preprint arXiv:2204.03044, 2022.
Cimpoi, M., Maji, S., Kokkinos, I., Mohamed, S., and Vedaldi, A. Describing Textures in the Wild. In Pro-
10

---

Update Your Transformer to the Latest Release: Re-Basin of Task Vectors
Ihaacoo, G., Wortsman, M., Wightman, R., Gordon, C.,, Car- lini, N., Taori, R., Dave, A., Shankar, V., Namkoong, H., Miller, J., Hajishirzi, H., Farhadi, A., and Schmidt, L. OpenCLIP, 2021.
McCloskey, M. and Cohen, N. J. Catastrophic Interfer- ence in Connectionist Networks: The Sequential Learn- ing Problem. In Psychology of learning and motivation, volume 24, pp. (09-165. Academic Press, 1989.
Ilharco, G., Wortsman, M., Gadre, S. Y., Song. S., Hajishirzi, H., Kornblith, S., Farhadi, A., and Schmidt, L. Patching open- open-vocabulary models by interpolating weights. Ad- vances Neural Information Processing Systems, 2022.
Navon, A., Shamsian, A., Achituve, I., Fetaya, E.,, Chechik, G., and Maron, H. Equivariant Architectures for Learning in Deep Weight Spaces. In International Conference on Machine Learning. 2023.
Ilharco, G., Ribeiro, M. T., Wortsman, M., Schmidt, L., Hajishirzi, H., and Farhadi, A. Editing Models with Task Arithmetic. In International Conference on Learming Representations, 2023.
Netzer, Y., Wang, T., Coates, A., Bissacco, A., Wu, B., Ng, A. Y., et al. Reading Digits Natural Images with Unsupervised Feature Learning. In Neural Information Processing Systems Workshops. Granada, 2011.
Imfeld, M., Graldi, J., Giordano, M., Hofmann, T., Anag- nostidis, S., and Singh, S. P. Transformer Fusion with Optimal Transport. In International Conference Learn- ing Representations, 2024.
Ortiz-Jimenez, G., Favero, A., and Frossard, P. Task Arith- metic in the Tangent Space: Improved Editing of Pre- Trained Models. Advances in Neural Information Pro- cessing Systems, 2024.
Izmailov, P., Podoprikhin, D., Garipov. T., Vetrov, D., and Wilson, A. G. Averaging Weights Leads to Wider Optima and Better Generalization. arXiv preprint arXiv:1803.05407, 2018.
Peña, F. A. G., Medeiros, H. R., Dubail, T., Aminbeidokhti, M.. Granger, E., and Pedersoli, M. Re-basin via implicit Sinkhorn differentiation. Proceedings of the IEEE conference on Computer Vision and Pattern Recognition, 2023.
Jonker, R. and Volgenant, T. A shortest augmenting path algorithm for dense and sparse linear assignment prob- lems. In Papers of the 16th Annual Meeting of DGOR in Cooperation with NSOR/Vorträge der 16. Jahrestagung der DGOR zusammen mit der NSOR, 1988.
Porrello, A., Bonicelli, L.,, Buzzega, P., Millunzi, M., Calder- ara, S., and Cucchiara, R. A Recond-Ordes Perspective Model Compositionaliiy and Incremental Learning. In International Conference Learning Representations, 2025.
Jordan, K., Sedghi, H., Saukh, O., Entezari, R., and Neyshabur, REPAIR: REnormalizing Permuted Acti- vations for Interpolation Repair. In International Confer- on Learning Representations, 2023.
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al. Learning Transferable Visual Models From Natural Language Supervision. In International Conference on Machine Learning, 2021.
Jovanović, I. and Stanić, Z. Spectral distances of graphs. Linear Algebra and its Applications, 436(5), 2012.
Krizhevsky, A., Hinton, G., et al. Learning Multiple Layers of Features from Tiny Images. Technical Report, Univer- sity of Toronto, 2009.
Rame, A., Kirchmeyer, M., Rahier, T., Rakotomamonjy, A., patrick gallinari, and Cord, M. Diverse Weight Averaging Generalization. In Advances in Neural Information Processing Systems, 2022.
for
Lu, M. Y., Chen, B., Williamson, D. F., Chen, R. J., Liang, I., Ding, T., Jaume, G., Odintsov, I., Le, L. P., Gerber, G., et al. A visual-language foundation model for computa- tional pathology. Nature Medicine, 30, 2024.
Ramé, A., Ahuja. K., Zhang, J., Cord, M., Bottou, L., and Lopez-Paz, D. Model Ratatouille: Recycling Diverse Models for Out-of- Distribution Generalization. In Inter- national Conference on Machine Learning, 2023.
Mall, U., Phoo, C. P., Liu, M. K., Vondrick, C., Hariharan, B., and Bala, K. Remote Sensing Vision-Language Foun- dation Models without Annotations via Ground Remote Alignment. In International Conference Learning Representations, 2024.
Singh, S. P. and Jaggi, M. Model Fusion via Optimal Trans- port. Advances in Neural Information Processing Systems, 2020.
Stallkamp. J., Schlipsing, M.. Salmen, J., and Igel, C. The German Traffic Sign Recognition Benchmark: A multi- class classification competition. In The 2011 interna- tional joint conference on neural networks. IEFE, 2011.
Matena, M. S. and Raffel, C. A. Merging Models with Fisher- Weighted Averaging. Advances in Neural Infor- mation Processing Systems, 35, 2022.

---

Update Your Transformer to the Latest Release: Re-Basin of Task Vectors
Stoica, G., Bolya, D., Bjorner, J., Ramesh, Hearn, T., and Hoffman, Ziplt! Merging Models from Different Tasks without Training. In International Conference Learning Representations, 2024.
Wang, A., Singh, A., Michael, J., Hill, F., Levy, O., and Bowman, S. R. GLUE A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding. In International Conference on Learning Representations, 2019.
Wortsman, M., Ilharco, G., Gadre, S. Y., Roelofs, R., GontijoLopes,, R., Morcos, A. S., Namkoong, H., Farhadi, A., Carmon, Y., Kornblith, S., et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. International Conference on Machine Learning, 2022a.
Wortsman, M., Ilharco, G., Kim, J. W., Li, M., Kornblith, S., Roelofs, R., Lopes, R. G., Hajishirzi, H., Farhadi, A., Namkoong, H., et al. Robust fine-tuning of zero- shot models. In Proceedings of the IEEE conference on Computer Vision and Pattern Recognition, 2022b.
Yadav, P., Tam, D., Choshen, L., Raffel, C. A., and Bansal, M. TIES -erg-nggg Resolving Interference When Merg- ing Models. Advances in Neural Information Processing Systems, 36, 2024.
Zhang, F. Z., Albert, P., Rodriguez-Opazo. C., van den Hengel, A., and Abbasnejad, E. Knowledge Composition using Task Vectors with Learned Anisotropic Scaling. Ad- vances Neural Information Processing Systems, 2024.
12

---

Update Your Transformer to the Latest Release: Re-Basin of Task Vectors
A. Appendix
Define:
A.1. On the Invariance to Permutations of our Metric for Inter-head Alignment
Q= X k= = XW
= XWn,
(22)
and split them by head:
Proposition A.1. Let E R be arbi- Irary. For any ht denote its singular values by ((h) (2) mmin(m,))) (h)), where (1(h) (2)h) ... O. For two matrices hi, h2 of the same shape, define
Q = [Q1, 1Q1, Q2.
Qul,
Qi E E Rxde,
(23)
and similarly for K and V. Let Pinter be inter-head per- mutation in SH, with induced permutation vector т, and let p Sd, be the intra-head permutation for head i. We form the block-penee pocckppemttationn matrix:
dp(h),he)) (10(b)) lp
(15)
where -1 is the usual p-norm for vectors. Then, for any permutation matrices P, E Rmxm and Pe E RTx", the row- and column-permucte matrix
H
P Patn
gta(t)
pl) intra!
(24)
where (f a binary x matrix with single at (1, x1)), and @ denotes the Kronecker product.
N' = P.hPo
(16)
has exactly the same singular values as h. In particular,
Step 1: Permuting the projection weights
1p) =0
(17)
Applying Paun to the query projections gives:
for every p. making d invariant under row- and column- permutations of h.
Q' = X W, Patm QQPat
11
11
Proof. If a permutation matrix, then PTp i.e.
Σ Qj atn Li, d)
is orthogonal. Furthermore, the singular values of any matrix h are given by the square root of the cigenvalues of h. h. If
i=1
N' P. Pc, then
Qu 100 Pi 110
11
(h') (PAPPJ (PrNP)
(18)
i=1
P. KT PAPPE
(19)
Hence,
ATP......
(20)
Q = Q+ 1(x)
intra
(1)
(25)
Since P. TTh P. a similarity transform not change the eigenvalues, & l and (f') h have the same eigen- values, and in turn h and h' share the same singular values. hence o(h') (h), and therefore
which does
where the new head Q corresponds to the head designated by the inter-head permutation '(i), modified according
to p ((0) The same applies to:
(()
dr, h = (lh - ahh')) II, 0,
(21)
K = K, intra
V = (1)
(26)
intra
proving that, for any row or column permutation of h, the distance d(h, h²) remains unchanged.
D
Step 2: Permuting the attention scores
Because each p()
is orthogonal (PPT = T), the attention intra
A.2. Proof of Equivariance of Multi-Head Attention to Structured Permutations 3.1
scores satisfy:
Q.KTT
Proof. We provide detailed, step-by-step proof showing that our two-stage alignment procedure sereduure-interheeee hnter-head reorder- ing followed by intra-head permutations preserves the functionality of a multi-head self-attention layer. Let:
A = softmax
Vdl.
Q. 1(6) intra
1000
KT" T (1))
= softmax
intra
dd
X E S be the input sequence.
Qx-1 KT
x
softmax
= Am-1(() 1())
Wg. We. W. E Rd Xdmast be the query, key, and value projection matrices.
Vdl.
Thanks to the orthogonality of the intra-head permutation blocks, the attention scores are only influenced by the inter- head permutation.
be the number of attention heads, each of dimen- sionality d; = d H.
13

---

Update Your Transformer to the Latest Release: Re-Basin of Task Vectors
Step 3: Permuting the value outputs
GTSRB
DTD
5
Task Gain (%)
For each head,
0
O A.V. (1) (o) Pintra xx-1(1) P
(()
(27)
-10
Task Vector
Ours
Step 4: Reconstructing the final output
-15
2
0
2
Concatenating all heads yields:
Supp. - Gain (%)
O' = 10 OH1== OP
(28)
-5
Conclusion. Applying Paun to the projection matrices is thus equivalent to permuting the output of multi-head attention. The self-attention layer remains functionally equivalent, and the original output can be recovered via O O'PT attn
-10
-15
0
2
0
2
Alpha
Alpha
Figure 6: Zero-shot gain/drop relative 0B of naive 0g + oT (blue) and strategy On + am((T) (red) varying a.
A.3. Full Procedure to Manage Residual Connections
Proof. We begin with the standard formulation of transformer block, ignoring LayerNorm for simplicity:
Battn W, MH).
A.4. Extended Comparison on the Application the Task Vector
Battn + X,
(29)
2 W.ReI.U((1z.).
In Fig. 6, we extend the sensitivity analysis of the scaling co- efficient a originally presented in Fig. 4 4- to additional visual classification tasks.
Zont
Ky 2;
Ignoring the ReLU activation function as well, we examine the impact of applying permutation to one layer within transformer block and then reversing it in the subsequent layer. This transfor- mation leads
A.5. Proof of Proposition 3.2
Proposition A.2. Proof. To assess how computational complexity scales with model size, define:
Zum Pwo WoPat (Patm MHA PIT (Pax)).
Pwo + Pinx,
(30)
2 Pw., W2W. Fw,,WiW, (z.)).
L: number of layers, evenly divided into MLP ( and self-attention (
= Pw.yy + Pw.z..
To ensure consistency in the permutation applicd to both addends within each residual block, replace the identity mapping with permutation composition where I, Pwo P. and Iam
H: number of attention heads.
Each MLP layer contains two lincar projections with dimension (dm: dz. and (d,, dm.). assuming d,, = d,-
Pw. Pw. This results in:
= PWo Kut + I.Pax = Pwr Zatn + Pwo 2.Z Pw,2s Pwn Pw,a, Pw,2..
Self-attention layers have Q, K, and V matrices, each of size (dm, dm).
(31)
After incorporating these compositions, the permutations remain consistent across each residual path, simplifying the block equa- tions to:
We now estimate the complexity for a single iteration of the weight-matching algorithm.
Zamn Pw, Wo (MHH(()))
MLP Layers. The main computational cost comes from computing a pairwise similarity matrix between rows of pro- jection matrices (O(d))) )), and solving (dm, dm) assign- ment via Hungarian algorithm (O(d),)) )). Hence, per-layer cost is:
2, Pw, Zum + Pwox,
(32)
2, Pw: W2 (i(zi)).
Zmm Pw325 + Pw.2,
With Pw, serving the input permutation for the subsequent layer.
(
(33)
14

---

Update Your Transformer to the Latest Release: Re-Basin of Task Vectors
Self-Attention Layers. Split into steps: inter-head
and intra-head permutation.
Inter-head permutation:
6/4 SVDs over matrices of size
(4r dm)
(
6d
(34)
H
Distance matrix over heads:
3
(35)
2
Hungarian algorithm over (H, H) matrix:
0(H3)
(36)
Intra-head permutation:
Per-head similarity:
2 O
dm
d
dm
(37)
H
Hungarian algorithm per head:
((
dm
(38)
H
Summed over heads:
d3
dm
)
d3
d3
H
=0
(39)
H
H
Total Self-Attention Cost per Layer.
2Hd
d3
d3
o
3
(40)
H
H
H²
Final Complexity. Summing across & MLP and & atten- tion layers:
L
6d
3H dm
d3 (41)
)
3
H
H
H
112
This expression can be algebraically simplified to more compact equivalent form:
o
(1.
3 + H22 H2dmm) )) H
3
(42)
H
So, the complexity scales polynomially with dm and H, and remains significantly lower than data-based fine-tuning.
15


Reviews:
[{"review_id": "hduvHpOoQR", "reviewer": "Reviewer_S7w3", "summary": "This paper explores how to update a fine-tuned downstream model from an older version of a pre-trained model to a newer pre-trained model without requiring re-finetuning. The paper proposes the TransFusion method, which incorporates attention head matching, alignment, and residual connection handling. Experiments on CV and NLP tasks demonstrate the effectiveness of the proposed approach.\n\n### Post-rebuttal\nI will maintain my score of 4.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Backbones used in NLP experiments\n\n> ... For example, which pre-trained models were used in the NLP experiments?\n\nWe sincerely apologize for the omission. In our experiments, we used two variants of OpenCLIP's ViT-B-16 text encoder: \"ViT-B-16/commonpool-l-s1b-b8k\" for $\\theta\\_A$ and \"ViT-B-16/datacomp-l-s1b-b8k\" for $\\theta\\_B$. We specifically chose these versions over models like BERT or T5 because we needed two variants of the same base model trained in different ways to represent ours $\\theta\\_A$ and $\\theta\\_B$.\n\n## Application to LLMs\n\n> ... The paper only conducts experiments on classification tasks. However, the problem it addresses is clearly more relevant to large language models. Can this method be applied to models at the 7B scale?\n\nOur method can be seamless applied to scenarios beyond classification (e.g., detection, segmentation, VQA). This flexibility arises from our approach not relying on classification-specific techniques nor imposing any assumptions regarding the type of loss function used during optimization.\n\nRegarding the application to 7B models such as LLaMA 2 and Mistral7b, this represents an intriguing area that we plan to explore in future work. We consider it feasible; indeed, as discussed in the response to Reviewer gzgf, the complexity of weight matching scales **linearly** with the number of layers. We also advocate for restricting fine-tuning operations to self-attention layers only, for efficiency reasons. Indeed, our method exhibits appealing computational complexity for these layers, as noted again in the response to Reviewer gzgf. Notably, this selective fine-tuning approach aligns conceptually with established state-of-the-art techniques such as QLORA [a], which fine-tune only the self-attention layers while excluding feed-forward layers. Consequently, re-basin operations could similarly be confined exclusively to self-attention layers, given that no information transfer is necessary for the remaining layers (as the task vector would be null).\n\n[a]: Dettmers et al., QLoRA: Efficient Finetuning of Quantized LLMs, NeurIPS 2023.", "claims_and_evidence": "The contributions claimed in the paper are supported by the experimental results.", "methods_and_evaluation_criteria": "The proposed method meets the requirements of the intended application scenario.", "theoretical_claims": "The theoretical proofs in the paper seem correct, but since I am not a researcher in this specific field, I cannot be completely certain.", "experimental_designs_or_analyses": "The experiments are designed with pretrained models and datasets from both the language and vision domains, demonstrating the effectiveness of the proposed method.", "relation_to_broader_scientific_literature": "If the pre-training and fine-tuning processes are viewed as two different tasks, the problem being studied may be related to continual learning, which aims to mitigate catastrophic forgetting when learning new knowledge.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths:\n+ The problem studied in the paper is practically meaningful, and the motivation behind the proposed method is clearly articulated.\n+ The proposed method is well-supported by theoretical foundations.\n\nWeaknesses:\n+ Some experimental details are not clearly specified. For example, which pre-trained models were used in the NLP experiments?\n+ The paper only conducts experiments on classification tasks. However, the problem it addresses is clearly more relevant to large language models. Can this method be applied to models at the 7B scale?", "other_comments_or_suggestions": "N/A", "questions_for_authors": "see above", "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": "3YZLNApMoM", "reviewer": "Reviewer_yjkC", "summary": "The introduces a new rebasin method for models for keeping models up-to-date as their underlying pretrained backbones evolve, focusing on transformers in particular. The author's method involves a “transport” the fine-tuning modifications—captured as a task vector—from the older base model to a new checkpoint​ without re-training or using data. The paper introduces TransFusion, a data-free and training-free re-basin procedure that realigns the task vector to be applicable on a new model release. The paper applies their method to visual and NLP tasks with ViT-B and different NLP classifiers.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Clarification\nWe clarify that our paper does not assume similarity between $\\theta\\_A$ and $\\theta\\_B$. As stated in Sec. 1 (line 67) and 3, these checkpoints may result from training on distinct data distributions or techniques. The re-basin mechanism is indeed designed to align models with significant differences in parameter space.\n\n## Using a small dataset to enhance transport\nWe agree; if even a small amount of data were available for each class, it would certainly be beneficial to use it. However, there are real-world scenarios where retaining data is not possible. That said, if these constraints do not apply, our method can be effectively utilized in conjunction with fine-tuning. To show this, starting from a small subset of $10$ shots per class, we follow \\[a\\] and learn one scaling coefficient per layer $\\alpha=\\[\\alpha\\_1,...,\\alpha\\_{|L|}\\]$. We then examine the performance after fine-tuning these coefficients.\n\nMethod|EuroSAT|DTD|GTSRB|SVHN\n-|-|-|-|-\n$\\theta\\_{B}+\\alpha\\tau$|+7.93|-1.44|+4.70|-15.98\n$\\theta\\_{B}+\\alpha\\pi(\\tau)$|**+10.00**|**+1.21**|**+6.80**|**+10.52**\n\nThere is a significant gain when fine-tuning a model that has undergone re-basin using our approach, represented as $\\theta_{B}+\\alpha \\pi(\\tau)$. In contrast, fine-tuning from $\\theta_{B}+\\alpha \\tau$ (which does not involve permutation) produces inferior results. This underscores that rebasing and fine-tuning should not be viewed as mutually exclusive but as complementary strategies.\n\n\\[a\\] Knowledge composition using task vectors with learned anisotropic scaling. NeurIPS 2024.\n## More experiments on different model sizes\nWe repeated the tests of Tab.1 using the CLIP ViT-L/14 models pretrained on \"laion400m\\_e31\" ($\\theta\\_A$) and \"commonpool\\_xl\\_clip\\_s13b\\_b90k\" ($\\theta\\_B$).\n\n||EuroSAT|EuroSAT-Supp|DTD|DTD-Supp|GTSRB|GTSRB-Supp|SVHN|SVHN-Supp|\n|-|-|-|-|-|-|-|-|-|\n|$\\theta\\_{B}$|67.37|88.66|64.94|88.66|60.15|88.66|68.53|88.66|\n|$\\theta\\_{B}+\\tau$|-4.07|-0.33|**+0.74**|-0.36|+0.70|-0.93|+2.04|-3.73|\n|**Ours**|**+1.15**|**-0.01**|+0.37|**-0.20**|**+2.29**|**-0.13**|**+5.33**|**-0.73**|\n\nThe results align with those reported in the paper.\n## Computational Overhead\nFor an extensive complexity analysis, please refer to our response to gzgf. We highlight two additional aspects:\n\n- By separating multi-head attention alignment into inter- and intra-head, we significantly reduce scaling issues related to width. For example, in CLIP’s ViT-B/16 (12 heads), the search space is reduced from $768!$ to $12 \\times 64!$. For a ViT-L/14, GiT Re-basin’s space grows to $1024!$, while TransFusion’s remains much smaller at $16 \\times 64!$.\n- The weight-matching procedure depends only on $\\theta\\_A$ and $\\theta\\_B$, and not on the downstream tasks to be transferred. This ensures that associated costs can be incurred only once. In practice, the entity responsible for releasing the updated model can precompute and distribute the permutation matrices, eliminating the computational burden for end users.\n## Functional role\nIntuitively, encoding semantic similarity would require optimizing permutations based on the activations of each layer, which is hard in a data-free scenario like ours. However, if few data are available, one can leverage our decoupled inter- and intra-head approach and replace our energy-based metric with one crafted on activations discrepancy.\n\nNevertheless, while singular values do not directly encode high-level semantics, they capture structural properties. The largest singular value of a linear transformation corresponds to its Lipschitz constant; hence, aligning attention heads with similar distributions of singular values implies aligning heads with comparable sensitivity to input perturbations. Moreover, the spectrum of singular values indicates how information is propagated: a low-rank attention head typically focuses on a narrow subspace of the input, whereas a full-rank head affects a broader range of directions.\n## Hyperparameters\nConsidering $\\tau$, it is optionally scaled by a hyperparameter $\\alpha$ fixed at $1$ for all experiments, due to our assumption of no access to data, which precluded further tuning. To assess sensitivity to $\\alpha$, we kindly refer to Fig. 4 (main paper) -- applying TransFusion leads to increased robustness across various choices of $\\alpha$.\n## Residual connections\nWe remark that our approach to handling residual connections is not specific to transformers but is applicable to any residual block. Indeed, residual connections assume an implicit identity mapping \n$z=I\\_n z\\_{\\text{out}}+I\\_n x$, where $I\\_n$ represents the identity matrix. However, once the permutations are applied to all layers, this assumption no longer holds, resulting in $z=P\\_{\\text{out}} z\\_{\\text{out}}+P\\_{\\text{input}} x$. As shown in Fig. 5, this deviation is sufficient to break the functional equivalence of the model.\n## Backbones for NLP tests\nWe apologize for the omission; please refer to our response to S7w3.", "claims_and_evidence": "The core concern I have about the claims of the paper are as follows\n\n* Assumption of structural similarity of models: One core assumption of the method is that the new checkpoint $\\theta_B$ is largely similar to tthe original $\\theta_A$. Is this always true? For instance, if $\\theta_B$ is trained under different conditions such as a new data distribution or slight architectural tweaks, then the permutation alignment may not capture necessary transformations. Moreover, when the representation spaces of your two networks diverge, then the transport of $\\tau$ could be suboptimal or even detrimental.\n* Scalability/Computational Overhead: The procedure requires solving a series of linear assignment problems for every transformer layer and for each attention head pair. Although the Hungarian algorithm is polynomial in complexity, the cumulative cost might become significant for very deep or wide models.\n* Hyperparameter Tuning: Could the authors discuss more on how to choose $\\alpha$ in practice? I worry that sensitivity may make this difficult to tune. \n* Disadvantages of a data free transfer: One major disadvantage and limitation I would have liked to see discussed is about cases where domain-specific adjustments are necessary. In many real-world cases, even a minimal dataset could help fine-tune the alignment further, suggesting that a semi-supervised or data-augmented variant might offer better performance. It would be helpful for the authors to discuss _when_ we should use TransFusion.", "methods_and_evaluation_criteria": "I see no issues with the methods and evaluation criteria. The authors evaluate on well-known NLP tasks which seem reasonable for this paper although I would have like to see pretrained decoder models. I would be curious about evaluating on a dataset like GSM after finetuning to see capabilities of the model on a dataset that is (reasonably) potentially out of distribution.", "theoretical_claims": "I checked proofs in the appendix and didn't see any issues. However, I do have one concern about theoretical aspects of this work. \n\n* Spectral Metric for Head Matching: While using singular values to form a permutation-invariant distance seems justified, it's not clear if this metric always aligns heads based on functional roles. Two heads with similar singular value profiles might perform very different roles in the network. This could lead to cases where heads that are similar in “energy” are paired, but the semantic content of their learned representations is not matched, potentially reducing the efficacy of the transport.\n* Residual Connections: The method decides redefine the identity mapping in residual connections. To me, this may be sensitive to the specific ordering and scale of the permutation matrices. And, this also keeps this method from applying to other non-transformer based methods.", "experimental_designs_or_analyses": "* The authors should expand on what encoder language model they use. Is it BERT, T5, etc.\n* It would be great for the authors to discuss how their method scales to deeper and wider models. If the authors could run different models sizes, that would be very useful for contextualizing their approach.", "relation_to_broader_scientific_literature": "The most significant paper that is similar to this approach is obviously Git Rebasin (ICLR 2023). The authors cite and discuss this paper at length to explain the similarities and differences of their approach. The authors also discuss task arithmetic and weight interpolation at length, as this is relevant to their approach.", "essential_references_not_discussed": "None that I can find.", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "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": "yoB7OFsW80", "reviewer": "Reviewer_gzgf", "summary": "This paper addresses a critical challenge for foundation models, i.e., fine-tuned models becoming obsolete when their base models are updated. The authors introduce TransFusion, a data-free method to transfer task vectors from an old base model to a new one. By leveraging a structured permutation strategy tailored for Transformers, i.e., addressing multi-head attention alignment and residual connections through spectral theory, the method ensures functional equivalence while preserving generalization. Extensive experiments on vision (e.g., EuroSAT, SVHN) and NLP tasks (e.g., GLUE) demonstrate significant improvements over baselines including Git Re-Basin, validating the effectiveness and robustness of the proposed approach.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Complexity analysis\nTo assess how computational complexity scale with model size, we define:\n\n- $|L|$: number of layers, evenly divided into MLP ($\\frac{|L|}{2}$) and self-attention ($\\frac{|L|}{2}$).\n- $|H|$: number of attention heads.\n- Each MLP layer contains two linear projections with dimension $(d\\_m, d\\_h)$ and $(d\\_h, d\\_m)$. We assume $d\\_m = d\\_h$ for simplicity.\n- Self-attention layers have Q, K, and V matrices, each with dimensions $(d\\_m, d\\_m)$.\n\nWe now estimate the complexity for MLP and self-attention layers for a single iteration of the weight-matching algorithm.\n### MLP Layers\nPermutation alignment for MLP layers resembles GiT Rebasin. The main computational cost involves computing a matrix of pairwise dot products between rows in the projection matrix, with complexity $O(d\\_m^3\\)$. Subsequently, applying the Hungarian algorithm to a $(d\\_m, d\\_m)$ matrix also has complexity $O(d\\_m^3)$. Thus, each MLP layer incurs $O(d\\_m^3)$.\n### Self-Attention Layers\nThe analysis is split into two steps: inter-head and intra-head permutations.\n\n**Inter-head permutation**:\n\n- Computing singular value decompositions (SVDs) for matrices Q, K, V across $2$ networks (A,B) and $|H|$ heads results in $6|H|$ SVDs. Each SVD on head-level matrices sized $(\\frac{d\\_m}{|H|},d\\_m)$ has complexity $O(\\frac{d\\_m^3}{|H|^2})$. Hence, total complexity for all SVDs is $O(\\frac{6 d\\_m^3}{|H|})$.\n- Computing distance matrices for Q, K, V, each sized $(|H|,|H|)$, involves $d\\_m$ operations per element, hence the complexity is $O(\\frac{3|H|^2 d\\_m}{2})$.\n- Applying the Hungarian algorithm to the resulting $(|H|,|H|)$ matrix incurs complexity: $O(|H|^3)$.\nThus, inter-head permutation complexity is: $O(\\frac{6 d\\_m^3}{|H|}+\\frac{3|H|^2 d\\_m}{2}+|H|^3)$.\n\n**Intra-head permutation**:\n\n- Computing similarity matrices for intra-head permutations involves cost matrices of dimensions $(\\frac{d\\_m}{|H|},\\frac{d\\_m}{|H|})$, incurring complexity per head: $O((\\frac{d\\_m}{|H|})^2 \\times d_m)$=$O(\\frac{d\\_m^3}{|H|^2})$.\n- Applying the Hungarian algorithm separately for each head results in a per-head computational complexity of $O((\\frac{d\\_m}{|H|})^3)$.\nSumming across all heads, total intra-head permutation complexity is $O(|H| \\times ((\\frac{d\\_m^3}{|H|^2})+(\\frac{d\\_m}{|H|})^3)=O(\\frac{d\\_m^3}{|H|}+\\frac{d\\_m^3}{|H|^2})$.\n\nTherefore, total complexity per self-attention layer is: $\nO(\\frac{6 d\\_m^3}{|H|}+\\frac{3|H|^2 d\\_m}{2}+|H|^3+\\frac{d\\_m^3}{|H|}+\\frac{d\\_m^3}{|H|^2})$.\n### Overall Complexity\nConsidering all layers, the total complexity combines the contributions from MLP and self-attention layers, yielding:  \n$O\\big(\\frac{|L|}{2}d\\_m^3+\\frac{|L|}{2} (\\frac{6 d\\_m^3}{|H|}+\\frac{3|H|^2 d\\_m}{2}+|H|^3+\\frac{d\\_m^3}{|H|}+\\frac{d\\_m^3}{|H|^2})\\big).$\n\nThis complexity scales **polynomially**, dominated by terms involving $d\\_m^3$ and $|H|^3$. The method remains substantially more efficient than full retraining, avoiding costly gradient computations or data iterations.\n## Aligning models with different architectures\nThe idea of aligning models with minor variations in architecture, such as different layer counts, is worth exploring. To address this, one simple approach could involve selectively pruning layers from the model with more layers to match its counterpart. For instance, one could remove redundant and unimportant layers. Alternatively, one could adopt the opposite strategy and replicate the last block of the smaller network multiple times to achieve a match. We plan to explore these aspects further in future work.\n## Related works\nBoth our work and that mentioned by the reviewer aim to improve transfer during fine-tuning. However, while DR-Tune and AdaRand assume access to downstream task data during fine-tuning, our method requires no training data to transfer fine-tuning from one model to another. Moreover, these solutions employ data-driven regularization during optimization; instead, our approach uses a parameter alignment technique that eliminates the need for training steps.\n## Quality of task vectors\n> ... How does the quality of the task vectors affect the performance of the proposed method? ... What does the limited gains on DTD mean? Is it related to task vector quality or dataset characteristics?\n\nOur approach follows a twofold procedure, i.e., align the two models $\\theta\\_A$ and $\\theta\\_B$, and then transport the task vector $\\theta\\_B + \\pi(\\tau)$. Regarding the quality of the task vector, the first stage is unaffected as it solely depends on $\\theta\\_A$ and $\\theta\\_B$. The second step, instead, employs $\\tau$ and thus it could be affected by low-quality fine-tuning $\\tau$ (after all, garbage in, garbage out). For instance, when considering DTD, we argue that the lower quality of the task vector is directly attributable to the challenging characteristics of the dataset, which features only 40 examples per class on average (for comparison, the other datasets we use have at least around 1000 examples per class).", "claims_and_evidence": "The core claims are well-supported. As shown in Table 1, the proposed method consistently outperforms baselines in zero-shot task accuracy (e.g., +4.95% on EuroSAT, Table 1) while maintaining generalization (e.g., minimal drops on ImageNet-R). Figure 5 and Table 3 also show the advantages of using Functional Equivalence and Transformer-Specific Design.", "methods_and_evaluation_criteria": "The proposed methods are reasonable, building upon established neural network permutation invariances and extending them thoughtfully to Transformer architectures. Their innovative two-step spectral-based attention head alignment and residual connection permutation strategy effectively address previously unresolved challenges (head contamination and residual mismatch). The chosen evaluation metrics, including zero-shot accuracy on downstream tasks and generalization to support sets, are appropriate and directly relevant to assessing the effectiveness and practical applicability of the proposed solution.", "theoretical_claims": "The authors provide rigorous theoretical analyses, particularly regarding their proposed permutation-invariant spectral distance metric for attention head alignment. Their mathematical proof in Appendix A.1 confirms that singular value-based distance metrics are indeed invariant to permutations, logically and rigorously justifying the attention alignment strategy. Additionally, their handling of residual connections is clearly explained and mathematically sound, ensuring consistency across permutations and preserving functional equivalence.", "experimental_designs_or_analyses": "Experimental design is thorough, comprehensive, and rigorous. Evaluations span multiple tasks, including image classification and NLP benchmarks, demonstrating cross-modal applicability. Baselines (e.g., Git Re-Basin, Optimal Transport, and vanilla transport) are effectively selected to illustrate clear performance distinctions. The results consistently show substantial improvements from TransFusion, with sufficient statistical robustness inferred from the consistently large margins observed.", "relation_to_broader_scientific_literature": "The paper is well-connected to existing literature. The following works are also suggested to be discussed:\n[1] Zhou et al. DR-Tune: Improving Fine-tuning of Pretrained Visual Models by Distribution Regularization with Semantic Calibration. In ICCV, 2023.\n[2] Yamaguchi et al. Adaptive Random Feature Regularization on Fine-tuning Deep Neural Networks. In CVPR, 2023.", "essential_references_not_discussed": "No critical missing references were identified. The existing cited literature covers the majority of essential works thoroughly, and any minor omissions do not detract from understanding or evaluating the paper’s core contributions.", "other_strengths_and_weaknesses": "Strengths:\n+This paper proposes a novel solution to Transformer re-basin with theoretical guarantees.\n+The proposed method enables cost-effective model updates without data access.\n+This paper provides strong empirical validation across modalities and tasks.\n\nWeaknesses:\n-The propsoed method requires identical model architectures. Will slight changes (e.g., layer counts) deteriorate the alignment?\n-How is the scalability of the permutation alignment for large models?\n-How does the quality of the task vectors affet the performance of the proposed method?.", "other_comments_or_suggestions": "No.", "questions_for_authors": "1.Could the method handle minor architectural changes (e.g., added layers) via adaptive permutations?\n2.What does the limited gains on DTD mean? Is it related to task vector quality or dataset characteristics?\n3.How does the computational complexity scale with model size (e.g., ViT-B vs. ViT-L)?", "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": "vpSK3wM4aa", "reviewer": "Reviewer_ejBz", "summary": "This paper introduces TransFusion, a method that re-basins task vectors by aligning model weights to base models in the parameter space, aiming at adapting task vectors to a later version of the model. In particular, TransFusion employs a two-step permutation process: inter-head alignment using spectral distance and intra-head matching to align each pair of head matrices. Experiments on vision and NLP tasks show that TransFusion outperforms existing re-basin methods, preserving fine-tuning effects with improved zero-shot generalization.\n\n## update after rebuttal\nThe rebuttal helped address my concerns on functional equivalence, so I raised my score to 3.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Functional equivalence\nThe central criticism of the reviewer revolves around the preservation of functional equivalence in self-attention layers. To address this justifiable concern, we provide a proof showing that our two-stage method **ensures** functional equivalence.\n\n**Theorem** Let $P\\_{inter}$ be a permutation of attention heads, and $P\\_{intra}$ be a set of permutations applied independently within each head. Then, applying $P\\_{inter}$ followed by $P\\_{intra}$ preserves the functional equivalence of the multi-head self-attention layer.\n\nNotation:\n\n- $X$: input sequence of shape $(S,d\\_m)$\n- $W_q,W_k,W_v$: weight matrices each of shape $(d\\_m,d\\_m)$\n- $|H|$: number of heads\n- $d\\_k=d\\_m/|H|$: number of units within each head\n\nThe self-attention layer:\n\n- computes $Q,K,V$ matrices as $Q=X W\\_q$, $K=X W\\_k$, $V=X W\\_v$\n- splits Q into $|H|$ heads, $Q=\\[Q\\_1,Q\\_2,...,Q\\_{|H|}\\]$ (the same for $K$ and $V$).\n\n**Step 1** The two-stage permutation procedure involving inter- and intra-head swaps can be summarized by a single, composed permutation matrix $P\\_{attn}$. Such a matrix of shape $(d\\_m,d\\_m)$ has a notable form, namely a grid $|H|\\times|H|$ of block matrices. Each block has shape $(d\\_k,d\\_k)$ and is either filled with zeros or a valid matrix permutation of $d\\_k$ units. Moreover, in each row and column, there can be only one permutation block. An example with three heads:\n```\nP_{attn}=\nP_{intra}^0, 0, 0\n0, 0, P_{intra}^1\n0, P_{intra}^2, 0\n```\nwith heads swapped according to $P\\_{inter}$:\n```\nP_{inter}=\n1, 0, 0\n0, 0, 1\n0, 1, 0\n```\nFor the sake of notation, we will use the permutation vector $\\pi$ to denote how $P\\_{inter}$ swaps heads. In this example $\\pi=[0,2,1]$.\n\nNotably, a block matrix $P\\_{attn}$ with the features mentioned above can be written as:\n\n$P\\_{attn}=\\sum\\_i^{|H|} E^{i,\\pi(i)} \\otimes P\\_{intra}^i$\n\nwhere $\\otimes$ is the Kronecker product and $E^{i,\\pi(i)}$ is a matrix filled with zeros except at position $i,\\pi(i)$, where it contains a one.\n\n**Step 2** When applying the block matrix $P_{attn}$ to the query weight $W\\_q$, we get the following query tokens:\n\n$Q^{'}=X W\\_q P\\_{attn}=Q P\\_{attn}=\\big[ \\sum\\_{j=1}^{|H|} Q\\_{j} P\\_{attn}[j,i] \\big]\\_{i=1}^{|H|}=\\[ Q\\_{\\pi^{-1}(i)} P\\_{intra}^{\\pi^{-1}(i)} \\]\\_{i=1}^{|H|}$\n\nwhere $P\\_{attn}[j,i]$ refers the block at position $(j,i)$ (which can be either zero-filled or not). The relation is relevant as it entails:\n\n$Q\\_i^{'}=Q\\_{\\pi^{-1}(i)} P\\_{intra}^{\\pi^{-1}(i)}$\n\nThe new head $Q\\_i^{'}$ corresponds to the head designated by the inter-head permutation $\\pi^{-1}(i)$, modified according to the permutation $P\\_{intra}^{\\pi^{-1}(i)}$ applied to its units. Note that the same result applies to $K\\_i^{'}$ and $V\\_i^{'}$.\n\n**Step 3: attention** For each head, the attention matrix (after permutation) is:\n\n$A\\_i^{'}=\\text{softmax}(Q\\_i'{K\\_i'}^T/\\sqrt{d\\_k})=\\text{softmax}(Q\\_{\\pi^{-1}(i)} P\\_{intra}^{\\pi^{-1}(i)} {P\\_{intra}^{\\pi^{-1}(i)}}^T {K\\_{\\pi^{-1}(i)}}^T /\\sqrt{d\\_k})$\n\n$\\quad \\ =\\text{softmax}(Q\\_{\\pi^{-1}(i)}{K\\_{\\pi^{-1}(i)}}^T/\\sqrt{d\\_k})$\n\n$\\quad \\ =A\\_{\\pi^{-1}(i)}$\n\nThanks to the orthogonality of the intra-head permutation blocks, the attention scores are only influenced by the inter-head permutations.\n\n**Step 4: output** For each head, the output is:\n\n$O\\_i^{'}=A\\_i^{'}V\\_i^{'}=A\\_{\\pi^{-1}(i)}V\\_{\\pi^{-1}(i)}P\\_{intra}^{\\pi^{-1}(i)}=O\\_{\\pi^{-1}(i)}P\\_{intra}^{\\pi^{-1}(i)}$\n\nThe final output is obtained by concatenating all heads:\n\n$O'=\\[O\\_1^{'},O\\_2^{'},...,O\\_{|H|}^{'} \\]=\\[ O\\_{\\pi^{-1}(i)} P\\_{intra}^{\\pi^{-1}(i)} \\]\\_{i=1}^{|H|}=\\big[ \\sum\\_{j=1}^{|H|} O\\_{j} P\\_{attn}[j,i] \\big]\\_{i=1}^{|H|}=O P\\_{attn}\n$\n\n**Result** To sum up, applying a block permutation $P\\_{\\text{attn}}$ to each projection matrix is equivalent to permuting the output of the self-attention mechanism $O'=O P\\_{\\text{attn}}$. This demonstrates functional equivalence in the context of our approach. When the output is fed to the next layer, we have to multiply it by ${P\\_{\\text{attn}}}^T$ to recover the original output.\n\n> In my opinion, if the rows of different head matrices align better than rows in the same head matrix, permuting the rows across different head matrices is acceptable as long as the functional equivalence is preserved.\n\nDue to space constraints, a detailed discussion is not possible. It is important to note that permuting rows across heads means that the blocks within $P\\_{\\text{attn}}$ at each head level are not necessarily orthogonal. Hence, the associated $P\\_{\\text{intra}}$ term does not cancel out when calculating the input to the softmax function. This creates a challenge in reversing the effects of the permutation, which is necessary to ensure that functional equivalence is maintained after applying the softmax function.\n## Other aspects\nFor details on the backbones used in our NLP experiments, as well as additional tests across different model sizes and a complexity analysis, please see our response to yjkC.", "claims_and_evidence": "While the paper presents feasible empirical results, some claims lack theoretical guarantees or deeper analysis. For example, It is not explicitly shown how well TransFusion aligns two self-attention layers.", "methods_and_evaluation_criteria": "While the paper presents promising results, it is unclear whether the proposed two-step permutation process fully preserves the functional equivalence of the self-attention mechanism. In objective functions (8) and (9), each permutation matrix is optimized separately. The approach appears to lack explicit constraints to guarantee that the transformed model is functionally equivalent to the original model. Further clarification or theoretical verification would strengthen the evidence for the effectiveness of TransFusion.\n\nIn addition, complexity is an important metric for evaluating the efficiency of re-basin [1]. However, this paper does not include the complexity of TransFusion theoretically or empirically.\n\n[1] Ainsworth, Samuel, Jonathan Hayase, and Siddhartha Srinivasa. \"Git Re-Basin: Merging Models modulo Permutation Symmetries.\" The Eleventh International Conference on Learning Representations.", "theoretical_claims": "I checked the correctness of Theorem A.1, i.e., permutation preserves the spectral distance. This theorem provides motivation for choosing the spectral distance to measure the alignment of different attention head matrices.", "experimental_designs_or_analyses": "- The paper does not clearly specify the details of the transformer models used for NLP experiments\n- The paper evaluates only one transformer model backbone, limiting the generalizability of the findings. It would be helpful to include more experiments on different architectures to ensure the robustness and scalability of the proposed TransFusion.", "relation_to_broader_scientific_literature": "This paper proposes an important application of model re-basin (parameter matching): to adapt task vectors to different versions of the base model. Additionally, it proposes a novel permutation for the self-attention mechanism to align the multi-head attention parameters while preserving the functional equivalence. This can benefit further studies on permutation invariance and mode connectivity of transformers.", "essential_references_not_discussed": "I don't find any missing references.", "other_strengths_and_weaknesses": "Strengths:\n- This study extends parameter matching techniques to transformer models, improving the alignment of self-attention layers by permutation.\n- This paper is overall easy to follow and well-organized.\n\nWeaknesses:\n- The definition of symbols is a bit confusing. $H$ is first defined as the number of heads and then defined as the head matrix. Additionally, the head matrix $H_i^A$ should have a subscript (q, k, or v) to indicate the role of the matrix. \n- It is mentioned that the proposed head alignment techniques address the head contamination problem. However, it is unclear to me why we need to tackle this issue. In my opinion, if the rows of different head matrices align better than rows in the same head matrix, permuting the rows across different head matrices is acceptable as long as the functional equivalence is preserved.\n- It is strange to separately optimize $P_{inter}$ and $P_{intra}$ as the result permutation might not preserve the functional equivalence of a self-attention layer, which is a prerequisite for re-basin. It would be better to explicitly show how the joint application of $P_{inter}$ and $P_{intra}$ leads to functional equivalence.\n- Although I like the novel problem setting of adapting task vectors to different base models, this new problem setting does not bring distinct technical challenges. The novel part of two-step attention head alignment only solves the head contamination problem (with an unclear motivation to me). It would be helpful to highlight the technical challenges and contributions of TransFusion for better significance.", "other_comments_or_suggestions": "Please refer to Strengths And Weaknesses", "questions_for_authors": "Please refer to Strengths And Weaknesses", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper investigates the question of whether it possible to “upgrade” a fine-tuning of a base model to a new upgraded base model, without having to fine-tune again, by relying instead on task vectors. The authors show that this is possible by adapting ideas from model [re-basin](https://arxiv.org/abs/2209.04836) and improving it specifically for transformers by carefully handling attention heads and residual connections, showing promising empirical results on vision and language tasks.

The reviewers all agreed that the problem and proposed solution are of significant interest and I thus recommend acceptance.
Decision: Accept (poster)