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

Joshua Steier

cs.LGMar 3, 2026

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

Authors: Joshua Steier

Organizations: Independent Researcher

Abstract

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\times$ in geometric information accessible after composition removal ( $R^2_{\mathrm{geom}}$ from 0.081 to 0.533 for HOMO-LUMO gap). Three factors explain this gradient. Task alignment dominates: models trained on HOMO-LUMO gap ( $R^2_{\mathrm{geom}}$ 0.44--0.53) outscore energy-trained models by $\sim$ 0.25 $R^2$ 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 ( $R^2 = 0.59$ vs.\ 0.38 in $L{=}0$ ), while $L{=}0$ (scalar) channels encode HOMO-LUMO gap ( $R^2 = 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 $R^2 = 0.68$ -- $0.95$ on a purely compositional target, and recommend linear probes for this setting.

Figures & tables

Figure 1: The CPD pipeline. Representations

\mathbf{X}

are extracted from a frozen model, then OLS projection (fold-wise) removes the linear composition component

\mathbf{Z}\hat{\boldsymbol{\beta}}

, producing a geometric residual

\mathbf{X}_{\mathrm{geom}}

. Ridge regression on the residual yields

R^{2}_{\mathrm{geom}}

(faithful). Gradient boosted trees on the same residual produce inflated scores by reconstructing the projected-out composition signal.

Model	Equiv.	TP	Params	Training Objective	Training Data
PaiNN	✓	×	1.17M	HL gap	QM9 (30 ep)
ViSNet	✓	×	1.67M	HL gap	QM9 (50 ep)
MACE-HL	✓	✓	0.73M	HL gap	QM9 (30 ep)
MACE pretrained	✓	✓	3.70M ∗	E + F	MPTraj (millions)
ANI-2x	×	×	13.7M	E	ANI-1x (diverse)
PaiNN-energy	✓	×	1.17M	E (U 0 )	QM9 (30 ep)

Table 1: Models evaluated. “Equiv.” indicates whether the architecture produces equivariant intermediate representations. “TP” indicates tensor product message passing. PaiNN and MACE each appear in both HOMO-LUMO and energy-trained variants, enabling two independent

2\times 2

factorial tests of task alignment.

Model	Equiv.	Target	$\boldsymbol{R^{2}_{\mathrm{full}}}$	$\boldsymbol{R^{2}_{\mathrm{geom}}}$	$\boldsymbol{R^{2}_{\mathrm{comp}}}$
PaiNN	✓	HL gap	0.949	0.533 $\pm$ 0.004	0.389
ViSNet	✓	HL gap	0.877	0.513 $\pm$ 0.010	0.389
MACE-HL	✓	HL gap	0.715	0.439 $\pm$ 0.005	0.389
MACE pretrained	✓	Energy	0.786	0.364 $\pm$ 0.007	0.389
ANI-2x	×	Energy	0.771	0.331 $\pm$ 0.006	0.389
PaiNN-energy	✓	Energy	0.656	0.310 $\pm$ 0.005	0.389

Table 2: Fold-wise CPD results on HOMO-LUMO gap. All values are mean

R^{2}\pm

standard deviation across 30 seeds

\times

5 folds. Models are ordered by

R^{2}_{\mathrm{geom}}

Figure 2: The linear accessibility gradient.

R^{2}_{\mathrm{geom}}

for HOMO-LUMO gap across ten models, color-coded by training regime. Three clusters emerge: HOMO-LUMO-trained (blue), energy-trained with diverse data (green), and energy-trained on QM9 only (orange/red). Inset: the

2\times 2

factorial crossing PaiNN’s architecture with training objective, showing the task alignment effect (

\Delta=0.223

) is

4.6\times

larger than the architecture effect (

\Delta=0.048

). Dashed line shows the composition

R^{2}

ceiling.

Model	HL Gap	Dipole	Polariz.	ZPVE
PaiNN	0.533	0.171	0.051	0.188
PaiNN-energy	0.310	0.201	0.081	0.248
ViSNet	0.513	0.141	0.056	0.190
MACE-HL	0.439	0.144	0.068	0.175
MACE pretrained	0.364	0.298	0.060	0.247
ANI-2x	0.331	0.266	0.063	0.247

Table 3:

R^{2}_{\mathrm{geom}}

across four QM9 properties (fold-wise CPD, 30-seed). PaiNN-HL and PaiNN-energy (PaiNN-E) show how task alignment reshapes geometric accessibility across properties: HOMO-LUMO gap drops by 0.223 when switching to energy training, while ZPVE increases by 0.060. Bottom row shows the composition

R^{2}

ceiling.

Channel	HL Gap $R^{2}$	Dipole $R^{2}$
Full representation	0.786	0.499
$L{=}0$ (scalar)	0.756	0.379
$L{=}1$ (vector norm)	0.337	0.586

Table 4: Information routing in MACE pretrained. Bold indicates the dominant channel for each property. MACE routes scalar properties through

L{=}0

and vector properties through

L{=}1

. Wilcoxon signed-rank test confirms the dipole preference for

L{=}1

(

p=1.86\times 10^{-9}

, 30 seeds).

Figure 3: Information routing by symmetry channel. Left: MACE routes scalar properties (HL gap) through $L{=}0$ channels and vector properties (dipole) through $L{=}1$ channels (black boxes highlight the dominant diagonal). Right: ViSNet concentrates all information in its scalar stream; vector channels carry $R^{2}=0.018$ .

Model	Ridge	GBT
MACE pretrained	$-$ 0.003	0.95
ViSNet	$-$ 0.003	0.89
ANI-2x	$-$ 0.002	0.93
SchNet	$-$ 0.003	0.85
DimeNet++	$-$ 0.003	0.68
MACE QM9	$-$ 0.003	0.71

Table 5: Probe comparison on average atomic mass (purely compositional). Ridge correctly returns

\approx 0

. GBTs produce inflated scores.

Figure 4: GBT inflation on a purely compositional target (average atomic mass). Ridge regression correctly returns

R^{2}\approx 0

on the geometric residual (blue). Gradient boosted trees recover

R^{2}=0.68

–

0.95

from the same residuals (red) by reconstructing the projected-out composition signal through nonlinear feature interactions.

Model	$\mathbf{X}_{\mathrm{geom}}$	$\mathbf{X}_{\mathrm{comp}}$
PaiNN	94.6% $\pm$ 0.1	52.5% $\pm$ 0.5
ViSNet	88.1% $\pm$ 0.4	52.5% $\pm$ 0.5
ANI-2x	87.6% $\pm$ 0.2	52.5% $\pm$ 0.5
MACE pretrained	84.3% $\pm$ 0.3	52.5% $\pm$ 0.5
DimeNet++	78.3% $\pm$ 0.5	52.5% $\pm$ 0.5
SchNet	78.1% $\pm$ 0.5	52.5% $\pm$ 0.5

Table 6: Structural isomer ordering accuracy (38,784 pairs, logistic regression, 5-fold CV).

\mathbf{X}_{\mathrm{comp}}

scores at chance by construction.

Figure 5: Structural isomer ordering accuracy (38,784 pairs). For each model, the geometric residual (colored bars) distinguishes isomers well above chance, while the compositional component (gray bars) scores exactly at 52.5% for every model, as required since isomers share identical formulas. PaiNN’s 94.6% means correct ordering 19 out of 20 times.

Figure 6: Sample efficiency of the linear accessibility gradient. The gradient emerges at

N=50

(Spearman

\rho=0.964

) and stabilizes by

N=500

(

\rho=1.0

). PaiNN at

N=50

(

R^{2}_{\mathrm{geom}}=0.296

) already exceeds SchNet at

N=2{,}000

(0.262), demonstrating that linearly disentangled representations are dramatically more sample-efficient.

Property	Baseline	MACE	SevenNet	CHGNet	$\boldsymbol{\Delta_{\mathrm{best}}}$
Band gap	0.530	0.723	0.669	0.666	+0.193
Formation E	0.789	0.986	0.953	0.926	+0.197
Is_metal (clf)	0.889	0.897	0.895	0.899	+0.010
Density	0.922	0.954	0.931	0.935	+0.032
Vol./atom	0.924	0.968	0.968	0.960	+0.044

Table 7: Materials Project crystals: best-layer

R^{2}

per model and property. Baseline is linear regression on composition features alone.

\Delta_{\mathrm{best}}

is the improvement of the best model over the composition baseline.

Robustness Check	Spearman $\rho$	Section
Fold-wise vs. global CPD	1.000	5.1
FWL partial $R^{2}$	0.943 ∗	5.2
Z1 vs. Z2 (raw counts)	1.000	5.3
Z1 vs. Z3 (fractions only)	1.000	5.3
Z1 vs. Z4 (binary presence)	1.000	5.3
CPD vs. LEACE (HL gap)	1.000	5.4

Table 8: Consolidated robustness summary. Spearman

\rho

is computed between the default model ranking and the ranking under each alternative. All checks preserve the gradient.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Fold-wise	Global	$\boldsymbol{\Delta}$
ViSNet	0.513	0.444	+0.069
MACE pretrained	0.364	0.295	+0.047
ANI-2x	0.331	0.276	+0.055
SchNet	0.262	0.238	+0.024
DimeNet++	0.240	0.176	+0.065
MACE QM9 97ep	0.101	0.090	+0.011

Appendix

Table 9: Fold-wise vs. global CPD on HOMO-LUMO gap (30-seed). Global QR over-removes geometric signal by projecting out a wider subspace informed by test-fold composition vectors.

Model	Standard $R^{2}_{\mathrm{geom}}$	FWL Partial $R^{2}$	$\boldsymbol{\Delta}$
ViSNet	0.444	0.826	+0.383
ANI-2x	0.276	0.691	+0.415
MACE pretrained	0.295	0.644	+0.349
SchNet	0.238	0.538	+0.300
DimeNet++	0.176	0.440	+0.264
MACE QM9 97ep	0.090	0.167	+0.076

Appendix

Table 10: Frisch-Waugh-Lovell partial

R^{2}

on HOMO-LUMO gap. FWL residualizes both

\mathbf{X}

and

y

, measuring the fraction of non-compositional target variance explained by the geometric residual.

Model	Z1 (default)	Z2 (raw)	Z3 (frac only)	Z4 (binary)
ViSNet	0.467	0.448	0.466	0.527
MACE pretrained	0.364	0.325	0.362	0.430
ANI-2x	0.331	0.310	0.331	0.402
SchNet	0.262	0.253	0.266	0.336
DimeNet++	0.240	0.211	0.238	0.292
MACE QM9 97ep	0.101	0.108	0.099	0.119

Appendix

Table 11: Composition feature sensitivity on HOMO-LUMO gap (

R^{2}_{\mathrm{geom}}

). Z1: element fractions + normalized atom count (default). Z2: raw element counts + raw atom count. Z3: element fractions only, no atom count. Z4: binary element presence + normalized atom count. Spearman

\rho=1.000

for all pairwise comparisons.

Model	CPD	LEACE	$\boldsymbol{\Delta}$
ViSNet	0.467	0.472	+0.005
MACE pretrained	0.364	0.362	$-$ 0.002
ANI-2x	0.331	0.319	$-$ 0.012
SchNet	0.262	0.280	+0.018
DimeNet++	0.240	0.243	+0.003
MACE QM9 97ep	0.101	0.068	$-$ 0.033

Appendix

Table 12: CPD vs. LEACE concept erasure on HOMO-LUMO gap (

R^{2}_{\mathrm{geom}}

). LEACE finds the theoretically optimal linear projection for removing composition. Mean

|\Delta|=0.012

. Spearman

\rho=1.000

Model	Geom $\rho$	Geom $p$	Comp $\rho$	Comp $p$
PaiNN	0.768	0.00	$-$ 0.004	0.385
ViSNet	0.458	0.00	$-$ 0.004	0.390
MACE pretrained	0.234	0.00	$-$ 0.004	0.403
ANI-2x	0.136	7.2e-160	$-$ 0.004	0.406
MACE QM9 97ep	0.138	6.4e-165	$-$ 0.004	0.392
DimeNet++	0.089	1.8e-68	0.001	0.802

Appendix

Table 13: Structural isomer distance correlations. Spearman correlation between pairwise

\mathbf{X}_{\mathrm{geom}}

distances and HOMO-LUMO gap differences for 38,784 isomer pairs. All

\mathbf{X}_{\mathrm{comp}}

correlations are near zero (

|\rho|<0.005

p>0.38

), confirming CPD isolates geometric from compositional information.

Model	Original dims	Original $R^{2}_{\mathrm{geom}}$
PaiNN	128	0.533
ViSNet	128	0.513
MACE pretrained	256	0.364
ANI-2x	1008	0.331
SchNet	128	0.262
DimeNet++	128	0.240

Appendix

Table 14: PCA dimensionality matching. All representations projected to 128 dimensions via PCA before applying CPD. The ranking is preserved, confirming the gradient reflects representation quality rather than dimensionality.

Data and representations
Symbol	Type	Definition
$n$	scalar	Number of molecules in the probe set
$d$	scalar	Dimensionality of model representation
$k$	scalar	Number of composition features (default 6)
$\mathbf{X}\in\mathbb{R}^{n\times d}$	matrix	Molecular representations from frozen model
$\mathbf{Z}\in\mathbb{R}^{n\times k}$	matrix	Composition feature matrix

Appendix

Table 15: Summary of notation.

Constant	Description	Value
Probe set construction
$n$	Number of QM9 molecules sampled	2,000
Molecule seed	np.random.RandomState(seed)	42
QM9 elements	Elements present in dataset	C, H, N, O, F
Cross-validation
$K$	Number of CV folds	5

Appendix

Table 16: Fixed experimental constants.

Model	Ridge	MLP	NL Gap	Category
PaiNN	0.534	0.833	+0.299	High linear + high NL
ViSNet	0.468	0.789	+0.321	High linear + high NL
MACE pretrained	0.363	0.753	+0.390	Moderate + high NL
ANI-2x (PCA)	0.327	0.607	+0.279	Moderate all
MACE QM9	0.100	0.439	+0.339	Low linear + moderate NL
SchNet	0.260	0.470	+0.210	Moderate all

Appendix

Table 17: Linear vs. nonlinear probing on

\mathbf{X}_{\mathrm{geom}}

for HOMO-LUMO gap. “NL Gap” is the difference between MLP and Ridge

R^{2}

, measuring how much geometric information is nonlinearly but not linearly accessible. ANI-2x was PCA-projected to 128 dimensions (99.94% variance retained) to prevent MLP divergence on 1008-dimensional input.

Model	Glob+noInt	Glob+Int	FW+noInt	FW+Int
MACE pre.	$-$ 27.89	0.295	$-$ 27.82	0.364
ViSNet	$-$ 27.77	0.444	$-$ 27.67	0.467
SchNet	$-$ 27.95	0.238	$-$ 27.92	0.262
ANI-2x	$-$ 27.94	0.276	$-$ 27.91	0.331

Appendix

Table 18: Intercept diagnostic on HOMO-LUMO gap. Without intercept, Ridge regression produces catastrophic

R^{2}\approx-28

because the residualized

\mathbf{X}_{\mathrm{geom}}

has near-zero mean and Ridge without intercept cannot model the target’s mean.

Explore similar work

Convex Compositional Reasoning Models
May 22, 2026Meir Roketlishvili, Semyon Semenov, Maksim Bobrin +7Compositional Reasoning Energy-Based Models
Mohamed bin Zayed University of Artificial Intelligence, Abu Dhabi, UAE · Applied AI Institute, Computational Imaging Lab · AXXX
GLACIER: A Multimodal Student-Teacher Foundation Model for Molecular Property Prediction
Jun 9, 2026Emily Nguyen, Yongchan Hong, Harsh Toshniwal +2Multimodal Pretraining Molecular Property Prediction
Department of Medical Biochemistry and Biophysics Science for Life Laboratory Karolinska Institutet Stockholm, Sweden · Amazon
QT-Net: Rethinking Evaluation of AI Models in Atomic Chemical Space
May 11, 2026Pablo Martínez Crespo, Stefano Ribes, Martin Rahm +6Graph Neural Networks Quantum Chemistry
Department of Computer Science and Engineering Chalmers University of Technology and University of Gothenburg · Department of Chemistry and Chemical Engineering Chalmers University of Technology · Technology Research Intel Corporation +3

May 22, 2026J/K move · Enter open · S save

Convex Compositional Reasoning Models

Meir Roketlishvili, Semyon Semenov, Maksim Bobrin, Viktor Kovalchuk+6

Mohamed bin Zayed University of Artificial Intelligence, Abu Dhabi, UAE · Applied AI Institute, Computational Imaging Lab · AXXX

Compositional Reasoning Energy-Based Models Nonconvex Optimization Input Convex Neural Networks Neural Combinatorial Optimization

Open paper AskPDF arXiv

Compositional energy-based models can generalize to larger combinatorial reasoning problems by reusing a learned factor energy across many local constraints. In our paper, we show that a key bottleneck in compositional reasoning is not composition itself, but the non-convex geometry of the learned energy landscape. To solve this problem, we introduce Convex Compositional Energy Minimization (CCEM), a framework that parameterizes each factor with an input-convex neural network and optimizes the composed energy over a tight convex relaxation of the feasible set. Because convexity is preserved under summation, the global relaxed objective remains convex, enabling deterministic projected first-order optimization. CCEM is trained in two stages: factor-level contrastive learning to shape local energy basins, followed by end-to-end refinement through an unrolled projected solver. Our experiments show that our models trained on small subproblems or a single problem size transfer to larger instances without retraining.

Abstract

Explore similar work

Convex Compositional Reasoning Models

GLACIER: A Multimodal Student-Teacher Foundation Model for Molecular Property Prediction

QT-Net: Rethinking Evaluation of AI Models in Atomic Chemical Space