🤖 machine learning

Representational Alignment Across Model Layers and Brain Regions with Multi-Level Optimal Transport

This paper introduces Multi-Level Optimal Transport (MOT), a unified framework that overcomes the limitations of standard layer-wise alignment by jointly inferring globally consistent, soft transport plans between neural network layers and brain regions, thereby enabling robust, interpretable comparisons across models of varying depths and architectures.

Original authors: Shaan Shah, Meenakshi Khosla

Published 2026-04-23

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Shaan Shah, Meenakshi Khosla

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to compare two different cities: City A (a small, cozy town) and City B (a massive, sprawling metropolis).

Your goal is to understand how similar their "neighborhoods" are. In the world of Artificial Intelligence (AI) and neuroscience, these "neighborhoods" are layers of a computer network or parts of the human brain, and the "people" living there are neurons.

The Old Way: The "Best Match" Game

For a long time, scientists compared these cities using a method called Greedy Pairwise Matching.

Imagine you are a tour guide in City A. You look at your first neighborhood (Layer 1) and ask, "Which neighborhood in City B looks most like me?" You find the best match, say, Neighborhood 5. You write that down.

Then you move to your second neighborhood (Layer 2) and ask the same question. You find its best match, maybe Neighborhood 6.

The Problem:

The "One-to-One" Trap: This method forces a strict rule: One neighborhood in City A must match exactly one in City B. But what if City A's "Layer 1" is actually a mix of features found in City B's "Layer 5, 6, and 7"? The old method can't see that. It forces a square peg into a round hole.
The "Selfish" Tour Guide: Because each neighborhood picks its own best match independently, you might end up with 10 neighborhoods from City A all claiming they match Neighborhood 5 in City B, while Neighborhoods 1 through 4 in City B are completely ignored. It's chaotic and unbalanced.
No Big Picture: You get a list of matches, but you don't get a single score that tells you, "Overall, how similar are these two cities?"

The New Way: Multi-Level Optimal Transport (MOT)

The authors of this paper propose a new framework called Multi-Level Optimal Transport (MOT). Think of this as a Master Logistics Planner who looks at the entire map of both cities at once.

Instead of asking, "Who is my best friend?" the planner asks, "How can we distribute the entire population of City A across City B so that everyone is happy, no one is left out, and the total travel cost is minimized?"

Here is how it works, using our city analogy:

1. The "Mass Distribution" (Soft Matching)

In the old method, a neighborhood had to pick one partner. In MOT, a neighborhood can split its "mass" (its importance or features).

Example: Neighborhood 1 in City A might send 40% of its "traffic" to Neighborhood 5 in City B, 30% to Neighborhood 6, and 30% to Neighborhood 7.
Why it helps: This perfectly handles the problem of different city sizes. If City A is small and City B is huge, City A's features can naturally spread out over several layers of City B without breaking the rules.

2. The "Global Harmony" (Consistency)

The planner ensures that the whole system is balanced.

Every neighborhood in City A must send out 100% of its traffic (nothing is lost).
Every neighborhood in City B must receive a fair share of traffic (nothing is ignored or overloaded).
The Result: You get a single, clear score for the whole city comparison, and a map showing exactly how the two cities relate to each other globally.

3. The "Rotation" Trick

Sometimes, two neighborhoods might be doing the exact same thing, but they are just "rotated" (like a map turned sideways). The old methods might say they are totally different because the coordinates don't match.
The authors added a special feature to MOT (called MOT+R) that acts like a compass. It realizes, "Oh, this neighborhood is just the same as that one, but turned 90 degrees." It rotates the map to align them perfectly, revealing hidden similarities.

What Did They Discover?

The researchers tested this on three very different things:

AI Models: Comparing different sizes of Large Language Models (like LLaMA and Qwen).
Brains: Comparing the visual cortex of different humans looking at the same pictures.
AI vs. Brains: Comparing computer vision models to human brains.

The Surprising Findings:

Natural Order: Without being told to do so, MOT naturally figured out that "Early layers match early layers, and deep layers match deep layers." It found the hidden hierarchy that the old methods missed.
Depth Differences: When comparing a small AI to a big AI, MOT showed that the small AI's single layer was actually "spread out" across three layers of the big AI. The old method just couldn't see this connection.
Better Scores: In almost every test, MOT gave a more accurate and meaningful comparison than the old "Best Match" method.

The Big Takeaway

Imagine you are trying to compare two different recipes for a cake.

The Old Way: You compare the flour in Recipe A to the flour in Recipe B, then the sugar to the sugar. If Recipe A has one bowl of flour and Recipe B has three bowls, you get confused.
The MOT Way: You look at the entire process. You realize that the single bowl of flour in Recipe A is actually distributed across the mixing, baking, and cooling stages of Recipe B. You see the whole story of how the ingredients flow through the system.

MOT allows scientists to finally compare AI and brains (and different AIs) in a way that respects their complexity, handles different sizes, and reveals the true, hidden structure of how they think. It turns a messy, one-to-one guessing game into a harmonious, global alignment.

1. Problem Statement

Current methods for comparing neural representations (both between artificial neural networks and between AI models and the human brain) rely on greedy, pairwise layer-wise matching. In this standard approach, each layer in a source network is independently matched to the single "best" layer in a target network based on a similarity metric (e.g., RSA, CKA, or linear predictivity).

The authors identify three fundamental limitations in this paradigm:

Rigid One-to-One Correspondence: It fails when networks have different depths ( $L \neq M$ ) or when features in one layer are distributed across multiple layers in another.
Asymmetry and Lack of Global Consistency: Matching is performed independently for each layer, ignoring the global activation structure. This leads to asymmetric results (matching $A \to B$ differs from $B \to A$ ) and no single unified score for the entire network alignment.
Noise Sensitivity: By optimizing matches locally, these methods risk overfitting to noise in individual layers and fail to capture hierarchical organizational principles (e.g., early layers mapping to early layers).

2. Methodology: Multi-Level Optimal Transport (MOT)

The paper proposes Multi-Level Optimal Transport (MOT), a unified framework that jointly infers soft, globally consistent couplings at two levels: neuron-to-neuron and layer-to-layer.

A. Core Framework

MOT operates by solving a hierarchical Optimal Transport (OT) problem:

Inner Level (Neuron-to-Neuron Transport):
- For every candidate pair of layers $(\ell, m)$ , MOT computes a soft coupling matrix $Q_{\ell m}$ .
- This is done by minimizing the transport cost between the tuning profiles of neurons in layer $\ell$ and layer $m$ (using a dissimilarity metric like correlation distance).
- Unlike hard permutation methods, $Q_{\ell m}$ allows neurons in one layer to map to a mixture of neurons in another, handling layers of unequal width.
Outer Level (Layer-to-Layer Transport):
- The inner costs form a layer-level cost matrix $C \in \mathbb{R}^{L \times M}$ .
- MOT solves a second OT problem to find a global layer coupling matrix $P \in \mathbb{R}^{L \times M}$ .
- Constraints: The problem is constrained by a transportation polytope where:
  - Each source layer distributes exactly 100% of its representational "mass" across target layers.
  - The total mass received by each target layer is balanced.
- Result: Instead of forcing a hard $1:1$ match, $P_{\ell m}$ represents the fraction of layer $\ell$ 's representation explained by layer $m$ . This naturally handles depth mismatches (e.g., one shallow layer distributing mass across three deep layers).
Reconstruction and Scoring:
- The alignment quality is evaluated by reconstructing source layer activations ( $\hat{X}_\ell$ ) using the target layers weighted by $P$ and $Q$ .
- The global score is the mean correlation between original and reconstructed activations on held-out data.

B. Rotation-Invariant Extension (MOT + R)

Since many neural representations (especially in Vision Transformers) are invariant to basis rotations, the authors introduce MOT + R.

This variant introduces orthogonal rotation matrices $R_{\ell m}$ for each layer pair.
It minimizes the reconstruction error $\|X_\ell R_{\ell m} - Y_m Q_{\ell m}^\top\|_F^2$ via alternating minimization (updating $Q$ and $R$ iteratively).
This ensures that geometric equivalences are captured even if the coordinate bases differ.

C. Three-Level MOT (Training Trajectories)

The framework is extended to a Three-Level MOT to align training trajectories. It adds a checkpoint-level OT problem on top of the layer-level problem, allowing for the alignment of two models across their entire training history, not just at a single snapshot.

3. Key Contributions

Theoretical Framework: A principled global alignment method that produces symmetric, globally consistent assignments with a single network-level score, avoiding the pitfalls of greedy local optimization.
Handling Depth Mismatches: By allowing soft, many-to-many mappings, MOT aligns networks of different depths without forcing inappropriate one-to-one correspondences.
Rotation Invariance: The MOT + R extension captures geometric similarities in rotated subspaces, crucial for Transformer architectures.
Emergent Hierarchical Structure: Without explicit constraints, MOT naturally recovers hierarchical organization (early-to-early, deep-to-deep) and reveals how deeper networks decompose computations that shallower networks compress into single layers.
Proof-of-Concept for Trajectories: Demonstrates the ability to align entire training trajectories, uncovering checkpoint-wise correspondences missed by standard methods.

4. Experimental Results

The authors evaluated MOT across three domains:

Large Language Models (LLMs):
- Setup: Compared LLaMA-3.2 and Qwen-2.5 models of varying sizes.
- Result: MOT achieved significantly higher reconstruction scores than pairwise baselines.
- Insight: MOT revealed smooth, diagonal transport plans (early layers $\to$ early layers). It showed that a single layer in a smaller model distributes its mass across multiple consecutive layers in a larger model, reflecting the refinement of features.
Human Visual Cortex (fMRI):
- Setup: Aligned fMRI responses from 4 subjects viewing natural scenes (V1–V4 regions).
- Result: MOT achieved reconstruction scores comparable to pairwise OT but with structured region-to-region correspondences (e.g., V1 aligning to V1) that pairwise methods failed to recover.
- Significance: Demonstrates that global alignment captures biological homology better than local matching.
Vision Transformers (ViTs):
- Setup: Compared DINOv2 and ViT-MAE models.
- Result: Vanilla MOT showed mixed results due to rotation sensitivity. However, MOT + R substantially outperformed all baselines (including pairwise + R), achieving much higher reconstruction accuracy and clear hierarchical alignment.
- Insight: Rotation invariance is critical for aligning Transformer representations.
Training Trajectories:
- Setup: Aligned Pythia-14M models across 154 checkpoints.
- Result: Three-level MOT revealed a coherent, diagonal correspondence between training checkpoints of two different seeds, whereas pairwise matching produced noisy, unstructured maps.

5. Significance and Conclusion

This paper addresses a critical gap in representational analysis: the inability of current methods to compare networks with different architectures or depths in a globally consistent manner.

Interpretability: MOT provides a "soft transport plan" that is more interpretable than a list of best-matched layers. It reveals how representations are distributed across depth, offering new insights into how depth affects feature learning.
Universality: The method works across artificial models (LLMs, ViTs) and biological data (fMRI), suggesting a universal principle of representational alignment.
Future Impact: By providing a rigorous measurement foundation for alignment, MOT enables future research into why different systems converge on similar representations (e.g., testing the "Platonic Representation Hypothesis" or "Contravariance Principle").

In summary, MOT replaces the greedy, local view of network comparison with a global, optimal transport perspective, enabling richer, more robust, and biologically plausible comparisons of neural systems.