Information Routing in Atomistic Foundation Models: How Task Alignment and Equivariance Shape Linear Disentanglement
Not on the map. Why?
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.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 12 of its 19 references that are in the map. First posted on arXiv 3 Mar 2026.
Builds on
12 of 19 references.
References · 12
- 2019Similarity of Neural Network Representations Revisitedcited by 1,695
- 2016Understanding intermediate layers using linear classifier probescited by 1,279
- 2018Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representationscited by 929
- 2019Designing and Interpreting Probes with Control Taskscited by 608
- 2017SchNet: A continuous-filter convolutional neural network for modeling quantum interactionscited by 569
- 2022MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fieldscited by 467
- 2018What you can cram into a single vector: Probing sentence embeddings for linguistic propertiescited by 452
- 2021Equivariant message passing for the prediction of tensorial properties and molecular spectracited by 438
- 2024The Platonic Representation Hypothesiscited by 362
- 2020Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Moleculescited by 254
- 2023LEACE: Perfect linear concept erasure in closed formcited by 254
- 2022The Design Space of E(3)-Equivariant Atom-Centered Interatomic Potentialscited by 100
Built on by
Not cited yet.