# Information Routing in Atomistic Foundation Models: How Task Alignment and Equivariance Shape Linear Disentanglement

arXiv: https://arxiv.org/abs/2603.03155 · HTML: https://mlsanity.com/paper/2603.03155/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Interpreting biomedical and science models, one of 13 directions in Mechanistic interpretability.
- Builds on 12 of its 19 references that are in the map.
- First posted on arXiv 3 Mar 2026.

## Placed in

- Mechanistic interpretability / [Interpreting biomedical and science models](https://mlsanity.com/mechinterp/interpreting-biomedical-and-science-models/)

## Builds on (most cited first, 12 of 12)

- [Similarity of Neural Network Representations Revisited](https://mlsanity.com/paper/1905.00414/) (2019)
- [Understanding intermediate layers using linear classifier probes](https://mlsanity.com/paper/1610.01644/) (2016)
- [Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations](https://mlsanity.com/paper/1811.12359/) (2018)
- [Designing and Interpreting Probes with Control Tasks](https://mlsanity.com/paper/1909.03368/) (2019)
- [SchNet: A continuous-filter convolutional neural network for modeling quantum interactions](https://mlsanity.com/paper/1706.08566/) (2017)
- [MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields](https://mlsanity.com/paper/2206.07697/) (2022)
- [What you can cram into a single vector: Probing sentence embeddings for linguistic properties](https://mlsanity.com/paper/1805.01070/) (2018)
- [Equivariant message passing for the prediction of tensorial properties and molecular spectra](https://mlsanity.com/paper/2102.03150/) (2021)
- [The Platonic Representation Hypothesis](https://mlsanity.com/paper/2405.07987/) (2024)
- [LEACE: Perfect linear concept erasure in closed form](https://mlsanity.com/paper/2306.03819/) (2023)
- [Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules](https://mlsanity.com/paper/2011.14115/) (2020)
- [The Design Space of E(3)-Equivariant Atom-Centered Interatomic Potentials](https://mlsanity.com/paper/2205.06643/) (2022)

## Abstract (arXiv)

What determines whether a molecular property prediction model organizes its representations so that geometric and compositional information can be cleanly separated? We introduce Compositional Probe Decomposition (CPD), which linearly projects out composition signal and measures how much geometric information remains accessible to a Ridge probe. We validate CPD with four independent checks, including a structural isomer benchmark where compositional projections score at chance while geometric residuals reach 94.6% pairwise classification accuracy. Across ten models from five architectural families on QM9, we find a linear accessibility gradient: models differ by 6.6× in geometric information accessible after composition removal (R2geom from 0.081 to 0.533 for HOMO-LUMO gap). Three factors explain this gradient. Task alignment dominates: models trained on HOMO-LUMO gap (R2geom 0.44--0.53) outscore energy-trained models by ~0.25 R2 regardless of architecture. Within-architecture ablations on two independent architectures confirm this: PaiNN drops from 0.53 to 0.31 when retrained on energy, and MACE drops from 0.44 to 0.08. Data diversity partially compensates for misaligned objectives, with MACE pretrained on MPTraj (0.36) outperforming QM9-only energy models. Inside MACE's representations, information routes by symmetry type: L=1 (vector) channels preferentially encode dipole moment (R2 = 0.59 vs. 0.38 in L=0), while L=0 (scalar) channels encode HOMO-LUMO gap (R2 = 0.76 vs. 0.34 in L=1). This pattern is absent in ViSNet. We also show that nonlinear probes produce misleading results on residualized representations, recovering R2 = 0.68--0.95 on a purely compositional target, and recommend linear probes for this setting.