원자 수준 기반 모델에서의 정보 라우팅: 등변성이 어떻게 선형적으로 분리된 표현을 생성하는가

Information Routing in Atomistic Foundation Models: How Equivariance Creates Linearly Disentangled Representations

원자 수준 기반 모델은 중간 표현에서 어떤 정보를 담고 있으며, 이러한 정보는 어떻게 구성되어 있는가? 본 연구에서는 Composition Projection Decomposition (CPD)라는 방법을 소개하며, 이는 QR 투영을 사용하여 학습된 표현에서 합성 신호를 선형적으로 제거하고, 기하학적 잔류량을 분석한다. QM9 분자와 Materials Project의 결정 구조에 대한 5가지 아키텍처 패밀리의 8개 모델을 분석한 결과, 분리도(disentanglement)의 경향성이 나타났다. 텐서 곱 등변성 아키텍처(MACE)는 합성 신호를 제거한 후 기하학적 정보가 거의 완전히 선형적으로 접근 가능하다는 것을 보여주었다 ($R^2_{ ext{geom}} = 0.782$, HOMO-LUMO gap 기준). 반면, 수작업으로 설계된 설명자(ANI-2x)는 동일한 정보를 비선형적으로 얽히게 만든다 ($R^2_{ ext{geom}} = -0.792$, Ridge 회귀 기준; $R^2 = +0.784$, MLP 기준). MACE는 목표에 특화된 신호를 불변 표현 채널을 통해 전달한다. 예를 들어, 쌍극자 모멘트는 $L = 1$ 채널을 통해, HOMO-LUMO gap은 $L = 0$ 채널을 통해 전달되며, 이는 ViSNet의 벡터-스칼라 아키텍처에서는 관찰되지 않는 패턴이다. 투영된 잔류량에 대한 그래디언트 부스팅 트리 기반 분석은 체계적으로 과장된 결과를 보여주며, 순수 합성 목표에 대해 $R^2 = 0.68$에서 $0.95$ 사이의 값을 나타낸다. 따라서, 본 연구에서는 선형적인 분석 방법을 주요 지표로 권장한다. 선형적으로 분리된 표현은 선형 분석 시 더 효율적인 샘플링을 가능하게 하며, 이는 예측 정확도 외에도 등변성 아키텍처의 실용적인 장점을 시사한다.

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.