On Privileged and Convergent Bases in Neural Network Representations
In this study, we investigate whether the representations learned by neural
networks possess a privileged and convergent basis. Specifically, we examine
the significance of feature directions represented by individual neurons.
First, we establish that arbitrary rotations of neural representations cannot
be inverted (unlike linear networks), indicating that they do not exhibit
complete rotational invariance. Subsequently, we explore the possibility of
multiple bases achieving identical performance. To do this, we compare the
bases of networks trained with the same parameters but with varying random
initializations. Our study reveals two findings: (1) Even in wide networks such
as WideResNets, neural networks do not converge to a unique basis; (2) Basis
correlation increases significantly when a few early layers of the network are
frozen identically.
  Furthermore, we analyze Linear Mode Connectivity, which has been studied as a
measure of basis correlation. Our findings give evidence that while Linear Mode
Connectivity improves with increased network width, this improvement is not due
to an increase in basis correlation.
Davis Brown, Nikhil Vyas, Yamini Bansal
# On Privileged and Convergent Bases in Neural Network Representations

###### Abstract

In this study, we investigate whether the representations learned by neural networks possess a privileged and convergent basis. Specifically, we examine the significance of feature directions represented by individual neurons. First, we establish that arbitrary rotations of neural representations cannot be inverted (unlike linear networks), indicating that they do not exhibit complete rotational invariance. Subsequently, we explore the possibility of multiple bases achieving identical performance. To do this, we compare the bases of networks trained with the same parameters but with varying random initializations. Our study reveals two findings: (1) Even in wide networks such as WideResNets, neural networks do not converge to a unique basis; (2) Basis correlation increases significantly when a few early layers of the network are frozen identically.

Furthermore, we analyze Linear Mode Connectivity, which has been studied as a measure of basis correlation. Our findings give evidence that while Linear Mode Connectivity improves with increased network width, this improvement is not due to an increase in basis correlation.

## 1 Introduction

While neural networks are black-box function approximators that are trained end-to-end to optimize a loss objective, their emergent internal layer-wise representations are important objects for both _understanding deep learning_ and _direct downstream use_. Internal representations of neural networks can be useful tools for interpretability [4, 5, 24], teaching us how neural networks perform the computations they do, as well as for understanding the implicit biases of gradient-based neural network training [1, 2]. Moreover, representations are often directly used for downstream tasks that the network was not originally trained for, like in transfer learning or representation learning. Thus, we would like to develop a better understanding of the mathematical properties of neural network representations.

One such property is whether neural networks representations have a _privileged basis_[10]. That is, are the features represented by each individual neuron significant, or is information stored only at a population level in neurons? This question is important, for instance, in interpretability, where attempts have been made to interpret features represented by individual neurons (such as edge or curve detectors in convolutional networks [5]). This question is also closely related to
that of _invariances_ exhibited by neural representations -- what are the set of transformations that can be applied to representations that keep the final network accuracy unchanged? In particular, if the representations are rotation invariant, then an individual neuron does not carry significant information.

To understand this further, consider the simple case of a two-layer neural network without any non-linear activation functions. That is, the function output of the network is \(f(x)=W_{2}W_{1}x\) with weights \(W_{2},W_{1}\) and inputs \(x\). Here, the first layer representations are \(W_{1}x\). This representation exhibits rotation invariance - we can rotate the first layer representations by an arbitrary orthonormal matrix \(O\) (giving us rotated first layer representations \(OW_{1}x\)) but the subsequent layer can invert this rotation and recover the original function \(f(x)=W_{2}O^{-1}OW_{1}x\). Thus, an individual neuron could represent any arbitrary feature for the same functional outputs in the network.

**Our Contributions:** In our work, we start by showing that the perceived permutation invariance of representations at high width is actually a result of a of kind of noise-averaging -- the correlation between the activities of neurons after accounting for permutations remains _nearly_ constant, as we scale the width (Section 3). This shows that while metrics like linear mode connectivity may suggest permutation invariance, the effect disappears when examined at a neuron level. Since this casts some doubt on the presence of a privileged basis, in Section 4 we ask if _any_ basis of neural representations is likely to work equally well. To do so, we consider a random rotation of a layer, and ask if it can be inverted by the later layers with training. We find that this is not the case. Thus, combined these results suggest that while the basis of neural representations matter, there is no one unique basis that is required to achieve the same functional accuracy. Finally, in Section 5, we ask what kinds of constraints can be imposed on the network to obtain a unique neural basis consistent across different training runs.

## 2 Related Work

**Convergent learning.** Also referred to as _universality_, convergent learning is the conjecture that different deep learning models learn very similar representations when trained on similar data [20, 24]. Much of the work in mechanistic interpretability [9, 23, 24, 25] has leveraged the universality conjecture to motivate research for toy models, with the hope that the methods and interpretations developed for these more tractable models will scale to larger and more capable models. Recently, [6] examined universality by reverse-engineering a toy transformer model for a group composition task. Attempts to test for convergent learning include _representation dissimilarity_ comparisons, notably neuron alignment [14, 19, 20] and correlation analysis / centered kernel alignment [17, 22].

Model stitching [2, 18] extracts features from the early layers of model \(f\) and inserts them into the later layers of model \(g\) (usually via a learned, low-capacity connecting layer \(\varphi\)). If the representations between these models can be combined such that the resulting'stitched' model, \(g_{>l}\circ\varphi\circ f_{\leq l}\), achieves a low loss on a downstream task, the models are called'stitching connected' for the layer \(l\) for that task.

**Linear Mode Connectivity:** It has been conjectured in [1, 12, 15] that, for different models learned by SGD with equal loss, once the permutation symmetries of neural networks are taken into consideration, linear interpolations between them of the form \(\theta_{\alpha}=(1-\alpha)\theta_{1}+\alpha\theta_{2}\) for \(0<\alpha<1\) have at least constant loss.

**Privileged Basis:** It is often taken for granted in the interpretability literature that the activation basis is _privileged_, at least in layers with elementwise operations (namely nonlinearities)
[3, 13, 28, 29, 30]. On the other hand, while the residual stream of transformer models has no obvious elementwise operation and is thus not an obviously privileged basis, [7] provides evidence for outlier basis-aligned features. To study this phenomenon, [11] demonstrate that transformers can learn rotationally invariant representations in their residual stream using a similar procedure to what we describe in Section 4. We ask the complementary question of whether layers with nonlinearities can learn in an arbitrary basis.

## 3 SGD Basis and Linear Mode Connectivity

In this section we explore whether two neural networks with different random initializations converge to the same basis. In other words, are the two neural networks the same up to a permutation of neurons per layer? It is clear that this will hold for an infinite width network since at infinite width there is no 'randomness' in initialization. We are interested in answering whether it holds for large yet feasible widths. Recently [1, 12, 15] have studied the weaker claim of if different neural networks are linear mode connected after an appropriate permutation to the neurons per layer. Specifically they study the loss/error barrier when interpolating1 between the two networks (after permutation). We will call this barrier as perm-LMC barrier. In [1, 15] it was found that for networks trained with SGD on CIFAR-10, the perm-LMC barrier become very small (\(<2\%\) error barrier) for large yet feasible widths. On the other hand, for ImageNet trained models, while the barrier decreases with width, it remains quite large. There are at least two reasons for why perm-LMC barrier could decrease with width:

Footnote 1: Following Jordan et al. [15], we reset batch norm statistics in all of the experiments.

1. The bases of two trained networks becomes closer with larger width.
2. LMC improves with width even if the two networks do not become closer in their basis with width.

To answer this question, we need a measure of how close two networks are to being permutations of each other. Let \(\text{Perm}_{n}\) be the set of permutations over \(n\) elements. For a layer with \(n\) neurons and activations \(X^{1}\) and \(X^{2}\) of the two networks, we use \(\max_{p\in\text{Perm}_{n}}\sum_{i\in[n]}\operatorname{Cov}(X^{1}_{i},X^{2}_{ p(i)})\) which we call permutation-correlation (perm-corr)2. This is also the measure used by Li et al. [19] and Jordan et al. [15] to find an appropriate permutation for calculating the permutation-LMC barrier.

Footnote 2: For all experiments we will report values of perm-corr averaged across all the layers.

**perm-corr is nearly constant across width:** In Figure 1 (left) we use the exact setup of Jordan et al. [15] and plot perm-corr and perm-LMC barrier (between two networks trained with independent initializations and batch orderings) as a function of width where we use Resnet20 as the base network and Resnet20-32x as the widest one. perm-corr remains almost constant with width. This suggests that we are very far from widths needed for neurons basis to be close i.e. perm-corr to be near 1. As in prior works we do find that perm-LMC barrier goes down with width. But combined these results suggest that **perm-LMC barrier goes down with width because linear mode connectivity becomes more feasible at high width rather than due to a better match between the basis of neurons**. We now explore this further.

The previous experiment already verified that basis correlation (operationalized by perm-corr) does not improve with width. We now verify that linear mode connectivity, even without an increase
in similarity of basis, improves with width. To do so, we consider the LMC barrier3 between two networks without permuting the neurons.

Footnote 3: Jordan et al. [15] did not reset batch norm statistics for this experiment and hence did not observe this improvement with width.

In Figure 1 (middle) we find that this quantity (green) improves significantly with width, adding support to our hypothesis. We can now ask why width has this effect. Intuitively, the LMC barrier for two uncorrelated networks (or even networks with some fixed amount of correlation) improves if the networks are robust to noise. This is because being robust to noise means that we can treat the other network as noise and maintain the performance of the first network. Figure 1 (middle) we plot (orange) the performance of a network formed by averaging a trained and a random network (with same weight norms). We find that the accuracy drop (compared to the trained network) of this network also improves with width adding support to our hypothesis. Together these experiments suggest that that while perm-LMC barrier going down with width is an interesting empirical finding in its own right, it overestimates the similarity of neurons across different networks and in particular is driven by other factors such as robustness to noise to a large extent.

### Change in Permutation-Correlation due to Training

Another approach to think about the arising of a unique basis is to see how to training process affects perm-corr i.e., does perm-corr improve due to training? In Figure 1 (right) we plot perm-corr for a trained and random network and find that two _random_ networks have higher perm-corr then for two trained network, for all but the narrowest width networks. This suggests that there is a large variance in the basis of neurons that are found in a trained network. This argues against convergent learning in neural networks from the perspective of the basis of neurons.

## 4 Can Networks Learn Rotationally Invariant Representations?

In the last section we saw the neuron basis for networks trained from independent initialization are not aligned (compared to the baseline of two randomly initialized networks), i.e. there is no unique

Figure 1: Do randomly initialized neural networks converge to the same basis? (**left**) perm-LMC barrier drops with width but perm-corr is nearly constant. (**middle**) LMC barrier (without permutation) also drops with width. (**right**) perm-corr does not improve through training for most widths.
basis. This raises the following natural question: Do all bases lead to good representations? Specifically, in this section, we ask if the representations learnt by neural networks can be made rotation invariant. We build on the linear example considered in the introduction, where any orthonormal transformation of the representation can be inverted by the following layer successfully.

To test this in non-linear networks trained with gradient descent, we take a Myrtle-CNN trained with SGD on CIFAR-10 trained to 92% accuracy. Let's say the network is \(f(x)=\sigma(A_{l}(\sigma(A_{l-1}(...))))\), where \(A_{l}\) denotes the pre-activation of layer \(l\) and \(\sigma\) denotes the non-linearity. Then, we perform the following procedure: We sample a random orthonormal matrix \(O_{1}\) and multiply it with the _pre-activations_ of the first layer \(O_{1}A_{1}\). Then, we freeze this rotated layer, and retrain all the remaining layers of the network on the same training dataset. Next, we apply a random orthonormal matrix to the second layer \(O_{2}A_{2}\), freeze the first two layers and retrain all layers \(l>2\). We repeat this procedure successively for all the layers in the network. Thus, \(f_{l,\text{rotated}}(x)=\sigma(A_{l}(\sigma(O_{l-1}A_{l-1}(\sigma(O_{l-2}A_{l- 2}(...)))))\).

Figure 4 shows the resulting error when we freeze up to \(l\) layers. We find that retraining _cannot_ invert this random rotation and that the network accuracy degrades significantly. We also observe that the error gets worse for later layers in the network. For comparison, we also plot the error of a network with random weights (with the same distribution as the initialization) for the first \(l\) layers and show that the increase in error from performing random rotations is similar to _random features_-not training the layers at all! For reference, we repeat the above procedure but we only freeze one layer at a time (Figure 4). We also observe a significant increase in the error of the network, which gets worse as we go deeper into the network. Our findings highlight that the basis of the network are not in fact rotation invariant. This suggests that the directions represented by individual neurons are in fact signifcant, and we cannot use an arbitrary basis for the network.

## 5 The residual stream might enable convergent bases

Finally, we ask if there are other empirically common choices that may enable the basis to be fixed in certain ways. We highlight the candidate of a _residual stream_. In transformer networks, for instance,

Figure 2: _Random feature_ vs _recovering random rotation_ performance for Myrtle CNN models trained on CIFAR-10. **(left)** successive freezing, rotating, and retraining; **(right)** single layer freezing and rotating.
there is a largely linear residual connection that runs from the input (after the embedding layer) to the remainder of the network [9]. If the residual stream is indeed important to the computations performed by the network, it may impose a basis on the network that roughly tries to align with the input embedding. For models that share the same tokenizer or embedding layer, this might suggest that we can combine networks with residual streams _as is_ -- without even having to compute permutations to match the neurons to perform symmetry correction [1].

However, prior work [2] has shown that the early layers of a vision network can be replaced with random features without significant loss in performance. We see in Figure 2 that the first few layers of the network are relatively more resistant to random rotations. This suggests that even if the residual connections were providing a force for the basis to align, the residual connections would align to the basis computed by these random transformations of the early layers.

Thus, we experiment with freezing early blocks of the network. First, we train a model \(A\) with \(n\) layers. Then, we train a new model \(B\) with the first \(l\) frozen layers of \(A\), training the \(n-l\) top-layers from scratch with a new random initialization. The perm-corr is computed between layers, Figure 3 shows our results for a ConvNeXt [21] model trained on CIFAR-10 (left) and a Vision Transformer (ViT) [8] trained on ImageNet (right). We also measure the identity stitching penalty, where we evaluate the performance of pairs models stitched with the _identity function_\(\varphi=\mathrm{id}\), i.e. models of the form \(g_{>l}\circ f_{\leq l}\).

We find that freezing only 2 blocks is sufficient to have a significantly higher perm-corr (like-wise, lower error for identity stitching Figure 3 (middle)) in ConvNeXt models and ViTs. This convergence phenomenon is considered across different model widths and for different residual stream structures in Appendix A.

Together, these results suggest a relatively cheap procedure for _fixing the basis_ when training different neural networks, an (as noted) desirable property for interpretability research and also relevant for modular and distributed neural network training [16, 26]. For example, a related phenomenon for networks fine-tuned from the same base model was used in [27] to create a combined model more accurate than its constituent models.

Figure 3: Enforcing a consistent basis by freezing \(l\) early layers. (**left**) perm-corr for ConvNeXt models trained on CIFAR-10 (**middle**) identity stitching CIFAR-10 ConvNeXt models (**right**) perm-corr for ViT models trained on ImageNet
## 6 Discussion and conclusion

We find that, in some important respects, linear mode connectivity overstates the similarity of neurons across runs. Our results suggest that while the LMC barrier improves with network width, it can in part be explained by factors beyond similarity. On the other hand, we find strong evidence that for neural networks, only certain bases (namely, those that are not rotation invariant) lead to good representations. Finally, we provide a straightforward procedure to enable a convergent basis; this is desirable for both interpretability and modular training.

## Acknowledgements

NV is supported by a Simons Investigator Fellowship, NSF grant DMS-2134157, DARPA grant W911NF2010021, and DOE grant DE-SC0022199.
