Transferable SCF-Acceleration through Solver-Aligned Initialization Learning¶
会议: NeurIPS 2026
arXiv: 2604.21657
OpenReview: https://openreview.net/forum?id=Xjaap3pzlS
论文: https://arxiv.org/abs/2604.21657
领域: 其他
关键词: 自洽场加速, 密度泛函理论, 求解器对齐学习, 有效相对迭代数, 等变图神经网络
一句话总结¶
针对机器学习矩阵初猜模型在跨分子尺度外推时反而导致自洽场（SCF）迭代减速的问题，本文揭示该缺陷源于基态监督与非线性求解器动力学的不对齐，提出端到端反向穿透 SCF 求解器的训练框架 SAIL 与计入隐性开销的 ERIC 指标，在 QM40 和 QMugs 大分子上首次实现杂化泛函自洽场的稳健外推加速。
研究背景与动机¶
Kohn-Sham 密度泛函理论（KS-DFT）是现代计算化学和电子结构计算中应用最广泛的基石方法，其庞大的计算量常年占据全球高性能计算中心巨额的算力资源。KS-DFT 的核心瓶颈在于求解 Roothaan-Hall 广义特征值方程时所需的自洽场（SCF）非线性迭代，每轮迭代构建 Fock 矩阵的计算开销极大（半局域泛函为 \(\mathcal{O}(B^3)\)，杂化泛函高达 \(\mathcal{O}(B^4)\)）。为了降低迭代轮数，近年来学术界提出了基于等变图神经网络预测哈密顿量矩阵或密度矩阵的机器学习初猜方法。然而，Liu 等人在 2025 年的里程碑式发现给该领域蒙上了阴影：基于矩阵预测的初猜模型在推断比训练集稍大的分子时，SCF 迭代收敛速度急剧退化，甚至显著慢于传统的原子密度重叠（SAD/MINAO）初猜，出现严重的外推减速现象。
这种外推失败在过去被普遍归咎于网络在拓扑和尺寸分布外的泛化能力缺陷，导致研究转向利用局部原子中心辅助基展开系数（coefficient-based ansatz）来规避矩阵维度的增长。然而，系数预测范式本质上仅严格适用于局域泛函（如 PBE），在有机化学与药物发现所必需的 meta-GGA（如 SCAN）和杂化泛函（如 B3LYP）中，电子动能密度和非局域精确交换项无法仅由电子密度单值决定，必须依赖启发式人工近似或静态退化回退，在更高阶泛函上的外推加速效益大打折扣。更为关键的矛盾在于：深入检验表明，常规矩阵模型在分布外测试集上的基态代理损失（surrogate loss）依然很小，模型并没有丢失拟合基态目标的能力，但产出的初始猜测却将非线性优化轨迹推入了极其劣质的吸引盆。
本文明确指出，矩阵初猜模型的泛化退化并非表征维度的外推瓶颈，而是一个经典的「监督目标失配（supervision misalignment）」问题——直接在收敛基态张量上计算 Frobenius 范数残差，平等惩罚了不改变物理能量的冗余旋转，却忽视了非线性求解器对占据-虚轨道混合方向的极端敏感性。核心 idea：放弃拟合静态基态收敛解，直接将不可微或可微的 SCF 求解器完整纳入端到端反向传播回路，提出求解器对齐初猜学习（SAIL），利用占据-虚轨道流形能量梯度监督求解器迭代轨迹，引导初猜自发学习最利于求解器快速收敛的动态表征。
方法详解¶
整体框架¶
SAIL 框架将分子构型映射到量子化学自洽场求解器的整个流水线统一为端到端可微分系统。模型首先输入分子的原子核电荷数 \(\mathbf{Z}\) 与空间三维坐标 \(\mathbf{R}\)，利用高阶 \(\mathrm{SE}(3)\) 等变图网络提取几何与环境表征；根据输出方案，分别经过辅助基系数读出头或基于张量积的矩阵块读出头，构造出初始猜测密度矩阵 \(\mathbf{P}^{(0)}\)；随后，\(\mathbf{P}^{(0)}\) 被送入具备 DIIS 加速特性的完全可微自洽场求解器中，展开前向迭代 \(T\) 步，计算每步密度矩阵在能量流形上的切空间梯度作为标量损失反向传播，更新神经网络权重。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
In["输入：分子几何构型 (Z, R)"] --> Enc["EquiformerV2 等变图网络<br/>参数化输出初猜张量"]
Enc --> Init["初始密度矩阵构建 P(0)<br/>系数展开或原子对矩阵块拼接"]
Init --> SCF["完全可微自洽场迭代 (SCF)<br/>含 DIIS 加速器的 T 轮前向求解"]
SCF --> D1["求解器轨迹对齐优化<br/>端到端反向传播穿透整个求解过程"]
SCF --> D2["密度矩阵流形能量梯度损失<br/>提取占据-虚轨道流形切空间梯度"]
D1 -.->|梯度更新| Enc
D2 --> D1
SCF --> D3["有效相对迭代数 (ERIC)<br/>统筹初猜 Fock 构建与迭代耗时"]关键设计¶
1. 求解器轨迹对齐优化：以自洽动力学替代静态基态监督 传统方法将初猜学习视为纯粹的张量回归任务，利用预计算的收敛基态密度矩阵 \(\mathbf{P}^*\) 或 Fock 矩阵 \(\mathbf{F}^*\) 进行有监督拟合。然而，自洽场本质上是一个高度非线性的不动点迭代过程，微小的张量预测误差在外推大分子时可能产生恶性的谱失真，导致收敛加速算法（如 DIIS）发生剧烈震荡甚至数值停滞。求解器轨迹对齐优化通过将整个 SCF 迭代过程视为循环神经网络式的动态系统，在前向传播中直接执行 \(T\) 轮自洽场循环（包括 Fock 矩阵构建、广义特征值分解 gEVP 以及 DIIS 历史残差子空间外推），并在每一步施加收敛损失，沿时间步反向穿透整个求解器计算图。该方法彻底摆脱了昂贵的基态标签生成需求，无需预先对大分子运行耗时的完整 DFT，即可直接在分子构型上自监督学习求解器收敛动态。
2. 密度矩阵流形能量梯度损失：隔离占据-虚轨道真实物理自由度 在自洽场收敛性判定中，传统矩阵 Frobenius 范数 \(\|\Delta \mathbf{P}\|_F\) 存在严重理论缺陷：秩为 \(N_{\text{occ}}\) 的有效密度矩阵构成光滑的 Grassmann 流形，其真实的能量变分自由度仅由占据轨道（occ）与虚轨道（virt）之间的旋转决定（共 \(N_{\text{occ}} \times N_{\text{virt}}\) 个自由度），而占据-占据（OO）和虚-虚（VV）内部的旋转完全不改变电子密度和体系总能量。Frobenius 范数无差别地惩罚整个 \(B \times B\) 空间，往往被不相关的内部旋转误差所主导。为此，SAIL 设计了基于密度矩阵切空间的能量梯度损失，在每轮迭代 \(t\) 中，将当前 Fock 矩阵投影到占据-虚轨道正交块： $\(G_{ia}^{(t)} = \left(\mathbf{C}^{\text{occ},(t)\top} \mathbf{F}^{(t)} \mathbf{C}^{\text{virt},(t)}\right)_{ia}\)$ 在自洽收敛点处，由于 Fock 矩阵在自身本征基下对角化，非对角块严格为零（\(\mathbf{G} = \mathbf{0}\)）。SAIL 定义每步损失为均方根变分梯度： $\(\mathcal{L}_{\nabla}^{(t)} = \sqrt{\frac{1}{N_{\text{occ}} N_{\text{virt}}} \sum_{i=1}^{N_{\text{occ}}} \sum_{a=1}^{N_{\text{virt}}} \left|G_{ia}^{(t)}\right|^2}\)$ 该标量损失严格反映了当前轨道偏离驻点的距离，并通过轨道数归一化使得不同原子规模的分子损失具备严格可比性。
3. 有效相对迭代数（ERIC）：显式纠正初猜隐藏计算开销 以往文献广泛采用相对迭代次数（RIC）衡量初猜质量：\(\text{RIC} = N_{\text{iter}}^{\text{ML}} / N_{\text{iter}}^{\text{ref}}\)。然而在实际落地中，RIC 与真实的 GPU 墙钟时间（Wall-time）严重脱节。在 \(\Delta\)-残差预测范式中，预测哈密顿量残差 \(\Delta \mathbf{F}\) 之前，必须先在 MINAO 初猜上调用一次昂贵的 Fock 矩阵构建算子 \(\mathcal{M}_{\mathbf{P} \to \mathbf{F}}\)；对于辅助基系数模型，从系数构建初始 Fock 矩阵同样存在显著计算开销。根据基准分析，杂化泛函下单次 Fock 构建耗费单步总时间的 97.8% 以上。传统统计将初猜阶段的 Fock 构建视作「免费准备工作」，人为隐瞒了 1 轮完整计算，导致例如由 10 步压减到 6 步时宣称加速 40%，实际却进行了 7 次 Fock 构建。为此，本文定义有效相对迭代数： $\(\text{ERIC} = \frac{\#\,\text{total Fock builds required with ML-initialization}}{\#\,\text{Fock builds starting from }\mathbf{P}_{\text{ref}}^{(0)}}\)$ ERIC 统一统筹模型前处理、隐式残差基构建与后续自洽迭代的全部 Fock 算子调用，为跨方法、跨泛函的加速评估建立了客观真实的物理基准。
损失函数 / 训练策略¶
SAIL 采用「代理预训练 + 求解器微调」的两阶段优化策略。在预训练阶段，模型首先利用 QM9 子集分子在基态代理损失上收敛（系数模型采用解析重叠 \(L^2\) 密度误差，矩阵模型采用混合 Frobenius-\(L^1\) 损失），快速建立起合理的空间电荷分布雏形。在微调阶段，固定预训练权重为起点，以全轨迹平均梯度损失为优化目标： $\(\mathcal{L}_{\text{tot}} = \frac{1}{T} \sum_{t=1}^{T} \mathcal{L}_{\nabla}^{(t)}\)$ 训练过程采用适用于神经网络隐层正交特征优化的 Muon 优化器配合余弦退火学习率策略（基础学习率设为 \(5 \times 10^{-4} \sim 1 \times 10^{-3}\)），并引入衰减率为 0.995 的参数指数滑动平均（EMA）与范数为 1.0 的梯度截断以平抑 DIIS 奇异点扰动。虽然默认展开 \(T=10\) 步，但消融实验证明截断至 \(T=4\) 步即可达到无损性能，显著降低显存开销。
实验关键数据¶
主实验¶
模型在 QM9 小分子（重原子数 \(\le 9\)，总原子数 \(\le 20\)）上训练，并在显著超出训练分布的 QM40 数据集（重原子数高达 37，分子尺寸达训练集 4 倍）上进行严格的外推评估。基线参考初猜统一为 PySCF 的工业级默认方案 MINAO/SAD，基组选用 def2-SVP，收敛阈值固定为 \(10^{-9}\,\text{Ha}\)。
| 输入表征 | 泛函类别 | 代理损失 (Base) | 代理损失 (SAIL) | RIC (Base) | RIC (SAIL) | ERIC (Base) | ERIC (SAIL) | 相比Base提升 |
|---|---|---|---|---|---|---|---|---|
| 辅助基系数 \(\mathbf{c}\) | PBE (GGA) | 0.003 | 0.011 | 0.56 | 0.58 | 0.63 | 0.65 | -0.02 |
| 辅助基系数 \(\mathbf{c}\) | SCAN (mGGA) | 0.012 | 8.558 | 0.87 | 0.66 | 0.96 | 0.73 | +0.23 |
| 辅助基系数 \(\mathbf{c}\) | B3LYP (Hybrid) | 0.003 | 6.706 | 0.82 | 0.67 | 0.90 | 0.75 | +0.15 |
| 密度矩阵 \(\mathbf{P}\) | PBE (GGA) | 0.320 | 1.108 | 1.04 | 0.63 | 1.04 | 0.63 | +0.41 (恢复加速) |
| 密度矩阵 \(\mathbf{P}\) | SCAN (mGGA) | 0.292 | 1.076 | 1.04 | 0.67 | 1.04 | 0.67 | +0.37 (恢复加速) |
| 密度矩阵 \(\mathbf{P}\) | B3LYP (Hybrid) | 0.308 | 0.931 | 1.05 | 0.72 | 1.05 | 0.72 | +0.33 (恢复加速) |
| Fock 矩阵 \(\mathbf{F}\) | PBE (GGA) | 0.019 | 0.294 | 1.32 | 0.61 | 1.39 | 0.68 | +0.71 (恢复加速) |
| Fock 矩阵 \(\mathbf{F}\) | SCAN (mGGA) | 0.024 | 0.292 | 1.28 | 0.69 | 1.35 | 0.77 | +0.58 (恢复加速) |
| Fock 矩阵 \(\mathbf{F}\) | B3LYP (Hybrid) | 0.032 | 0.413 | 1.57 | 0.66 | 1.65 | 0.74 | +0.91 (恢复加速) |
在 QMugs 数据集（重原子数高达 90，尺度达到训练集的 10 倍）上针对 B3LYP 泛函进行的极限外推测试中，\(\Delta\)-密度矩阵模型全程维持约 0.78 的平坦 ERIC，并在 GPU4PySCF 上取得端到端 \(1.35\times\) 真实墙钟加速比。
消融实验¶
基于 QM40 数据集对 \(\mathbf{P}\) 模型（SCAN 泛函）的关键架构与训练设计进行系统消融分析：
| 实验配置 | 代理损失 \(\mathcal{L}_{\mathbf{P}}\) | RIC (↓) | ERIC (↓) | 核心设计结论说明 |
|---|---|---|---|---|
| SAIL 完整模型 (\(T=10\)) | 1.076 | 0.67 | 0.67 | 默认基准，加速比最高且极为稳定 |
| 替换损失为 DIIS 对易子残差 \(\mathcal{L}_{\text{comm}}\) | 1.100 | 0.67 | 0.67 | 轨道能量梯度与对易子残差效果等价 |
| 求解器迭代截断 \(T=8\) | 1.074 | 0.67 | 0.67 | 步数削减对收敛引导无损 |
| 求解器迭代截断 \(T=4\) | 1.148 | 0.67 | 0.67 | \(T=4\) 保持最佳加速，显著节省微调开销与显存 |
| 求解器迭代截断 \(T=2\) | 0.982 | 0.68 | 0.68 | 早期动力学仍可提供足够梯度信号 |
| 求解器迭代截断 \(T=1\) | 1.163 | 0.69 | 0.69 | 仅穿透单步即可基本扭转外推减速 |
| 非残差学习（无 \(\Delta\)-MINAO） | 0.845 | 0.72 | 0.72 | 缺失底层原子对称基线使 ERIC 退化 5 个百分点 |
| 单阶段端到端从头训练（无预训练） | 1.473 | 0.75 | 0.75 | 虽优于传统 Base (1.04)，但不如两阶段方案充分 |
关键发现¶
- 反直觉的代理损失与迭代加速权衡：所有 SAIL 微调模型的基态代理损失均比预训练基线高出数倍甚至两个数量级（如系数模型在 SCAN 上代理损失由 0.012 升至 8.558），但 ERIC 却大幅骤降（0.96 降至 0.73）。这强力证明了拟合收敛基态并非好的初猜训练目标；强行压缩基态的无关多余自由度反而损害了求解器的瞬态吸引子结构。
- 矩阵预测模型的全面复活：常规基态训练下，\(\mathbf{P}\) 和 \(\mathbf{F}\) 预测在 QM40 上全面失效（\(\text{ERIC} > 1.0\)），随分子规模增大呈线性恶化。经 SAIL 优化后，密度矩阵模型一跃成为最具普适性的最优方案，在 PBE、SCAN 和 B3LYP 上分别斩获 0.63、0.67 和 0.72 的极佳 ERIC，彻底打破了「矩阵模型无法尺寸外推」的固有认知。
- Fock 预测的「隐性税收」被 ERIC 揭穿：在 B3LYP 下，\(\Delta\)-Fock 模型表面上取得了 0.66 的最低 RIC（看似优于密度矩阵模型的 0.72），但计入初猜所需的额外 Fock 构建后，其实际 ERIC 变为 0.74，综合性能逊于无需额外构建的密度矩阵方案。
亮点与洞察¶
- 发现并矫正核心科研认知偏差：推翻了以往文献关于「矩阵模型不能尺寸泛化」的唯拓扑论假设，从优化几何与能量驻点流形切空间的角度，确凿指出了基态回归目标与非线性迭代动力学失配的本质病因。
- 构建科学的无偏评估基准 ERIC：精准指出了过去 Hamiltonian 和系数学习文献中普遍忽视的初猜预构建隐藏耗时，杜绝了通过转嫁算力至数据准备阶段来刷高 RIC 指标的现象。
- 真正实现药物化学尺度杂化泛函加速：成功将基于 AI 的自洽场无损加速推广至具备 90 个重原子、包含真实药物骨架的 QMugs 大分子，并在工业界重度依赖的 B3LYP 杂化理论级别下兑现了端到端 \(1.35\times\) 的 GPU 物理加速。
局限与展望¶
- 基组与泛函耦合性：目前矩阵预测模型仍与特定的原子轨道（AO）基组紧密绑定，更换基组需重新微调，计算成本需在后续的大规模计算中摊销。
- 化学元素与构型受限：目前验证主要局限在稳定有机分子的常规元素（H, C, N, O, F），在开壳层过渡金属配合物、周期性晶体系统以及高能非平衡态过渡态构型上的迁移能力尚待进一步拓展验证。
- 计算资源门槛：反向穿透 SCF 迭代需要维持较深的前向激活值计算图，尽管截断至 \(T=4\) 可大幅缓解，微调过程对 GPU 显存的需求依然高于浅层图神经网络。
相关工作与启发¶
- vs. Liu et al. (2025) 辅助基系数模型：Liu 等人因观察到矩阵模型外推退化而放弃矩阵路线、转用辅助基标量系数预测；本文通过 SAIL 修复了矩阵模型的根源缺陷，在更广泛的 meta-GGA 和杂化泛函上超越了系数退化方案。
- vs. QHFlow / QHNet / Graph2Mat 等基态回归模型：传统方法均以拟合稳态基态张量为目标，在分布外不可避免地陷入非物理的求解阻滞；SAIL 则是面向求解器本身的目标导向式微调范式。
- vs. 神经泛函替代模型：完全取代 DFT 的端到端代理网络不可避免地承受着未知的泛化误差漂移；SAIL 作为纯初猜加速器，完全保留了 Kohn-Sham DFT 原生的物理严格性与数值精度。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 颠覆了基态拟合传统范式，首创面向非线性自洽求解器动力学端到端微分的初猜对齐学习框架。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 3 种输入表征、3 级 Jacob 天梯泛函，跨越 QM9、QM40、QMugs 达 10 倍分子尺度，并对开销进行了毫厘不差的软硬件耗时分解。
- 写作质量: ⭐⭐⭐⭐⭐ 理论剖析透彻，数学流形物理意义清晰，直面领域固有认知盲区，论证有力。
- 价值: ⭐⭐⭐⭐⭐ 真实缩减高精度杂化 DFT 计算耗时达 25%~35%，为加速高通量第一性原理计算与机器学习力场数据集构建提供了无损利器。