大原子模型的“巴别塔”：深度解析 DPA-3 与 MACE 的多任务/多头机制异同
在构建 大原子模型（Large Atomistic Models, LAMs） 的征途中，研究人员面临着一个被称为“巴别塔”的终极难题：数据不兼容。
现有的开源数据集（如 Materials Project, OC20, SPICE）虽然庞大，但它们在生成时采用了不同的量子力学近似方法（如 PBE, B3LYP, \(\omega\)B97M 等不同的 DFT 泛函）。如果把这些数据直接混在一起训练，模型就会因为标签的系统性偏差而“精神分裂”。
为了解决这个问题，剑桥大学的 MACE 团队和深势科技的 DPA-3 团队分别给出了自己的答案。虽然两者的目的都是实现多任务学习（Multi-task Learning）和跨领域知识迁移（Cross-learning），但它们在底层架构设计上却有着本质的区别。
本文将深入剖析这两种机制的数学原理、架构差异以及各自的优劣势。
1. MACE 的解法：多头回放微调 (Multi-Head Replay)
MACE 团队在构建统一基础力场时，采用的是经典的 “一干多枝”策略，即多头读出机制（Multi-Head Readout）。
1.1 数学表达
在 MACE 中，模型被分为两部分：共享的骨干网络（Backbone）和独立的数据集读出头（Readout Head）。对于属于特定数据集（Head）的原子 \(i\)，其能量预测公式为：
$$ E_i^{(\text{head})} = \sum_s \mathcal{R}^{(\text{head}, s)}(\mathbf{h}_i^{(s)}) + E_{0, z_i}^{(\text{head})} $$- \(\mathbf{h}_i^{(s)}\)：骨干网络提取的、与数据集无关的共享节点特征。
- \(\mathcal{R}^{(\text{head}, s)}\)：专属于该数据集的独立读出网络（通常是一个浅层 MLP）。
- \(E_{0, z_i}^{(\text{head})}\)：该精度下特定元素的参考孤立原子能量。
1.2 核心思想
MACE 的逻辑是：底层的物理几何规律（如键长、键角、多体相互作用）是普适的，因此用一个强大的 Backbone 来学习；而不同 DFT 泛函带来的能量绝对值差异，则交给最后一层的不同 Head 来分别拟合。 为了防止在微调新 Head 时忘记旧知识，MACE 引入了回放机制（Replay Buffer），在训练时按比例混入预训练数据。
2. DPA-3 的解法：数据集编码 (Dataset Encoding)
DPA-3 团队（以及之前的 DPA-2）敏锐地察觉到了传统多头机制的局限性，转而采用了一种 “条件生成” 的策略，即数据集编码机制。
2.1 数学表达
在 DPA-3 中，不存在物理隔离的多个读出头。整个网络只有一个统一的拟合网络 \(\mathscr{F}\)。原子 \(i\) 的能量预测公式为：
$$ E_i = \mathscr{F} \left( v_i^{(1,L)}, c(\mathscr{D}_m) \right) + e_m(Z_i) $$- \(v_i^{(1,L)}\)：经过线图序列（LiGS）提取的共享节点特征。
- \(c(\mathscr{D}_m)\)：数据集编码（Dataset Encoding），通常是一个代表数据集 \(\mathscr{D}_m\) 的 One-hot 向量。
- \(\mathscr{F}\)：全局唯一的能量拟合网络（MLP）。
- \(e_m(Z_i)\)：针对该数据集拟合的单原子能量偏置。
2.2 核心思想
DPA-3 的逻辑是：把“这是什么精度的数据”当作一种特征（Condition/Prompt），直接喂给神经网络。网络不仅共享底层的几何特征提取，连顶层的能量拟合逻辑也是共享的。网络内部会自动学习如何根据 \(c(\mathscr{D}_m)\) 的提示，对输出的能量进行微调。
3. 核心差异深度对比
虽然两者都成功实现了跨数据集的联合训练，但架构上的差异带来了截然不同的特性：
差异一：参数量与扩展性 (Parameter Scaling)
- MACE (多头机制)： 参数量随数据集数量线性增长 \(\mathcal{O}(N)\)。如果你要融合 31 个数据集，你就必须在模型末端挂载 31 个 MLP。这在数据集较少时没问题，但面对 OpenLAM 这种海量异构数据集时，会导致模型体积臃肿。
- DPA-3 (数据集编码)： 参数量是恒定的 \(\mathcal{O}(1)\)。无论融合 3 个还是 31 个数据集，拟合网络 \(\mathscr{F}\) 始终只有一个，仅仅是输入的 One-hot 向量维度增加了几位。这使得 DPA-3 在处理极多任务时具有无与伦比的扩展性。
差异二：跨领域知识迁移的深度 (Depth of Cross-learning)
- MACE： 知识迁移是隐式且受限的。不同数据集的知识只能在 Backbone 中交汇，一旦进入 Readout Head，信息就彻底隔离了。这意味着模型无法在顶层学习到“PBE 泛函和 B3LYP 泛函在能量映射上的系统性关联”。
- DPA-3： 知识迁移是显式且全局的。因为拟合网络 \(\mathscr{F}\) 也是共享的，模型不仅在底层共享几何特征，在顶层也能学习到不同 DFT 精度之间的非线性映射关系。这使得 DPA-3 在 Zero-shot（零样本）泛化测试中表现出极强的鲁棒性。
差异三：训练策略与工程实现
- MACE： 采用两阶段训练。先在最大规模的数据集（如 OMAT）上预训练 Backbone，然后冻结或微调 Backbone，同时训练多个 Head。这种方式工程上容易控制，且可以通过 Replay 机制灵活调整对特定领域的偏好。
- DPA-3： 倾向于端到端的并发多任务训练。将所有数据集混合，通过 Dataset ID 区分，直接从头开始训练（Concurrent multi-task training）。这种方式对数据采样采样率的平衡（Data Balancing）要求极高。
4. 总结与评价
我们可以用一个通俗的比喻来总结两者的区别：
- MACE 的多头机制 就像是一家 “联合诊所”：前台（Backbone）负责给病人做全身检查（提取特征），然后根据病人的病种（数据集），分发给不同的专科医生（Readout Head）开处方。
- DPA-3 的数据集编码 就像是一位 “全科神医”：医生（统一网络）不仅看体检报告，还会看病人的医保类型（Dataset Encoding），然后在大脑中综合处理，直接开出处方。
| 对比维度 | MACE (Multi-Head) | DPA-3 (Dataset Encoding) |
|---|---|---|
| 网络架构 | 共享 Backbone + 多个独立 Readout MLP | 共享 Backbone + 统一 Readout MLP |
| 区分数据集方式 | 路由到不同的网络分支 | 作为特征向量（One-hot）输入 |
| 参数量扩展性 | 随数据集数量线性增加（较差） | 几乎不随数据集数量增加（极佳） |
| 知识迁移级别 | 仅限底层特征表示 | 贯穿特征表示与顶层能量拟合 |
| 适用场景 | 数据集数量较少，且各数据集物理差异极大时 | 融合海量（数十个）异构数据集，追求极致泛化时 |
未来趋势： 从大语言模型（LLM）的发展史来看，通过 Prompt（条件编码）来统一多任务（如 DPA-3 的做法）往往比物理隔离的多头架构具有更高的上限和更好的 Scaling Law。DPA-3 的设计显然更加面向未来“万物互联”的 OpenLAM 时代；而 MACE 的多头回放机制，则在现阶段提供了一种极其稳定、可解释性强的工程解决方案。两者都在 AI for Science 的历史上留下了浓墨重彩的一笔。