Representational Alignment Across Model Layers and Brain Regions with Hierarchical Optimal Transport
This provides a more interpretable framework for comparing representations in AI and neuroscience, particularly for networks with different architectures, but it is incremental as it builds on optimal transport techniques.
The paper tackles the problem of aligning representations across neural network layers and brain regions, which standard methods handle poorly due to asymmetric results and depth mismatches, by proposing Hierarchical Optimal Transport (HOT) that matches or surpasses existing methods in alignment quality across vision models, language models, and brain data.
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.