Similarity and Matching of Neural Network Representations

1 Introduction . A central topic in the analysis of deep neural networks is the investigation of the learned representations . A good understanding of the features learned on inner layers provides means to analyse and advance deep learning systems . A fundamental question in this line of inquiry is “ when are two representations similar ? ” We contribute to this line of research by introducing a novel approach to study the above question . The key idea of our work is to ask the following , somewhat different question : “ in what way are two representations similar , once we know that they are similar ? ” In this paper , we present a conceptual framework and a methodology that allows to meaningfully pose and study this question . For this purpose , we will distinguish two approaches to define similarity : representational similarity and functional similarity . Representational similarity Representational similarity notions define similarity via computing statistics on two different data embeddings that capture appropriate geometrical or statistical properties . There exist several statistical measures , [ Raghu et al. , 2017 , Morcos et al. , 2018 , Kornblith et al. , 2019 ] each of which serves as a different notion of similarity . These measures have been used successfully to obtain valuable insights about the inner workings of deep neural networks [ Nguyen et al. , 2021 , Mirzadeh et al. , 2021 , Neyshabur et al. , 2020 , Wu et al. , 2020 ] . Functional similarity While representational similarity concerns the data embeddings , in contrast , functional similarity concerns the functions that produce and process these embeddings , i.e. , parts of the function compositions that the neural networks realize . With these elements , we can ask new types of questions that we were not able to ask using solely the data embeddings . For example , we can ask the following : “ can network B achieve its task using the representations of network A ? ” 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Our paper focuses on this specific theme . In particular , we investigate this by taking the activations of network A at a given layer , transform it with an affine map , then use the result as an input of the same layer in network B . In other words , we stitch together the two networks with an affine stitching layer . This technique first appeared in Lenc and Vedaldi [ 2019 ] to study the invariance and equivariance properties of convolutional networks . Evaluating the performance of this combined network provides an alternative viewpoint on similarity of representations , the viewpoint of functional similarity . We investigate this novel perspective of reflecting on representational similarity through the lens of functional similarity . A brief outline of this work1 is the following : • After a detailed introduction to model stitching ( Section 3 ) and presenting two types of stitching methods ( Section 4 ) , we empirically demonstrate that trained convolutional networks with the same architecture but different initializations can be stitched together , even with a single , affine stitching layer in many cases without significant performance loss ( Section 5 ) . We will refer to this compatibility property as the ‘ matchability ’ of representations . • Observing the matchability of representations as a quite robust property of common vision models , we have a wide range of examples in our hands when we know that representations are functionally similar . This permits us to study the relation between representational similarity measures and functional similarity . In our experiments , we show that the values of similarity indices are not necessarily indicative of task performance on a stitched network , and perhaps more surprisingly , high-performance stitchings can have differences in the values of similarity indices such as Centered Kernel Alignment ( CKA ) [ Cortes et al. , 2012 , Kornblith et al. , 2019 ] . This also reflects on phenomena experienced by the ‘ end-users ’ of these indices . For instance , in the context of continual learning Mirzadeh et al . [ 2021 ] observe that CKA remains constant , while the accuracy on previous tasks drops drastically ( Section 6 ) . • Constraining the space of transformations in the stitching layer provides means to analyse how the information is organised in the representations . As an example of this probing methodology , we study bottleneck stitching layers , and show that when doing SVD in the representational space , the magnitudes of the principal components are not in direct correspondence with the information content of the latent directions ( Section 7 ) . As SVD is commonly used with embeddings , and there are representational similarity measures , for example SVCCA [ Raghu et al. , 2017 ] , that use SVD as an ingredient , these results might be of interest to both theorists and practitioners . • The transformation matrices of well-performing stitching layers capture how representations are functionally similar . This provides the opportunity to empirically investigate “ in what way are two representations similar , once we know that they are similar ? ” . As a first step on this avenue , we present results regarding the uniqueness and sparsity properties of the stitching matrices ( Section 8 ) . 2 Related work . Matching representations The idea of 1x1 convolution stitching layers connecting the lower half of a network with the upper half of another network first appears in Lenc and Vedaldi [ 2019 ] . They demonstrate that such hybrid “ Franken-CNNs ” can approach the accuracy of the original models in many cases . For shallower layers , this phenomenon can appear even when the first network is trained on a different task , suggesting that lower level representations are sufficiently generic . Mapping layers as introduced by Li et al . [ 2016 ] are single layer linear 1x1 convolutional networks mapping between two representations . They do not incorporate information from the second half of the network . The authors train mapping layers with L1 regularized MSE , given their main goal of identifying sparse mappings between the filters of two different networks . As Li et al . [ 2016 ] notes , learning mapping layers can get stuck at poor local minima , even without L1 regularization . This forces the authors to restrict their experiments to the shallowest layers of the networks they inspect . We manage to overcome the issue of poor local optima by initializing from the optimal least squares matching ( Section 8 ) , and as our experiments demonstrate , this alleviates the issue , resulting in stitched networks that often achieve the performance of their constituents . Concurrently and independently of our work , Bansal et al . [ 2021 ] also apply model stitching to investigate similarities between deep learning models . They observe that for a given task , even very different architectures can be successfully stitched , e.g. , a transformer with a residual network . 1Code is available at the project website : https : //bit.ly/similarity-and-matching . Similarity indices Canonical Correlation Analysis ( CCA ) provides a statistical measure of similarity . CCA itself has a tendency to give too much emphasis to the behavior of low-magnitude high-noise latent directions . Similarity indices were designed to make CCA more robust in this sense , namely Singular Value CCA [ Raghu et al. , 2017 ] and Projection Weighted CCA [ Morcos et al. , 2018 ] . The current state of the art in measuring similarity is Centered Kernel Alignment ( CKA ) [ Cristianini et al. , 2001 , Cortes et al. , 2012 , Kornblith et al. , 2019 ] . Kornblith et al . [ 2019 ] conducted a comparative study of the measures CCA , SVCCA , PWCCA as similarity indices , and they observed that the other similarity indices do not succeed in identifying similar layers of identical networks trained from different initializations , whereas CKA manages to do so , therefore in this paper we chose to focus on the latter . Kornblith et al . [ 2019 ] also consider nonlinear CKA kernels , but as they demonstrate this does not lead to meaningful improvements , so we focus only on the linear variant ( also commonly known as RV coefficient Robert and Escoufier [ 1976 ] ) , in line with other literature employing CKA similarity [ Nguyen et al. , 2021 , Mirzadeh et al. , 2021 , Neyshabur et al. , 2020 , Wu et al. , 2020 ] . Functional similarity The general framework of knowledge distillation [ Hinton et al. , 2015 , Romero et al. , 2014 ] can be seen as fostering end-to-end functional similarity of a student network to a teacher network . This relation and the versatile utilisation of this technique underpins the importance of the notion of functional similarity . Yosinski et al . [ 2014 ] studies transfer learning from a related perspective . 3 Introducing Dr. Frankenstein . 3.1 Preliminaries and notation . Let fθ : X → Y denote the input-output function of a feedforward artificial neural network with m ∈ N layers , thus , fθ = fm ◦ · · · ◦ f1 , where X is the input space , Y is the output space , fi : Ai−1 → Ai are maps between activation spaces Ai−1 and Ai for i ∈ [ m ] with A0 = X , and θ are the parameters of the network . We consider models that are trained on a datasetD = { ( xi , yi ) } ni=1 in a supervised manner with inputs xi ∈ X and labels yi ∈ Y , where n ∈ N is the dataset size . The central tool of our investigation involves splitting up the network at a given layer to a representation map and a task map . The representation map at layer L is a function RL : X → AL that maps each data point xi to its activation vector ai , L at layer L : RL ( xi ) = fL ◦ · · · ◦ f1 ( xi ) = ai , L , where AL is the activation space at layer L. The activation vectors for the dataset D at layer L are simply denoted by AL = ( ai , L ) ni=1 . The task map at layer L is a function TL : AL → Y that maps an activation vector at layer L to the final output of the network : TL ( ai , L ) = fm ◦ · · · ◦ fL+1 ( ai , L ) . Thus , the input-output map of the network fθ is simply the composition fθ = TL ◦RL : X → Y for all L = 1 , . . . , m. We will omit the index of the layer L when this does not hurt clarity . 3.2 The Frankenstein learning task . Definition 3.1 ( Frankenstein network ) . Fix two ( pretrained ) neural networks fθ and fφ . We will call a transformation S : Aθ , L → Aφ , M between the two activation spaces at layers L and M a stitching layer . Given a stitching layer S , the Frankenstein or stitched network corresponding to S is the composition F = FS : X → Y , F ( x ) = Tφ , M ◦ S ◦Rθ , L ( x ) , where Tφ , M is the task map at layer M for the network fφ and Rθ , L is the representation map at layer L for network fθ . The stitching layer thus realizes a correspondence between the representations of two neural networks : it transforms the activations of one network at a particular layer to be a suitable input to the corresponding layer in the other network . For better readability , we will often call fθ as Model 1 and fφ as Model 2 , and index the corresponding data accordingly ( f1 = fθ , f2 = fφ , etc. ) . In the experiments of this paper , Model 1 and Model 2 have the same architecture and L = M . Definition 3.2 ( Frankenstein learning task ) . Given two pretrained neural networks fθ and fφ , a task given by a labeled dataset D = { ( xi , yi ) ∈ X × Y : i = 1 , . . . , n } and a loss function L : Y ×Y → R , we will call the Frankenstein learning task at layer L the task of finding the stitching layer S : Aθ , L → Aφ , L which minimizes the total value of the loss L on the dataset D. Frankenstein is an evaluation framework The Frankenstein learning task primarily provides an evaluation framework for matching representations , and it does not tie our hands on how the transformation function S is produced . As an evaluation metric , we can either use the mean loss value of the Frankenstein learning task on an evaluation set , or any other appropriate metric ( e.g. , accuracy in the case of classification tasks ) . Stitching convolutional representations Our toolset is generic , but here we will only consider convolutional architectures . Thus , the activation spaces have the structure of rank-3 tensors Rw×h×c , where w and h are the spatial width and height and c is the number of feature maps . We will only consider linear stitching layers of the form M : Rc → Rc ( 1 × 1 convolutional layers ) . When formulating least squares problems on n many activation tensors , we always mean the least squares problem in Rc , which in practice means reshaping the n× w × h× c tensor to a nwh× c matrix .

Summary: The paper develops and applies a set of methods to determine the functional equivalence of architecturally-identical neural networks trained with different initializations. These methods involve taking the activations from a given layer l in one network, inputting them into a "stitching layer" (either a linear projection or a 1x1 convolutional layer), then taking the activations from the stitching layer and inputting them into layer l+1 in a second network. The stitching layer is optimized either to minimize the models' loss on a task, or to minimize the distance between its own outputs and the activations of the second network. The authors find that stitching layers can typically be learned that have minimal impact on model accuracy and loss. They also demonstrate some results that they claim undermine the utility of CKA, a popular representational similarity metric.