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: A Tale of Two Swiss Rolls

Imagine two Swiss roll cakes. One lies flat on a plate as a flattened spiral; the other stands up, rolled into a cylinder. Measuring their surfaces externally, their shapes are completely different. But to an ant crawling on the cake's surface, the two rolls feel identical—the local curvature, distances, and angles are the same.

This is the difference between intrinsic and extrinsic geometry.

Traditionally we compare neural networks the way we compare the external shapes of two Swiss rolls: looking at angles, distances, and correlations between vectors. Cayco Gajic and Pellegrino pose a fundamental question: if two networks look entirely different in state space yet solve problems in strikingly similar ways, how do we detect it?

Their answer: look at intrinsic geometry, not extrinsic geometry.

Part 1: A Neural Network Is a Distorted Cloth

The Manifold Hypothesis

The manifold hypothesis states that although data may appear high-dimensional (e.g., a 1024×1024 image has millions of pixels), it actually lies on a low-dimensional "surface." Imagine a 2D sheet of cloth crumpled into 3D space—externally a mess, but intrinsically still a 2D plane, merely distorted.

A neural network essentially distorts this cloth: each layer twists, stretches, and folds it until entangled classes are separated enough for a linear classifier.

The Problem with Traditional Methods

Conventional comparison methods—CKA, CCA, RSA, Procrustes—compare the extrinsic shape of the distorted cloth. But the same intrinsic structure can be embedded into entirely different extrinsic spaces. A spiral can be rolled into a cylinder or laid flat; Procrustes says they are very different, yet intrinsically they are the same.

For neural networks this means: two networks may share the same "solution strategy" yet appear very different due to weight initialization, rotations, or dimensionality—and vice versa.

Part 2: The Pullback Metric—Capturing the Essence of the Distortion

Given an input manifold ℳ and network φ, the pullback metric works as follows: for nearby points p and q on the manifold with tangent vector v, the network maps the displacement to J·v (J is the Jacobian). The dot product in hidden space becomes:

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

The matrix G(p) = JᵀJ is the pullback metric. It tells us whether two directions on the original manifold become closer, farther, orthogonal, or parallel after distortion.

Key insight:

> It fully describes *how* the manifold is warped, independent of *where* the result sits in the embedding space.

Two networks with completely different weights have identical pullback metrics if they distort the input manifold in the same way.

Part 3: Spectral Ratio—Measuring the Distance Between Two Metrics

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

> G₁v = λG₂v

The λ values are independent of the coordinate system chosen to describe the manifold. The authors define the Spectral Ratio distance:

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

• If the metrics are identical, all λ = 1 and the distance is 0.
• If they differ drastically in some directions, the distance approaches 1.
• It is naturally bounded in [0, 1], directly interpretable as a similarity.

The authors prove it is a pseudometric, satisfying separation, symmetry, and the triangle inequality—endowing the space of SPD matrices with a genuine metric structure.

Part 4: MSA—A Metric on Riemannian Metric Spaces

MSA extends the comparison from single points to the entire input manifold:

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

In words: traverse every point on the input manifold, compute the spectral-ratio distance between the two networks' pullback metrics at that point, and average.

Core properties:

1. State-space rotation invariance: any orthogonal transformation of hidden layers leaves MSA unchanged. 2. Coordinate invariance: reparameterizing the input manifold leaves MSA unchanged.

MSA thus captures a network's intrinsic computational style rather than incidental implementation details.

Part 5: Experiments—What Can MSA See That Traditional Methods Cannot?

Experiment 1: Rich vs. Lazy Learning

Network training has two famous regimes:

• Rich regime: the network learns structured features of the data, producing compact, meaningful representations.
• Lazy regime: the network memorizes training samples and interpolates with random features in high-dimensional space.

Traditional methods: Procrustes and RSA report rich and lazy networks as very similar—their PCA projections do look alike, both spiraling the input space.

MSA: reports near-zero similarity. Rich networks learn smooth, structured warps capturing the data's essential structure; lazy networks apply essentially random distortions to squeeze training points into correct classes. Like a carefully designed Swiss roll versus a crumpled sheet of paper—both look like "a ball" from outside, but their intrinsic geometries differ completely.

Experiment 2: Comparing Dynamical Systems (RNN vs. SSM)

Task: a sequential working-memory task where the network receives two angle inputs and must recall and output them after a delay. Two architectures:

• RNN: fully connected recurrent network
• SSM: structured state-space model (with HiPPO initialization)

MSA findings:

1. MSA clearly distinguishes RNNs from SSMs even when solving the same task. 2. MSA identifies training status—trained vs. untrained models show clearly different similarity. 3. Geometric dynamics over time: RNNs maintain a relatively stable geometry throughout the delay, while SSM geometry changes over time.

Compared to RSA and DSA: RSA distinguishes architectures but is insensitive to training status; DSA is weak on both counts; MSA captures both geometric structure and dynamical evolution.

Experiment 3: Information Geometry of Diffusion Models

The most exciting part—applying MSA to Stable Diffusion XL. The setup defines a 2D manifold ℳ from bilinear interpolation of four text embeddings; each point π ∈ ℳ corresponds to a distribution over generated images.

Findings:

1. Information-geometry evolution during diffusion: early timesteps (near t=1, Gaussian noise) have markedly different geometry from late ones (near t=0, generated images). 2. Effect of guidance: for classifier-free guidance strength γ, MSA finds that beyond a certain threshold, the manifold geometry becomes *closer* again to the no-guidance case. This suggests an optimal guidance level where the information geometry deviates most from the unguided model—potentially a new angle for hyperparameter tuning.

Part 6: Deeper Reflections—Why Does This Matter?

For Interpretability

Current interpretability largely "dissects" networks—which neurons activate, which connections matter. MSA offers a geometric perspective: instead of examining parts, examine how the network warps the data manifold. Like understanding a person by studying neurons (anatomy) versus studying their thinking (geometry)—the latter may be closer to the essence of understanding.

Connection to Geometric Algebra Transformers

MSA and GATr (Geometric Algebra Transformer) share an intuition: neural computation is fundamentally geometric.

• GATr uses geometric algebra tools to build networks (rotors operating on vectors).
• MSA uses Riemannian geometry tools to analyze networks.

Limitations and Future Directions

The authors honestly note MSA's limitations:

1. Requires an explicit manifold representation: for data without a clear manifold (e.g., neuroscience experiments), the manifold may need to be learned first. 2. Ignores downstream use: MSA only examines hidden-layer geometry, not how later layers consume it—if the decoder is low-rank, some representational directions may be irrelevant. 3. Correlational, not causal: similarity metrics reveal associations, not causal relationships.

Future possibilities:

• Guiding network design: can manipulating geometry create more robust or efficient models?
• Combining causal inference: not just "are these networks similar?" but "what happens if we change this geometric feature?"
• Broader applications: reinforcement learning, graph neural networks, population coding in neuroscience.

Conclusion: The Geometric Eye

The paper's core message in one sentence:

> To compare two neural networks, don't ask *where they are*—ask *how they warp the world*.

It is a shift of perspective—from extrinsic to intrinsic, static to dynamic, discrete to continuous. We once treated networks as black boxes, then statistical models, then interpretable feature extractors. Now we are beginning to treat them as geometric objects—mappings that warp manifolds of data.

This may be one path toward truly understanding the nature of intelligence. After all, the history of physics teaches us: 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

Key points

• MSA compares neural representations via intrinsic geometry (pullback metrics JᵀJ) rather than extrinsic state-space coordinates.
• The spectral ratio d_SR(G₁,G₂) = 1 − √(λ_min/λ_max) on generalized eigenvalues is a coordinate-invariant pseudometric on SPD matrices.
• MSA is invariant to hidden-state rotations and input-manifold reparameterization.
• It separates rich vs. lazy learning regimes that Procrustes/RSA conflate, and distinguishes RNN vs. SSM architectures and training status better than RSA or DSA.
• Applied to Stable Diffusion XL, MSA reveals how information geometry evolves across diffusion timesteps and suggests an optimal classifier-free guidance strength.

*Note: the arXiv identifier above is quoted from the original post and has not been independently verified.*