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

📄 arXiv: 2603.03155v1 📥 PDF

作者: Joshua Steier

分类: cs.LG, cs.AI, physics.chem-ph

发布日期: 2026-03-03

💡 一句话要点

提出CPD方法，揭示原子级模型中等变性如何解耦线性表示，提升模型性能。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 原子级模型 等变性 表示解耦 线性探测 组合投影分解

📋 核心要点

  1. 现有原子级模型中间表示的编码方式和信息组织方式尚不明确，阻碍了模型的可解释性和优化。
  2. 提出CPD方法，通过QR投影线性移除组合信号，并分析几何残差，从而揭示模型表示的解耦程度。
  3. 实验表明，等变架构(MACE)能产生线性解耦的表示，提升了样本效率，优于手工设计的描述符(ANI-2x)。

📝 摘要（中文）

本文提出了一种名为组合投影分解(CPD)的方法，该方法利用QR投影从学习到的表示中线性移除组合信号，并探测几何残差。通过对QM9分子和Materials Project晶体上五个架构系列的八个模型进行分析，发现了一个解耦梯度：张量积等变架构(MACE)产生的表示在移除组合信息后，几何信息几乎完全线性可访问（HOMO-LUMO能隙的R^2_geom = 0.782），而手工设计的描述符(ANI-2x)则以非线性的方式纠缠相同的信息（Ridge回归下的R^2_geom = -0.792；MLP下的R^2 = +0.784）。MACE将特定于目标的信号通过不可约表示通道进行路由——偶极子到L=1，HOMO-LUMO能隙到L=0——在ViSNet的矢量-标量架构中未观察到这种模式。研究表明，投影残差上的梯度提升树探针被系统性地膨胀，在一个纯粹的组合目标上恢复了R^2 = 0.68--0.95，并建议线性探针作为主要指标。线性解耦的表示在线性探测下更具样本效率，这表明等变架构除了原始预测精度之外，还具有实际优势。

🔬 方法详解

问题定义：原子级模型在中间层如何编码原子信息，特别是几何信息和组合信息之间的关系？现有方法难以有效解耦这些信息，导致模型的可解释性差，且可能影响模型的泛化能力和样本效率。现有方法难以区分模型学习到的表示中哪些是真正有用的几何信息，哪些是仅仅依赖于原子组成的信息。

核心思路：核心思想是通过线性投影的方式，将原子表示中的组合信息（例如原子类型和数量）移除，然后分析剩余的几何信息。如果模型能够很好地解耦组合信息和几何信息，那么在移除组合信息后，剩余的表示应该仍然包含丰富的几何信息，并且可以通过简单的线性模型进行预测。

技术框架：论文提出的CPD方法包含以下几个步骤：1) 使用原子级模型提取原子结构的中间层表示。2) 使用QR分解，将表示投影到一个与组合信息正交的子空间。3) 使用线性模型（例如Ridge回归）或非线性模型（例如MLP）来预测目标属性（例如HOMO-LUMO能隙）。4) 通过比较原始表示和投影后的表示的预测性能，评估模型的解耦程度。

关键创新：最重要的创新点在于提出了CPD方法，能够定量地评估原子级模型中组合信息和几何信息的解耦程度。通过分析不同架构的模型，发现等变架构能够更好地解耦这些信息，从而提升模型的样本效率。与现有方法相比，CPD方法提供了一种更直接、更可解释的方式来理解原子级模型的内部表示。

关键设计：CPD方法的关键设计包括：1) 使用QR分解进行投影，保证移除的是线性相关的组合信息。2) 使用多种线性/非线性模型进行预测，评估解耦后的表示的线性可分性。3) 针对不同的目标属性，分析模型在不同不可约表示通道上的信息路由情况。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，MACE等张量积等变架构能够产生线性解耦的表示，在移除组合信息后，几何信息几乎完全线性可访问（HOMO-LUMO能隙的R^2_geom = 0.782）。相比之下，手工设计的描述符（ANI-2x）则以非线性的方式纠缠相同的信息（Ridge回归下的R^2_geom = -0.792）。线性解耦的表示在线性探测下更具样本效率，表明等变架构具有实际优势。

🎯 应用场景

该研究成果可应用于原子级建模领域，帮助设计更高效、更可解释的原子级模型。通过理解模型如何编码原子信息，可以更好地优化模型架构和训练策略，提升模型在材料发现、药物设计等领域的性能。此外，该方法还可以用于评估不同模型的表示学习能力，为模型选择提供依据。

📄 摘要（原文）

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.