SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhanpeng Zhou $^{1}$ , Yongyi Yang $^{2}$ , Xiaojiang Yang $^{1}$ , Junchi Yan $^{1*}$ , Wei Hu $^{2\dagger}$
$^{1}$  Dept. of Computer Science and Engineering & MoE Key Lab of AI, Shanghai Jiao Tong University
$^{2}$  Dept. of Electrical Engineering & Computer Science, University of Michigan
{zzp1012, yangxiaojiang, yanjunchi}@sjtu.edu.cn
{yongyi,vvh}@umich.edu
Recent work has revealed many intriguing empirical phenomena in neural network training, despite the poorly understood and highly complex loss landscapes and training dynamics. One of these phenomena, Linear Mode Connectivity (LMC), has gained considerable attention due to the intriguing observation that different solutions can be connected by a linear path in the parameter space while maintaining near-constant training and test losses. In this work, we introduce a stronger notion of linear connectivity, Layerwise Linear Feature Connectivity (LLFC), which says that the feature maps of every layer in different trained networks are also linearly connected. We provide comprehensive empirical evidence for LLFC across a wide range of settings, demonstrating that whenever two trained networks satisfy LMC (via either spawning or permutation methods), they also satisfy LLFC in nearly all the layers. Furthermore, we delve deeper into the underlying factors contributing to LLFC, which reveal new insights into the permutation approaches. The study of LLFC transcends and advances our understanding of LMC by adopting a feature-learning perspective. We released our source code at https://github.com/zzp1012/LLFC.
Despite the successes of modern deep neural networks, theoretical understanding of them still lags behind. Efforts to understand the mechanisms behind deep learning have led to significant interest in exploring the loss landscapes and training dynamics. While the loss functions used in deep learning are often regarded as complex black-box functions in high dimensions, it is believed that these functions, particularly the parts encountered in practical training trajectories, contain intricate benign structures that play a role in facilitating the effectiveness of gradient-based training [36, 4, 39]. Just like in many other scientific disciplines, a crucial step toward formulating a comprehensive theory of deep learning lies in meticulous empirical investigations of the learning pipeline, intending to uncover quantitative and reproducible nontrivial phenomena that shed light on the underlying mechanisms.
One intriguing phenomenon discovered in recent work is Mode Connectivity [10, 5, 11]: Different optima found by independent runs of gradient-based optimization are connected by a simple path in the parameter space, on which the loss or accuracy is nearly constant. This is surprising as different optima of a non-convex function can very well reside in different and isolated "valleys" and yet this does not happen for optima found in practice. More recently, an even stronger form of mode
*Corresponding author. SJTU authors are partly supported by NSFC(62222607, 61972250, U19B2035) and Shanghai Municipal Science and Technology Major Project (2021SHZDZX0102).
$^{\dagger}$  Wei Hu acknowledges support from the Google Research Scholar Program.
Figure 1: Illustration of Layerwise Linear Feature Connectivity (LLFC). Each tensor comprises three feature maps when provided with three input images.
connectivity called Linear Mode Connectivity (LMC) was discovered [24, 9]. It depicts that different optima can be connected by a linear path of constant loss/accuracy. Although LMC typically does not happen for two independently trained networks, it has been consistently observed in the following scenarios:
- Spawning [9, 7]: A network is randomly initialized, trained for a small number of epochs (e.g. 5 epochs for both ResNet-20 and VGG-16 on the CIFAR-10 dataset), and then spawned into two copies which continue to be independently trained using different SGD randomnesses (i.e., for mini-batch order and data augmentation).
- Permutation [6, 1]: Two networks are independently trained, and the neurons of one model are permuted to match the neurons of the other model while maintaining a functionally equivalent network.
The study of LMC is highly motivated due to its ability to unveil nontrivial structural properties of loss landscapes and training dynamics. Furthermore, LMC has significant relevance to various practical applications, such as pruning and weight-averaging methods.
On the other hand, the success of deep neural networks is related to their ability to learn useful features, or representations, of the data [27], and recent work has highlighted the importance of analyzing not only the final outputs of a network but also its intermediate features [20]. However, this crucial perspective is absent in the existing literature on LMC and weight averaging. These studies typically focus on interpolating the weights of different models and examining the final loss and accuracy, without delving into the internal layers of the network.
In this work, we take a feature-learning perspective on LMC and pose the question: what happens to the internal features when we linearly interpolate the weights of two trained networks? Our main discovery, referred to as Layerwise Linear Feature Connectivity (LLFC), is that the features in almost all the layers also satisfy a strong form of linear connectivity: the feature map in the weight-interpolated network is approximately the same as the linear interpolation of the feature maps in the two original networks. More precisely, let  $\theta_{A}$  and  $\theta_{B}$  be the weights of two trained networks, and let  $f^{(\ell)}(\theta)$  be the feature map in the  $\ell$ -th layer of the network with weights  $\theta$ . Then we say that  $\theta_{A}$  and  $\theta_{B}$  satisfy LLFC if
$$
f ^ {(\ell)} \left(\alpha \boldsymbol {\theta} _ {A} + (1 - \alpha) \boldsymbol {\theta} _ {B}\right) \propto \alpha f ^ {(\ell)} \left(\boldsymbol {\theta} _ {A}\right) + (1 - \alpha) f ^ {(\ell)} \left(\boldsymbol {\theta} _ {B}\right), \quad \forall \alpha \in [ 0, 1 ], \forall \ell . \tag {1}
$$
See Figure 1 for an illustration. While LLFC certainly cannot hold for two arbitrary  $\theta_{A}$  and  $\theta_{B}$ , we find that it is satisfied whenever  $\theta_{A}$  and  $\theta_{B}$  satisfy LMC. We confirm this across a wide range of settings, as well as for both spawning and permutation methods that give rise to LMC.
LLFC is a much finer-grained characterization of linearity than LMC. While LMC only concerns loss or accuracy, which is a single scalar value, LLFC (1) establishes a relation for all intermediate feature maps, which are high-dimensional objects. Furthermore, it is not difficult to see that LLFC applied to the output layer implies LMC when the two networks have small errors (see Lemma 1); hence, LLFC can be viewed as a strictly stronger property than LMC. The consistent co-occurrence of LLFC and LMC suggests that studying LLFC may play a crucial role in enhancing our understanding of LMC.
Subsequently, we delve deeper into the underlying factors contributing to LLFC. We identify two critical conditions, weak additivity for ReLU function and a commutativity property between two trained networks. We prove that these two conditions collectively imply LLFC in ReLU networks, and provide empirical verification of these conditions. Furthermore, our investigation yields novel insights into the permutation approaches: we interpret both the activation matching and weight matching objectives in Git Re-Basin [1] as ways to ensure the satisfaction of commutativity property.
In summary, our work unveils a richer set of phenomena that go significantly beyond the scope of LMC, and our further investigation provides valuable insights into the underlying mechanism behind LMC. Our results demonstrate the value of opening the black box of neural networks and taking a feature-centric viewpoint in studying questions related to loss landscapes and training dynamics.
(Linear) Mode Connectivity. Freeman and Bruna [10], Draxler et al. [5], Garipov et al. [11] observed Mode Connectivity, i.e., different optima/modes of the loss function can be connected through a non-linear path with nearly constant loss. Nagarajan and Kolter [24] first observed Linear Mode Connectivity (LMC), i.e., the near-constant-loss connecting path can be linear, on models trained on MNIST starting from the same random initialization. Later, Frankle et al. [9] first formally defined and thoroughly investigate the LMC problem. Frankle et al. [9] observed LMC on harder datasets, for networks that are jointly trained for a short amount of time before going through independent training (we refer to this as the spawning method). Frankle et al. [9] also demonstrated a connection between LMC and the Lottery Ticket Hypothesis [8]. Fort et al. [7] used the spawning method to explore the connection between LMC and the Neural Tangent Kernel dynamics. Lubana et al. [23] studied the mechanisms of DNNs from mode connectivity and verified the mechanistically dissimilar models cannot be linearly connected. Yunis et al. [38] showed that LMC also extends beyond two optima and identified a high-dimensional convex hull of low loss between multiple optima. On the theory side, several papers [10, 21, 32, 26, 25, 16] were able to prove non-linear mode connectivity under various settings, but there has not been a theoretical explanation of LMC to our knowledge.
Permutation Invariance. Neural network architectures contain permutation symmetries [13]: one can permute the weights in different layers while not changing the function computed by the network. Ashmore and Gashler [2] utilized the permutation invariance of DNNs to align the topological structure of two neural networks. Tatro et al. [31] used permutation invariance to align the neurons of two neural networks, resulting in improved non-linear mode connectivity. In the context of LMC, Entezari et al. [6], Ainsworth et al. [1] showed that even independently trained networks can be linearly connected when permutation invariance is taken into account. In particular, Ainsworth et al. [1] approached the neuron alignment problem by formulating it as bipartite graph matching and proposed two matching methods: activation matching and weight matching. Notably, Ainsworth et al. [1] achieved the LMC between independently trained ResNet models on the CIFAR-10 dataset using weight matching.
Model Averaging Methods. LMC also has direct implications for model averaging methods, which are further related to federated learning and ensemble methods. Wang et al. [33] introduced a novel federated learning algorithm that incorporates unit permutation before model averaging. Singh and Jaggi [30], Liu et al. [22] approached the neuron alignment problem in model averaging by formulating it as an optimal transport and graph matching problem, respectively. Wortzman et al. [35] averaged the weights of multiple fine-tuned models trained with different hyper-parameters and obtained improved performance.
3 Background and Preliminaries
Notation and Setup. Denote  $[k] = \{1,2,\dots ,k\}$ . We consider a classification dataset  $\mathcal{D} = \{(\pmb {x}_i,y_i)\}_{i = 1}^n$ , where  $\pmb {x}_i\in \mathbb{R}^{d_0}$  represents the input and  $y_{i}\in [c]$  represents the label of the  $i$ -th data point. Here,  $n$  is the dataset size,  $d_0$  is the input dimension and  $c$  is the number of classes. Moreover, we use  $\pmb {X}\in \mathbb{R}^{d_0\times n}$  to stack all the input data into a matrix.
We consider an  $L$ -layer neural network of the form  $f(\pmb{\theta}; \pmb{x})$ , where  $\pmb{\theta}$  represents the model parameters,  $\pmb{x}$  is the input, and  $f(\pmb{\theta}; \pmb{x}) \in \mathbb{R}^c$ . Let the  $\ell$ -th layer feature (post-activation) of the network be  $f^{(\ell)}(\pmb{\theta}; \pmb{x}) \in \mathbb{R}^{d_\ell}$ , where  $d_\ell$  is the dimension of the  $\ell$ -th layer ( $0 \leq \ell \leq L$ ) and  $d_L = c$ . Note that
$f^{(0)}(\pmb {\theta};\pmb {x}) = \pmb{x}$  and  $f^{(L)}(\pmb {\theta};\pmb {x}) = f(\pmb {\theta};\pmb {x})$  . For an input data matrix  $\pmb{X}$  , we also use  $f(\pmb {\theta};\pmb {X})\in \mathbb{R}^{c\times n}$  and  $f^{(\ell)}(\pmb {\theta};\pmb {X})\in \mathbb{R}^{d_{\ell}\times n}$  to denote the collection of the network outputs and features on all the datapoints, respectively. When  $\pmb{X}$  is clear from the context, we simply write  $f^{(\ell)}(\pmb {\theta}) = f^{(\ell)}(\pmb {\theta};\pmb {X})$  and  $f(\pmb {\theta}) = f(\pmb {\theta};\pmb {X})$  . Unless otherwise specified, in this paper we consider models trained on a training set, and then all the investigations are evaluated on a test set.
We use  $\mathrm{Err}_{\mathcal{D}}(\pmb{\theta})$  to denote the classification error of the network  $f(\pmb{\theta}; \cdot)$  on the dataset  $\mathcal{D}$ .
Linear Mode Connectivity (LMC). We recall the notion of LMC in Definition 1.
Definition 1 (Linear Mode Connectivity). Given a test dataset  $\mathcal{D}$  and two modes $^3$ $\pmb{\theta}_{A}$  and  $\pmb{\theta}_{B}$  such that  $\mathrm{Err}_{\mathcal{D}}(\pmb{\theta}_{A})\approx \mathrm{Err}_{\mathcal{D}}(\pmb{\theta}_{B})$ , we say  $\pmb{\theta}_{A}$  and  $\pmb{\theta}_{B}$  are linearly connected if they satisfy
$$
\operatorname {E r r} _ {\mathcal {D}} \left(\alpha \boldsymbol {\theta} _ {A} + (1 - \alpha) \boldsymbol {\theta} _ {B}\right) \approx \operatorname {E r r} _ {\mathcal {D}} \left(\boldsymbol {\theta} _ {A}\right), \quad \forall \alpha \in [ 0, 1 ]. \tag {2}
$$
As Definition 1 shows,  $\theta_{A}$  and  $\theta_{B}$  satisfy LMC if the error metric on the linear path connecting their weights is nearly constant. There are two known methods to obtain linearly connected modes, the spawning method [9, 7] and the permutation method [6, 1].
Spawning Method. We start from random initialization  $\pmb{\theta}^{(0)}$  and train the model for  $k$  steps to obtain  $\pmb{\theta}^{(k)}$ . Then we create two copies of  $\pmb{\theta}^{(k)}$  and continue training the two models separately using independent SGD randomnesses (mini-batch order and data augmentations) until convergence. By selecting a proper value of  $k$  (usually a small fraction of the total training steps), we can obtain two linearly connected modes.
Permutation Method. Due to permutation symmetry, it is possible to permute the weights in a neural network appropriately while not changing the function being computed. Given two modes  $\theta_{A}$  and  $\theta_{B}$  which are independently trained (and not linearly connected), the permutation method aims to find a permutation  $\pi$  such that the permuted mode  $\theta_B^\prime = \pi (\theta_B)$  is functionally equivalent to  $\theta_{B}$  and that  $\theta_B^{\prime}$  and  $\theta_{A}$  are linearly connected. In other words, even if two modes are not linearly connected in the parameter space, they might still be linearly connected if permutation invariance is taken into account.
Among existing permutation methods, Git Re-Basin [1] is a representative one, which successfully achieved linear connectivity between two independently trained ResNet models on CIFAR-10. Specifically, a permutation  $\pi$  that maintains functionally equivalent network can be formulated by a set of per-layer permutations  $\pi = \{P^{(\ell)}\}_{\ell=1}^{L-1}$  where  $P^{(\ell)} \in \mathbb{R}^{d_{\ell} \times d_{\ell}}$  is a permutation matrix. Ainsworth et al. [1] proposed two distinct matching objectives for aligning the neurons of independently trained models via permutation: weight matching and activation matching:
$$
\text {W e i g h t m a t c h i n g :} \quad \min  _ {\pi} \| \boldsymbol {\theta} _ {A} - \pi (\boldsymbol {\theta} _ {B}) \| ^ {2}. \tag {3}
$$
$$
\text {A c t i v a t i o n m a t c h i n g :} \quad \min  _ {\pi} \sum_ {\ell = 1} ^ {L - 1} \| \boldsymbol {H} _ {A} ^ {(\ell)} - \boldsymbol {P} ^ {(\ell)} \boldsymbol {H} _ {B} ^ {(\ell)} \| _ {F} ^ {2}. \tag {4}
$$
Here,  $\pmb{H}_A^{(\ell)} = f^{(\ell)}(\pmb{\theta}_A; \pmb{X})$  is the  $\ell$ -th layer feature matrix, and  $\pmb{H}_B^{(\ell)}$  is defined similarly.
Main Experimental Setup. In this paper, we use both the spawning method and the permutation method to obtain linearly connected modes. Following Frankle et al. [9], Ainsworth et al. [1], we perform our experiments on commonly used image classification datasets MNIST [18], CIFAR-10 [15], and Tiny-ImageNet[17], and with the standard network architectures ResNet-20/50 [12], VGG-16 [29], and MLP. We follow the same training procedures and hyper-parameters as in Frankle et al. [9], Ainsworth et al. [1]. Due to space limit, we defer some of the experimental results to the appendix. Notice that [1] increased the width of ResNet by 32 times in order to achieve zero barrier, and we also followed this setting in the experiments of the permutation method. The detailed settings and hyper-parameters are also described in Appendix B.1.
4 Layerwise Linear Feature Connectivity (LLFC)
In this section, we formally describe Layerwise Linear Feature Connectivity (LLFC) and provide empirical evidence of its consistent co-occurrence with LMC. We also show that LLFC applied to the last layer directly implies LMC.
3Following the terminology in literature, a mode refers to an optimal solution obtained at the end of training.
Figure 2: Comparison of  $\mathbb{E}_{\mathcal{D}}[1 - \mathrm{cosine}_{\alpha}(\pmb{x}_i)]$  and  $\mathbb{E}_{\mathcal{D}}[1 - \mathrm{cosine}_{A,B}(\pmb{x}_i)]$ . The spawning method is used to obtain two linearly connected modes  $\pmb{\theta}_A$  and  $\pmb{\theta}_B$ . Results are presented for different layers of various model architectures on different datasets, with  $\alpha \in \{0.25, 0.5, 0.75\}$ . Standard deviations across the dataset are reported by error bars. More results are in Appendix B.2.
Figure 3: Comparison of  $\mathbb{E}_{\mathcal{D}}[1 - \mathrm{cosine}_{\alpha}(\pmb{x}_i)]$  and  $\mathbb{E}_{\mathcal{D}}[1 - \mathrm{cosine}_{A,B}(\pmb{x}_i)]$ . The permutation method is used to obtain two linearly connected modes  $\pmb{\theta}_A$  and  $\pmb{\theta}_B$ . Results are presented for different layers of various model architectures on different datasets, with  $\alpha \in \{0.25, 0.5, 0.75\}$ . Standard deviations across the dataset are reported by error bars. More results are in Appendix B.2.
Definition 2 (Layerwise Linear Feature Connectivity). Given dataset  $\mathcal{D}$  and two modes  $\theta_{A},\theta_{B}$  of an  $L$  -layer neural network  $f$  , the modes  $\theta_{A}$  and  $\theta_{B}$  are said to be layerwise linearly feature connected if they satisfy
$$
\forall \ell \in [ L ], \forall \alpha \in [ 0, 1 ], \exists c > 0, s. t. c f ^ {(\ell)} (\alpha \boldsymbol {\theta} _ {A} + (1 - \alpha) \boldsymbol {\theta} _ {B}) = \alpha f ^ {(\ell)} (\boldsymbol {\theta} _ {A}) + (1 - \alpha) f ^ {(\ell)} (\boldsymbol {\theta} _ {B}). \tag {5}
$$
LLFC states that the per-layer feature of the interpolated model  $\theta_{\alpha} = \alpha \theta_{A} + (1 - \alpha)\theta_{B}$  has the same direction as the linear interpolation of the features of  $\theta_{A}$  and  $\theta_{B}$ . This means that the feature map  $f^{(\ell)}$  behaves similarly to a linear map (up to a scaling factor) on the line segment between  $\theta_{A}$  and  $\theta_{B}$ , even though it is a nonlinear map globally.
LLFC Co-occurs with LMC. We now verify that LLFC consistently co-occurs with LMC across different architectures and datasets. We use the spawning method and the permutation method described in Section 3 to obtain linearly connected modes  $\theta_{A}$  and  $\theta_{B}$ . On each data point  $x_{i}$  in the test set  $\mathcal{D}$ , we measure the cosine similarity between the feature of the interpolated model  $\theta_{\alpha}$  and
Figure 4: Comparison of  $\mathbb{E}_{\mathcal{D}}[1 - \mathrm{cosine}_{LMC}(\pmb{x}_i)]$  and  $\mathbb{E}_{\mathcal{D}}[1 - \mathrm{cosine}_{notLMC}(\pmb{x}_i)]$ . Both the spawning and permutation methods are used to obtain two linearly connected modes. Standard deviations across the dataset are reported by error bars.
linear interpolations of the features of  $\theta_{A}$  and  $\theta_{B}$  in each layer  $\ell$ , as expressed as  $\mathrm{cosine}_{\alpha}(\pmb{x}_i) = \cos [f^{(\ell)}(\alpha \pmb{\theta}_A + (1 - \alpha)\pmb{\theta}_B;\pmb{x}_i),\alpha f^{(\ell)}(\pmb{\theta}_A;\pmb{x}_i) + (1 - \alpha)f^{(\ell)}(\pmb{\theta}_B;\pmb{x}_i)]$ . We compare this to the baseline cosine similarity between the features of  $\theta_{A}$  and  $\theta_{B}$  in the corresponding layer, namely  $\mathrm{cosine}_{A,B}(\pmb{x}_i) = \cos [f^{(\ell)}(\pmb{\theta}_A;\pmb{x}_i),f^{(\ell)}(\pmb{\theta}_B;\pmb{x}_i)]$ . The results for the spawning method and the permutation method are presented in Figures 2 and 3, respectively. They show that the values of  $\mathbb{E}_{\mathcal{D}}[1 - \mathrm{cosine}_{\alpha}(\pmb{x}_i)]$  are close to 0 across different layers, architectures, datasets, and different values of  $\alpha$ , which verifies the LLFC property. The presence of small error bars indicates consistent behavior for each data point. Moreover, the values of  $\mathbb{E}_{\mathcal{D}}[1 - \mathrm{cosine}_{A,B}(\pmb{x}_i)]$  are not close to 0, which rules out the trivial case that  $f^{(\ell)}(\pmb{\theta}_A)$  and  $f^{(\ell)}(\pmb{\theta}_B)$  are already perfectly aligned.
To further verify, we also compare the values of  $\mathbb{E}_{\mathcal{D}}[1 - \mathrm{cosine}_{\alpha}(\pmb{x}_i)]$  of two linearly connected modes with those of two modes that are independently trained (not satisfying LMC). We measure  $\mathrm{cosine}_{0.5}$  of the features of two modes that are linearly connected and two modes that are independently trained, denoted as  $\mathrm{cosine}_{LMC}$  and  $\mathrm{cosine}_{notLMC}$  correspondingly. In Figure 4, the values of  $\mathbb{E}_{\mathcal{D}}[1 - \mathrm{cosine}_{LMC}(\pmb{x}_i)]$  are negligible compared to  $\mathbb{E}_{\mathcal{D}}[1 - \mathrm{cosine}_{notLMC}(\pmb{x}_i)]$ . The experimental results align with Figures 2 and 3 and thus we firmly verify the claim that LLFC co-occurs with LMC.
LLFC Implies LMC. Intuitively, LLFC is a much stronger characterization than LMC since it establishes a linearity property in the high-dimensional feature map in every layer, rather than just for the final error. Lemma 1 below formally establishes that LMC is a consequence of LLFC by applying LLFC on the output layer.
Lemma 1 (Proof in Appendix A.1). Suppose two modes  $\pmb{\theta}_{A}$ ,  $\pmb{\theta}_{B}$  satisfy LLFC on a dataset  $\mathcal{D}$  and  $\max \{\mathrm{Err}_{\mathcal{D}}(\pmb{\theta}_{A}), \mathrm{Err}_{\mathcal{D}}(\pmb{\theta}_{B})\} \leq \epsilon$ , then we have
$$
\forall \alpha \in [ 0, 1 ], \operatorname {E r r} _ {\mathcal {D}} \left(\alpha \boldsymbol {\theta} _ {A} + (1 - \alpha) \boldsymbol {\theta} _ {B}\right) \leq 2 \epsilon . \tag {6}
$$
In summary, we see that the LMC property, which was used to study the loss landscapes in the entire parameter space, extends to the internal features in almost all the layers. LLFC offers much richer structural properties than LMC. In Section 5, we will dig deeper into the contributing factors to LLFC and leverage the insights to gain new understanding of the spawning and the permutation methods.
5 Why Does LLFC Emerge?
We have seen that LLFC is a prevalent phenomenon that co-occurs with LMC, and it establishes a broader notion of linear connectivity than LMC. In this section, we investigate the root cause of LLFC, and identify two key conditions, weak additivity for ReLU activations and commutativity. We verify these conditions empirically and prove that they collectively imply LLFC. From there, we provide an explanation for the effectiveness of the permutation method, offering new insights into LMC.
5.1 Underlying Factors of LLFC
For convenience, we consider a multi-layer perceptron (MLP) in Section 5, though the results can be easily adapted to any feed-forward structure, e.g., a convolutional neural network $^{4}$  (CNN). For