Representational Dissimilarity Metric Spaces for Stochastic Neural Networks
Authors: Lyndon Duong
Jingyang Zhou
Josue Nassar
Jules Berman
Jeroen Olieslagers
Alex H Williams
Venue: ICLR 2023 poster (2023)
Keywords: stochastic neural networks
noise
representational geometry
dissimilarity metrics

Abstract: Quantifying similarity between neural representations---e.g. hidden layer activation vectors---is a perennial problem in deep learning and neuroscience research. Existing methods compare deterministic responses (e.g. artificial networks that lack stochastic layers) or averaged responses (e.g., trial-averaged firing rates in biological data). However, these measures of _deterministic_ representational similarity ignore the scale and geometric structure of noise, both of which play important roles in neural computation. To rectify this, we generalize previously proposed shape metrics (Williams et al. 2021) to quantify differences in _stochastic_ representations. These new distances satisfy the triangle inequality, and thus can be used as a rigorous basis for many supervised and unsupervised analyses. Leveraging this novel framework, we find that the stochastic geometries of neurobiological representations of oriented visual gratings and naturalistic scenes respectively resemble untrained and trained deep network representations. Further, we are able to more accurately predict certain network attributes (e.g. training hyperparameters) from its position in stochastic (versus deterministic) shape space.

REPRESENTATIONAL DISSIMILARITY METRIC SPACES FOR STOCHASTIC NEURAL NETWORKS
Lyndon R. Duong, Jingyang Zhou, Josue Nassar, Jules Berman, Jeroen Olieslagers, Alex H. Williams Center for Computational Neuroscience, Flatiron Institute, New York, NY, 10010
Center for Neural Science, New York University, New York, NY, 10003 {lyndon, duong, jingyang. zhou, alex. h. williams j@nyu.
ABSTRACT
Quantifying similarity between neural representations c.g. hidden layer activa- tion vectors-is a perennial problem in deep learning and neuroscience research. Existing methods compare deterministic responses (e.g. artificial networks that lack stochastic layers) or averaged responses (e.g., trial-averaged firing rates in bi- ological data). However, these measures of deterministic representational similar- ity ignore the scale and geometric structure of noise, both of which play important roles in neural computation. To rectify this, we generalize previously proposed shape metrics (Williams et al., 2021) to quantify differences in stochastic repre- sentations. These new distances satisfy the triangle inequality, and thus can be used as rigorous basis for many supervised and unsupervised analyses. Leverag- ing this novel framework, we find that the stochastic geometries of neurobiological representations of oriented visual gratings and naturalistic scenes respectively re- semble untrained and trained deep network representations. Further, we are able to more accurately predict certain network attributes (e.g. training hyperparameters) from its position in stochastic (versus deterministic) shape space.
INTRODUCTION
Comparing high-dimensional neural pepponsee--nurobioolon firing rates or hidden layer tivations in artificial networks-is a fundamental problem in neuroscience and machine learn- ing (Dwivedi & Roig, 2019; Chung & Abbott, 2021). There representational dissimilarity including canonical correlations analysis (CCA; Raghu et al., 2017), centered kernel alignment (CKA; Kornblith al., 2019), representational similarity analysis (RSA; Kriegeskorte et al., 2008a), shape metrics (Williams et al., 2021), and Riemannian distance (Shah- bazi al., 2021). Intuitively, these measures quantify similarity the geometry of neural responses while removing expected forms of invariance, such as permutations over arbitrary neuron labels.
many methods for quantifying
However, these methods only compare deterministic representations-iic. networks that can be represented as a function f Z + I where n denotes the number of neurons and Z denotes the space of network inputs. For example, each z Z could correspond an image, and f=) is the response evoked by this image across population of neurons (Fig. 1A). Biological networks are essentially never deterministic this fashion. In fact, the variance of a stimulus-evoles response often larger than its mean (Goris et al., 2014). Stochastic responses also arise in the deep learning literature in many contexts, such as in deep generative modeling (Kingma & Welling, 2019), Bayesian neural networks (Wilson, 2020), or to provide regularization (Srivastava et al., 2014).
neural
Stochastic networks may be conceptualized as functions mapping each z E Z to a probability distribution, F(- 2), over RR (Fig. IB, Kriegeskorte & Wei 2021). Although it is easier to study the representational geometry of the average response, it is well understood that this provides an incomplete and potentially misleading picture (Kriegeskorte & Douglas, 2019). For instance, the ability to discriminate between two inputs 2, z' E Z depends on the overlap of F(z) and F(z'), and not simply the separation of their means (Fig. IC-D). A rich literature neuroscience has built on top of this insight (Shadlen et al., 1996; Abbott & Dayan, 1999; Rumyantsev et al., 2020). However,
"qqual contribution.

---

A
Generalized shape metrics are an alternative approach to quantifying representational dissimilarity. The idea is to compute the distance after minimizing deterministic network
B
stochastic network
C
D
class
class 1
E
rotate
Figure 1: (A) Illustration of deterministic network mapping inputs, (color-coded images) into points IR.. (B) Illustration of a stochastic network, where each input, z E Z, maps onto a dis- tribution, F- =), over IR., (C) Example where noise correlations impair discriminability between two image classes. (D) Example where noise correlations improve discriminability (see Abbott & Dayan, 1999). (E) Illustration of two stochastic networks with equivalent representational geometry.
to our knowledge, no studies have compared noise correlation structure across animal subjects or species, as has been done with trial-averaged responses. In machine learning, many studies have characterized the effects of noise model predictions (Sietsma & Dow, 1991; An, 1996), but only a handful have begun to characterize the geometry of stochastic hidden layers (Dapello et al., 2021).
To address these gaps, we formulate a novel class of metric spaces over stochastic neural represen- tations. That is, given two stochastic networks F; and Fj, we construct distance functions d(F,, F2) that are symmetric, satisfy the triangle inequality, and are equal zero if and only if F; and F3 are equivalent according to pre-defined criterion. In the deterministic limit-i.. as F; and F, map onto Dirac delta functions approach converges to well-studied metries over shape spaces (Dry- den & Mardia, 1993; Srivastava & Klassen, 2016), which were proposed by Williams et al. (2021) to measure distances between deterministic networks. The triangle inequality required to derive thcoretical guarantees for many downstream analyses (e.g. nonparametric regression, Cover & Hart 1967, and clustering, Dasgupta Long 2005). Thus, we lay an important foundation for analyzing stochastic representations, akin to results shown by Williams et al. (2021) in the deterministic case.
2 METHODS
2.1
DETERMINISTIC SHIAP METRICS
We begin by reviewing how shape metries quantify representational dissimilarity in the deterministic case. In the Discussion (sec. 4), we review other related prior work.
Let (fi..... : f& } denote deterministic neural networks, each given by function fx : Z R= Representational similarity between networks is typically defined with respect to a set of M inputs, (.,,..., 2M) E 2M We can collect the representations of each network into matrix:
fx(21)
Xx=
(I)
fa (*M)
A naive dissimilarity measure would be the Euclidean distance, || -X,-Xp XXIFF. This is nearly always useless. Since neurons typically labelled in arbitrary order, the very least-bo be invariant to permutations. Intuitively, we desire a notion of distance such that d(X,, X,) d(X,, , X III for any permutation matrix, II E Ruxn Linear CKA and RSA achieve this by computing the dissimilarity between X.X! and X,X instead of the representations.
notion of distance should-2 at
nuisance transformations (e.g. permuta- tions or rotations in R²), ("), Let Φv IR x + IR M xn be a fixed, "preprocessing function" for each network and let G be set of nuisance transformations on Ra. Williams al. (2021) showed that:
d d(X, x,
min | ( (X,)T|| TEG
(2)
2

---

is metric over equivalent neural representations provided two technical conditions are met. The first is that G is a group of linear transformations. This means that: (a) the identity is in the set of nuisance transformations (I E 9). (b) every nuisance transformation is invertible by another nuisance transformation (if T E g then T-1 E 9). and (c) nuisance transformations are closed under composition (T,T2 transformation is an isometry, meaning that || - X, = ||x; XýIIF for every T E G. Several choices of satisfy these conditions including the permutation group, P, and the orthogonal group, O, which respectively correspond to the set of all permutations and rotations on R"
if T E 9 and T2 E 9). The second condition is that every nuisance
Equation (2) provides recipe to construct many notions of distance. To enumerate some ex- amples, we will assume for simplicity that Φ1 = = ΦK Φ and all networks have n neu- rons. Then, to obtain a metric that invariant to translations and permutations, we can set <(x) (I/n) - ) and G = Pnn). If we instead set G (n)) we recover the well- known Procrustes distance, which is invariant to rotations. Finally, if we choose 6(-) to whiten the covariance of X, we obtain notions of distance that are invariant to linear transformations and are closely related to CCA. Williams et al. (2021) provides further discussion and examples.
An attractive property of equation (2) is that it establishes metric space over deterministic represen- tations. That is, distances are symmetric (x X,) = d( 5. and satisfy the triangle inequality d(X,,, Xx) I d(X,, X,) + d(X, j. X&). Further, the distance is zero if and only if there exists a T E 9 such that o; = Φا X,, ,, These fundamental properties are needed to rigorously establish many statistical analyses (Cover& & Hart, 1967; Dasgupta & Long, 2005).
2.2 STOCHASTIC SHAPE METRICS
Let (Fi,..., } denote a collection of K stochastic networks. That is, each F& is a function that maps each input = E Z to a conditional probability distribution Fs(- =). How can equation (2) be generalized to measure representational distances in this case? In particular, the minimization in equation (2) is over a Euclidean "ground metric," and we would like to choose compatible metric over probability distributions. Concretely, let D(P, Q) be a chosen "ground metric" between two distributions P and Q. Let 6n and 6, denote Dirac masses at x. y E IR"" and consider the limit that P + 8n and Q - by. In this limit, we seek ground metric for which D(6m, 6y) is related to ||æ - yll. Many probability metrics and divergences fail to meet this criterion. For example, if y, then the Kullback-Leibler (KL) divergence approaches infinity and the total variation distance and Hellinger distance approach a constant that does not depend on ||= - y/|-
In this work, explored two ground metrics. First, the
p- Wasserstein distance (Villani, 2009):
(P,) = (inf [|| Y|P
where p 2 1, and the infimum is taken over all random variables (X, Y) whose marginal distribu- tions coincide with P and Q. Second, the energy distance (Székely & Rizzo, 2013):
Ep Q) (E [||X Y | - X' - E l - 1/2
(4)
O 0 < < and X, X' i.d. P and Y, Y Lid. Q. As desired, we have Wp, (82, 6,) = || yl for any p. and (ge. 5y) = x - y| Thus, when q = for example, the energy distance converges to the square root of Fuclidean distance in the deterministic limit. Interestingly, when g 2, the energy distance produces a deterministic metric on trial-averaged responses (see appendix F.1).
The Wasserstein and energy distances are intuitive generalizations of Euclidean distance. Both can be understood as being proportional to the amount of energy it costs to transport a pile of dirt (a probability density P) to different configuration (the other density Q). Wasserstein distance is based on the cost of the optimal transport plan, while energy distance is based on the the cost of a random (i.e. maximum entropy) transport plan (see Supp. Fig. 1, and Feydy et al. 2019).
Our main proposition shows that these two ground metrics can be used to generalize equation (2). Proposition 1 (Stochastic Shape Metries). Let Q be distribution on the inpul space. Let $1....,4K ØK be measurable functions mapping onto R.". and let Fee = F; 1-11 Let D2 denote the squared "ground metric," and let G be group of isometries with respect to D. Then, 1/2
d(F,, F2)
-min TEg
E $~QQ
D2
F.
2), (711))0)))
(5)
3

---

defines metric over equivalence classes, where F; is equivalent to F3 if only if there is T g such that F (- 2) and Fÿ 12) T-1 are equal for all z E supp(Q)-
Above, we use the notation P o o- to denote the pushforward measure-ice ic. the measure defined by the function composition, (A)- for a measurable set A, where P is a distribution and ø measurable function. A proof is provided in appendix C. Intuitively, T plays the same role as in equation (2). which is to remove nuisance transformations (e.g. rotations or permutations; see Fig. IE). The functions , ΦK also play the same role as "preprocessing functions," " imple- menting steps such as whitening, normalizing by isotropic scaling, or projecting data onto principal subspace. For example, to obtain a cranslaton-iivvaiian distance, we can subtract the grand mean response from each conditional distribution. That is, 0% (x) = x Eg a [ (ær |=) [z]]
2.3 PRACTICAL ESTIMATION OF STOCHASTIC SHAPE METRICS
Stochastic shape distances (eq. 5) are generally more difficult to estimate than deterministic distances (eq. 2). In the deterministic case, the minimization over T E G is often a well-studied problem, such as linear assignment (Burkard et al., 2012) or the orthogonal Procrustes problem (Gower & Dijksterhuis, 2004). In the stochastic case, the conditional distributions F- x) often do not even have a parametric form, and can only be accessed by drawing samples- e.g. by repeated forward passes in artificial network. Moreover, Wasserstein distances suffer well-known curse of dimensionality: in n-dimensional spaces, the plug-in estimator converges at a very slow rate proportional to s -1/n,, where is the number of samples (Niles- Weed & Rigollet, 2022).
Thus, to estimate shape distances with Wasserstein ground metrics, we assume that, (. 2). is well-approximated by a Gaussian for each z E Z. The 2-Wasserstein distance has a closed form expression in this case (Remark 2.31 in Peyré & Cuturi 2019 and Theorem 1 in Bhatia et al. 2019):
W2(N E,), E,)) (1m, My 13,l
min UEO(m)
|| -//""> 1/2
(6)
where (N, Σ) denotes a Gaussian density. It is important not to confuse the minimization over U E (n) in this equation with the minimization over nuisance transformations, T G, in the shape metric (eq. 5). These two minimizations arise for entirely different reasons, and the Wasser- stein distance is not invariant to rotations. Intuitively, we can estimate the Wasserstein- shape metric by minimizing over U E (n) and T E G in alternation (for full details, see appendix D.I).
Approximating (. 2) as Gaussian is common in neuroscience and deep learning (Kriegeskorte & Wei, 2021; Wu et al., 2019). In biological data, we often only have enough trials to estimate the first two moments of neural response, and one may loosely appeal to the principle of maximum entropy to justify this approximation (Uffink, 1995). In certain artificial networks the Gaussian assumption is satisfied exactly, such as in variational autoencoders (see sec. 3.3). Finally, even if the Gaussian assumption is violated, equation (6) can still be a reasonable ground metric that is only sensitive to the first two moments (mean and covariance) of neural responses (see appendix E.3).
The Gaussian assumption is also unnecessary if we use the energy distance (eq. 4) as the ground metric instead of Wasserstein distance. Plug-in estimates of this distance converge at much faster rate in high-dimensional spaces (Gretton et al., 2012; Sejdinovic et al., 2013). In this case, we propose two-stage estimation procedure using iteratively reweighted least squares (Kuhn, 1973), followed by a "metric repair" step (Brickell et al., 2008) which resolves small triangle inequality violations due distance estimation error (see appendix D.2 for full details).
We discuss computational complexity in Appendix 1.1.1 and provide user-friendly implementations of stochastic shape metrics at: gi thub. com/ahwil/ llia/netrep.
2.4 INTERPOLATING BETWEEN
AND COVARIANCE- -ENSSTIVVE METRICS
An appealing feature of the 2- 2-Wasserstein distance for Gaussian measures (eq. 6) is its decomposi- tion into two terms that respectively depend on the mean and covariance. We reasoned that it would be useful to isolate the relative contributions of these two terms. Thus, we considered the following generalization of the 2- Wasserstein distance parameterized by scalar, 0 S a S 2:
W (P,, P3) = (allm 1nl12 + (2 - a)) min UEO(())
|| Σ 1 U/ 1/2
(7)
4

---

A toy representations
B
randomly rotated
MDS embedding alpha
D
alpha 2
perrelation
2
scale
alpha
E
correlation
Toy Dataset. (A) 16 out of 99 "toy networks" with different correlation structure (horizon- tal axis) and covariance scale (vertical axis). Colors indicate distributions conditioned on different network inputs (as in Fig. IB). (B) Same as A, with random rotations applied in neural activation space. These rotated representations are used in subsequent panels. (C) 2D embedding of networks in stochastic shape space (a ground metric, G = O). Numbered points correspond to labeled representations in panel B. Colormap indicates ground truth covariance parameters. (D) Same as C, but with
Figure 2:
2 (covariance-insensttte)) (E) Same as C, but with 0 (mean-insensitive).
where P., Py are distributions with means μà J and covariances E,, E,. Appendix E we show that this defines a metric and, by extension, shape metric when plugged into equation (5).
We can a interpolate between a Euclidean metric on the mean responses and a metric on covariances known as the Bures metric (Bhatia et al., 2019). When a and the distributions are Gaussian, we recover the 2-Wasserstein distance. Thus, by sweeping a, we can utilize a spectrum of stochastic shape metrics ranging from distance that isolates differences in trial-average geometry (a = 2) to a distance that isolates differences noise covariance geometry (a 0). Distances along this spectrum can all be understood as generalizations of the usual "earth mover" interpretation of Wasserstein distance-the covariance-insensitee differences in the mean while the mean-insensitive metric (a 0) only penalizes transport due to differences in the orientation and scale of covariance. Simulation results in Supp. Figure 2 provide additional intuition for the behavior of these shape distances as is adjusted between 0 to 2.
metric (a = 2) only penalizes transport due to
3
RESULTS AND APPLICATIONS
3.1 TOY DATASET
We begin by building intuition on a synthetic dataset in 2 dimensions with M = inputs. Each response distribution was chosen to be Gaussian, and the mean responses were spaced linearly along the identity line. We independently varied the scale and correlation of the covariance, producing 2D space "toy networks." Figure 2A shows a sub-sampled 4 x grid of toy networks. To demonstrate that stochastic shape metrics are invariant nuisance transformations, random rotation to each network's activation space (Fig. 2B). The remaining panels show analyses for 99 randomly rotated toy networks spaced over a 11 grid (11 correlations and scales).
applied
Because the mean neural responses were constructed to be identical (up to rotation) across net- works, existing measures of representational dissimilarity (CKA, RSA, CCA, etc.) all fail to capture the underlying structure of this toy dataset (Supp. Fig. 3). In contrast, stochastic shape distances can elegantly recover the 2D space of networks we constructed. In particular, we computed the 99 99 pairwise distance matrix between all networks (2- WWssersttein ground metric and rotation invariance,
O) and then performing multi-dimensional scaling (MDS; Borg Groenen, 2005) to obtain 2D embedding. This reveals 2D grid of networks that maps onto our constructed arrangement (Fig. 2C). Again, since the toy networks have equivalent geometries on average, a deterministic metric obtained by setting a in eq. 7 (covariance-insensitee
metric) fails to recover this struc- ture (Fig. 2D). Setting a = 0 in eq. 7 inseesiissesitivee stochastic metric) also fails to recover a sensible 2D embedding (Fig. 2E), since covariance ellipses of opposite correlation by a 9 rotation. Thus, we are only able to fully distinguish the toy networks in Figure 2A by taking both the mean and covariance into account when computing shape distances.
be aligned
5

---

Example višp rnpresenttion for oriented gratings
Exampie ViSp repesertation natural
comelated mean and representations static gratings
Gross animal aitferences in VISp representations scenes
20
00
337
10.0
50"
102"
1507
0m0
sta60 gratings
stnon befineen covtrnacco
-20
PC 1
20
PC 1
13.07
natural soones
22.5
E
Differanco benween indvicual pair of arimals
means
Dovorianoee
maans
cova nanoss
PO2
p-0.57
20
10
20
ditan01 brtwoen consrtinons 10 5)
PC1
0
Figure 3: (A) In an example mouse VISp, neuronal responses form different means and covari- ances for six grating orientations. (B) In an example mouse VISp, neuronal responses form different means and covariances for six different natural scenes. (C) Covariance distances dominate dif- ferences between recording sessions for artificial grating stimuli, but not for natural scenes. (D) Averaged distance (across sessions) between mean responses for natural scenes is larger compared to for gratings. (E) Observations in (C) generally hold for individual pairs of recording sessions.
In Supp. Fig. 4 we show that using energy distance (eq. 4, = 1) as the ground metric produces similar result. Similar to Fig. 2C, MDS visualizations reveal the expected 2D manifold of toy networks. Indeed, these alternative distances correlate-but do not coincide exactly-with the dis- tances shown in Figure 2 that were computed with a Wasserstein ground metric (Supp. Fig, 4D).
3.2 BIOLOGICAL DATA
Quantifying representational similarity is common in visual neuroscience (Shi et al., 2019; Kriegeskorte et al., 2008b). To our knowledge, past work has only quantified similarity in the geom- etry of trial-averaged responses and has not explored how the population geometry of noise varies across animals or brain regions (e.g. how the scale and shape of the response covariance changes). We leveraged stochastic shape metrics perform a preliminary study on primary visual cortical recordings (VISp) from K 31 mice in the Allen Brain Observatory.' The results we present below suggest: (a) anioss-animal variability in covariance geometry comparable in magnitude to variability in trial-average geometry, (b) across-anima distances in covariance and trial-average geometry are not redundant statisties as they are only weakly correlated, and (c) the relative contri- butions of mean and covariance geometry to inter-animal shape distances are stimulus-dependent. Together, these results suggest that neural response distributions contain nontrivial geometric struc- ture in their higher-order moments, and that stochastic shape metrics can help dissect this structure.
We studied population responses (evoked spike counts, see appendix B.2) to two stimulus sets: a set of 6 static oriented grating stimuli and set of 119 natural scene images. Figure 3A shows neural responses from one animal to the oriented gratings within a principal component subspace (top), and the isolated mean and covariance geometries (bottom). Figure 3B similarly summarizes neu- ral responses to six different natural scenes. In both cases, the scale and orientation of covariance within the first two PCs varies across stimuli. Furthermore, the scale of trial-to-trial variance was comparable to across-condition variance in the response means. These observations can be made individually within each animal, but a stochastic shape metric (2-Wasserstein ground metric, and tation invariance G = O) enables us to quantify differences in covariance geometry across animals. We observed that the overall shape distance between two animals reflected mixture of differences in trial-average and covariance geometry. Specifically, by leveraging equation (7), we observed that mean-insensitive (a 0) and oovariance-insenstiiie (a = 2) distances between animals have similar magnitudes and are weakly correlated (Fig. 3C-E).
See appendix B.2 for full details. Data are available at: observatory wrw.ttto...briimmmp...m arrainmmpp.org/visuuu wtw..o//vvsaallclonngg//////////

---

A
a=0
a=2
0.3
factor
dip-
10
Clase
Shuffed
anneal
Recondarztion
1.0
Untrakned Traned
0.0
10 10
10
307
0.6
Disentange Gassificatoon
accuracy
0.0
beta to
beta hi
factor la
factor hi
dii to
dip hl
dip lo
cip-lthi
beta to lo
beta t:
annazled lo annealed hi
1.0
0.5-
0.--
N
Figure 4: (A) Dissimilarity matrices with varying a (top) and corresponding 2D embeddings (bot- tom) for 1800 VAFs trained dSpri tes. Six different VAE objectives (color hue) were used, each with six possible regularization strengths (tint) repeated over 50 random seeds. (B) Untrained net- works were farther apart in mean-insensitive distance (a = while trained networks were largely separated by covariance-insentn
distance (a 2). (C) kNN prediction of network's s random seed. (D) Predicting reconstruction loss using &NN regression. (E) kNN prediction accuracy of
VAE objective and regularization strength (hi vs lo) for metrics with different a (color scale). (F)
Regression performance predicting factor disentanglement scores (see Supp. B.3 for details).
Interestingly, the ratio of these distances was reversed across the two stimulus sets differrnces in covariance geometry across animals were larger relative to differences in average for oriented gratings, while the opposite was true for natural scenes (Fig. 3C-E). Later we will show an intrigu- ingly similar trend when comparing representations between trained and untrained deep networks.
3.3 VARIATIONAL AUTOENCODERS AND LATENT FACTOR DISENTANGLEMENT
Variational autoencoders (VAEs; Kingma Welling, 2019) are well-known class of deep gener- ative models that map inputs, E 2 (e.g. images), onto conditional latent distributions, F.. =), which are typically parameterized as Gaussian. Thus, for each high-dimensional input z,, the en- coder network produces distribution N N(,,, Σ,) in relatively low-dimensional latent space ("bot- tleneck layer"). Because of this, VAEs are a popular tool for unsupervised, nonlinear dimensionality reduction (Higgins et al., 2021; Seninge et al., 2021; Goffinet et al., 2021; Batty et al., 2019). How- ever, the vast majority of papers only visualize and analyze the means, {ur,... pMr ). and ignore the covariances, EE,,... ΣM), ). generated by these models. Stochastic shape metries enable to compare both the mean and covariance structure learned by different VAEs. Such comparisons can help us understand how modeling choices impact learned representations (Locatello et al., 2019) and how reproducible or identifiable learned representations are in practice (Khemakhem et al., 2020).
We studied collection of 1800 trained networks spanning six variants of the VAE framework six regularization strengths and 50 random seeds (Locatello et al., 2019). Networks were trained on synthetic image dataset called dSprites (Matthey al., 2017). which mark within the VAE disentanglement literature. Each image has set of ground truth latent factors which Locatello et al. (2019) used compute various disentanglement scores for all networks.
well-established bench-
We computed stochastic shape distances between over 1.6 million network pairs, demonstrating the scalability of our framework. We computed rotation-invariant distances (g = O) for the generalized Wasserstein ground metric (a 0, 0.5, 1, 1.5, in equation 7; Fig. 4A) and for energy distance (q 1 in equation 4; Supp. Fig. 7A, Supp. B.3.3). In all cases, different VAE variants visibly clustered together in different regions of the stochastic shape space (Fig. 4B, Supp. Fig. 7B). Interestingly, the covariance-in -innenceiiinnniiiveeeee (a = 2) shape distance tended to be larger than the mean-inss -rcan-insensitve (a = 0) shape distance (Fig. 4B), in agreement with the biological data natural images (Fig. 3C, bottom).
7

---

A
W 0,0
B
T=01 7=03 0.5
e
W R.a 0.06
W 8,0 0.1
Figure 5: (A) Left, Patch-Gaussian augmented images for different patch width, W, and train-time noise level, a. Right, MDS embedded activation vectors from Patch-Gaussian trained networks. Col- ors correspond to different images, points correspond to independent samples, T = 0.1. (B) Dis- tance matrices (left) and low-dimensional MDS embedding (right) of networks for different shape distances parameterized by and different levels of Gaussian corruption at test time,
Even more interestingly, this relationship was reversed in untrained VAEs (Fig. 4B), similar to the biological data on artificial gratings (Fig. 3B, top). We trained several hundred VAEs on MNIST and CIFAR-10 10 to confirm these results persisted across more complex datasets (Supp. Fig. 9: Supp. B.3). Overall, this suggests that the ratio of a and a shape distances may be a useful summary statistic of representational complexity. We leave detailed investigation of this to future work.
Since stochastic shape distances define proper metric spaces without triangle inequality violations, identify the k-nearest neighbors (&NN) of each network within this space, and these neighborhoods to perform nonparametric classification and regression (Cover & Hart, 1967). This simple approach sufficient predict most characteristics of a network, including its random seed (Fig. 4C), average training reconstruction loss (Fig. 4D), its variant of the VAE objective in- cluding regularization strength (Fig. 4E), and various disentanglement scores (Fig. 4F). Detailed pro- cedures for these analyses are provided in Appendix B.3. Notably, many of these predictions about network identity (Fig. 4E) were more accurate for the novel stochastic shape metrics < 2), compared to existing shape metrics (a 2, deterministic metric on the mean responses; Williams al. 2021). Similarly, many disentanglement score predictions (Fig. 4F) improved when consider- ing both covariance and means together (0 < 2).
The fact that we can often infer a network's random seed from its position stochastic shape space suggests that VAE features may have limited interpretability on this dataset. These limitations appear to apply both to the mean (a = 2) and the covariance (a = 0) representational geometries, as well as to intermediate interpolations. Future work that aims to assess the identifiability of VAE representations may find it useful to use stochastic shape metrics to perform similar analyses.
3.4 EFFECTS OF
DATA AUGMENTATION ON ARTIFICIAL DEEP NETWORKS
Despite the success of artificial neural networks on vision tasks, they are still susceptible to small input perturbations (Hendrycks & Dietterich, 2019). simple and popular approach induce robustness in deep networks is Patch-Gaussian augmentation (Lopes et al., 2019). which adds Gaus- sian noise drawn from (0,0 random image patches of width W during training (Fig. 5A, left column). At test time, network robustness is assessed with images with spatially uniform noise drawn from (00 r²). Importantly, the magnitude of noise at test time, т, may be distinct from noise magnitude during training, From Fig. 5A (right column), augmentation reccond-fortte rows) qualitatively leads to more robust hidden layer representations on noisy data compared to networks trained without it (first row). While tion is empirically successful (for quantitative details, see Lopes et al., 2019), how W and change hidden layer representations to confer robustness remains poorly understood.
that using
augmenta-
To investigate, we trained a collection of 339 ResNet-18 networks (He et al., 2016) on CIFAR- 10 (Krizhevsky, 2009), sweeping 16 values of W, 7 values of a, and 3 random seeds (see
8

---

appendix B.4 for details). While the architecture is deterministic, can consider to be stochastic mapping by absorbing the random Gaussian perturbation parametrriied by T r-into the first layer of the network and allowing the stochasticity to percolate through the network. Representations from a fully connected layer following the final average pooling layer were used for this analysis. We computed stochastic shape distances across all 57,291 pairs of networks across six values of and three shape metrics parameterized by E {0, 2 defining the ground metric in equation (7).
Sweeping across and revealed a rich set of relationships across these networks (Fig. 5B and Supp. Fig. 11). While a complete investigation is beyond the scope of this paper, several points are worthy of mention. First, the mean-insensitive (a = 0. top row) and (a = 2. bottom row) metrics produce clearly distinct MDS embeddings. Thus, the new notions of stochastic representational geometry developed in this paper (corresponding to a = 0) provide new information to existing distance measures (corresponding to a 2). Second, the arrangement of networks in stochastic shape space reflects both W and Ø, sometimes in 2D grid layout that maps nicely onto the hyperparameter sweep (e.g. hyperparameters but different random seeds tend to be close together in shape space. Third, the test-time noise, т, also intricately impacts the structure revealed by all metries. Finally, embeddings based on 2- Wasserstein metric (a = 1) qualitatively resemble embeddings based on the covariance- insensitive metric (a 2) rather than the mean-insensitive metric (a = 0).
0 and 0.01). Networks with the same
4 DISCUSSION AND RELATION TO PRIOR WORK
We have proposed stochastic shape metrics as a novel framework to quantify representational dis- similarity across networks that respond probabilistically to fixed inputs. Very few prior works have investigated this issue. To our knowledge, methods within the deep learning literature like CKA (Ko- mblith et al., 2019) have been exclusively applied to deterministic networks. Of course, the broader concept of measuring distances between probability distributions appears frequently. For example, to quantify distance between two distributions over natural images, Fréchet inception distance (FID; Heusel et al., 2017) computes the 2- 2-Wasserstein distance within hidden layer representation space. While FID utilizes similar concepts our work, it addresses very different problem-i.e., how to compare two stimulus sets within the deterministic feature space of single neural network, rather than how to compare the feature spaces of two stochastic networks over a single stimulus set.
A select number of reports in neuroscience, particularly within the fMRI literature, have addressed how measurement noise impacts RSA (an approach very similar to CKA). Diedrichsen et al. (2020) discuss how measurement noise induces positive bias in RSA distances, and propose approaches to correct for this bias. Similarly. Cai et (2016) propose a Bayesian approach to RSA that performs well in low signal-to-noise regimes. These papers essentially aim to develop methods that are robust to noise, while we were motivated to directly quantify differences noise scale and geometric structure across networks. It is also common to use Mahalanobis distances weighted by inverse noise covariance to compute intra-network representation distances (Walther et al., 2016). This procedure does not appear to quantify differences in noise structure between networks, which we verified on a simple "toy dataset" (compare Fig. 2C to Supp. Fig. 3). Furthermore, the Mahalanobis variant of RSA typically only accounts for a single, stimulus-indepe indeeedddind noise covariance. In contrast, stochastic shape metrics account for noise statistics that change across stimuli.
Overall, our work meaningfully broadens the toolbox of representational geometry to quantify stochastic neural responses. The strengths and limitations of our work are similar to other ap- proaches within this toolbox. A limitation in neurobiological recordings is that we only observe a subset of the total neurons in each network. Shi et al. (2019) document the effect of subsampling neurons on representational geometry. Intuitively, when the number of recorded neurons is large relative to the representational complexity, geometric features are not badly distorted (Kriegeskorte & Diedrichsen, 2016; Trautmann et al., 2019). We show that our results are not qualitatively affected by subsampling neurons in Supp. Fig. 6. Another limitation is that representational geometry does not directly shed light on the algorithmic principles of neural computation (Maheswaranathar et al., 2019). Despite these challenges, representational dissimilarity measures are one of the few quanti- tative tools available to compare activations across large collections of complex, blackbbox models, and will be mainstay of artificial and biological network analysis for the foreseeable future.
9

---

REFERENCES
F Abbott and P Dayan. The effect of correlated variability on the accuracy of a population code. Neural Comput. 11(1):11-11,, January 1999.
Guozhong An. The Effects of Adding Noise During Backpropagation Training Generalization Performance. Neural Computation, 8(3):643-674, 04 1996. ISSN 0899-7667. doi: 10.1 10.1162/neco. 1996.8.3.643. URL httphh://doi..rrg//1.... 1162//eco.. 1996.
Eleanor Batty, Matthew Whiteway, Shreya Saxena, Dan Biderman, Taiga Abe, Simon Musall, Winthrop Gillis, Jeffrey Markowitz, Anne Churchland, John P Cunningham, Sandeep R Datta, Scott Linderman, and Liam Paninski. Behavenet: nonlinear embedding and bayesian neural de- coding of behavioral videos. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https: /tttt::://preee http//occeenngs....o neurips cc/paper/2819/ffiee 210443dF8965228837727244447 papee pdf.
Rajendra Bhatia, Tanvi Jain, Yongdo Lim. the buses-wasserstiin distance between pos- itive definite matrices. Expositiones Mathematicae, 37(2):165-191, 2019. ISSN 0723-0869. doi: https//dboi.rgg/10.100.... cxw..6....mmma........ 000.0.... URL httt::://ww..scieneerreei sccemme//orrcceee//////773338813
Ingwer Borg and Patrick JF Groenen. Modern multidimensional scaling: Theory and applicationn. Springer Science & Business Media, 2005.
Justin Brickell, Inderjit Dhillon, Suvrit Sra, and Joel A Tropp. The metric nearness problem. SIAM
J. Matrix Anal. Appl. 30(1):375-333 -366,999,, January 2008
Rainer Burkard, Mauro Dell' Amico, and Silvano Martello. Assignment Problems. Society for Industrial and Applied Mathematics, 2012. doi: 10.1137/1 9917//997781619722338..8 URL https:// epubs. siam.
1137/1 1137/1.9781111333
Mingho Cai, Nicolas W Schuck, Jonathan W Pillow, and Yael Niv. A bayesian method for ducing bias in neural representational similarity analysis. In D. Lee, M. Sugiyama, U. Luxburg, Guyon, and Garnett (eds.), Advances Neural Information Processing Systems, volume 29, Curran Associates, Inc., 2016. URL https::/ptocceeeeeeeeee wwwppporeeennsseeeeeee. neurips /c/pppper/216/fflee b06f50d1 hotpa//0000000001111a11cccceccea
.df..
YueY Chung and L F Abbott. Neural population geometry: An approach for understanding biological and artificial neural networks. Curr: Opin. Neurobiol., 70:137-144, October 2021.
T. Cover and P. Hart. Nearest neighbor pattern classification. IEEE Transactions on Information Theory, 13(1):21-27, 1967. doi: 10. 1999/TTT.197710539666
Joel Dapello, Jenelle Feather, Hang Le, Tiago Marques, David Cox, Josh McDermott, James J DiCarlo, and Sueyeon Chung. Neural population geometry reveals the role of stochasticity in robust perception. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 15595- 15607. Curran Associates, Inc., 2021. URL -Paper.
www.diinnnnnnnuiisseeeeeeeee www.npps.ccppppee/e....... cc/paper/202V/
Sanjoy Dasgupta and Philip M Long. Performance guarantees for hierarchical clustering, Journal of Computer and System Sciences, 70(4):555-569 2005.
Jörn Diedrichsen, Eva Berlot, Maricke Mur, Heiko H. Schütt, Mahdiyar Shahbazi, and Nikolaus Kriegeskorte. Comparing representational geometries using whitened unbiased-diisance-mme similarity, 2020.
Ian Dryden and Kanti V Mardia. Multivariate shape analysis. Sankhyā: The Indian Journal of Statistics, Series A, pp. 480- 1993.
Kshitij Dwivedi and Gemma Roig. Representation similarity analysis for efficient task taxonomy & transfer learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2019.
10

---

Jean Feydy, Thibault Séjourné. François-Xavia Xavier Vialard, Shun-Ichi Amari, Alain Trouve, and Gabriel Peyré. Interpolating between optimal transport and MMD using sinkhorn divergences. volume 89 of Proceedings of Machine Learning Research, pp. 2681 16900. PMLR, 2019.
Jack Goffinet, Samuel Brudner, Richard Mooney, and John Pearson. Low-dimensional learned feature spaces quantify individual and group differences in vocal repertoires. Elife, 10, May 2021.
Robbe L T Goris, Anthony Movshon, and Eero Simoncelli. Partitioning neuronal variability. Nat. Neurosci., 17(6):858 865, June 2014.
C. Gower and Garmt B. Dijksterhuis. Procrustes problems. Oxford University Press, Oxford New York, 2004. ISBN 9780198510581.
Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola. A kemel two-sample test. Journal of Machine Learning Research, 13(25):723-77,, 2012. URL http://jmmr. wwwplrr.orgppppppeeeeeeee pttpp//p/t/////eett//e// 2a. html.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Dan Hendrycks and Thomas Dictterich. Benchmarking neural network robustness to common cor- ruptions and perturbations. In International Conference on Learning Representations, 2019. URL https://openreviens met//foruumiictttttiiCMM
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochre- iter. Gans trained by two time-scale update rule converge to a local nash equilibrium. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Cur- ran Associates, Inc., 2017. URL https: http://oooeceiieeeeeeeeeee www.ig.....i...... cc/paaprr/211//fi22 415447771000fff55887766000888880 .ad
Irina Higgins, Le Chang, Victoria Langston, Demis Hassabis, Christopher Summerfield, Doris Tsao, and Matthew Botvinick. Unsupervised deep learning identifies semantic disentanglement in single inferotemporal face patch neurons. Nat. Commun., 12(1):6456, November 2021.
Norman R Howes. Modern Analysis and Topology. Springer, New York, NY, 1995 edition, June 1995.
Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoen- coders and nonlinear ica: A unifying framework. In Silvia Chiappa and Roberto Calandra (eds.), Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pp. 2207-2217, PMLR, 26-28 Aug 2020. URL https: httpss////.ppceeess/.eeeeee mlr pwwss/eeeslv /108// ttt//khhmaae220..... html.
Diederik P Kingma and Max Welling. An introduction variational autoencoders. Foundations
and Trends in Machine Learning, 12(4):307-392, 2019.
Simon Komblith, Mohammad Norouzi. Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. volume 97 of Proceedings of Machine Learning Research, pp- Long Beach, California, USA, 2019, PMLR.
Nikolaus Kriegeskorte and Jöm Diedrichsen. Inferring brain- computational mechanisms with mod- els of activity measurements. Philosophical Transactions of the Royal Society B: Biological Sciences, 371((1055) 2016. doi: (t09///sb..016.0777
Nikolaus Kriegeskorte and Pamela K Douglas. Interpreting encoding and decoding models. Curr.
Opin. N 55:167-179, April 2019.
Nikolaus Kriegeskorte and Xue-Xin Wei. Neural tuning and representational geometry. Nature Reviews Neuroscience. 2021. 2021, ISSN 1471-0048, doi: 10.1039/4441533-0110000001 URL https:

---

Nikolaus Kriegeskorte, Maricke Mur, and Peter Bandettini. Representational similarity analysis connecting the branches of systems neuroscience. Frontiers in Systems Neuroscience, 2:4, 2008a. ISSN 1662-5137. doi: 01.338//neuro.06.004...200 URL https://ww..footierrssi froniiernniii... org/article/10 3389/neuro.. 06. 1.6.000.22 2008.
Nikolaus Kriegeskorte, Maricke Mur, Douglas Ruff, Roozbeh Kiani, Jerzy Bodurka, Hossein Esteky, Keiji Tanaka, and Peter A Bandettini. Matching categorical object representations in inferior temporal cortex of man and monkey. Neuron, 60(6): 12:0::14211111111 December 2008b.
Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009.
Harold W Kuhn. A note on fermat's problem. Math. Program., 4(1):98-107, December 1973.
Kenneth Lange. MM Optimization Algorithms. Society for Industrial and Applied Mathematics, Philadelphia, PA, 2016. doi: 11113//1971111197440000 URL https: siam. org/doi/ abs/10. 1137/1 9781611974409.
Olivier Ledoit and Michael Wolf. A well-conditioned estimator for large-dimensional covariance matrices. Journal of Multivariate Analysis, 88:365-411 2004.
Francesco Locatello, Gabriele Abbati, Thomas Rainforth, Stefan Bauer, Bernhard Schölkopf, and Olivier Bachem. On the fairness of disentangled representations. In H. Wallach, H. Larochelle, A. Beygelzimer, F. ('Aché-Buc, E. Fox, and Garnett (eds.), Advances in Neurul Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https: http://edoocc......///o///////ep .ddf.
Raphael Gontijo Lopes, Dong Yin, Ben Poole, Justin Gilmer, and Ekin D. Cubuk. Improving bustness without sacrificing accuracy with patch gaussian augmentation, 2019.
Niru Maheswaranathan, Alex Williams, Matthew Golub, Surya Ganguli, and David Sussillo. Uni- versality and individuality in neural dynamics across large populations of recurrent networks. In H. Wallach, H. Larochelle, Beygelzimer, d'Alché- Buc, E. Fox, and Garnett (eds.), Ad- vances in Neural Information Processing Systems 32, pp. 15629-15641. Curran Associates, Inc., 2019.
Loic Matthey, Irina Higgins, Demis Hassabis, and Alexander Lerchner. dsprites: Disentanglement testing sprites dataset.
2017.
Jonathan Niles-Weed and Philippe Rigollet. Estimation of wasserstein distances in the spiked trans- port model. Bernoulli, to appear, 2022.
Gabriel Peyré and Marco Cuturi. Computational optimal transport: With applications to data sci- ence. Foundations and Trends@ in Machine Learning, 11(5-6):355- 607, 2019.
Maithra Raghu, Justin Gilmer, Jason Yosinski, and Jascha Sohl- Sohl-Dickstein. Sveca: Singular vector Guyon, U. V.
canonical correlation analysis for deep learning dynamics and interpretability. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems 30, pp. 6076 6085. Curran Associates, Inc., 2017.
Oleg Rumyantsev, Jérôme A Lecoq, Oscar Hernandez, Yanping Zhang, Joan Savall, Radostaw Chrapkiewicz, Jane Li, Hongkui Zeng, Surya Ganguli, and Mark J Schnitzer. Fundamental bounds the fidelity of sensory cortical coding. Nature, 580(7801):100-100 April 2020.
Dino Sejdinovic, Bharath Sriperumbudur, Arthur Gretton, and Kenji Fukumizu. Equivalence of distance-base and rkhs-based statisties in hypothesis testing. The Annals of Statistics, 41(5): 2263-2291, 2022/08/01/ 2013. URL http: //www//www///// tion date: October 2013.
http://wwww////t//t//t///o///ttt jstor. org/ htgg/saable/2356555 Full publica-
Lucas Seninge, Ioannis Anastopoulos, Hongxu Ding, and Joshua Stuart. VEGA is interpretable generative model for inferring biological network activity in single-cell transcriptomics. Nat. Commun., 12(1):5688 September 2021.
12

---

M N Shadlen, K H Britten, W T Newsome, and J A Movshon. A computational analysis of the relationship between neuronal and behavioral responses to visual motion. J. Neurosci. 16(4): 1486-1510, February 1996.
Mahdiyar Shahbazi, Ali Shirali, Hamid Aghajan, and Hamed Nili. Using distance on the riemannian manifold to compare representations in brain and in models. Neurolmage, 239:118271 2021. ISSN 1053-8119 19. doi: https://doi.org/10..0000 rcietcedirect.. wtt.sccccc/cccccccc wttp 2021. 2021.118271 URL https: /www.
9200444..
Jianghong Shi, Eric
and Michael Buice. Comparison against task driven artificial neu- ral networks reveals functional properties in mouse visual cortex. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buu E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 5764-5774, Curran Associates, Inc., 2019.
Jocelyn Sietsma and Robert J.F. Dow. Creating artificial neural networks that generalize. Neural Net- works, 4(1):67-79, 1991. ISSN 0893-6080. doi: https://doi.org/111100666000 99000000000000000332222 J90033-2. URL https: //wl::::://w..... wwt.://ciinncdrrre.......m. ctm/ http://c.....cccceeeearre//////a
Anuj Srivastava and Eric P Klassen. Functional and shape data analysis. Springer Series in Statis-
tics. Springer, New York, NY, edition, October 2016.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: simple way to prevent neural networks from overfitting. Journal of Machine Learn- ing Research, 15(56):929-11558 2014. URL http: wwtggpaprs//v5///// srivastvva1a html.
B. Stellato, G. Banjac, P. Goulart, A. Bemporad, and S. Boyd. OSQP: an operator splitting solver for quadratic programs. Mathematical Programming Compulation, 12(4):637-672, 2020. doi: 01007//1223200200011777
00./ URL https: 10t0////..0//11133..3 1000///1122220222-00111333
Gâbor J Székely and Maria L Rizzo. Energy statistics: A class of statistics based distances. J. Stat. Plan. Inference, 143(6):1249-1277) 1272, August 2013.
Gábor J. Székely and Maria L. Rizzo. The energy of data. Annual Review of Statistics and Its Application, 4(1):447-479, 2017. doi: 10. 14 16-054026.
Jun Tong, Rui Hu, Jiangtao Xi, Zhitao Xiao, Qinghua Guo, and Yanguang Yu. Linear shrinkage estimation of covariance matrices using low-complexity cross-validation. Signal Processing, 148: 223-233, 2018
Eric M. Trautmann, Sergey D. Stavisky, Subhaneil Lahiri, Katherine C. Ames, Matthew T. Kauf- man, Daniel O'Shea, Saurabh Vyas, Xulu Sun, Stephen I. Ryu, Surya Ganguli, and Kr- ishna V. Shenoy, Accurate estimation of neural population dynamics without spike sorting, Neu- ron, 103(2):292-308 e4, 2022/11/15 2019. doi: 10. (000666.eeroon2211..5.000 URL https: //doi. org/10 10. 1016/j neuron. 2019.00 05.003.
Jos Uffink. Can the maximum entropy principle be explained consistency require- ment? Studies in History and Philosophy of Science Part B: Studies in History and Phi- losophy of Modern Physics, 26(3):223 261, 1995. ISSN 1355-2198. doi: https:/doi.oo org//
(0.10066/355-2118850000155 URL https: bii/135522 9895000151.
Cédric Villani. Optimal Transport. Springer Berlin Heidelberg. 2009.
Alexander Walther, Hamed Nili, Naveed Ejaz, Arjen Alink, Nikolaus Kriegeskorte, and Jörn Diedrichsen. Reliability of dissimilarity measures for multi-voxel pattern analysis. Neuroimage, 137:188-200, August 2016.
Alex H. Williams, Erin Kunz, Simon Korblithh, and Scott W. Linderman. Generalized shape metrics on neural representations. Advances Neural Information Processing Systems, volume 34, 2021.
Andrew Gordon Wilson. The case for bayesian deep learning, 2020
13

---

Anqi Wu, Sebastian Nowozin, Edward Meeds, Richard E. Turner, Jose Miguel Hernandez-Lobato, and Alexander L. Gaunt. Deterministic variational inference for robust bayesian neural networks. In International Conference on Learning Representations, 2019. URL https://openreeiie. for
14

---

APPENDICES
A
SUPPLEMENTARY FIGURES
A
D
C
W
align
B
E
Supplementary Figure 1: Proposed method intuition using distances based on either W2 or E ground metrics. (A) and (B) Two example stochastic network representations five stimuli (colors). (C) The optimal alignment of the representations over nuisance transformations (e.g. rotations, O). (D) Intuitively, the 2-Wasserstein distance (W2) the minimum cost of turning one density (pile of dirt) to another (Villani, 2009; Peyré & Cuturi, 2019). Here highlight the distance between the two green densities to reduce clutter. (E) Energy distance is based on the maximum- entropy transport map between the two distributions (Feydy et al., 2019).
15

---

A
representation
representation
mean alignment
covariance alignment
reprosomatooal dssance
3.0
0.5
1.5
1.5
20
noozured reprosentational distanoes
lower bound reprecetaatooal distanoss
D
Representation 2 is rotated to different extent for different a values
0.0
0,22
0.44
0.67
0.89
1.11
1.33
1.56
1.78
2.0
covariance alignment
alignment
Supplementary Figure 2: Simulated example showing how varying the a parameter in W2 induces different rotational alignments between neural representations (g O). (A) Two stimulated stochas- tic representations for three stimulus inputs. Colors represent different input conditions (M 3), hollow points represent sampled representations from the first network and filled points represent sampled representations from the second network. The example constructed so that no rotation simultaneously align both the means and covariances.
(B) If the stochastic metric only takes means into account (a = 2), after rotating one of the representations by 90%, two sets of represen- tational means completely overlap, and the distance becomes O. If the stochastic metric only takes covariances into account (a = 0), the optimal alignment between the two sets of covariances is either 09 180", and after this rotation, distance between representations again is O. (C) When both and a = 2, distance between the two representations is O, so the lower bound for the distance for a in the range between 0 and 2 also We computed the stochastic metric within this range of a, and the final distance is generally above the lower bound. (D) Optimal rotation between the two representations at different values of
RSA (Euclidean, Spearman]
RSA (Eaciidaa, Pearson]
RSA Maahlannttis, Spearman)
RSA (Mahalanobis Naason))
CKA (Linear Kemel)
Supplementary Figure Associated with Figure 2 from the main text. Embeddings of "toy dataset" networks (see Fig. 2A-B) visualized by multi-dimensional scaling of existing dissimilarity measures. Each point represents network, the color scheme the same as Fig. 2C. All methods fail to recover reasonable embedding which captures representational differences (compare with stochastic shape metric embeddings in Fig. 2C and Supp. Fig. 4C). Starting from the left, the first two plots representational similarity analysis (RSA; Kriegeskorte al. 2008a) with two forms of correlation distance (Spearman and Pearson) applied to Euclidean representational similarity ma- trices. The two plots use Mahalanobis distance re-weighted by the noise covariance (Walther et al., 2016) rather than Euclidean distance. The final plot shows an embedding by centered kernel alignment with linear filter (Kornblith al., 2019).
16

---

A
B
C
D
raw distance estimates (energy dist, qF1)
distances after L2 metric repair
MDS repaired distances
@Wasserstei dist)
1.0
shape distance
6.
63
shape
(energy dist, q=1)
20
Supplementary Figure 4: Associated with Figure from the main text. Stochastic shape metries with energy distance also recover the "ground truth" structure of synthetic "toy data". (A) Matrix of estimated pairwise distances computed with & ground metric on the synthetic data shown in Fig. 2A. (B) Matrix of pairwise distances after quadratic metric repair (see appendix F.3) was performed to correct for minor triangle inequality violations. (C) Multidimensional scaling embedding of the distance matrix in panel B into 2D Euclidean space. Compare with Fig. 2C. (D) Linear correlation between stochastic shape distances with energy distance ground metric (i.e. off-diagonal entries of panel B) and 2- 2-Wasserstein ground metric (i.e. off-diagonal entries of Fig. 2B). Red dashed line denotes the best linear model according to least-squares criterion.
Covariance distance dominates representation- al difference for drifted gratings
distance between means
20
60
distance between covariances
Supplementary Figure 5: Associated with Figure 3 from the main text. Drifted gratings (4 drifting directions, 75 repeats each) were presented different set of experimental sessions. Like (artifi- cial) static gratings, representational distances across sessions for drifted gratings dominated by covariance differences.
17

---

Sampling without replacement
static gratings
natural scenes
100reamm
15-
35-
dilance
10
0.0
2.0
1.0
2.0
cislanoe beween
dislance beween means
dialance beween covariances
dislance beweon means
covarianoes
Sampling with replacement
static gratings
natural scenes
15
15
distance
10
10
5-
5-
0.0 distance beween
0.0 distance
a
distance
baween
beween
a
distance
beween means
Dovarisnces
means
Dovanamines
Supplementary Figure 6: We use a simulation to explore how the size of the neural population recording affects conclusions about representational distances. In particular, does the ratio of mean-insensitive to covariance covariance-inseesittvee across-animal distances (a 0 vs. a = 2) change when sub-sample neurons? For this simulation, recorded from their respective VISps. We randomly sample a subset of n neurons among these recorded neurons (n 10, 30, 50, 70, 90, 100), and computed representational distance (a 0, 1) using only the subset. For panel A and B, we sampled the neurons without replacement, and for panel C and we sample with replacement (bootstrapping). We observe that for all tested a, representational distance increases with the number of neurons within the subset. This is expected because distances will generically increase with the dimension (e.g. the Euclidean distance between two random vectors in high dimensions will tend to be large, relative to low dimensions). However, ratio a = 0 and a shape distances is preserved when subsampling neurons (all lines are trending upward as function of a).. Error bars capture how the computed distances vary across 15 random draws of n neurons from the recorded population.
chose two mice that have 102 and 110 neurons
18

---

A
300
803
B
Normalized Vg qrep. 0e
Normalized 1R 472.00 2.00
0.8
0.8
0.8
0.8
0.6
0.5
Normalized vRg
0.6
0.6
0.6
0.4
Normaltog
0.2
0.0
0.0
0.0
0.0
1.00
1.00
Normalized Energy
Normalized Energy
Normalized Energy
Normallzzed Fnergy
Normalized Energy
D
Clossifiation
accuracy
1.0
000900
0.5
Energy
eeato
0-2.0
beta
factorr
factor
dip-i
dip-illo
dipei hi
betastz
annealed annealed
0=1.5
E
Disentangle
1.0
1-1.0
Score
0.
0-0.
0.0
M
C
Supplementary Figure 7: Associated with VAE analyses (Figure 4) in the main text. (A) Dissimi- larity matrices measured for 1800 Seprites-trriing VAEs from Locatello et al. (2019) using gener- alized interpolated 2- Wasserstein (Equation 7) with varying (first five columns), and using energy distance (Equation 4) with 64 samples for each unique input (right-most column). Row/column ordering of each matrix is the same in Figure 4. (B) 21) embeddings corresponding distance matrices in (A). Colors are the same as in Figure 4. We aligned to the left-most panel using Pro- crustes analysis, allowing for scaling and rotations/reflections. (C) Re-scaling W distances and energy distances such that they lie between [o, 1] reveals that the distribution energy distances agrees best with W2 10= (middle column). (D) Predicting objective and regularization strength us- ing distance matrices in (A). (E) Predicting disentanglement scores using distance matrices in (A). See Supp. B.3 for more details.
== 0.00
0.50
2 1=00
a. 1.50
Q 2.00
4.0
Distortion (ratio)
2.5
2.0
1.5
10
20
10
20
10
Embedding dim
Supplementary Figure 8: Associated with VAE analyses (Figure 4) in the main text. Distortion induced by multidimensional scaling of 1800 x 1800 dissimilarity matrices with varying embedding dimensionality. Different shading represents (10th-90th) and (25th-75th) percentiles. See Supp. B.3 for details.
19

---

MNIST
CIFAR-10
10P
10:1
10-1
Untrained
Untrained Trained
Trained
10
10-2
10-2
10"
10
10-1
10P
10 a
Supplementary Figure 9: Associated with VAE analyses (Figure 4) in the main text. We initialized 350 3-VAEs and trained them on MNIST (left) and CIFAR-10 (right) with different values of B in the loss function. Training led to inter-network distances being dominated by (a = 2) dissimilarity, in agreement with Figure 4B of the main text. See Supp. B.3 for training details.
20

---

0.01
0.05
0.1
0.2
0.3
0.5
2
Supplementary Figure 10: Distance matrices for different values of a and т.
0.01
0.05
0.1
0.2
0.3
0.5
2
Supplementary Figure 11: Two dimensional embedding of the distance matrices in Fig. 10 for different values of a (row) and (column).
21

---

B SUPPLEMENTAL METHODS
B.I CODE AND REPRODUCIBILITY
We have attached zip file containing self-contained Python script with our implementation of the interpolated generalized Wasserstein distance (equation 7) and energy distance (equation 4) in file named met metric. setric.p.... py. Additional analysis and source code will be located at gi thub. www.ccom//aaa.cae ahwillia/ netrep.
B.2 ALLEN BRAIN OBSERVATORY DATA
B.2.1 DATA PRE-PROCESSING.
In each recording session, gratings (6 orientations) and natural scenes (119 images) were presented to one mouse, and between 19 to 110 neurons in VISp were recorded by extracellular microelectrode arrays (Neuropixels Visual Coding dataset). Each neuron's response the sum of action potentials (spikes) emitted within a 250 millisecond time window (the duration of the stimulus presentation in this dataset). To compare how similar single set of stimuli were represented across sessions, we Gaussaan-aprroximated the data recorded for each stimulus. Then for each stimulus class (gratings and scenes), compared between sets of Gaussians (one per stimulus).
image was measured as
Our metrics compared between sets of Gaussians that have the same dimensionality, so we per- formed PCA to equalize dimensionality across all sessions. We concatenated data recorded for different images within single stimulus class, and extracted the first 19 principal components in replacement of the total number of recorded neurons for further analysis. On average, the extracted 19 principal components explained 83% of the data variance in response to gratings, and 76% the variance to natural scenes.
B.2.2 ESTIMATING RESPONSE MEAN AND COVARIANCE.
To compare between two neural representations, our stochastic metrics take two sets of Gaussian means and covariances as inputs, where each of which is estimated from the principal components (PCs) extracted from the data.
In each session, a stimulus (either grating or a scene) was presented over 50 repeats. The number of repeats is large compared to conventional neuroscience experiments, but it is still small compared to the total number of recorded neurons (e.g. 110 neurons in one session), or the total number PCs, which introduces challenges to covariance estimation. For the mean of each stimulus representation, we used the sample mean from the PCs. When number of samples is relatively small, sample covariance has known bias: it tends over-estimate large eigenvalues, and under-estimate small eigenvalues the population covariance. One standard and effective fix in the literature is to shrinkage estimator (S*) linear interpolation between identity matrix (I) and the sample covariance (S) (e.g. Ledoit & Wolf (2004); Tong et al. (2018)):
S = 7 + 71-
(8)
This interpolation reduces the eigenvalue bias by balancing between eigenvalues of the sample co- variance (overly skewed eigenvalue spectrum), and that of the identity matrix (flat eigenvalue spec- trum). 7 of the shrinkage covariance estimator was chosen using cooss-validation. To obtain the cross-validation training set, randomly sample half of the epochs from each data trial, and for test set, we used the remaining half.
B.3 VARIATIONAL. AUTOENCODERS AND LATENT FACTOR DISENTANGLEMENNT (SUPPLEMENT TO SEC. 3.3)
B.3.1 VAE OBJECTIVES AND ARCHITECTUREE USED IN THIS STUDY
Because conventional VAE encoders output a latent Gaussian conditional mean and covariance, this makes them an ideal framework with which to apply stochastic shape metrics. In particular, the interpolated Wasserstein distance (equation 7) is exact in this case. We used set of 1800 VAEs
22

---

trained on dSprites from the extensive study by Locatello al. (2019). These include six variants of the VAE objective (3, Factor, 3-TC, DIP-I, DIP-II, Annealed), each with six different levels of regularization strength and 50 repetitions at different random seeds. The dimensionality of the latent representation, from which we obtained activations used in this study, was 10D. We refer the reader to their supplemental document for more details about each architecture and training scheme. The authors provided metadata associated with each network such as training hyperparameters as well as factor disentanglement scores (see below).
In addition to the VAES trained on dSprites, we trained 350 3 VAES on MNIST and CIFAR-IO. To remain consistent with the study by Locatello al. (2019), we trained 3-VAEs at 6-8 different levels of regularization strength (1 < 3 < 16) and 50 random initialization seeds. We used the standard VAE symmetric encoder-decoder architecture with layers (I=
= 3 for MNIST and for CIFAR-10) 10), each with 64 4 x 4 convolutional filters with stride 2, followed by fully-connected layer with 256 hidden units and ReLU activations. The latent representations of these networks were diagonal Gaussians, and were 10D (MNIST) or 50D (CIFAR-10). The final 2D convolution- transpose layer of the decoder used a sigmoid nolinearity to ensure outputs were between [0, 1].
We zero-padded the height and width of MNIST images from 28 x 28 + 32 x 32. Model training used batch sizes of 64 images, and up to 1000 training epochs. We used the Adam optimizer with IE-4 learning rate and model checkpoints at each epoch. Models used in this study were from checkpoints corresponding to the lowest validation loss during training. Latent activations used for shape metric analysis in this study were obtained using held-out test set of 3500 images.
B.3.2 W2 -
vs. W3 W BEFORE AND AFTER TRAINING
Fig. 4B of the main text shows that, prior to training, VAEs are primarily separated by mean- insensitive distance (W2,a covariance-insensitiv distance (a = 2). We sought to confirm whether this effect persisted across different datasets using VAES trained on MNIST and CIFAR-10 (described above). SuppFig. shows that pairwise network W2 distances before and after training indeed exhibit this effect on these more complex datasets. We reproduced these effects using both default PyTorch weight ini- tialization and Kaiming weight initialization.
0, equation 7), whereas after training they are separated by
B.3.3 COMPUTING ENERGY DISTANCE BETWEEN TRAINED VAES
In addition measuring interpolated Wasserstein distances (equation 7), we also repeated our anal- yses using energy distance (equation 4). Rather than requiring computing means and covariances, this method operates directly on samples. Since VAF latents are parameterized as Gaussian, we generated data by randomly sampling from the Gaussian defined by the model's conditional mean and covariance for a given input. We sampled 64 samples for 2048 images and computed pair- wise energy distances between all 1800 networks in the (Locatello et al., 2019) dSprite tes dataset (SuppFig. 7). Interestingly, the energy dissimilarity matrix was qualitatively different than all of the W2 dissimilarity matrices (SuppFig. 7A). The geometry of the embedded points was accordingly different than embeddings derived from WE distances (SuppFig. 7B).
The energy dissimilarity matrix seemed to correlate with those derived from W distances (Supp- Fig. 7C). We noted, however, that after re-scaling the dissimilarity matrices such that they lie be- tween [0. 1]. the distribution of pairwise energy distances was most in line with interpolated Wasser- stein distances when a = 1 (SuppFig. 7C middle panel).
We repeated the classification and disentanglement kNN analyses done in the main text using neigh- borhoods defined by energy distance (SuppFig. 7D,E). In most cases, using energy distance per-
formed as well as, but sometimes worse than W the covariance-insensitine Wasserstein metric. is possible that computing energy distance using a higher number of samples per image than 64 would improve estimates and downstream regression/classifcattin performance. In general it would be interesting to examine the effects of sample size and empirical energy distance estimate convergence. We leave deeper investigation into this for future work.
23

---

B.3.4 LOW-DIMENSIONAL PROJECTIONS
To determine a reasonable embedding dimensionality for networks, we performed the following analysis. Given symmetric K K distance matrix D, with elements d(i,j), used multidimen- sional scaling to embed networks into a low M-dimensional space. Networks in this embedded space can be encoded by a new, Euclidean distance matrix D with elements d(i,j). For each element on the upper-triangle of these matrices, computed distortion ratio,
(ii,) d(i.j) Distortion (i.o) = max max (i. 1/)(1,
(9)
(10)
By sweeping embedding dimensionality M from 1-20, determined that using MDDS embedding dimensionality of M=15 produced reasonably minimal distortions (SuppFig. 8) for all the distance matrices. After embedding the networks into 15D, we then performed principal components analysis to obtain the scatterplots in Fig. 4AA and SuppFig. 7. In the main text, we used orthogonal Procrustes align the principal components of each subpanel to the left-most panel. For SuppFig. 7B, we again aligned all panels to the left-most panel using Procrustes analysis, but allowed for re-scaning in order to compensate for energy distances being on an arbitrary scale compared with W2 distances.
B.3.5 VAE DISENTANGLEMENT METRICS
For each of the 1800 VAES trained on dSpri tes. Locatello et al. (2019) computed large array of factor disentanglement scores proposed by previous studies. The scores abbreviated Fig. 4F are listed in the below table. We list the different disentangement scores using the same naming convention as the work of Locatello et (2019). We refer the reader their supplement for more details on each of these scores.
Disentanglement score
A
B-VAE eval accuracy
B
Disentanglement, Informativeness, Completeness (DCI) disentanglement
C
DCI completeness
DCI informativeness
Factor VAE eval accuracy Logistic regression mean test accuracy
G
Boosted trees mean test accuracy
Discrete mutual information gap (MIG)
Modularity score
Explicitness test score
K
Separated Attribute Predictability (SAP) score
Gaussian total correlation
M
Gaussian Wasserstein correlation
Gaussian Wasserstein normalized correlation
Mutual information score
B.3.6 k-NEAREST NEIGHBORS ANALYSES
Because the stochastic metries used in this study satisfy the triangle inequality, this permitted parapamettric analyses using k-nearest neighbors (kNN) to determine whether network similar- ity carried information about model hyper- parammerrsssrrs and task performance. We used scikit-learn' #MeighborsClasn soClassif assifier and KNei KNeighbor gherrr sRegr essor for classification and regression analyses, re- spectively. We withheld a test set and performed 6-fold cross-validation on the remaining data to determine k. the number of neighbors use for clasiffication/rerrsssoo We reported final perfor- mance using the average score on the held-out test set.
For classification analyses, trained models to decode random initial seed (1/50 chance, Fig. 4C), and model objective along with regularization strength (6 objective x 6 regularization strengths in the Locatello study, i.e. 1/36 chance, Fig. 4E). In terms of regression analyses, trained models to predict training reconstruction loss Fig. 4D and disentanglement scores Fig. 4F and reported average R² on the held-out test set.
24

---

B.4 ADDITIONAL DETAILS FOR PATCH-GAUSSEN AUGMENTATION EXPERIMENTS
B.4.1 TRAINING AND ARCHITECTURE DETAILS
We use the ResNet- 18 architecture (He et al., 2016) where an intermediate fully-connected layer of dimension 100 added after the final average pooling layer, followed by a linear readout layer. All analyses were done the representations produced of this intermediate fully-connected layer.
Following standard practice, images were randomly cropped, followed by a random horizontal flip. augmentation was applied, where the entire noisy patch
A modified version of the constrained to reside in the image. Lastly, we subtract off the per-channel mean and divide by the per-channel standard deviation. For the Patch-Gaussian augmentation,
16 values of patch width, W E {2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32} and different values of noise scale, a E {0.05, 0.1, 0.3, 0.5, 0.8, 1.). leading 17 6 112 possible (W, o) combinations. For each (W. σ) pair, we trained 3 networks, each with a different random seed, leading 16 x 7 x 3 336 networks. As a baseline, we also trained networks with no Patch- Gaussian augmentation over three random seeds, giving us 339 total networks.
swept
We used stochastic gradient descent with momentum of 0.9, batch size of 128 and weight decay of 1E-4. Networks were trained for 200 epochs where the learning rate was initially set to 0.1 and halved every 60 epochs.
B.5 VISUALIZATION OF HIDDEN LAYER REPRESENTATIONS
To visualize the effect Gaah-aaa Syyeer-parametee on hidden layer representations as shown in Fig. 5A, we randomly selected one image from each of the 10 classes, e.g. 21...., 210- For each image i-and a given value 7-we drew 100 samples from (2,,,) T) and collected the hidden layer representations, leading to 1,000 points total. nutli-dimensional scaling then applied to embed the representations into two dimensions.
B.5.1 STOCHASTIC SHAPE METRIC COMPUTATION
2,000 images were used for computing the stochastic shape metric. To estimate the conditional mean and covariance for each image, 1,000 samples were first drawn from N(z,, 7). The conditional mean was estimated via a Monte Carlo estimator. The conditional covariance was computed by first computing the Monte Carlo estimator and then adding 0.0001 the diagonal ensure the covariance is well-conditioned.
To visualize the metric shape induced by the stochastic shape metric, multi-dimensionall scaling was used to embed the networks into 20 dimensions. Principal component analysis was then done to linearly project the MDS embeddings onto the top principal components.
We used three different values for the interpolated Wasserstein distance, {0. 1,2] and 6 values for the magnitude of the input perturbation, (0.01, 0.05, 0.1, 0.2, 0.3, 0.5). All distance matrices are shown in Fig. 10. The corresponding two-dimensional embedding are shown in Fig. 11.
C PROOF OF PROPOSITION 1
We first prove two lemmas, from which the main proposition immediately follows. Lemma 1. If is group of isometries metric space (d1, then
d(x,
min di (z, T()))
(11)
TEG
is pseudometric which can be used to define metric over equivalence classes [x] = {8) xz where the equivalence relation is defined as:
+ BTTg such that r T(y)
(12)
Proof. This proof is more or less reproduced from Williams et al. (2021), and similar arguments can be found elsewhere within the statistical shape analysis literature.
25

---

The equivalence relation in equation (12) is self-evident. This simply states that d(x, and only if = T(y) for some alignment transformation T E G. Then we define our equivalence relation as: ~ y if and only if d(æ; y) O. In other words, although d technically only defines pseudometric on S. is easily associated to a proper metric on a set of equivalence classes, i.e. the quotient space (5/ >) See Howes (1995) for more background details (page 27, in particular).
if
Now we prove that d is symmetric. Let Tæy denote the optimal transformation from Y to X. That is, Tay - argminreg di (xr, T(y)) and yz argminnreg di (y, T())). Then, using the fact that di is symmetric and that G defines group of isometries, we have
d(z, 9) = d1 (z, (v)) d1 (T (Tey (1), d1 Try (x))) s di (y. ggn = d(y, xx)
but also
d(y,z) = di (v. (x) = di ( (y))) S di (z,, Tny(w))) d(x, v).
di
The only way for both inequalities to hold is for d(z,y) = d σ) Also, we see that Tzy T Typ which we will exploit below.
remains to prove the triangle inequality. This is done as follows:
d(x, v) (z, (y)
(13)
< (z. (T.
(14)
(1, Tux (=)) + (2),
(15)
di (z, Tus (x)) (1 (2
(16)
d(x, 2) + d(z.g)
(17)
The first inequality follows from replacing the optimal alignment, Tay. with sub-optimal align- ment, given by function composition Tę= Tzy- (Recall that 9 is a group and so is closed under function compositions.) The second inequality follows from the triangle inequality d1, after
choosing Tax (2) as the midpoint. The penultimate step follows from T being an isometry on di and since Tus G, we have T G by the group properties of G.
Lemma 2. Let S2) be a metric space, let f()) and g(-) be functions mapping Z + S2, and let Q be probability distribution supported on Z. Then,
E {((()))())))) 1/2 2~Q
(f.gg)
(18)
is a metric over the set of functions mapping 2 + S2.
Proof. Since d2 is a metric, we have d2(2, 3) > o * y. Recall our assumption that the support of Q equals Z. Thus, if there exists E Z for which f(z) # g(2), the expectation will evaluate to positive number and we have d1 (f. g) > 0. So we conclude d, (f(.. g)) = if and only if f and define the exact same mapping from 2 > S2.
is also obvious that dj (f. g) = d1 (g. f), due to the symmetry of d2. Thus, it only remains to prove the triangle inequality.
Fix any function h: Z + S. Due to the triangle inequality on d2, we have:
d((f.) ( x~Q
(B(f(z). 1/2
E. (de(f()). dc((()))((())))) 1/2 ENQ
(19)
Now let X = dz(f(2), h(z)) and Y = da(h(z), 9(2)). Note that is a random variable (sam- pled from Q), and so X and Y are also random variables. We now recall two elementary facts: |||||== (2 defines a norm over random variables, and ||X + Y S | + ||Y|| for any two random variables (Minkowski's 's inequality). Our definitions of X and Y imply that the right hand side of equation (19) can re-written ||X + Yll2. And we can therefore conclude the proof since:
d((f.g) < |X + 21l2 < ||||| + |Y 2 = (f, + di (h g)
(20)
26

---

Main proof. Let restate and then prove proposition 1. We want to show that the following:
F) 2) T ) )
1/2
d(F;;
min TED
E $~Q
D2
Fi
(21)
is pseudometric over stochastic networks-i.e, a pseudometric over functions F that map inputs E onto probability distributions. Recall that (- z) is shorthand notation for F( F() where o=1 is the pre-image of Φ.
Our key assumptions are that -(.,)) ) is a metric over probability distributions and that a group of isometry transformations with respect to this metric-i.e. for any pair of probability distributions F and G. we have that:
D(F.G) D(FoT G
(22)
for any G. is well-known that the Wasserstein distance (Villani. 2009) and energy dis- tance (Sejdinovic et al., 2013; Székely & Rizzo, 2017) are probability metrics. Further, it is easy to show that orthogonal pushforward transformations are isometries for both metrics. For example, we have for the 2- Wasserstein distance that:
(P. (P. Q) inf E ||x - Y1I² = inf E IXX - TYI (P
(23)
for any orthogonal transformation T. Thus, for our purposes we can think of G as being any sub- group of the orthogonal group.
Now that we have reminded ourselves of the main proposition, let us turn to the proof.
Proof. Let us define:
1/2
d1 (F (F xna
D2
(7, (2)1=
(24)
Plugging this into equation (21), we have:
d(Fi, F2) mmin min di (F. 139 T T)
(25)
Lemma 2 tells us that d, is a metric. Thus, lemma applies to equation (25). This permits us to conclude that is pseudometric and defines a metric over sets of equivalent neural representations, as claimed.
D PRACTICAL ESTIMATION OF STOCHASTIC SHAPE METRICS
In both biological and artificial networks, do not have parametric forms for the conditional distributions over neural population responses. Instead, we can only draw samples from these distributions-e.. e.g., by feeding input into artificial network and performing stochastic for- ward pass, or by recording evoked spike counts to a sensory stimulus in biological data. We consider simple experimental setup: we given K stochastic neural networks {Fi...., Fx) ). M network ). and L. repeated observations or measurements of the neural re-
inputs or conditions (z1....;= ²M) sponses each input. For example, in an artificial network that ingests image data, Af would denote (km) R
the number of images in a test set and denotes the number of samples per image. Let x; denote sample é, from network k; to condition m. That is,
(km) Fe
F (x =m) i.i.d. for (L,m, m,,k k) E (1.... L.) x (I....,M) M} (1....,y) K}.
(26)
D. METRICS BASED 2- WASSERSTEIN DISTANCE AND GAUSSIAN ASSUMPTION
Our main assumption in this section is that distributions over neural activations are multi- variate Gaussians. That is, for each stochastic network and every input z Z, we have (z) = (x), E,()),, where μ4 Z + R2 and Σ, : Z + gaxn, T : R + R is linear pushforward map, then the pushforward measure is still Gaussian and is defined by (z) T (TT (2)TT)
27

---

The 2- 2-Wasserstein distance between two multivariate Gaussian distributions has well-known closed form expression (Peyré Cuturi, 2019, Remark 2.31):
(2 (u, 2,) (1l - 14 B(E, )
(27)
where ((.,)) is the Bures metric between positive definite matrices. Typically one sees the Bures metric defined
B = [Tr Tr[E,] 2 -,,2//11/3 1/2
(28)
we use this expression, minimizing B TE,T" over nuisance transformations T E 9 is not straightforward. However, an equivalent formulation of the Bures metric is:
BE, Ej) min U
1 - U1I
(29)
where the minimization is over U E O()). The equivalence between equations (28) and (29) is already established in the literature (see Theorem of Bhatia al. 2019). For the sake of complete- ness we have included a proof in appendix F.4.
(km) K,M.L. Recall that are given sampled neural responses {-/ Jk,m. as specified in equation (26). Using these, we can estimate the mean and covariance of each distribution:
fi (Zmm) =
1 Σ a
(km) and [x(zm) = I
Σ
(km) (km)
(a (x (30)
Here, we've used the typical maximum likelihood estimators. However, any consistent estimator will suffice.
The proposition below summarizes the main result of this section. Using this proposition, we pro- duce an estimate of the distance between two stochastic networks by alternating minimization (i.e. block coordinate descent) over T, .,,..., UM- Each parameter update can often be solved exactly. For example, we typically consider the case of orthogonal nuisance transformations, i.c. G Onn), in which case all parameter updates correspond to solving an orthogonal Procrustes problem (Gower & Dijksterhuis, 2004). Further, the minimizations over (U1...., Un} can be done in parallel.
Proposition 2. If F. (z) and (z) are both Gaussian for all z E Z, then:
Σ TT. (g 1/2 TÉ, (zm (zm)) 1/2 1/2
Ff) min T.U
A
consistent estimator of stochastic shape distance (eq. 5) with the 2-Wasserstein distance used as the "ground metric." The minimization in the above equation is performed over T E G and Um =(n)) for all E M}.......
Proof. Plugging equations (27) and (29) into our definition of stochastic distance (eq. 5 from propo-
sition 1), we have:
( T2)(2)]1 min (z) -T 1/2
d(F,, min E 2
[1((z))
(31)
T
Given M i.id. samples 2m Q for E (I,...,M,, estimate the expectation with an empirical average:
1/2 (2mm))) mm. 1/2 (32)
min
(1((2m) (2m
min U.
TJ
12=1
2 Althoouhh there are certain tricks one can exploit to compute the gradient (Newton- (Newton-Schulz Schulz iterations), the constraint that E G is non-trive trivirival.. When ς is continuous manifold (e.g. the orthogonal or special orthogonal group), one can resort to manifold optimization algorithms. These algorithum are somewhat cumbersome but nonetheless plausible approach. However, even this would not cover the case where G is discrete set (c.g., the permutation group).
28

---

Next, pull out the minimization over each Un E O()) outside the sum. Additionally, since G a group of isometries on IR", we know that G is a subgroup of the orthogonal group. Thus, every T E G is an orthogonal matrix, so T Um is also an orthogonal matrix. Thus we introduce a change of variables Um T Um and minimize over Um E (n),, as this attains the same minimum value. In summary, we have:
M M
(2)( 1/2
min T.U1..., UM
[mm((m)) Tuy (Zm (21))l
(33)
The only remaining step is to replace every (u(z) and [() with some consistent estimator, such as the empirical mean and covariance (see eq. 30). In the limit as M + oo and - oc, we have convergence to the true distance due to the law of large numbers.
D.1.1 ALGORIITHMIC COMPLEXITY AND COMPUTATIONAL CONSIDERATIONS
To compute distances using the Wasserstein ground metric between two coussian-distribure stochastic network representations (Eq. 6), used closed form updates of the orthogonal pro- crustes problem (Gower & Dijksterhuis, 2004) for T E O and (Um)m in alternation, using $ iterations. Importantly, T and each Un can be solved exactly each alternation (described above). For K stochastic networks, we must consider (K32) total pairwise comparisons.
Computing the optimal T at each step involves aligning two stochastic representations, each com- prising a stacked matrix of M n-dimensional means and M n X covariances using Procrustes alignment (Gower& & Dijksterhuis, 2004). This involves a matrix multiplication and singular value decomposition with O combined complexity, assuming < M. Similarly, at each step, com- puting the Bures metric requires solving for M n x n orthogonal matrices, (Um M m== with total complexity ((M Thus, the total algorithmic worst-case time complexity is O(K SMn").
Notably, this computation is highly parallelizable over the K2 pairwise comparisons. For the VAE and aatch-Gaussial results in the main text (Figures 4 and 5), we distributed the distance matrix cal- culation by distributing pairwise comparisons over single CPU cores, with each comparison taking few seconds complete.
D.2 METRICS BASED ON ENERGY DISTANCE
This section outlines an alternative measure of stochastic representational distance that does not require any parametric assumption (e.g. Gaussian) on stochastic neural responses. We use energy distance, & defined in equation (4), as the ground metric appearing in proposition 1. As explained in the main text, this distance has favorable estimation properties in high dimensional spaces relative to the Wasserstein distances. It remains an open problem to develop estimation procedures for Wasserstein- stochastic shape distances without the assumption of Gaussianity.
To compute the stochastic shape distance, we need to solve the following optimization problem:
argmin E TEg z~Q
(F)
(34)
Note that we have squared the expression occurring in the main proposition-i.e. we have dropped the (-) operation as this does not effect the optimal alignment transformation.
First, let's focus on the innermost term inside the expectation of equation (34). Let X ~ (- 2) and Xi ~ P 2), independently and treating the input z as fixed for now. Now, since G is a group of isometries with respect to the Euclidean norm, we have:
(-12, T- = E X; - TX, 1 - E - XX11 E TX, - TX
E X; -X,
X|11 J X111
The final two terms are constant with respect to T, so we can drop them from the objective function without effecting the result. Thus, equation (34) can be simplified to:
argmin E Teg z~Q
E| TX1
(35)
29

---

Here, the outer expectation is over network inputs & and the inner expectation is over conditional distributions, X; Pi(( z) and X, (. (1).
Recall again that are given sampled responses {ar (a= K.M.L. specified in equation (26). To Tk,, construct consistent estimator, evoke the law of large numbers to replace the expectations with empirical averages. Equation (35) becomes:
(km))
M
L
L
argmin TEG
ΣΣΣ W: E=pp-
|
(8m)
Tarpmm
(36)
MT2
When y - 2, the optimal T E 9 can often be identified efficiently-e.g.. by solving Procrustes problem when G O (Gower Dijksterhuis, 2004) or a linear assignment problem when 9 = P (Burkard al., 2012). However, when q 2 the stochastic shape distance only depends on the mean and is insensitive to higher-order moments of the neural response (see appendix F.1) In the more interesting case where q * 2, we can iteratively re-weighted least squares (Kuhn, 1973) to identify the solution. Details of this well-known algorithm are provided in appendix F.2.
Now, let T" E G be the solution to equation (36). Using this is straightforward estimate the desired stochastic shape distance. Our estimate of d(Fi, Pÿ) is:
(m) 1 LIL-I) (im) -T" x
Σ E>p
|
(im) (im)
2 (im)
(mm)
||
T
where the sums over >> p are over all L(T. - 1)/2 pairwise combinations between sampled activations. Each of the three terms in the expression above is consistent (though not unbiased) estimator of its corresponding term in definition of energy distance (eq. 4). However, if the final two terms above are over-estimated in magnitude and the first term is under-estimated, the overall estimate of d(F,, Fy) may be negative. This violates perhaps the most important property of a metric space that distances should be nonnegative. Triangle inequality violations are also possible.
We propose a simple fix using basic ideas from the literature on metric repair (Brickell et al., 2008). Given collection of K networks, we the procedure above compute estimate of the K K distance matrix D where Dij = d(F,, Fj). We then find the matrix D that closest to our estimate
D according quadratic loss, and which satisfies all the axioms of a metric space. This amounts to solving quadratic program, detailed in appendix F.3.
E INTERPOLATED METRICS BETWEEN DISTANCES
AND
E. PROOF THAT W2 IS A METRIC
We start by proving well-known and basic lemma, which states that the tp norm of collection of metries also defines a metric.
Lemma 3. Let d di...., dn be collection of metrics set S. Then, for any p > 1,
dx, v) = Vd (x,,)) +... dd(aay))
(37)
metric S.
Proof. is obvious that d(x, y) = 0 if and only if dj (x, 2) = .... = dn (I, 3) = and that d(x, y) = d(y, x.). So is only non-trivial to prove the triangle inequality.
Let d(x, y) denote the vector in IR" holding each distance. That is:
d (x. y)
d(x, v)
(38)
dn (z. v)
'One could alternatively consider using manifold optimization methods when & is a continuous manifold. However, these methods are somewhat cumbersome and aren't easy to extend to the case where & is a discrete set, such the set of all permutations.
30

---

The triangle inequality now follows from:
d(r.v) = (di(z, +.... dn 1), = -I((x,,)))
(39)
I [|((r, m) + d(m, f)lle
(40)
S || m) + || v)ll
(41)
= d(x, m) + d(m, y)
(42)
for all x E S, y E S, and m E S. The first inequality follows from the triangle inequality on each di...., d2 and then from I|- being a monotonically incressing function of each vector coordinate. The second inequality follows from the sub-additivity property of all norms.
This lemma provides further perspective on the closed form expression we gave in appendix D.1 for the 2- 2-sasserttein distance between Gaussian distributions. In particular, if P. = NW,, Σ;) and Pj = (N(,,, we saw in equation (27) that:
(P.,,, =
+ (E,, E,))
(43)
where d is the squared Euclidean distance between the means and dg is the squared Bures metric between the covariances. Thus, the 2- 2-Wasserstein distance between Gaussians can intuitively thought of as the &a norm of this pair of metrics.
It trivial verify that all the properties of a metric are preserved the distance is multiplied sealar > 0. That is, if g is a metric, then d(x, y) ag(z, y) is also metric. Combining this with lemma is obvious that,
W (P) = (a Mz) + (2- ax - E,)) 1/2
(44)
which is simply re-statement of equation (7), is a metric for any 0 < 2.
E.2 LOWER BOUND ON INTERPOLATED SHAPE DISTANCES
We now derive a simple lower bound on stochastic shape distances when W2 is used as ground metric. Let dz denote the squared shape distance of interest, for any chosen S 2. We have:
(F. = min Eg
(W2) (1)( T-1)
TEG
min Eg a: (H, THy(2)) + (2 (2- aa) (2), TEg
2 a min E, d' (uH)))) + +22 min Eg (x), TzEg
T.EG
The inequality here follows from the linearity of expectation and then from separately minimizing the two terms. The inequality is tight if the optimal value of Tę equals the optimal value of Tz- Furthermore, the two minimized terms in the final expression proportional to the squared shape distance when a = and a = 0, respectively:
1
min Eg ((( TEG
(2 Ff) 2
(45)
1
min E, (x) Teg
-a (F,, Fi)
de
(46)
Thus, in summary we have:
2-a
(F(FFFF))
2 22 (F, +
d L (F,, F2)
(47)
2
2
2-a
dF,FF)))
1/2
d
L (F,
d
L (Fi,
(48)
2
31

---

E.3 INTERPRETATION OF W2 WHEN GAUSSIAN ASSUMPTION VIOLATED
then W2 can be interpreted as a natural extension of 2- 2-Wasserstein distance (see sec. 2.4). What if neural responses are not Gaussian Concretely, consider two distributions P; and Pj. which are not necessarily Gaussian. We can still define the first two moments (mean and covariance) of these distributions:
When neural responses are
μ P [z) and Σ; E, [(2 - pz)(zz - p4)
(49)
Using these, we can still compute (P,, Pj) as before.
However, it is no longer the case that this calculation will coincide with the 2-Wassesstein distance between P3 and Pj. That is, since the Gaussian assumption is unreliable, W (P,, H2 Pj). Because of this, can no longer conceptualize W2 the amount of energy taken to transport P. onto Pj with the parameter a differentially weighting the cost of transporting mass due to mismatches in the mean and covariance.
On the other hand, W2 may still be reasonable ground metric in many practical circumstances. In particular, it obvious that W2 (P,. P3) = if and only the mean and covariance of these distribu- tions match. Thus, it is a pseudometric over all probability distributions and a metric on equivalence classes defined by the equivalence relation P. ~ P, if and only fr; = μj and Σ; Σ, From this, it is easy to show that the stochastic shape metric (eq. 5) based on this ground metric also satisfies the metric space axioms, including the triangle inequality.
In high-dimensional datasets, it is often challenging to estimate and interpret the higher-order sta- tistical moments of a distribution. In the setting of comparing stochastic neural representations, one could argue that it is reasonable to settle for ground metric that is insensitive to these higher- order moments. From this perspective, W2 belongs to a larger family of ground metries that can be expressed:
(50) for some chosen metric on the means, du. and another chosen metric the covariances, dy. Again, no assumption on whether P. and Pj being Gaussian is strictly necessary. A more thorough explo- ration of these alternative ground metrics is potential direction of future research.
D(P;, P3) = (dh, (1i, 145) + (E, 2,))
F MISCELLANEOUS THEORY AND BACKGROUND
F.1 ENERGY DISTANCE AS A TRIAL- SHAPE METRIC WHEN q = 2
Performing representational dissimilarity analysis on trial-average activity measurements is already common practice in neuroscience. Here, we show that this approach arises as a special case of the stochastic shape distances explored in this manuscript. When g = 2, the energy distance is given by:
(2 (E - Y1I
- JELy 1/13)1//2
(51)
Since X and X' are independent and identically distributed random variables, we have:
_ - X'I22 = _ x] +
(52)
E| - IE| X
(53)
B[X x] - [X E[X]
(54)
Likewise,
E Y - Y'1 EE T -
(55)
Plugging these expressions into equation (51) and simplifying we that:
(P. = X - 2 - E x + E X E[x] - E + E|y] 1/2
-IE{XT-XT + Ey - X y| BXx E[ E|X
-
Yy
+
1 B
B
=[E[[X] T [[ + B X 2/2
=(IEE[X] - 1
= -||E[X] - E[Y]II
32

---

To summarize, have shown that the q 2 energy distance between distribution P and Q is equal to the Euclidean distance between the mean of P and the mean of Q. If we use this energy distance the ground metric, D, in proposition 1 to construct stochastic shape distance, essentially calculating a deterministic shape distance" on the mean response patterns.
F.2 ITERATIVELY REWFIGHTED LEAST SQUARES
Fix to be a value on the open interval (0,2) and consider the following optimization problem:
min
Σ
l T
(56)
TEg
It is easy to see that equation (36) is an instance of this problem for a particular choice of vectors
H2 E I2" and y E R".
Our key assumption will be that we can efficiently solve the following weighted least squares prob- lem:
min TEG
ws 1yy
Tx;
(57)
for any choice of weightings, Wi,.... WN. Again, this is possible when G is the orthogonal group (Procrustes problem) or the permutation group (linear assignment).
Iteratively re-weighted least squares algorithms minimization
family of methods that are encompassed by the even larger family of majorize-miniiization algorithms (Lange, 2016). The specific method we extension to Weiszfeld's 's algorithm (Kuhn, 1973). Our starting point is
deploy can be viewed to recognize that the function + ga/2 is concave for 0 < q < 2 and 2 0. Thus, can derive an upper bound using the first-order Taylor expansion:
(s + 870/22 s 54/2 +
99/2) = 54/2 4 ds
d
(58)
2 8(1-g/2)
for any such that s + 8
We will now use this fact to derive upper bound the objective function in equation (56). Let T() E G represent our estimate of the optimal T E G after iterations of our algorithm, and let E G denote any feasible transformation. Then, E (I,..., NJ,, define:
(t)
lly, T()x;1
(59)
((0) b1 Tx - (1)
(60)
Notice that these definitions imply 8 + 510 2 O. Now, plugging into equation (58), we have:
(0()+
( + 6(!1)) 9/2 =11y: Tx <
+
(t)
9/2
!
8
+
(61)
2
(1)
g/2)
This an upper bound for each term in the sum of the original objective function. Therefore, plugging in the definitions of a( and g(0) we have:
S(T) Σ 1
T <
Σ |. T(O mr mm:11) 9/2
+
ạ
(1ly: T(t) (1 q//2)
11y: Tw, 12 8/
(e)
(62)
Σ ll - T1:0 |1y:
Tx; 1|2 - |/ (1)) xxl
(63)
y: (() 2-8
A Q(T ((())
(64)
"Specifically, the distances covered under Proposition in Williams al. (2021).
33

---

Here, view Q(T T function of -ic.. T(t) Q(T T provides an upper bound the objective function for any T E G. Furthermore, it is easy to check that f(T( Q(T()) - the upper bound is tight at T = TT().
is fixed. The calculations above show that
We now have all the necessary ingredients to derive an algorithm. We start by initializing T(1) g
by some method. Then we compute (T((2, TT(),.... iteratively according to:
y T Tx, y: T x,
T(+ argmin Q(T T argmin TEg
(65)
TEg
The last equality here follows from dropping terms from equation (63) that are constant." Intuitively, at each step we are minimizing a surrogate function Q(T) | (+)) that upper bounds the true objective
function. This surrogate function is easy to optimize since the minimization is a special case of equation (57) with weightings:
W3 y,
T xz 12-9
(66)
Furthermore, because we showed that the upper bound is tight T T have:
f(T((+))) < Q(T(++))
(())
min
Q(T ( s Q(T()) ( f(T
(67)
TeS
which shows that the the objective function never increases as the algorithm progresses.
F.3 QUADRATIC METRIC REPAIR
We are given symmetric estimate of distance matrix D RK*K which may contain nega- tive entries and triangle inequality violations. Let d E 1)) be vector holding the upper triangular entries of D, excluding the diagonal. Then, consider the following optimization problem:
minimize subject to
|| 2
2 0 2;
K(K - 1)/2}
+ xj - 20,
V(i.i,)) E TK
where Tk the ((k) directed triples of indices corresponding triangle inequality con- straint. This is quadratic program-i.c. i.e., a convex optimization problem with quadratic objective and linear inequality constraints. To solve this problem, we use the open-source and highly opti- mized OSQP solver (Stellato et al., 2020). The number of inequality constraints grows cubically as K increases, so finding an exact solution may be computationally expensive for analyses of large collections of stochastic neural networks.
F,4 REFORMULATING THE BURES METRIC
Here will prove that equations (28) and (29) are equivalent. A similar statement is proved in Theorem of Bhatia et al. (2019). Our proof relies on the following lemma.
Lemma 4. Let X E Raxx be a matrix with singular value decomposition X USVT. Then VUT (x
Proof. This follows from the construction of the singular value decomposition. First, recognize that X can be written as the product of an orthogonal Q and symmetric positive semidefinite matrix, P. as follows:
X X(X X 1/2 [
1/2
(68)
=P
'It is important to understand that we are treating for the minimization. We also further simplified by rescaling Q(T TT(t)) by factor of 2/g, which doesn't affect the value at which the minimum is attained.
( as a constant. Only terms that depend T matter
34

---

It is easy to check that Q'Q QQ" I. Now, since P is positive semidefinite, we have P VSV for some orthogonal matrix V and nonnegative diagonal matrix S. Defining = QV, we arrive at the SVD of X QP USV Now we see that:
U QV X XX
UV = (1Xx))11 x)-1/2
(69)
Taking the transpose of this prove the lemma.
Now we proceed to prove the main result.
Proposition 3. Let A and B be two positive definite matrices. Then
| - Tr(A B 2 QEO
min
(70)
Proof. The minimization over Q is an instance of the well-known orthogonal procrustes problem (Gower & Dijksterhuis, 2004). This has a closed form solution. Specifically, denoting the singular value decomposition of 1/2 A 1/2 as USV USVT, we have:
Q
argmin ||A A/2 =B VUT
(71)
QEO
Now, by Lemma above, we have:
VUT = (( (B 12 T = (A 12 1/22/// (72)
Plugging this into the original problem, we have:
A 1/2 - Q² 21/2
Tr + B Q" 2
(73)
Tr B "A
(74)
Due to the eyclic trace property, this becomes:
+ 1L-B-(AAV2))
A " Tr(A+ B 2(A
(75)
claimed.
35


Reviews:
[{"review_id": "kmbL40Ss_E3", "reviewer": "Reviewer_yN72", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper presents a novel class of stochastic shape metrics, which are suitable for assessing similarities in the representations of both biological and artificial stochastic networks.  The paper presents a thorough analysis of the formal properties of their class of metrics, showing they satisfy the metric axioms, that they have other desirable formal properties, and their relationship to certain previous metrics.  In their experimentation, they focus on a generalized form of 2-Wasserstein distance, which allows a fine-tuning of its sensitivity to mean vs covariance structure, and has a tractable analytic estimator.  They provide a series of case studies demonstrating the use of this metric on synthetic, biological and artificial networks, leading to some intriguing observations. ", "summary_of_the_review": "I recommend acceptance of the paper, since it offers a general solution to an important problem that will be of interest both to the neuroscience and machine learning communities.  Further, the results will be of interest to both communities (although, as noted above, several of the observations could perhaps be followed up in more detail in making the case for the use of generalized forms of stochastic metrics).", "strength_and_weaknesses": "Strengths:\n\n- The problem of measuring distances between representations in stochastic networks is an important one, both in neuroscience and machine learning.  The formal properties investigated by the authors are well motivated, such as invariance to network permutations.  The authors offer a clean and general solution to this problem.\n- Some interesting observations are made in the experimentation, which I'm not aware of having being observed before.  Particularly, the observation that the response of biological networks (using Allen mouse brain observatory data) to different classes of responses can generate representations primarily focused on the mean or the covariance (noise) structure, and that this is mirrored in artificial networks pre and post training.  Also, the observation that many artificial network training properties, such as random seed, reconstruction loss and regularization can be predicted from a network's noise structure is unexpected, and may be relevant for instance in predicting which networks will generalize well (e.g. in statistical learning theory).\n\nWeaknesses:\n\n- It is unclear whether the 'minimization' operation over the group of nuisance transformations is necessarily the right way of assessing distances between stochastic representations (which is used since the work generalizes previous work on deterministic shape metrics).  Potentially, some form of uncertainty about the best alignment should also be preserved over the group of nuisance transformations.\n- The empirical work is somewhat preliminary, in the sense that many of the observations could be followed up in more detail with more substantial studies.  It is also limited in the sense that only covariance (pairwise) structure is considered when comparing 'noise' distributions empirically, while the framework would allow for higher-order properties of the distributions to be compared.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "XeuXjjkkwC", "reviewer": "Reviewer_pofg", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The submission entitled “Representational dissimilarity metric spaces for stochastic neural networks” proposed a new method to characterize the representational similarity of neural networks, taking into account the noise characteristics. The paper considered metrics over the stochastic shape using p-Wasserstein distance and some variations. Overall, this is a well-written paper. I enjoyed reading it. ", "summary_of_the_review": "Solid paper. The study is well executed with moderate innovation. \n", "strength_and_weaknesses": "\nStrength: \nNice combination of theory and experiments. \nThe numerical experiments are extensive. \nThe theoretical framework is elegant and will be useful in some applications.\n\nA few concerns:\n\nAlthough the numerical experiments covered several examples,  the insights revealed were quite limited. Some of the results could be obtained by using fairly standard methods to study the mean and the noise covariance structure (e.g., Section 3.2).\n\nThe Gaussian assumption needs to be better justified, perhaps through more simulation experiments.  Having a small number of trials in practice is not a good *justification* of the Gaussian assumption. The maximum entropy argument is also problematic as the maximum entropy distribution depends on the domain being studied. For variables that are non-negative, the maximum entropy distribution will not be Gaussian. \n\nThe innovation on the “stochastic representation” needs to be toned down. This is not a new concept. Every general model with noise can be naturally treated as a “stochastic representation”. The classic RSA approaches can also deal with stochastic presentations, under well-defined noise assumptions (e.g., reviewed recently in Kriegeskorte & Wei, Nature Reviews Neuroscience, 2021). \n\nThe following reference should be cited and discussed, as it is closely related to the current work:\nShahbazi, Mahdiyar, et al. \"Using distance on the Riemannian manifold to compare representations in brain and in models.\" NeuroImage 239 (2021): 118271.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Hj6fuLuufK", "reviewer": "Reviewer_Moyn", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper introduces proper distance metrics for quantifying the similarity of representations in stochastic networks, either biological or artificial. The paper builds on a recent article by Williams, which described metrics for quantifying similarity in deterministic networks by allowing for a set of preprocessing steps (e.g., whitening) and alignment operations (e.g., rotation) before applying a “ground metric” (e.g., Euclidean norm). This paper uses the same framework but applies it to a set of probabilistic ground metrics that were chosen to converge to the Euclidean norm in cases where the network is deterministic. The paper focuses on a Gaussian approximation to one of these metrics (Wassertein distance). The Gaussian approximation results in a metric that has two terms, one based on the similarity of the means, which might be thought of as the deterministic component, and one based on the similarity of the covariances after rotational alignment. They also describe a meta-metric that linearly interpolates between these terms in order to emphasize either the means or covariances. \n\nThe authors apply their framework to a set of toy networks that differ in their covariance scale and correlational structure and report that their metrics can uncover the latent structure (Fig 2C). They then apply their approach to neural recordings from mouse primary visual cortex to gratings and natural scenes (Fig 3). They find that their similarity metric is dominated by the noise covariance structure for static gratings but not for natural scenes, presumably because natural scenes drive more variation in the mean response. The authors then apply their method to a large collection of 1800 VAEs, and report that their distance metrics can be used to decode a variety of network properties with varying degrees of accuracy (Fig 4). Finally, the apply their model to a Patch-Gaussian network trained with varying amounts and sizes of Gaussian noise, and report that their metrics can recover this structure to some degree (Fig 5). \n", "summary_of_the_review": "The paper develops a framework to extend representational similarity to the stochastic regime, which I feel is an important contribution. The metric is solely motivated by how it behaves in the deterministic case which is a notable limitation in my opinion. There are no major novel findings from their analyses applying their method, but the analyses do serve to illustrate and validate the approach. The paper generally seems of high quality. \n\nPost-rebuttal comments:\n\nI apologize for my slow response. I thank the authors for taking so much time to consider and respond to my questions and comments. I was impressed with all of their responses including to other reviewers. I have upgraded my score from a 6 to an 8. I hope to see this paper accepted. ", "strength_and_weaknesses": "Strengths\n\nThe paper addresses an interesting and important topic: how to deal with stochasticity when quantifying large-scale similarity measures between biological or artificial networks. As noted, virtually all prior work has focused on the deterministic case, with comparatively superficial treatments of noise covariance structure (e.g., mahalanobis-based correction metrics). \n\nI found the writing clear, logical, and easy to follow. \n\nApplying their method to both biological and artificial networks is a strength in my opinion. \n\nThe method seems like it scales well to large datasets and many networks though not much information is reported about computation time as far as I can tell. \n\nWeaknesses\n\nAll the metrics are motivated by the desire to reproduce a sensible Euclidean metric in the deterministic limit. This seems like a potentially desirable property, but it doesn’t say anything about how metric behaves in the stochastic case, which is the whole point. In Figure 1, they illustrate how the covariance structure of a set of responses can impact discriminability, but it is not obvious how their metric relates to the stimulus discriminability. This strikes me as a major shortcoming, since it makes it hard to understand what the metric tells you. \n\nTheir Gaussian-approximated Wasserstein metric allows for a rotation of the covariance structure. This rotation appears to be separate from any rotation allowed to align the two networks. It is not clear to me why this is appropriate or desirable. Is this partly why the mean-insensitive version of their metric does not produce the desired result in their toy dataset?\n\nThe experiments are useful demonstrations of the method, but don’t seem to shed new light on brain or network function. The fact that stimulus-driven variation is greater for natural stimuli doesn’t seem particularly surprising and presumably explains the findings in Figure 3B/3C. It is also not clear that their analysis tells us much about noise structure. In Figure 3F, what is each datapoint? A session? Why would we expect distances to covary across sessions between natural scenes and gratings? Do the neurons differ across sessions?\n\nOne thing that would be interesting is to know is the extent to which the noise covariance structure of a visual region (e.g., V1) is stereotyped, in the sense of being similar across animals but distinct across regions, and whether this structure is predictable from the stimulus-driven component. This seems like the type of question their method might be able to answer. \n\nThe authors seem to make a big deal out of whether distances are above or below the diagonal when plotting the covariance- and mean-insensitive measures against each other (e.g., Fig 3B/C and 4B). It’s not clear to me that this has any particular significance. Why are squared distances between the means comparable to the Frobenius norm of the covariances?\n\nI can’t tell if the fact that you can decode network properties from their measures is notable (e.g., Fig 4). There is no baseline comparison against any alternative model. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster