Generalized Shape Metrics on Neural Representations

1 Introduction . The extent to which different deep networks or neurobiological systems use equivalent representations in support of similar task demands is a topic of persistent interest in machine learning and neuroscience [ 1–3 ] . Several methods including linear regression [ 4 , 5 ] , canonical correlation analysis ( CCA ; [ 6 , 7 ] ) , representational similarity analysis ( RSA ; [ 8 ] ) , and centered kernel alignment ( CKA ; [ 9 ] ) have been used to quantify the similarity of hidden layer activation patterns . These measures are often interpreted on an ordinal scale and are employed to compare a limited number of networks—e.g. , they can indicate whether networks A and B are more or less similar than networks A and C. While these comparisons have yielded many insights [ 4–12 ] , the underlying methodologies have not been extended to systematic analyses spanning thousands of networks . To unify existing approaches and enable more sophisticated analyses , we draw on ideas from statistical shape analysis [ 13–15 ] to develop dissimilarity measures that are proper metrics—i.e. , measures that are symmetric and respect the triangle inequality . This enables several off-the-shelf methods with theoretical guarantees for classification ( e.g . k-nearest neighbors , [ 16 ] ) and clustering ( e.g . hierarchical clustering [ 17 ] ) . Existing similarity measures can violate the triangle inequality , which complicates these downstream analyses [ 18–20 ] . However , we show that existing dissimilarity 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . measures can often be modified to satisfy the triangle inequality and viewed as special cases of the framework we outline . We also describe novel metrics within this broader family that are specialized to convolutional layers and have appealing properties for analyzing artificial networks . Moreover , we show empirically that these metric spaces on neural representations can be embedded with low distortion into Euclidean spaces , enabling an even broader variety of previously unconsidered supervised and unsupervised analyses . For example , we can use neural representations as the inputs to linear or nonlinear regression models . We demonstrate this approach on neural representations in mouse visual cortex ( Allen Brain Observatory ; [ 21 ] ) in order to predict each brain region ’ s anatomical hierarchy from its pattern of visual responses—i.e. , predicting a feature of brain structure from function . We demonstrate a similar approach to analyze hidden layer representations in a database of 432K deep artificial networks ( NAS-Bench-101 ; [ 22 ] ) and find a surprising degree of correlation between early and deep layer representations . Overall , we provide a theoretical grounding which explains why existing representational similarity measures are useful : they are often close to metric spaces , and can be modified to fulfill metric space axioms precisely . Further , we draw new conceptual connections between analyses of neural representations and established research areas [ 15 , 23 ] , utilize these insights to propose novel metrics , and demonstrate a general-purpose machine learning workflow that scales to datasets with thousands of networks . 2 Methods . This section outlines several workflows ( Fig . 1 ) to analyze representations across large collections of networks . After briefly summarizing prior approaches ( sec . 2.1 ) , we cover background material on metric spaces and discuss their theoretical advantages over existing dissimilarity measures ( sec . 2.2 ) . We then present a class of metrics that capture these advantages ( sec . 2.3 ) and cover a special case that is suited to convolutional layers ( sec . 2.4 ) . We then demonstrate the practical advantages of these methods in Section 3 , and demonstrate empirically that Euclidean feature spaces can approximate the metric structure of neural representations , enabling a broad set of novel analyses . 2.1 Prior work and problem setup . Neural network representations are often summarized over a set of m reference inputs ( e.g . test set images ) . Let Xi ∈ Rm×ni and Xj ∈ Rm×nj denote the responses of two networks ( with ni and nj neurons , respectively ) to a collection of these inputs . Quantifying the similarity between Xi and Xj is complicated by the fact that , while the m inputs are the same , there is no direct correspondence between the neurons . Even if ni = nj , the typical Frobenius inner product , 〈Xi , Xj〉 = Tr [ X > i Xj ] , and metric , ‖Xi−Xj‖ = 〈Xi−Xj , Xi−Xj〉1/2 , fail to capture the desired notion of dissimilarity . For instance , let Π denote some n × n permutation matrix and let Xi = XjΠ . Intuitively , we should consider Xi and Xj to be identical in this case since the ordering of neurons is arbitrary . Yet , clearly ‖Xi −Xj‖ 6= 0 , except in very special cases . One way to address this problem is to linearly regress over the neurons to predict Xi from Xj . Then , one can use the coefficient of determination ( R2 ) as a measure of similarity [ 4 , 5 ] . However , this similarity score is asymmetric—if one instead treats Xj as the dependent variable that is predicted from Xi , this will result in a different R2 . Canonical correlation analysis ( CCA ; [ 6 , 7 ] ) and linear centered kernel alignment ( linear CKA ; [ 9 , 24 ] ) also search for linear correspondences between neurons , but have the advantage of producing symmetric scores . Representational similarity analysis ( RSA ; [ 8 ] ) is yet another approach , which first computes an m×m matrix holding the dissimilarities between all pairs of representations for each network . These representational dissimilarity matrices ( RDMs ) , are very similar to the m ×m kernel matrices computed and compared by CKA . RSA traditionally quantifies the similarity between two neural networks by computing Spearman ’ s rank correlation between their RDMs . A very recent paper by Shahbazi et al . [ 25 ] , which was published while this manuscript was undergoing review , proposes to use the Riemannian metric between positive definite matrices instead of Spearman correlation . Similar to our results , this establishes a metric space that can be used to compare neural representations . Here , we leverage metric structure over shape spaces [ 13–15 ] instead of positive definite matrices , leading to complementary insights . In summary , there are a diversity of methods that one can use to compare neural representations . Without a unifying theoretical framework it is unclear how to choose among them , use their outputs for downstream tasks , or generalize them to new domains . 2.2 Feature space mapping , metrics , and equivalence relations . Our first contribution will be to establish formal notions of distance ( metrics ) between neural representations . To accommodate the common scenario when the number of neurons varies across networks ( i.e . when ni 6= nj ) , we first map the representations into a common feature space . For each set of representations , Xi , we suppose there is a mapping into a p-dimensional feature space , Xi 7→ Xφi , where X φ i ∈ Rm×p . In the special case where all networks have equal size , n1 = n2 = . . . = n , we can express the feature mapping as a single function φ : Rm×n 7→ Rm×p , so that Xφi = φ ( Xi ) . When networks have dissimilar sizes , we can map the representations into a common dimension using , for example , PCA [ 6 ] . Next , we seek to establish metrics within the feature space , which are distance functions that satisfy : Equivalence : d ( Xφi , X φ j ) = 0 ⇐⇒ X φ i ∼X φ j ( 1 ) Symmetry : d ( Xφi , X φ j ) = d ( X φ j , X φ i ) ( 2 ) Triangle Inequality : d ( Xφi , X φ j ) ≤ d ( X φ i , X φ k ) + d ( X φ k , X φ j ) ( 3 ) for all Xφi , X φ j , and X φ k in the feature space . The symbol ‘ ∼ ’ denotes an equivalence relation between two elements . That is , the expression Xφi ∼ X φ j means that “ X φ i is equivalent to X φ j . ” Formally , distance functions satisfying Eqs . ( 1 ) to ( 3 ) define a metric over a quotient space defined by the equivalence relation and a pseudometric over Rm×p ( see Supplement A ) . Intuitively , by specifying different equivalence relations we can account for symmetries in network representations , such as permutations over arbitrarily labeled neurons ( other options are discussed below in sec . 2.3 ) . Metrics quantify dissimilarity in a way that agrees with our intuitive notion of distance . For example , Eq . ( 2 ) ensures that the distance from Xφi to X φ j is the same as the distance from X φ j to X φ i . Linear regression is an approach that violates this condition : the similarity measured by R2 depends on which network is treated as the dependent variable . Further , Eq . ( 3 ) ensures that distances are self-consistent in the sense that if two elements ( Xφi and Xφj ) are both close to a third ( X φ k ) , then they are necessarily close to each other . Many machine learning models and algorithms rely on this triangle inequality condition . For example , in clustering , it ensures that if Xφi and X φ j are put into the same cluster as X φ k , then X φ i and X φ j can not be too far apart , thus implying that they too can be clustered together . Intuitively , this establishes an appealing transitive relation for clustering , which can be violated when the triangle inequality fails to hold . Existing measures based on CCA , RSA , and CKA , are symmetric , but do not satisfy the triangle inequality . By modifying these approaches to satisfy the triangle inequality , we avoid potential pitfalls and can leverage theoretical guarantees on learning in proper metric spaces [ 16–20 ] . 2.3 Generalized shape metrics and group invariance . In this section , we outline a new framework to quantify representational dissimilarity , which leverages a well-developed mathematical literature on shape spaces [ 13–15 ] . The key idea is to treat Xφi ∼X φ j if and only if there exists a linear transformation T within a set of allowable transformations G , such that Xφi = X φ j T . Although G only contains linear functions , nonlinear alignments between the raw representations can be achieved when the feature mappings Xi 7→ Xφi are chosen to be nonlinear . Much of shape analysis literature focuses on the special case where p = n and G is the special orthogonal group SO ( n ) = { R ∈ Rn×n | R > R = I , det ( R ) = 1 } , meaning that Xφi and X φ j are equivalent if there is a n-dimensional rotation ( without reflection ) that relates them . Standard shape analysis further considers each Xφi to be a mean-centered ( ( X φ i ) > 1 = 0 ) and normalized ( ‖Xφi ‖ = 1 ) version of the raw landmark locations held in Xi ∈ Rm×n ( an assumption that we will relax ) . That is , the feature map φ : Rm×n 7→ Sm×n transforms the raw landmarks onto the hypersphere , denoted Sm×n , of m × n matrices with unit Frobenius norm . In this context , Xφi ∈ Sm×n is called a “ pre-shape. ” By removing rotations from a pre-shape , [ Xφi ] = { S ∈ Sm×n | S ∼X φ i } for pre-shape X φ i , we recover its “ shape. ” To quantify dissimilarity in neural representations , we generalize this notion of shape to include other feature mappings and alignments . The minimal distance within the feature space , after optimizing over alignments , defines a metric under suitable conditions ( Fig . 2A ) . This results in a broad variety of generalized shape metrics ( see also , ch . 18 of [ 15 ] ) , which fall into two categories as formalized by the pair of propositions below . Proofs are provided in Supplement B . Proposition 1 . Let Xφi ∈ Rm×p , and let G be a group of linear isometries on Rm×p . Then , d ( Xφi , X φ j ) = min T∈G ‖Xφi −X φ j T ‖ ( 4 ) defines a metric , where Xφi ∼X φ j if and only if there is a T ∈ G such that X φ i = X φ j T . Proposition 2 . Let Xφi ∈ Sm×p , and let G be a group of linear isometries on Sm×p . Then , θ ( Xφi , X φ j ) = min T∈G arccos 〈 Xφi , X φ j T 〉 ( 5 ) defines a metric , where Xφi ∼X φ j if and only if there is a T ∈ G such that X φ i = X φ j T . Two key conditions appear in these propositions . First , G must be a group of functions . This means G is a set that contains the identity function , is closed under composition ( T 1T 2 ∈ G for any T 1 ∈ G and T 2 ∈ G ) , and whose elements are invertible by other members of the set ( if T ∈ G then T−1 ∈ G ) . Second , every T ∈ G must be an isometry , meaning that ‖Xφi −X φ j ‖ = ‖X φ i T −X φ j T ‖ for all T ∈ G and all elements of the feature space . On Rm×p and Sm×p , all linear isometries are orthogonal transformations . Further , the set of orthogonal transformations , O ( p ) = { Q ∈ Rp×p : Q > Q = I } , defines a well-known group . Thus , the condition that G is a group of isometries is equivalent to G being a subgroup of O ( p ) —i.e. , a subset of O ( p ) satisfying the group axioms . Intuitively , by requiring G to be a group of functions , we ensure that the alignment procedure is symmetric—i.e . it is equivalent to transform Xφi to match X φ j , or transform the latter to match the former . Further , by requiring each T ∈ G to be an isometry , we ensure that the underlying metric ( Euclidean distance for Proposition 1 ; angular distance for Proposition 2 ) preserves its key properties . Together , these propositions define a broad class of metrics as we enumerate below . For simplicity , we assume that ni = nj = n in the examples below , with the understanding that a PCA or zero-padding preprocessing step has been performed in the case of dissimilar network sizes . This enables us to express the metrics as functions of the raw activations , i.e . functions Rm×n × Rm×n 7→ R+ . Permutation invariance The most stringent notion of representational similarity is to demand that neurons are one-to-one matched across networks . If we set the feature map to be the identity function , i.e. , Xφi = Xi for all i , then : dP ( Xi , Xj ) = min Π∈P ( n ) ‖Xi −XjΠ‖ ( 6 ) defines a metric by Proposition 1 since the set of permutation matrices , P ( n ) , is a subgroup of O ( n ) . To evaluate this metric we must optimize over the set of neuron permutations to align the two networks . This can be reformulated ( see Supplement C ) as a fundamental problem in combinatorial optimization known as the linear assignment problem [ 26 ] . Exploiting an algorithm due to Jonker and Volgenant [ 27 , 28 ] we can solve this problem in O ( n3 ) time . The overall runtime for evaluating Eq . ( 6 ) is O ( mn2 + n3 ) , since we must evaluate X > i Xj to formulate the assignment problem . Rotation invariance Let C = Im − ( 1/m ) 11 > denote an m×m centering matrix , and consider the feature mapping φ1 which mean-centers the columns , φ1 ( Xi ) = X φ1 i = CXi . Then , d1 ( Xi , Xj ) = min Q∈O ‖Xφ1i −X φ1 j Q‖ ( 7 ) defines a metric by Proposition 1 , and is equivalent to the Procrustes size-and-shape distance with reflections [ 15 ] . Further , by Proposition 2 , θ1 ( Xi , Xj ) = min Q∈O arccos 〈Xφ1i , X φ1 j Q〉 ‖Xφ1i ‖‖X φ1 j ‖ ( 8 ) defines another metric , and is closely related to the Riemannian distance on Kendall ’ s shape space [ 15 ] . To evaluate Eqs . ( 7 ) and ( 8 ) , we must optimize over the set of orthogonal matrices to find the best alignment . This also maps onto a fundamental optimization problem known as the orthogonal Procrustes problem [ 29 , 30 ] , which can be solved in closed form inO ( n3 ) time . As in the permutationinvariant metric described above , the overall runtime is O ( mn2 + n3 ) . Linear invariance Consider a partial whitening transformation , parameterized by 0 ≤ α ≤ 1 : Xφα = CX ( αIn + ( 1− α ) ( X > CX ) −1/2 ) ( 9 ) Note that X > CX is the empirical covariance matrix of X . Thus , when α = 0 , Eq . ( 9 ) corresponds to ZCA whitening [ 31 ] , which intuitively removes invertible linear transformations from the representations . Thus , when α = 0 the metric outlined below treats Xi ∼ Xj if there exists an affine transformation that relates them : Xi = XjW + b for some W ∈ Rn×n and b ∈ Rn . When α = 1 , Eq . ( 9 ) reduces to the mean-centering feature map used above . Using orthogonal alignments within this feature space leads to a metric that is related to CCA . First , let ρ1 ≥ . . . ≥ ρn ≥ 0 denote the singular values of ( Xφαi ) > ( X φα j ) /‖X φα i ‖‖X φα j ‖ . One can show that θα ( Xi , Xj ) = min Q∈O arccos 〈Xφαi , X φα j Q〉 ‖Xφαi ‖‖X φα j ‖ = arccos ( ∑ ` ρ ` ) , ( 10 ) and we can see from Proposition 2 that this defines a metric for any 0 ≤ α ≤ 1 . When α = 0 , the values ρ1 , . . . , ρn are proportional to the canonical correlation coefficients , with 1/n being the factor of proportionality . When α > 0 , these values can be viewed as ridge regularized canonical correlation coefficients [ 32 ] . See Supplement C for further details . Past works [ 6 , 7 ] have used the average canonical correlation as a measure of representational similarity . When α = 0 , the average canonical correlation is given by ∑ ` ρ ` = cos θ0 ( Xi , Xj ) . Thus , if we apply arccos ( · ) to the average canonical correlation , we modify the calculation to produce a proper metric ( see Fig . 4A ) . Since the covariance is often ill-conditioned or singular in practice , setting α > 0 to regularize the calculation is also typically necessary . Nonlinear invariances We discuss feature maps that enable nonlinear notions of equivalence , and which relate to kernel CCA [ 33 ] and CKA [ 9 ] , in Supplement C .

Summary: The paper is concerned with issues that are fundamental to basic research on understanding natural and artificial neural networks. Such research often involves measuring the similarity of neural representations (patterns of neural activity). The authors propose desirable attributes for the methods used for such comparisons, e.g. that the similarity function should define a metric metric space: it should satisfy the constraints of equivariance, symmetry and the triangle inequality. They show that many commonly used methods do not satisfy these constraints but can be easily modified to do so. They also construct a metric especially suited to convolutional layers, which traditionally have not been straightforwardly analyzed by existing methods. The authors also discuss practical issues of general importance: They discuss what sample size is required for such metrics to be accurate/robust, which is highly important for neuroscience research where it can be difficult to obtain large sample sizes. They also talk about dimensionality reduction, showing that approximating neural embeddings with low dimensional Euclidean spaces, while suitable for visualization, can highly distort the embedding. They validate their new metrics by showing that it can recover known structure in large scale datasets.