Similarity measures are widely used to interpret the representational geometries used by neural networks to solve tasks. Yet, because existing methods compare the extrinsic geometry of representations in state space, rather than their intrinsic geometry, they may fail to capture subtle yet crucial distinctions between fundamentally different neural network solutions. Here, we introduce metric similarity analysis (MSA), a novel method which leverages tools from Riemannian geometry to compare the intrinsic geometry of neural representations under the manifold hypothesis. We show that MSA can be used to i) disentangle features of neural computations in deep networks with different learning regimes, ii) compare nonlinear dynamics, and iii) investigate diffusion models. Hence, we introduce a mathematically grounded and broadly applicable framework to understand the mechanisms behind neural computations by comparing their intrinsic geometries.
Figures & tables
Figure 1: Intrinsic vs. extrinsic geometric similarity.
Figure 2: The pullback metric of neural representations.
Figure 3: The spectral ratio is a distance over SPD matrices. a. The set of SPD matrices (here 2×2 ) is a solid cone in the Euclidean space of their entries. b. Slice of the SPD cone along matrices which are related by a scalar multiple. The slice is characterised by a variable r , defining the relative magnitude of the diagonal entries and θ defining the magnitude of the off-diagonal entries. The spectral ratio defines a similarity between pairs of matrices on the SPD cone, with different relative θ ( c ) or r ( d ).
Figure 4: Metric similarity analysis enables comparison of the intrinsic geometry of neural representations. Two neural networks receiving inputs from the same manifold with different hidden-layer geometries, corresponding to distinct Riemannian metrics.
Figure 5: Rich and lazy network representations have different intrinsic geometries. a. Network trained to map a 2D manifold to discrete classes. b. PCA applied to the activation of the rich and lazy network, coloured according to the target class. c. Mesh grid of the intrinsic geometry of the representations. d. Similarity between the rich and lazy representation; rich 2 is a different seed.
Figure 6: MSA captures hierarchical structure across neural representations. Networks with different learning regimes, tasks, architectures, and seeds. a. MSA, RSA and CCA applied to pairs of these networks. Note that RSA and CKA are equivalent when mean-centered ( Williams, 2024 ) . MSA captures the hierarchical structure across models while CCA and RSA do not. b. Instead of considering all combinations of such hyperparameters, we hierarchically vary them. All networks receive the same input on a 2-dimensional planar manifold. They have different learning regimes (rich vs. lazy), or different tasks (5 vs. 2 classes), or seeds (S1/S2, defining initial weights).
Figure 7: MSA enables comparing state-space models to recurrent neural networks. a. The network receives a sequence of angular inputs: ui(t)=fi(θ1)\mathds1t=t1+fi(θ2)\mathds1t=t2 for i=1,2 , with f1(x)=cos(x) and f2(x)=sin(x) , which it must keep in memory for a variable delay period. At the Go cue (signalled by u3(t)=\mathds1t=tgo ), it must output same sequence. b. PCA applied to the state of the RNN. c. MSA applied to the full input-time manifolds of different models. d. Eigenvectors of the doubly-centred pairwise MSA similarity matrix (i.e. MDS with MSA distance). e. Similarity of the geometry (RSA), dynamics (DSA) and geometry+dynamics (MSA) between the trained RNN and all other models. f. MSA applied to compare the representations within a single model at two timepoints (after the second input).
Figure 8: The latent space of diffusion models collapses without guidance. a. Schematic of StableDiffusionXL: a backward diffusion process transports a Gaussian distribution to a latent generative distribution. Samples are then mapped to image space via a variational auto-encoder (VAE). b. The diffusion process is conditioned on a text embedding. We generated a manifold of text embeddings via bilinear interpolation: z(π)=π1π2v1+(1-π1)π2v2+π1(1-π2)v3+(1-π1)(1-π2)v4. This defines a statistical manifold of generative distributions. c. Samples across the manifold for a fixed x(1) . d. Similarity between the statistical manifold geometry at any pair of timepoint. e. Visualization of the information geometry using MDS on the MSA distance. f. MSA similarity between the statistical manifold at different guidance levels vs. no guidance ( γ=1 ).