Geometry-aware similarity metrics for neural representations on Riemannian and statistical manifolds

Abstract

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 (8)

• 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\times 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 $\theta$ 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 $\theta$ (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.

Theorems & Definitions (10)

• Definition 2.1: Spectral ratio
• Proposition 2.2: The spectral ratio is a distance
• Definition 2.3: Metric similarity analysis
• Proposition 2.4: MSA is a distance over pullback Riemannian metrics
• Proposition 4.1: MSA is invariant to state-space rotations
• Proposition 4.2: MSA is invariant to changes in local coordinates
• proof : Proof of proposition \ref{['thm:sr_distance']}
• proof : Proof of proposition \ref{['thm:msa_distance']}
• proof : Proof of proposition \ref{['thm:rotation_invariance']}
• proof : Proof of proposition \ref{['thm:coordinate_invariance']}