Functional equivalence

The 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.

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.

Notation:

• $X$: input sequence of shape $(S,d\_m)$
• $W_q,W_k,W_v$: weight matrices each of shape $(d\_m,d\_m)$
• $|H|$: number of heads
• $d\_k=d\_m/|H|$: number of units within each head

The self-attention layer:

• computes $Q,K,V$ matrices as $Q=X W\_q$, $K=X W\_k$, $V=X W\_v$
• splits Q into $|H|$ heads, $Q=$ Q_1,Q_2,...,Q_{|H|} $$ (the same for $K$ and $V$).

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:

P_{attn}=
P_{intra}^0, 0, 0
0, 0, P_{intra}^1
0, P_{intra}^2, 0

with heads swapped according to $P\_{inter}$:

P_{inter}=
1, 0, 0
0, 0, 1
0, 1, 0

For 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]$.

Notably, a block matrix $P\_{attn}$ with the features mentioned above can be written as:

$P\_{attn}=\sum\_i^{|H|} E^{i,\pi(i)} \otimes P\_{intra}^i$

where $\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.

Step 2 When applying the block matrix $P_{attn}$ to the query weight $W\_q$, we get the following query tokens:

$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|}$

where $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:

$Q\_i^{'}=Q\_{\pi^{-1}(i)} P\_{intra}^{\pi^{-1}(i)}$

The 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^{'}$.

Step 3: attention For each head, the attention matrix (after permutation) is:

$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})$

$\quad \ =\text{softmax}(Q\_{\pi^{-1}(i)}{K\_{\pi^{-1}(i)}}^T/\sqrt{d\_k})$

$\quad \ =A\_{\pi^{-1}(i)}$

Thanks to the orthogonality of the intra-head permutation blocks, the attention scores are only influenced by the inter-head permutations.

Step 4: output For each head, the output is:

$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)}$

The final output is obtained by concatenating all heads:

$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} $

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.

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.

Due 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.

Other aspects

For 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.

Complexity analysis

To assess how computational complexity scale with model size, we define:

• $|L|$: number of layers, evenly divided into MLP ($\frac{|L|}{2}$) and self-attention ($\frac{|L|}{2}$).
• $|H|$: number of attention heads.
• 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.
• Self-attention layers have Q, K, and V matrices, each with dimensions $(d\_m, d\_m)$.

We now estimate the complexity for MLP and self-attention layers for a single iteration of the weight-matching algorithm.

MLP Layers

Permutation 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)$.

Self-Attention Layers

The analysis is split into two steps: inter-head and intra-head permutations.

Inter-head permutation:

• 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|})$.
• 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})$.
• Applying the Hungarian algorithm to the resulting $(|H|,|H|)$ matrix incurs complexity: $O(|H|^3)$. Thus, inter-head permutation complexity is: $O(\frac{6 d\_m^3}{|H|}+\frac{3|H|^2 d\_m}{2}+|H|^3)$.

Intra-head permutation:

• 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})$.
• Applying the Hungarian algorithm separately for each head results in a per-head computational complexity of $O((\frac{d\_m}{|H|})^3)$. Summing 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})$.

Therefore, total complexity per self-attention layer is: $O(\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})$.

Overall Complexity

Considering all layers, the total complexity combines the contributions from MLP and self-attention layers, yielding:
$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).$

This 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.

Aligning models with different architectures

The 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.

Related works

Both 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.

Quality of task vectors

... 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?

Our 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).

Backbones used in NLP experiments

... For example, which pre-trained models were used in the NLP experiments?

We 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$.

Application to LLMs

... 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?

Our 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.

Regarding 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).

[a]: Dettmers et al., QLoRA: Efficient Finetuning of Quantized LLMs, NeurIPS 2023.

Clarification

We 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.

Using a small dataset to enhance transport

We 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.

There 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.

$a$ Knowledge composition using task vectors with learned anisotropic scaling. NeurIPS 2024.

More experiments on different model sizes

We 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$).

The results align with those reported in the paper.

Computational Overhead

For an extensive complexity analysis, please refer to our response to gzgf. We highlight two additional aspects:

• 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!$.
• 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.

Functional role

Intuitively, 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.

Nevertheless, 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.

Hyperparameters

Considering $\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$.

Residual connections

We 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 $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.

Backbones for NLP tests

We apologize for the omission; please refer to our response to S7w3.