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

Each language version is independently generated for its own context, not a direct translation.

这篇论文就像是在给分子 AI 模型做了一次“深度体检”，目的是搞清楚：这些模型脑子里到底是怎么思考的？它们是把“分子由什么组成”和“分子长什么样”分得很清楚，还是混在一起乱成一团？

为了让你轻松理解，我们可以把分子想象成乐高积木，把 AI 模型想象成乐高大师。

1. 核心问题：大师是怎么看乐高的？

想象一下，你给 AI 看两个乐高模型：

模型 A：由 5 块红色积木和 3 块蓝色积木搭成的房子。
模型 B：由 5 块红色积木和 3 块蓝色积木搭成的船。

这两个模型的**成分（Composition）完全一样（都是 5 红 3 蓝），但形状（Geometry）**完全不同。

好的大师（理想模型）：能一眼看出“哦，虽然材料一样，但一个是房子，一个是船”。它的脑子里，关于“材料”的信息和关于“形状”的信息是分开的。
普通的大师（现实中的很多模型）：可能觉得“红色积木多，所以这个模型能量高”，它把材料和形状混在一起了。如果你问它关于形状的问题，它得先把材料的信息从脑子里“过滤”掉才能回答，这很费劲。

这篇论文就是想知道：哪些 AI 大师是“分得清”的？为什么有的分得清，有的分不清？

2. 新工具：CPD（成分剥离术）

以前的方法有个大漏洞：如果你直接问 AI“这个分子能量是多少”，AI 可能会说：“因为它是碳氢化合物，所以能量高。”它其实是在背成分表，而不是在分析形状。

作者发明了一个叫 CPD（成分探针分解） 的新工具，就像是一个**“成分过滤器”**：

第一步（过滤）：强行把 AI 脑子里关于“成分”（比如碳、氢、氧的比例）的信息全部抽走，就像把乐高积木的颜色标签撕掉，只留下形状。
第二步（测试）：剩下的信息里，还能不能看出分子的形状？

关键发现 1：非线性探针是个“骗子”
作者发现，如果用复杂的“非线性”工具（比如梯度提升树，GBT）去测试剩下的信息，它会作弊！即使你把成分信息抽走了，它也能通过复杂的曲线把成分信息“猜”回来，假装自己看懂了形状。

比喻：就像你让一个人背对着黑板，把黑板上的字擦掉（成分），然后问“黑板上原来写了什么”。如果让他用复杂的逻辑推理（非线性），他可能猜出“既然你擦掉了，那肯定写了字”，从而得高分。
结论：作者坚持只用简单的“线性探针”（Ridge 回归），因为它不会作弊，能真实反映 AI 脑子里到底剩下了多少关于形状的信息。

3. 三大发现：什么决定了大师的水平？

作者测试了 10 种不同的 AI 模型，发现它们对形状信息的“提取能力”差距巨大（有的模型能提取出 53% 的信息，有的只有 8%）。是什么导致了这种差距？

因素一：训练目标决定一切（Task Alignment）—— 最重要的因素

这是论文最惊人的发现。

比喻：
- 大师 A：专门训练去预测“房子的形状”（HOMO-LUMO 能隙，一种对形状很敏感的性质）。
- 大师 B：专门训练去预测“房子的总重量”（能量，主要由成分决定，形状影响很小）。
结果：即使大师 A 和大师 B 用的是同样的乐高积木（同样的架构），大师 A 的脑子里形状信息非常清晰，而大师 B 的脑子里形状信息模糊不清。
启示：如果你想要一个能理解分子形状的 AI，不要只看它架构多高级，要看它以前是学什么长大的。 让它学形状相关的知识，它才会把形状信息整理好。

因素二：数据多样性是“替补队员”

比喻：如果大师 B（学重量的）见过成千上万种不同形状的乐高（在大规模数据集 MPTraj 上预训练），它的水平会比只见过几种乐高（只在 QM9 数据集训练）的大师 B 好很多。
结论：虽然“学什么”最重要，但“见过多少”也能起到补救作用。见过世面多的模型，即使没专门学过形状，也能稍微懂一点。

因素三：对称性架构（Equivariance）是“锦上添花”

比喻：有些模型（如 MACE）天生就懂得“旋转不变性”（不管乐高怎么转，它都知道是同一个东西）。
结论：这种天赋很有用，但不是万能的。如果训练目标不对（比如只让它学重量），就算它有天赋，也学不好形状。只有“天赋 + 正确的训练目标”结合，才能达到最高水平。

4. 有趣的细节：MACE 模型的“内部交通”

作者还发现，MACE 这个模型内部像是一个分工明确的工厂：

L=0 通道（标量通道）：专门负责处理“数值大小”的信息（比如能隙）。
L=1 通道（向量通道）：专门负责处理“方向”的信息（比如偶极矩，像指南针一样有方向）。
比喻：就像工厂里，有的传送带专门运箱子（标量），有的专门运箭头（向量）。MACE 把不同类型的信息送到了不同的传送带上，互不干扰。
对比：另一个模型 ViSNet 虽然也有传送带，但最后所有信息都挤在“箱子传送带”上，箭头信息丢了。这说明 MACE 的“内部交通设计”更优秀。

5. 总结：这对我们意味着什么？

选模型看“出身”：如果你想用 AI 预测分子的电子性质（这很依赖形状），一定要选那些专门学过这类性质的模型，哪怕它的架构看起来简单点。不要盲目迷信“预训练大模型”，如果它只学过能量，它可能并不懂形状。
别被复杂工具骗了：在分析 AI 内部时，简单的线性测试往往比复杂的非线性测试更诚实。
形状信息很珍贵：那些能把“成分”和“形状”分得清清楚楚的模型，在需要少量数据就能学会新任务（Few-shot learning）时，表现会好得多。

一句话总结：
这篇论文告诉我们，AI 模型能不能“看清”分子的形状，主要取决于它以前“学过什么”（训练目标），而不是它“长得有多复杂”（架构）。 就像教孩子认字，如果你只让他背字典（成分），他可能认不出字怎么写（形状）；但如果你让他读故事（形状敏感的任务），他自然就能把字和形状对应起来。

Each language version is independently generated for its own context, not a direct translation.

1. 研究背景与核心问题 (Problem)

在分子性质预测领域，原子基础模型（如 MACE, SchNet, PaiNN 等）已能高精度预测能量、力和电子性质。然而，一个关键问题尚未解决：这些模型的中间表示（representations）是如何组织的？

具体而言，模型是否将分子的组成信息（Composition，即包含哪些元素及其比例）与几何信息（Geometry，即原子在空间中的排列）进行了清晰的分离（解耦）？

挑战：分子性质通常同时依赖于组成和几何，且两者高度相关。传统的探针（Probing）方法直接对原始表示进行训练，无法区分模型是真正学到了几何结构，还是仅仅利用了组成信息的捷径（Shortcut）。
现有方法的缺陷：作者发现，使用非线性探针（如梯度提升树 GBT）在去除组成信号后的残差上进行测试时，会出现严重的分数膨胀现象。GBT 能够利用高维残差中的非线性交互重建被投影掉的组成信号，导致在纯组成目标上也能获得极高的 $R^2$ 分数，从而产生误导性结论。

2. 方法论：组成探针分解 (CPD)

为了解决上述问题，作者提出了**组成探针分解（Compositional Probe Decomposition, CPD）**方法，旨在线性地分离组成信号并测量剩余几何信息的可访问性。

核心步骤：

定义组成特征 ( $Z$ )：构建包含元素分数（C, H, N, O, F）和标准化原子数的向量。
线性投影去除 ( $X_{geom}$ )：
- 在交叉验证的每个折叠内，使用普通最小二乘法（OLS）拟合表示矩阵 $X$ 对组成特征 $Z$ 的回归。
- 计算残差： $X_{geom} = X - Z\hat{\beta}$ 。
- 该残差 $X_{geom}$ 代表了与组成线性正交的部分（即几何、拓扑及非线性相互作用信息）。
- 关键点：采用**折叠内（Fold-wise）**投影以防止测试集信息泄露。
线性探针测量：
- 在残差 $X_{geom}$ 上使用**岭回归（Ridge Regression）**预测目标性质（如 HOMO-LUMO 能隙）。
- 记录 $R^2_{geom}$ 作为几何信息线性可访问性的度量。
- 严格限制：仅使用线性探针。作者证明非线性探针（如 GBT）会错误地重建被去除的组成信号，导致虚假的高分。

验证机制：

结构异构体基准（Structural Isomer Benchmark）：利用具有相同化学式但结构不同的异构体对。由于组成完全相同，组成探针的预测准确率应接近随机（50%），而几何残差应能区分异构体。
鲁棒性检查：包括与 LEACE（概念擦除）对比、不同组成特征定义、PCA 降维匹配等 12 项检查，确保梯度排序的稳定性。

3. 主要发现与结果 (Key Results)

作者对 5 种架构家族共 10 个模型（基于 QM9 数据集）进行了评估，发现了一个显著的线性可访问性梯度（Linear Accessibility Gradient）：在去除组成信号后，不同模型保留的几何信息量差异高达 6.6 倍（ $R^2_{geom}$ 从 0.081 到 0.533）。

这一梯度由三个相互作用的因素决定，其重要性排序如下：

因素 1：任务对齐（Task Alignment）占主导地位

发现：训练目标与探测目标的一致性对几何信息的可访问性影响最大。
数据：在 HOMO-LUMO 能隙任务上，专门训练该任务的模型（PaiNN, ViSNet, MACE-HL）的 $R^2_{geom}$ 在 0.44–0.53 之间；而仅训练能量（Energy）的模型，即使架构相同，分数也大幅下降（约低 0.25 $R^2$ ）。
消融实验：
- PaiNN：从 HOMO-LUMO 训练改为能量训练， $R^2_{geom}$ 从 0.533 降至 0.310 ( $\Delta = 0.223$ )。
- MACE：从 HOMO-LUMO 训练改为能量训练， $R^2_{geom}$ 从 0.439 降至 0.081 ( $\Delta = 0.338$ )。
结论：训练目标决定了模型是否被迫将几何信息编码为线性可访问的形式。

因素 2：等变性（Equivariance）的放大作用，但非替代作用

发现：等变架构（如 MACE, PaiNN）本身并不保证高几何可访问性。
反直觉结果：在能量训练组中，等变模型 MACE QM9 的分数（0.081）甚至低于不变量模型 SchNet（0.262）。
结论：等变性只有在与任务对齐结合时才能发挥最大效用。单独的等变性不足以产生解耦的表示。

因素 3：数据多样性（Data Diversity）的补偿作用

发现：在任务未对齐的情况下，大规模多样化数据可以部分弥补差距。
数据：在 MPTraj 上预训练的 MACE（ $R^2_{geom} = 0.364$ ）显著优于仅在 QM9 上训练的 MACE（0.081–0.101）。
结论：数据多样性可以拓宽几何信息的可访问性，但无法完全消除任务不对齐带来的损失。

信息路由机制（Information Routing）

MACE 模型：展示了基于不可约表示（Irreducible Representations）的清晰路由。
- L=0 (标量) 通道：主要编码 HOMO-LUMO 能隙（标量性质， $R^2=0.76$ ）。
- L=1 (矢量) 通道：主要编码偶极矩（矢量性质， $R^2=0.59$ ）。
ViSNet 模型：虽然也有标量和矢量流，但矢量流几乎不包含线性可提取信息（ $R^2 \approx 0.018$ ），表明其等变操作主要用于内部计算而非输出表示的显式路由。

4. 关键贡献 (Contributions)

CPD 方法论：提出并验证了一种新的探针方法，能够线性分离组成和几何信号，并揭示了非线性探针在残差分析中的系统性偏差。
三因素梯度理论：确立了“任务对齐 > 等变性 > 数据多样性”的层级关系，推翻了“等变架构自动产生更好几何表示”的假设。
信息路由洞察：揭示了 MACE 模型中基于对称性类型（标量/矢量）的信息路由机制，并指出这种机制并非所有等变架构的通用属性。
鲁棒性验证：通过 12 项独立检查（包括异构体分类、概念擦除、不同组成定义等），证明了模型排序的稳定性（Spearman $\rho = 1.0$ ）。

5. 意义与影响 (Significance)

对实践者的指导：
- 在选择预训练分子编码器时，训练目标比架构更重要。如果下游任务对几何敏感（如电子性质），应优先选择在该类目标上训练过的模型，而非仅仅选择等变模型。
- 大规模预训练（数据多样性）可以作为任务不对齐的补偿，但不能完全替代任务对齐。
方法论警示：
- 在分子建模及 NLP/CV 领域的表示探测中，严禁在去除概念信号后使用非线性探针（如 GBT），否则会导致虚假的“信息保留”结论。必须使用线性探针来衡量线性可访问性。
理论启示：
- 监督信号（任务目标）在塑造表示结构方面比归纳偏置（架构设计）更具决定性。这呼应了表示学习理论中关于无监督解耦需要归纳偏置，但监督信号能更有效引导特定因子分离的观点。
- 线性可访问性梯度反映了表示的组织方式而非信息总量。低分模型可能仍包含丰富的几何信息，只是以非线性纠缠的形式存在。

总结

该论文通过严谨的线性分解方法（CPD），揭示了原子基础模型中几何信息组织的核心规律：训练任务的对齐程度是决定几何信息是否线性可访问的首要因素。这一发现为分子模型的选择、架构设计以及表示探测方法论提供了重要的理论依据和实践指南。