What do atomistic foundation models encode in their intermediate representations, and how is that information organized? We introduce Composition Projection Decomposition (CPD), which uses QR projection to linearly remove composition signal from learned representations and probes the geometric residual. Across eight models from five architectural families on QM9 molecules and Materials Project crystals, we find a disentanglement gradient: tensor product equivariant architectures (MACE) produce representations where geometry is almost fully linearly accessible after composition removal ($R^2_{\text{geom}} = 0.782$ for HOMO-LUMO gap), while handcrafted descriptors (ANI-2x) entangle the same information nonlinearly ($R^2_{\text{geom}} = -0.792$ under Ridge; $R^2 = +0.784$ under MLP). MACE routes target-specific signal through irreducible representation channels -- dipole to $L = 1$, HOMO-LUMO gap to $L = 0$ -- a pattern not observed in ViSNet's vector-scalar architecture under the same probe. We show that gradient boosted tree probes on projected residuals are systematically inflated, recovering $R^2 = 0.68$--$0.95$ on a purely compositional target, and recommend linear probes as the primary metric. Linearly disentangled representations are more sample-efficient under linear probing, suggesting a practical advantage for equivariant architectures beyond raw prediction accuracy.
Alex: Welcome to another episode of ResearchPod. Sam, today we're diving into a paper on how AI models handle molecules—what's the key question here?
Sam: The paper, titled 'Information Routing in Atomistic Foundation Models: How Equivariance Creates Linearly Disentangled Representations' by Joshua Steier, looks at AI systems trained to predict properties of molecules and materials, like their energy levels or electronic structures. These atomistic foundation models take the positions of atoms as input and make accurate predictions. But the puzzle is whether their internal workings cleanly separate two things: the types and numbers of atoms—what researchers call composition—from how those atoms are arranged in 3D space, or geometry. Without that separation, or disentanglement, it's hard to tell if a prediction relies on the right factors.
Alex: So this paper is basically asking if these models mix up 'what atoms are there' with 'how they're shaped,' and whether that matters for using them reliably, say in drug design?
Sam: Yes, exactly—that confusion can lead models to predict properties based partly on atom types alone, masking whether the 3D shape truly drives the outcome. The study introduces a tool called Composition Projection Decomposition, or CPD, to test this. It mathematically strips away the composition signal from the model's internal representations, then checks what's left for pure geometric information. Across several models, they find some architectures keep these factors more separated than others.
Alex: Huh. So if a model predicts something like a molecule's energy gap well, it might still be cheating by leaning on atom counts instead of shape?
Sam: Precisely. They probe properties like the HOMO-LUMO gap—the energy difference between a molecule's highest occupied and lowest unoccupied orbitals, key for electronics and chemistry. In top-performing models using tensor product equivariant architectures like MACE, the geometric signal remains strongly accessible even after removing composition. This suggests equivariance—a property where the model respects 3D rotations—helps route information to specific internal channels.
Alex: Right, but why does that separation matter practically? Doesn't high accuracy already solve drug discovery problems?
Sam: Accuracy alone doesn't reveal the why behind a prediction, which is crucial for trust in fields like drug design—where you need to know if tweaking a molecule's shape will change its properties without composition side effects. The paper shows this disentanglement makes representations more efficient for new tasks with limited data.
Alex: So these equivariant models like MACE seem to handle the separation better. But how exactly do they keep geometry from mixing with composition inside?
Sam: Equivariant models build their internal features in layers that respect a molecule's 3D rotations—like keeping track of scalars that don't change when you spin the molecule, and vectors that point in specific directions. They split these into channels based on angular momentum, sort of like different radio frequencies carrying distinct signals: simple unchanging numbers go to one channel called L=0, while direction-sensitive info like arrows goes to L=1. Researchers label these building blocks irreducible representations, or irreps for short.
Alex: Okay, so channels like L=0 for steady properties, L=1 for ones with direction. Does the paper show certain properties stick to the right channels?
Sam: Yes—in MACE, they probe each channel separately after applying CPD to remove composition. For the dipole moment, which measures how charges separate like tiny magnets pointing one way, the signal routes strongly to the L=1 vector channels. The HOMO-LUMO gap, a scalar property without direction, stays in L=0. In MACE, the geometric part for HOMO-LUMO predicts nearly as well as the full representation—suggesting clean separation. Other models drop off more sharply.
Alex: Huh. And does that geometry signal hold up across models?
Alex: So MACE keeps geometry linearly accessible... But do all these equivariant models handle that separation equally well?
Sam: The paper sorts models into tiers based on how well the geometric signal survives CPD. Tensor product ones like MACE lead, with intermediate results for vector-scalar types like ViSNet or SchNet, and lower for others like DimeNet++ or ANI-2x. They check layer by layer on CPD residuals for HOMO-LUMO. MACE begins high and positive, ending strong. DimeNet++ starts deep negative—like the probe predicts worse than just guessing the average—then climbs slowly. These profiles act like fingerprints of how each builds disentanglement.
Alex: Huh, so MACE doesn't just end better—it starts from a stronger place... For models at the bottom like ANI-2x, is the geometry info really gone?
Sam: No—the info's there, just not linearly accessible. A simple nonlinear probe, like a two-layer network that can curve its fits, flips ANI-2x from negative to matching MACE's linear scores on HOMO-LUMO. The paper suggests the signal exists but encoded tangled, unlike MACE's cleaner paths.
Alex: Right... same underlying info, but one model's internals make it easier to pull out purely geometric effects without fancy probes. That distinction seems key for practical use.
Alex: Does the paper dig into how different equivariant architectures—like MACE versus ViSNet—handle that property-specific routing to L=0 or L=1 channels?
Sam: It does. In MACE, after summarizing the direction-sensitive L=1 channels by their lengths—making them rotation-proof like measuring arrow sizes without caring about point direction—the dipole magnitude signal stands out there, about one-and-a-half times stronger than in the steady L=0 channels. For the HOMO-LUMO gap, it's the reverse: L=0 carries the main load. ViSNet puts almost all usable signal in scalars for both properties—its vector channels add next to nothing. The paper finds this routing concentrated early in MACE, then shifting toward rotation-proof summaries in later layers.
Alex: Huh... so even among equivariant models, MACE routes dipole to vectors more cleanly... Does this hold beyond small molecules, say on crystal materials?
Sam: Yes, they test Materials Project crystals. MACE and ORB—a crystal specialist—match closely on CPD residuals for band gap and formation energy, both showing notable predictive power purely from geometry. Yet overall similarity metrics show their internals diverging deeper in, hinting matched endpoints from different paths.
Alex: Notable consistency on the geometry check across molecules and crystals.
Alex: But how reliable are these linear probes on the CPD leftovers? Do fancier ones change the picture?
Sam: Fancier probes can mislead. Imagine a tool that fits data by building a forest of decision trees, each splitting on patterns like branches—researchers call these gradient boosted trees, or GBT. On a test case that's purely about atom types, like average atomic mass with no shape involved, CPD should wipe out all signal, leaving probes scoring near zero. Linear ones do, but GBT falsely claims strong geometry scores. QR projection in CPD clears straight-line composition effects, but curved interactions linger—like echoes of atom counts in mixed features. Flexible GBT grabs those, inflating claims; linear probes stay honest.
Alex: Huh... so even after stripping composition, these tree probes sneak in leftover clues?
Alex: Does this probe difference show up in how models pack their info overall?
Sam: They use a fairness check called SHAP values, which measures how much each internal feature pushes predictions—like crediting players for a team's score. ANI-2x crams half its importance into just over 1% of dimensions—super tight—but that stays tangled nonlinearly after CPD. MACE spreads wider, yet linearly pulls geometry cleanly.
Alex: Right... divergent paths, same endpoint on pure shape in some cases... That's a grounded view of why architecture choices matter.
Alex: Does the paper look at how different properties use separate directions inside MACE's representations?
Sam: It does—they train simple linear predictors for multiple properties on MACE's full outputs and measure overlap by angle between those fits. Different properties point in mostly separate directions: no two align more than halfway except polarizability and zero-point vibrational energy, which push oppositely. This matches the channel routing we discussed.
Alex: Huh... so properties don't crowd the same space linearly... Does that separation help when you have very little training data?
Sam: Yes—they test by fitting those linear predictors with tiny datasets on full outputs and CPD leftovers for HOMO-LUMO. MACE reaches useful predictions quickest, giving a head start for scarce-label scenarios like new materials screening.
Alex: A practical edge on small data... But the paper flags some limits to keep this in perspective?
Sam: Several, yes. Tests use averaged-over-atoms summaries, not per-atom features—pooling can hide or create issues. Molecules are tiny, up to nine heavy atoms; bigger ones like proteins might shift the balance. MACE's edge mixes its design with larger pretraining data others lacked, and key models like EquiformerV2 weren't tested.
Alex: Fair cautions... So where does this leave practical next steps?
Sam: Two stand out: linearly separated geometry aids quick learning from few labels, favoring these models for discovery tasks. Overall, the work maps a spectrum of separation across architectures, urging linear checks over flexible ones that overclaim.
Alex: That's a solid diagnostic—clearer internals without the hype. Thanks, Sam, for breaking it down.
Sam: My pleasure, Alex. This advances how we inspect these models reliably. Thanks for listening to ResearchPod.