MALOQ: 用于量子输运的算符大规模加速学习

# MALOQ: 面向量子传输的算子加速学习

来源：https://arxiv.org/html/2606.28911

###### 摘要

机器学习 (ML) 算子模型可以通过训练来预测密度泛函理论 (DFT) 哈密顿量/密度矩阵，从而显著降低计算成本，将电子结构计算扩展到以前无法实现的规模。本文介绍MALOQ（面向量子传输的算子加速学习），这是一个用于训练和预测电子结构矩阵的应用，适用于由几个到10万个原子组成的系统，使用大型基组，并涵盖广泛的原子元素。基于最先进的 SO(2) 等变骨干架构，MALOQ 提供 (i) 用于处理高阶哈密顿矩阵数据的自定义数据处理内核，以及 (ii) 原子图的可扩展边向分布。在目前可用的大型分子哈密顿数据集上训练，与分子向分布框架相比，它将每轮训练时间减少了 30% 以上，并能够对任意大小的材料图进行推理。我们在 Alps 超级计算机上展示了针对 3,000-12,000 个原子的可扩展训练和推理，分别使用高达 192 个 GPU 和 256 个 GPU。

## I. 引言

机器学习 (ML) 模型通过将原子模拟的范围扩展到以前无法实现的规模，正在加速计算材料科学研究[52 (https://arxiv.org/html/2606.28911#bib.bib73)]。迄今为止，这些模型主要专注于学习分子和原子层面的性质。这包括一类 ML 原子间势 (MLIP)，它们预测分子动力学 (MD) 模拟中使用的分子能量和原子力。最先进的、基于图的“材料基础模型”现在能够以准 DFT 级别的精度为任意分子和复杂材料几何生成这些量[3 (https://arxiv.org/html/2606.28911#bib.bib9),46 (https://arxiv.org/html/2606.28911#bib.bib10),31 (https://arxiv.org/html/2606.28911#bib.bib76)]。

新出现的一类 ML 材料模型将这些架构扩展到由**算子**[14 (https://arxiv.org/html/2606.28911#bib.bib79),20 (https://arxiv.org/html/2606.28911#bib.bib12),33 (https://arxiv.org/html/2606.28911#bib.bib71),30 (https://arxiv.org/html/2606.28911#bib.bib11)] 定义的电子层面量。这些“算子模型”不学习以原子为中心的性质，而是直接训练于编码给定系统内原子间和原子内轨道相互作用的矩阵。最常见的算子是哈密顿量/福克算符 (H) 和密度算符 (P)，它们描述了所考虑原子结构内电子可以占据的状态的空间和能量分布。最近的工作将这些原理扩展到学习可以原子轨道基表示的其它算子，可以是直接的，如格林函数（H的逆）[43 (https://arxiv.org/html/2606.28911#bib.bib2)]，也可以是间接的，例如源自力的低秩 Hessian [8 (https://arxiv.org/html/2606.28911#bib.bib14)] 矩阵。

算子模型正在开发用于两个主要目的：

请看图注

图 1: 电子结构学习在分子（顶部）和材料（底部）中的应用。目标是学习原子结构（左）与其对应的电子哈密顿矩阵（中）之间的映射，并在 DFT 单独无法企及的规模上进行预测。从这些矩阵中，我们可以解析地构建电子密度或计算流经复杂材料的电子流（右）。

1. 1. 实现复杂材料中的量子传输：矩阵形式的哈密顿量是量子传输 (QT) 模拟的输入，旨在计算在不同电压配置下流经原子分辨器件的电子流。先进的 QT 求解器以密度泛函理论 (DFT) 级别的电子结构输入进行操作，DFT 是用于电子结构计算的最强大的第一性原理方法之一[6 (https://arxiv.org/html/2606.28911#bib.bib33)]。目前，QT 代码能够处理拥有超过 10k 个原子、存在复杂多体效应的系统[44 (https://arxiv.org/html/2606.28911#bib.bib7)]，与半导体公司制造的电子器件的物理尺寸相匹配[1 (https://arxiv.org/html/2606.28911#bib.bib152)]。然而，使用 DFT 代码生成 H 的复杂度为 O(N_orb^3)，其中 N_orb 是所关注结构中原子轨道的总数，这最终限制了可以处理的问题规模（原子数）。现有工作通常计算小周期系统的 H，然后将其拼接以构建设备尺度模型。同时，这些算子模型可以在 O(N_orb) 可扩展性下生成 H[41 (https://arxiv.org/html/2606.28911#bib.bib31)]。正如 MLIP 开启了 MD 模拟的新领域一样，算子模型有望解锁对真实复杂原子几何结构的 QT 建模[47 (https://arxiv.org/html/2606.28911#bib.bib38)]。

2. 2. 利用下游性质的缩放定律：ML 性质模型受限于数据——其性能随着训练数据的增多而提高[46 (https://arxiv.org/html/2606.28911#bib.bib10)]。然而，现有训练数据集中力和能量的生成已经在 DFT 级别消耗了数十亿核心小时[27 (https://arxiv.org/html/2606.28911#bib.bib8)]，因此目前不可能生成数量级更多的数据。在这些数据集中，每个 N 原子分子或材料晶胞的处理生成 O(1) 个能量标签，O(N) 个力标签，以及电子结构算符中的 O(N^2) 个轨道相互作用标签。通过预训练方案利用这些可用但至今被忽略的电子数据，为继续利用“神经缩放定律”（该定律定义了训练数据集大小与模型准确率之间的关系）提供了有效途径。在 DFT 中，电子结构算符内的轨道相互作用参数化了所有后续的原子级分子和材料性质。将这些信息纳入性质预测架构中，可以改善在原子和分子预测任务上的泛化能力[8 (https://arxiv.org/html/2606.28911#bib.bib14),20 (https://arxiv.org/html/2606.28911#bib.bib12),18 (https://arxiv.org/html/2606.28911#bib.bib5),40 (https://arxiv.org/html/2606.28911#bib.bib13),33 (https://arxiv.org/html/2606.28911#bib.bib71)]。小分子基准测试表明，从学习的电子结构算符解析计算力和能量，其准确性与使用 MLIP 的直接预测相当[23 (https://arxiv.org/html/2606.28911#bib.bib3)]。

电子结构矩阵的大型数据集最近变得可用，包含高达 10 PB[27 (https://arxiv.org/html/2606.28911#bib.bib8)] 的轨道相互作用。然而，在这种矩阵数据上进行训练，带来了与以前 ML 性质模型（如 MLIP）所遇到的不同计算挑战。首先，嵌入中使用的球谐特征必须匹配高阶轨道相互作用数据的维度，通常高达 f（有时是 g）轨道（需要角动量 l=6-8）。与 MLIP 中最大球谐系数的秩通常是超参数不同，电子结构预测需要高阶嵌入仅仅是为了产生正确维度的数据。其次，虽然 MLIP 通常对相互作用原子之间施加 ~6-8 Å 的截断，但处理轨道相互作用数据需要处理延伸到 12 Å 原子间距离的专用边嵌入[47 (https://arxiv.org/html/2606.28911#bib.bib38),20 (https://arxiv.org/html/2606.28911#bib.bib12)]。更大的截断会创建密集连接的图表示，节点度数超过 100，大大增加了需要处理的嵌入总量。最后，处理轨道相互作用数据需要在训练和推理期间频繁进行基变换，以在 H 的“标签”和“矩阵”表示之间进行转换。这三个关键差异共同显著增加了训练和推理的内存占用。

为了能够“大规模”进行算子预测，因此我们介绍了 MALOQ（面向量子传输的算子加速学习）包，这是一个加速的 ML 框架，用于训练电子结构矩阵并预测任意原子元素和结构尺寸（从大分子到超过 10k 个原子的材料系统，图 1 (https://arxiv.org/html/2606.28911#S1.F1)）。MALOQ 的核心是一个消息传递图神经网络 (GNN)，利用等变球面通道网络 (eSCN) 卷积[37 (https://arxiv.org/html/2606.28911#bib.bib62)]。其骨干架构在 MLIP（eSEN[11 (https://arxiv.org/html/2606.28911#bib.bib94)]）和哈密顿量预测任务（HELM[20 (https://arxiv.org/html/2606.28911#bib.bib12)]）上均取得了最先进的性能。在此之上，我们引入了两类计算创新，旨在实现电子结构矩阵的大规模并行训练和推理：

1. 1. 算子数据处理加速： (1a) 高效的“矩阵 ↔ 标签”转换内核，允许将哈密顿矩阵预处理为训练标签，速度可达 ~10^6 个原子/秒，并能从预测的球谐系数快速重建。 (1b) eSCN 卷积所需的广义旋转矩阵（Wigner-D 矩阵）的加速计算。

2. 2. 大型分子和材料原子图的自定义分布，以加速分子数据集的训练，并执行大型原子系统 H 的推理。我们将高效的逐点通信方案与边向分区相结合，以最小化通信并减少分区间的负载不平衡，实现了对高达 192 个 GH200 GPU 的强可扩展性，以及在包含 12,000 个原子（~1000 万个轨道相互作用）的结构尺寸上 >90% 的弱可扩展性效率。

这些贡献共同使我们能够 (1) 将不规则分子数据库上的训练速度提高 30-50%，以及 (2) 实现对包含超过 10k 个原子的大型材料图的哈密顿矩阵的快速推理。一方面，我们的工作使得可以模拟与工业相关的原子系统，现在可以在 DFT 级别以降低的计算复杂度轻松确定其量子传输性质。另一方面，它能够高效处理电子分辨率下可用的海量分子和材料性质数据，为 ML 性质模型在无需额外数据集生成的情况下继续受益于神经缩放定律创造了途径。

## II. 背景与相关工作

### II-A 电子结构问题

MLIP 和其他性质模型的训练数据通常使用 DFT 生成。这种第一性原理方法实现了 Kohn-Sham 方程，该方程自洽地耦合了所关注原子系统的电子密度和静电势，直到经过 N_iter 次迭代后达到收敛[24 (https://arxiv.org/html/2606.28911#bib.bib49)]。几种商业或免费可用的代码在原子结构上实现了这种自洽迭代的“密度 ↔ 势”方案[26 (https://arxiv.org/html/2606.28911#bib.bib149),32 (https://arxiv.org/html/2606.28911#bib.bib130)]。这些实现中许多在局域原子轨道基 |φ_i⟩ 上运行，这些轨道通常由收缩的高斯基函数构造[26 (https://arxiv.org/html/2606.28911#bib.bib149),32 (https://arxiv.org/html/2606.28911#bib.bib130)]，在旋转下像球谐函数 (Y_l^m(r̂′)) 一样变换。每个球谐函数由一个角动量度数 l 和阶数 m ∈ {-l,...,l} 指定，例如 s (l=0)，p (l=1)，d (l=2)，⋯。在这个基组中，Kohn-Sham 方程的形式为广义本征值问题：H ψ = ε S ψ。哈密顿矩阵 H 的元素为 H_ij = ⟨φ_i|Ĥ(r)|φ_j⟩，其中 Ĥ(r) 是哈密顿算符。重叠矩阵 S 由项 S_ij = ⟨φ_i|φ_j⟩ 组成，描述了局域轨道之间的有限重叠。它们可以从轨道基组预先计算。H 和 S 的大小均为 N_orb = ∑_{k=1}^{N_A} N_orb,k，其中 N_A 是原子总数，N_orb,k 是原子 k 的轨道（基元）数量，索引 k 遍历所有考虑的原子。Kohn-Sham 方程的解提供了能量本征值 (ε) 和波函数 (ψ)，由此可以导出电荷密度 ρ(r)。DFT 代码中的核心操作是重复求解上述广义本征值问题，直到 ρ(r) 和由此产生的静电势 V(r) 之间达到自洽收敛。该完整方法的计算复杂度每次迭代为 O(N_orb^3)。在这些自洽迭代结束时，可以从最终收敛的 H 计算原子和分子尺度的性质，例如力和能量。

### II-B 相关工作

应用 | 操作 | 任务 | 预测 | 分布式 | N_A (L_max) | 节点 | 边
--- | --- | --- | --- | --- | --- | --- | ---
MLIPs 
↳ DeepMD[17] | 线性 | F/E | ✓ | ✗ | ✓ | - | - | -
↳ Allegro[25] | TP | F/E | ✓ | ✗ | ✓ | - | - | -
↳ MACE[3] | TP | F/E | ✓ | ✗ | ✗ | - | - | -
↳ SevenNet[36] | TP | F/E | ✓ | ✗ | ✓ | - | - | -
↳ UMA[46] | eSCN | F/E | ✓ | ✗ | ✓ | - | - | -
哈密顿量 
↳ QHNet[51] | TP | H_ij | ✓ | ✓ | ✗ | 10 (4) | - | -
↳ WANet[30] | eSCN | H_ij | ✓ | ✓ | ✗ | 100 (6) | - | -
↳ HELM[20] | eSCN | H_ij | ✓ | ✓ | ✗ | 150 (6) | - | -
↳ DeepH[29] | eSCN | H_ij | ✓ | ✓ | ✗ | 216 (4) | - | -
**MALOQ** | **eSCN** | **H_ij** | **✓** | **✓** | **✓** | **3,000+ (4-8)** | - | -

表 I: 跨原子元素（近乎“通用”）的材料建模 GNN 应用选择，包括功能相似性以及在预测任务、用于更新节点/边特征的操作（Op.）以及分布式计算环境（Distr.）的可能性方面的差异。F/E = 力与能量，H_ij = 哈密顿量，TP = SO(3) 等变张量积，eSCN = SO(2) 等变球面通道网络。N_A = 最大训练结构尺寸。

GNN 现已成为计算材料科学的主流工具，它们被用于