Representational Alignment Across Model Layers and Brain Regions with Hierarchical Optimal Transport

TL;DR

Representational comparison across networks and brains is hampered by layer-wise, one-to-one alignments. HOT introduces a two-level, mass-conserving hierarchical OT framework that jointly infers inner neuron-to-neuron couplings and outer layer-to-layer couplings, yielding a single global alignment score and a soft, depth-agnostic transport plan with marginals $\sum_\ell P_{\ell m} = 1/M$ and $\sum_m P_{\ell m} = 1/L$ (and, in the inner level, $\sum_j Q_{ij} = 1/n_\ell$, $\sum_i Q_{ij} = 1/n'_m$). A rotation-invariant extension HOT+R further accounts for geometric transformations by optimizing orthogonal rotations $R_{\ell m}$ alongside transport plans. Across vision models, large language models, and human visual cortex data, HOT matches or surpasses greedy baselines, reveals coherent hierarchical correspondences (early-to-early, deep-to-deep), and shows how depth expands representations through soft, distributed layer mass. While computationally intensive, HOT provides richer, more interpretable mappings than previous methods and offers a principled path toward universal principles of representation in artificial and biological systems.

Abstract

Standard representational similarity methods align each layer of a network to its best match in another independently, producing asymmetric results, lacking a global alignment score, and struggling with networks of different depths. These limitations arise from ignoring global activation structure and restricting mappings to rigid one-to-one layer correspondences. We propose Hierarchical Optimal Transport (HOT), a unified framework that jointly infers soft, globally consistent layer-to-layer couplings and neuron-level transport plans. HOT allows source neurons to distribute mass across multiple target layers while minimizing total transport cost under marginal constraints. This yields both a single alignment score for the entire network comparison and a soft transport plan that naturally handles depth mismatches through mass distribution. We evaluate HOT on vision models, large language models, and human visual cortex recordings. Across all domains, HOT matches or surpasses standard pairwise matching in alignment quality. Moreover, it reveals smooth, fine-grained hierarchical correspondences: early layers map to early layers, deeper layers maintain relative positions, and depth mismatches are resolved by distributing representations across multiple layers. These structured patterns emerge naturally from global optimization without being imposed, yet are absent in greedy layer-wise methods. HOT thus enables richer, more interpretable comparisons between representations, particularly when networks differ in architecture or depth.

Figures (33)

• Figure 1: Left: Pairwise OT. Layers are matched independently, so multiple target layers can be mapped to the same source while other sources remain unused, yielding asymmetric, unbalanced mappings. Right: Hierarchical OT. HOT infers a globally consistent transport plan where each source layer distributes all its mass and each target layer receives exactly one unit, ensuring balanced, symmetric alignments that handle depth mismatches and reveal hierarchy.
• Figure 2: Transport plans for LLM alignment. Hierarchical OT (left) versus pairwise OT (right) for two cross-model comparisons: (a) Qwen-2.5 0.5B $\leftrightarrow$ LLaMA-3.2 3B and (b) LLaMA-3.2 1B $\leftrightarrow$ Qwen-2.5 3B. HOT reveals smooth, diagonal correspondences across layers, while pairwise OT produces noisier and less structured mappings.
• Figure 3: Transport plans for cross-subject brain alignment. Hierarchical OT (left) versus pairwise OT (right) for two randomly selected subject pairs: (a) Subject A $\leftrightarrow$ Subject C and (b) Subject B $\leftrightarrow$ Subject D. HOT recovers structured region-to-region correspondences that are absent in pairwise OT. Other subject pairs show similar trends (see Appendix \ref{['fig:fmris_all']}).
• Figure 4: Transport plans for vision model alignment. ViT-MAE Base $\leftrightarrow$ DINOv2 Giant (a) without rotation (HOT) and (b) with rotation augmentation (HOT+R). HOT+R captures geometric equivalences induced by rotations, yielding clearer correspondences than the rotation-sensitive variant.
• Figure A.1: Transport plans across LLM families and scales. Hierarchical OT (HOT) mappings are shown for six cross-model comparisons: (a) LLaMA-3.2 1B $\leftrightarrow$ LLaMA-3.2 3B, (b) Qwen-2.5 0.5B $\leftrightarrow$ Qwen-2.5 3B, (c) Qwen-2.5 0.5B $\leftrightarrow$ LLaMA-3.2 1B, (d) Qwen-2.5 0.5B $\leftrightarrow$ LLaMA-3.2 3B, (e) LLaMA-3.2 1B $\leftrightarrow$ Qwen-2.5 3B, and (f) LLaMA-3.2 3B $\leftrightarrow$ Qwen-2.5 3B. HOT uncovers structured, near-diagonal correspondences that persist across both intra-family (a,b) and cross-family (c-f) alignments, illustrating its robustness compared to pairwise OT.