Merging Text Transformer Models
from Different Initializations

Given two models trained on the same data but from from separate initializations, namely ${\bm{\theta}}_{A}$ and ${\bm{\theta}}_{B}$, we compute post-activation features for each layer or sublayer parameter ${\bm{W}}_{\ell}\subset{\bm{\theta}}$ in order to determine which features might correspond between models (Ainsworth et al., 2023; Stoica et al., 2024). In our setting, features are computed at the token level, and all special non-vocabulary and non-padding tokens are always included (such as [SEP], [CLS]). At a given layer or sublayer, we compute $d$-dimensional activations across $n$ tokens from both models ${\bm{X}}_{A},{\bm{X}}_{B}\in\mathbb{R}^{n\times d}$, and then determine feature relatedness via cross-correlation computed as the following:

where $\bm{\sigma}$ are feature standard deviations, and $\bm{\mu}$ are feature means. We choose to standardize and mean-center the features before comparing them because the magnitude of feature values can vary greatly in some pre-trained text Transformers (Puccetti et al., 2022). Then, given feature cross-correlations ${\bm{C}}\in\mathbb{R}^{d\times d}$, we compute the optimal permutation as the following:

where $\pi:\{1,2,...,d\}\rightarrow\{1,2,...,d\}$ is a permutation mapping. This optimization problem finds the feature correspondences between the two models that lead to the highest total correlation captured. This problem an instance of the assignment problem and can be solved using the Jonker-Volgenant algorithm (Crouse, 2016; Tatro et al., 2020; Ainsworth et al., 2023).

After converting the map $\pi^{*}$ to its corresponding permutation matrix ${\bm{P}}$, we can apply ${\bm{P}}$ to the original weight matrix ${\bm{W}}^{B}_{\ell}\subset{\bm{\theta}}_{B}$ so that the order of the layer’s features most closely resembles that of ${\bm{W}}^{A}_{\ell}\subset{\bm{\theta}}_{A}$:

We then apply ${\bm{P}}^{\mathrm{T}}={\bm{P}}^{-1}$ to the next layer in order to unpermute the new ordering in model $\theta_{B}$:

After applying all computed permutations to ${\bm{\theta}}_{B}$ to get ${\bm{\theta}}_{B^{\prime}}$, the final merged model is computed as $\lambda{\bm{\theta}}_{A}+(1-\lambda){\bm{\theta}}_{B^{\prime}}$ for some $\lambda\in[0,1]$.

In finding corresponding neurons in the Multi-Headed Attention (MHA) parameters, namely the key (${\bm{W}}_{K}$), query (${\bm{W}}_{Q}$), value (${\bm{W}}_{V}$), and linear layer (${\bm{W}}_{O}$) weights, we propose several methods to compute potential permutation matrices.

For each of the key, query, and value weights, the full parameter ${\bm{W}}\in\mathbb{R}^{d_{\text{model}}\times d_{\text{model}}}$ is logically partitioned into $h$ attention heads each of output dimension $d_{k}=d_{\text{model}}/h$ (Vaswani et al., 2017). In order to apply a permutation to these full weight matrices and maintain the functional equivalence of the overall model, permutations must operate on each attention head separately, and not permute features between attention heads. This is because the final hidden vector from MHA reflects a concatenation of the result from each head, which are computed separately with weights ${\bm{W}}_{K_{i}},{\bm{W}}_{Q_{i}},{\bm{W}}_{V_{i}}$ for head $i$.

Since our models are trained from distinct initializations, the correspondence of their attention heads may differ in addition to the correspondence of features within each head. This distinction is diagrammed in the first two permutations of Figure 2. We collect features from just after the attention computation and before the linear layer following attention. These features are the aforementioned concatenations. We first compute ${\bm{C}}=\text{corr}({\bm{X}}_{A},{\bm{X}}_{B})$ and then partition this correlation matrix by heads into $d_{k}\times d_{k}$ correlation matrices, for each potential attention head pair. Let ${\bm{C}}^{jk}$ be the block of the correlation matrix corresponding to head pair $(j,k)$. We then compute the optimal permutation for each unique head pair, and store its head-internal permutation and cost from the following:

We then compute the outer head correspondence permutation with a new assignment problem:

The outer permutation dictates the subset of previously computed inner permutations used in the final permutation, and the order in which to concatenate them, resulting in our 2-staged MHA permutation. We show this approach in Algorithm 1. We note that LinearSumAssignment refers to the specific assignment problem we encounter in our method.

The resulting permutation matrix ${\bm{P}}_{\text{MHA}}$ applies as following: ${\bm{P}}_{\text{MHA}}^{\mathrm{T}}$ can apply to ${\bm{W}}_{O}$ as described previously, but we apply ${\bm{P}}_{\text{MHA}}$ to each of ${\bm{W}}_{V},{\bm{W}}_{K}$, and ${\bm{W}}_{Q}$. We note that the $\langle{\bm{W}}_{Q_{i}}{\bm{x}},{\bm{W}}_{K_{i}}{\bm{x}}\rangle$ dot-product attention operation multiplies away the head dimension $d_{k}$, meaning we do not necessarily have to use the same permutation for ${\bm{W}}_{V}$ and $\left\{{\bm{W}}_{Q},{\bm{W}}_{K}\right\}$. However, we still have to use the same outer permutation for all three matrices, as the attention weights resulting from $\langle{\bm{W}}_{Q_{i}}{\bm{x}},{\bm{W}}_{K_{i}}{\bm{x}}\rangle$ are still head-specific. In experimentation, we do not find a notable difference between using a separate permutation for $\{{\bm{W}}^{Q},{\bm{W}}^{K}\}$²²2This permutation would be computed from features ${\bm{W}}_{Q}{\bm{x}}$ and ${\bm{W}}_{K}{\bm{x}}$, and using the same permutation for ${\bm{W}}_{Q},{\bm{W}}_{K}$ and ${\bm{W}}_{V}$. Therefore, we consider only the latter case for simplicity.

We show an example of an output from our proposed algorithm in Figure 2, as well as some alternative approaches to permuting MHA weights. The first permutation matrix shows the resulting block-diagonal structure of assuming that heads are aligned between different minima. The second matrix shows an example from our method, and the third matrix disregards head structure and allows permuting features across heads. We note that ignoring heads will not lead to a valid permutation $\pi$ where $f({\bm{x}};{\bm{\theta}})=f({\bm{x}};\pi({\bm{\theta}}))$, but we still include it for experimental comparisons.

Each Transformer layer, assuming no cross-attention, contains two residual connections. This means all layers, from the embedding layer to the output layer, is linked via consecutive residual connections. This diverges even from ResNets, which generally contain skip-connections every 2 or 3 layers (He et al., 2016).

We diagram the residual connections of a Transformer layer and their relationships to model parameters in Figure 1. The first connection skips past the multi-headed attention sublayer, and the second connection skips past the feed-forward sublayers, as diagrammed. The connections can be formulated as the following, where LN refers to LayerNorm:

As seen in the equations, the input and output of both sublayers are added to create a new output, which implies that if a permutation operation is applied to the output state, the permutation needs to be the same for both addends.

We note that the addends are normalized via the LayerNorm module, and any permutation to the output would need to permute the features of the LayerNorm module as well. We apply the permutation to the weights of LN. Because LN is not a full weight matrix, and maintains the same feature ordering as its input, we must apply the permutation to the addends of the residual connections as well (Stoica et al., 2024).

Now, ignoring the parameters from the LayerNorm module for a moment, we see that applying the permutation to the output of the second residual connection leads to the following:

We see that due to the residual structure, any permutation applied to second feed-forward weight parameter, ${\bm{W}}_{2}$, must also be applied to MHA, or more specifically the ${\bm{W}}_{O}$ matrix. To unpermute these features, we apply ${\bm{P}}^{\mathrm{T}}$ to where the permuted ${\bm{x}}_{a}^{r}$ and ${\bm{x}}_{f}^{r}$ states become inputs, which are ${\bm{W}}_{1}$ and {${\bm{W}}_{Q},{\bm{W}}_{V},{\bm{W}}_{K}\}$, respectively.

Because the input to each layer must be permuted ${\bm{P}}{\bm{x}}$, and the output of each layer is also permuted ${\bm{P}}{\bm{x}}_{f}^{r}$, we can see that the entire Transformer architecture uses the same $\{{\bm{P}},{\bm{P}}^{\mathrm{T}}\}$ matrices for all the weights involved in residual connections. This is unlike our other proposed permutations for multi-headed attention which are specific to each Transformer layer. Requiring the same permutation throughout the model reduces the degrees of freedom available in finding an optimal permutation, which is an important result demonstrating a potential difficulty of aligning Transformers.

At the ends of the models, namely the embedding layer(s) and output layer(s), we also apply these transformations, as the input to the first Transformer block and the output of the last Transformer block are permuted. We apply this ${\bm{P}}$ to the embedding weights, including positional and any special token embeddings, and at the final layer, we apply ${\bm{P}}^{\mathrm{T}}$ to the weight matrix immediately following the last Transformer block LayerNorm. This is usually a pooling or dense layer, depending on the model task.

Because of the multiple potential features that could contribute to the computation of the residual permutations, namely both ${\bm{x}}_{a}^{r}$ and ${\bm{x}}_{f}^{r}$ across all layers, we consider several strategies for learning these mappings. We test 4 approaches: First refers to only obtaining features from the features immediately following the embedding layer(s). Last refers to only obtaining features from just after the final Transformer layer’s LayerNorm. All refers to concatenating all features ${\bm{x}}_{f}^{r},$ and ${\bm{x}}_{a}^{r}$ from all Transformer layers, and separate refers to computing 2 $\{{\bm{P}},{\bm{P}}^{\mathrm{T}}\}$ pairs per Transformer layer, one for ${\bm{x}}_{f}^{r}$, and the other for ${\bm{x}}_{a}^{r}$. The separate approach also does not lead to a valid permutation like Ignore Heads in Section 3.2, but we also include it for experimental comparisons.

Unlike the residual stream or Multi-Headed attention, Feed-Forward sub-layers require no special attention in order to permute them into a new space. We simply compute correlations from the features after the computation of the first Feed-Forward layer (${\bm{W}}_{1}$), and compute ${\bm{P}},{\bm{P}}^{\mathrm{T}}$ separately for each Transformer layer. We apply these permutations as described in Section 3.1:

As many models have task-specific output layers like classification heads or masked-language modeling heads, there is also another permutation able to be computed between a pooling/dense layer and the actual model head, which tends to also be a linear layer. This permutation mapping would proceed as normal, like the feed-forward example, but we do not include it in our experimentation as we see no notable differences in its inclusion.

We investigate Transformer encoder-based masked language models in this work. Specifically, we consider 5 different BERT models, seeds 1 through 5, from the MultiBERTs reproductions (Devlin et al., 2019; Sellam et al., 2021). Each of these models is a bert-base-uncased checkpoint, but trained with a different random initialization and random batch ordering. These properties are the type of SGD-related variation we seek to study between different minima. All models use the same original BERT vocabulary and tokenizer. All of our reported experiments include a mean and standard error across the 10 unique pairings resulting from our 5 different models.

For classification tasks, we fine-tune each of the MultiBERTs models with a randomly initialized classification head, including pooling layer and classification layer weights. We keep the head initializations the same across models.

We report vanilla averaging as our main baseline for comparison, computed as ${\bm{\theta}}_{\text{avg}}=\frac{1}{2}({\bm{\theta}}_{A}+{\bm{\theta}}_{B})$.

We focus on two different tasks to test our method. We use the masked language modeling task to test our base method, as this is the main task that BERT and many other pretrained Transformer encoded models are trained with. We also consider fine-tuning these models for classification tasks, and then comparing them in their fine-tuned state. We use the General Language Understaning Evaluation (GLUE) benchmark (Wang et al., 2018) for our classification tasks, and exclude WNLI as in Devlin et al. (2019). GLUE is a set of 8 diverse natural language understanding classification tasks. We report scores on GLUE test sets from our reproductions in Table 4 in Appendix A.

For our experiments on masked language modeling, we use the validation set of the Wikitext-103 benchmark as our evaluation data Merity et al. (2016)³³3We obtain the wikitext-103-raw-v1 version, available from https://huggingface.co/datasets/wikitext. For computing model activations, we extract a random sample of just over 1 million sentences of the Books corpus Zhu et al. (2015). For a majority of our experiments, we sample 100k sentences and use this subset to compute features, unless stated otherwise. We take a diverse sample across different genres among the books available. We choose the Books corpus as it is part of the original pre-training data from BERT (Devlin et al., 2019).

For GLUE experiments, we use the full training data for each of the tasks to compute features, and the full validation sets to compute losses. The amount of data available for each task varies, and statistics are also reported in Table 4 in Appendix A.

To compute loss barriers, we compute several interpolations of ${\bm{\theta}}_{A}$ and ${\bm{\theta}}_{B}$, as $\lambda{\bm{\theta}}_{A}+(1-\lambda){\bm{\theta}}_{B}$. Specifically, we use 21 samples evenly spaced between $\lambda=0$ and $\lambda=1$, inclusive. We use the definition of loss-barrier as Frankle & Carbin (2018)⁴⁴4Referred to as linear interpolation instability in this work., defined as the maximum difference between the loss of an interpolation and the average loss of the base models:

To compute MLM loss/pseudo-perplexity, we use a masking probability of $p=0.15$ across block sizes of 128 tokens. For $N$ masked samples in text $\mathbf{W}$, we compute pseudo-perplexity as:

We report loss on GLUE tasks as normal, defined by each task, across the entire validation set.

We report results on the 10 MultiBERTs merges after merging different sets of Transformer components described in Section 3. We show pseudo-perplexity results across the range of interpolations, displayed Figure 3. We report vanilla averaging with no permutations, merging all feed-forward sublayers, merging all multi-headed attention sublayers, and merging all feed-forward sublayers and all multi-headed attention sublayers. Aligning and merging either the feed-forward sublayers or the attention sublayers clearly leads to a perplexity reduction over the baseline, and their combination leads to a stark reduction, of almost 7$\times$ the original perplexity at $\lambda=0.5$. We do not include permuting parameters involved in residual connections (Section 3.3) and the output projection (Section 3.4) here as they do not outperform merging only feed-forward and attention sublayers. We discuss the residual connections further in Section 5.3.

The consistently reduced barrier between minima indicates that these different models are connected with a lower loss path than seen without considering these models within a more similar loss space. We note that we do not observe a linear or convex loss path between these models, as sometimes observed in previous work on MLPs, ResNets, and VGGs (Tatro et al., 2020; Entezari et al., 2022; Ainsworth et al., 2023). In this line of prior work, the same data is generally used to train models, compute activations for alignment and merging, and compute loss barriers. Due to the extensive pretraining data of these masked language models, and the limited alignment and evaluation data we use, we do not test for linear mode connectivity in the same manner. Instead, the stark loss reduction seems to indicate that minima are connected with a barrier at least as high as what we report.

In understanding the extent to which these different Transformers learn similar representations throughout the model, we compute the average feature correlations between all 10 masked language model pairs. We report correlations for attention and feed-forward layer features before and after applying our method. Individual correlations are computed between features of the same index. We show these values in Figure 4. We find that the average feature correlations of aligned models are significantly higher than those of the original models, which demonstrates some success of our alignment procedure, but still no higher than 0.3. Previous work has reported higher average correlations on different architectures like ResNets (Tatro et al., 2020), but in the image domain. Some pre-trained transformers are also known to be sparsely activated and able to be pruned heavily (Li et al., 2023; Dalvi et al., 2020), which may also lead to lower average feature correlations.

We report loss barriers for our Head-Permutation multi-headed attention approach as compared to some alternatives also described in Section 3.2 in Table 1. These results reflect permuting both attention parameters as well as feed-forward parameters. We see that our proposed Head-Permutation approach for the attention sub-layer outperforms simple attention averaging, as well as approaches ignoring the multi-headed structure of the weight parameters (Ignore-Heads), and not allowing for different head correspondences across different models (Monotonic). We also show an example correlation matrix between the first multi-headed attention layer from 2 different MultiBERTs models in Figure 5. The correlation matrix shows clear attention head boundaries, as well as a scattered pattern that supports our proposed technique that does not assume any monotonically ordered head correspondence.

As described in Section 3.3, many of the permutation operations in the Transformer architecture are shared due to its repeated Add&Norm components. This linkage results in a huge reduction in the number of permutation symmetries of the Transformer architecture, and therefore the number of valid permutations of parameters involved in the residual stream. We report the loss barriers after applying our permutation alignment to only the parameters involved in the residual connections in Table 2. We find that an identity permutation performs significantly better than all other proposed approaches. Even the separate permutation approach does not outperform using the identity matrix $I_{d}$ as the residual merge. This approach does not even create a valid Transformer symmetry, but has many more degrees of freedom than the other proposed approaches.

We report loss barriers on the masked language modeling task while varying the amount of sentences used to compute correlations. All previous experiments are reported on $\sim$100k sentences. We combine feed-forward and attention layers in this setting. Results are shown in Figure 6. Values are fairly similar across data amounts, suggesting that there is no strong relationship between the amount of data used and the overall loss barrier. We do see, however, small variations between different data amounts, which are likely due to the content of the different sets of sentences used to create the correlation matrices. We believe that more than the quantity of tokens, the type and quality of text used for computing correlations likely matters more. We leave further investigation into specific data sources for aligning text-based Transformer models for future work.

We investigate loss barriers between fine-tuned BERT models across 8 different GLUE tasks. We compare loss barriers from vanilla averaging to those obtained after applying our method to combine the models. Different from our result on masked language modeling, we include residual permutations as we find it has a slight decrease in loss barriers over just feed-forward and attention permutations. We use the First approach to compute residual permutations. We report these results in Table 3. We additionally include loss barrier curves for all 8 tasks in Appendix B in Figure 7.

While we are able to observe lower loss barriers between minima, the trends of loss reduction across interpolations is inconsistent, especially compared to the masked language modeling setting. For example, while the maximum loss of vanilla averaging is higher than the maximum loss of our approach, there are still other interpolation weights where the loss of our approach is instead higher than vanilla averaging. We also note an interesting pattern of lower loss than either parent model for several tasks around $\lambda=0.15$ and again at $\lambda=0.85$ . Additionally, the loss barrier curves of some vanilla merges between these fine-tuned models have different behavior than their pretrained alternatives, as seen in Figure 7 with the presence of “M” like loss curve shapes between minima. While our method can extend to this setting to find a lower loss path between fine-tuned minima, it is unclear which kind of data is necessary to observe the lowest loss path. We leave further investigation into connecting fine-tuned models, and further understanding of their pre-merging connectivity in loss space, for future work.