MSA: When Neural Network Similarity Meets Riemannian Geometry

> Paper: *Geometry-aware similarity metrics for neural representations on Riemannian and statistical manifolds* > arXiv: 2603.28764 > Authors: N Alex Cayco Gajic, Arthur Pellegrino (ENS-PSL & UCL)

Opening: Two Swiss Rolls

Picture two Swiss roll cakes: one laid flat as a spiral, the other rolled into a cylinder. Externally, their shapes are completely different. But an ant crawling on the surface of either cake experiences the same local curvature, distances, and angles.

This is the distinction between intrinsic and extrinsic geometry. Traditional neural network comparison methods—measuring angles, distances, and correlations between vectors—compare extrinsic shape. The paper asks: if two networks look entirely different in state space yet solve problems in strikingly similar ways, how can we detect that? The answer: look at intrinsic geometry, not extrinsic geometry.

Part 1: A Neural Network Is a Warped Sheet of Fabric

The manifold hypothesis holds that high-dimensional data actually lies on a low-dimensional manifold—like a 2D sheet of fabric crumpled into 3D space. A neural network, layer by layer, warps this sheet: stretching, folding, and separating entangled classes until they become linearly separable.

Traditional methods (CKA, CCA, RSA, Procrustes) compare the crumpled shape of the fabric in ambient space. But the same intrinsic structure can be embedded in completely different ambient spaces. Two networks may share the same "solution strategy" yet look different due to initialization, rotation, or dimensionality—and vice versa.

Part 2: The Pullback Metric

Given an input manifold ℳ and network map φ, take two nearby points p, q with tangent displacement v. The network maps v to J·v (J = Jacobian) in hidden space. The inner product of mapped vectors becomes:

> (J·v₁) · (J·v₂) = v₁ᵀ (JᵀJ) v₂

The matrix G(p) = JᵀJ is the pullback metric. It tells us which input directions get stretched, compressed, or made orthogonal—capturing the *manner* of distortion, independent of where the result sits in ambient space.

Part 3: The Spectral Ratio

Given two SPD matrices G₁ and G₂, solve the generalized eigenvalue problem:

> G₁v = λG₂v

These λ values are coordinate-independent. The paper defines:

> d_SR(G₁, G₂) = 1 - √(λ_min / λ_max)

• If metrics are identical, all λ = 1 and distance = 0.
• If they differ drastically along some directions, distance approaches 1.
• The distance is bounded in [0, 1], and the authors prove it is a pseudometric satisfying identity, symmetry, and the triangle inequality.

Part 4: MSA

Metric Similarity Analysis averages the spectral ratio over the whole input manifold:

> d_MSA(φ₁, φ₂) = (1/Vol(ℳ)) ∫_ℳ d_SR(G₁(p), G₂(p)) dvol(p)

Key properties:

1. Invariance to orthogonal transformations of hidden layers. 2. Invariance to reparameterization of the input manifold.

MSA thus captures the network's intrinsic computation, not incidental implementation details.

Part 5: Experiments

Experiment 1: Rich vs Lazy Learning

• Rich regime: networks learn structured features; Lazy regime: networks memorize samples with random high-dimensional features.
• Procrustes and RSA report rich and lazy networks as very similar—their PCA projections both look like spirals.
• MSA reports near-zero similarity: rich networks warp the input smoothly and structurally; lazy networks distort it randomly. A carefully made Swiss roll vs. a crumpled sheet of paper.

Experiment 2: RNN vs SSM on a Working Memory Task

Networks receive two angle inputs and must recall them after a delay. Compared architectures: fully connected RNNs and structured state space models (SSMs, with HiPPO initialization).

Findings:

1. MSA clearly separates RNNs and SSMs even on the same task. 2. MSA distinguishes trained vs untrained models. 3. RNNs maintain relatively stable geometry through the delay; SSM geometry evolves over time. 4. RSA can separate architectures but is insensitive to training state; DSA (dynamic similarity analysis) is weak at both. MSA captures geometry *and* dynamics.

Experiment 3: Information Geometry of Diffusion Models

Applied to Stable Diffusion XL with a 2D manifold of text embeddings obtained by bilinear interpolation of four prompts.

Findings:

1. Geometry evolves during diffusion: early timesteps (near t=1, Gaussian noise) and late timesteps (near t=0, generated images) have clearly different information geometry. 2. Classifier-free guidance γ: above a threshold, manifold geometry becomes *closer* to the no-guidance case, implying an optimal guidance level where information geometry diverges most from the unguided model—a potential new angle for hyperparameter tuning.

Part 6: Why It Matters

For interpretability

Rather than dissecting which neurons fire, MSA offers a geometric view: study how the network warps the data manifold—analogous to studying someone's *way of thinking* rather than their anatomy.

Connection to Geometric Algebra Transformers

• GATr uses rotors to *build* networks.
• MSA uses Riemannian tools to *analyze* networks.

Both share the intuition that neural computation is fundamentally geometric.

Limitations

1. Requires an explicit manifold representation of input data; unclear for neuroscience data without learned manifolds. 2. Ignores downstream use: only hidden-layer geometry is considered; a low-rank decoder may ignore some representational directions. 3. Correlational, not causal.

Future directions: geometry-guided network design, combining MSA with causal inference, and applications to RL, graph neural networks, and neural population coding.

Conclusion

> To compare two neural networks, don't ask "where are they?"—ask "how do they warp the world?"

This is a shift of perspective: from extrinsic to intrinsic, static to dynamic, discrete to continuous. As physics history suggests, the deepest beauty often hides in geometry.

Further reading:

• Paper code (to be released)
• Related: GATr (Geometric Algebra Transformer)
• Feynman, *The Character of Physical Law*, on geometry and physics