Global Plane Waves from Local Gaussians: Periodic Charge Densities in a Blink¶

会议: ICML 2026
arXiv: 2601.19966
代码: https://github.com/Jotels/ELECTRAFI (有)
领域: 科学计算 / 机器学习+DFT / 电子结构
关键词: 电荷密度预测, 平面波基组, 各向异性高斯, 泊松求和, DFT 初猜加速

一句话总结¶

ELECTRAFI 在实空间预测一组各向异性高斯的参数，再利用高斯的解析傅里叶变换 + 泊松求和公式，在倒易空间一次性算出周期晶体的电荷密度的平面波系数，做一次逆 FFT 就拿到全场密度；在 NMAE 与 ChargE3Net 持平甚至更优的同时，推理快 \(463\times \sim 633\times\)，把端到端 DFT 总时间真正降下来 \(\sim 20\%\)。

研究背景与动机¶

领域现状：Kohn–Sham DFT 是物理/化学/材料里最常用的电子结构方法，主流加速思路分两条：一是机器学习势 (MLIP) 直接预测能量/力，绕开 SCF 循环；二是保留 DFT 流程、用 ML 预测一个好的初始电荷密度 \(\rho(\mathbf{r})\) 作为 SCF 起点，从而减少自洽迭代次数。后一条路线的好处是仍然由 DFT 收敛到目标泛函的严格自洽解。

现有痛点：周期体系的 ML 密度模型分两派——轨道模型用原子中心的球谐+径向基展开密度，对小分子很好但对金属/无机晶体所需的高 \(l\)、弥散基组爆炸；probe 模型（如 ChargE3Net、DeepDFT）把实空间网格点当作图节点跑 GNN，准确度强但每张图要查询 \(10^7\) 量级网格点，单结构推理动辄 \(30\sim80\) 秒，省下的 SCF 时间会被推理时间反吃掉，端到端反而变慢。

核心矛盾：周期性 + 长程库仑相互作用天然属于倒易空间，但已有模型要么在实空间做昂贵的周期镜像求和+球谐展开，要么直接用 NN 预测平面波系数（输出维度跟 FFT 网格走，scale 极差，仅作小补丁有效）。准确度和推理成本必须联合优化才能换来真实的 DFT 墙钟节省。

本文目标：构造一个 (1) 天然周期、(2) 与平面波 DFT 编码对齐、(3) 能解析处理长程结构、(4) 推理时间相对 DFT 可忽略的电荷密度模型，并且在真实 VASP 流水线里跑出净时间节省。

切入角度：高斯函数在傅里叶变换下保持高斯，且有闭式表达；同时泊松求和公式可以把"全空间高斯叠加"等价改写成"倒易格点上的傅里叶级数"。如果用一组浮动高斯做密度的局部展开，就能把"周期化"这一步整体让位给解析变换，不再做实空间镜像求和。

核心 idea：实空间预测局部各向异性高斯参数 \((w^{(j)},\boldsymbol{\mu}^{(j)},\boldsymbol{\Sigma}^{(j)})\) → 解析写出每个高斯在倒易格点 \(\mathbf{G}\) 处的傅里叶系数 → 求和得 \(\hat{\rho}(\mathbf{G})\) → 一次逆 FFT 还原周期实空间密度 \(\rho(\mathbf{r})\)。

方法详解¶

整体框架¶

输入：周期晶体的原子图（含 PBC 邻居）。骨干：改造自 EScAIP 的注意力 GNN，按 PBC 邻域更新标量/向量/张量流。读出：dyadic 邻域聚合，给每个原子输出三类通道 \(S\in\mathbb{R}^{N\times C}\)、\(V\in\mathbb{R}^{N\times C\times 3}\)、\(T\in\mathbb{R}^{N\times C\times 3\times 3}\)。分配：按"每个价电子分配 \(M\) 个高斯"的规则，原子 \(a\) 拿到 \(n_a=Mv_a\) 个高斯槽位，全结构共 \(N_\mathcal{N}=M\sum_a v_a\) 个高斯。参数化：每个高斯由 \((w^{(j)},\boldsymbol{\mu}^{(j)},\boldsymbol{\Sigma}^{(j)})\) 描述，签名权重 \(w^{(j)}=\tanh(\mathrm{MLP}(S^{(j)}))\)，中心 \(\boldsymbol{\mu}^{(j)}=\mathbf{R}_{a(j)}+\mathbf{d}^{(j)}\) 是相对宿主原子的偏移，协方差用 Gram 分解 \(\boldsymbol{\Sigma}^{(j)}=\gamma^{(j)}\mathbf{A}^{(j)}\mathbf{A}^{(j)\top}\) 保证正定。重建：把所有高斯解析变换到倒易空间，求和得 \(\hat{\rho}(\mathbf{G})\)，再用一次 IFFT 得到 \(\rho(\mathbf{r})\)，最后整体缩放权重让 \(\int_\Omega\rho\,d\mathbf{r}=N_e\)（电子数守恒）。

关键设计¶

解析傅里叶 + 泊松求和的周期化:
- 功能：把"实空间无穷镜像求和"换成"倒易空间有限格点求和"，一次性拿到天然周期的密度场。
- 核心思路：辅助非周期函数 \(\tilde\rho(\mathbf{r})=\sum_j w^{(j)}\mathcal{N}(\mathbf{r};\boldsymbol{\mu}^{(j)},\boldsymbol{\Sigma}^{(j)})\) 永远不在实空间显式求值。利用高斯的自互易性，每个分量在 \(\mathbf{G}\) 上有闭式 \(\hat{\mathcal{N}}(\mathbf{G})=\exp(-\tfrac{1}{2}\mathbf{G}^\top\boldsymbol{\Sigma}\mathbf{G})e^{-i\mathbf{G}\cdot\boldsymbol{\mu}}\)；求和得 \(\hat\rho(\mathbf{G})=\sum_j w^{(j)}\exp(-\tfrac{1}{2}\mathbf{G}^\top\boldsymbol{\Sigma}^{(j)}\mathbf{G})e^{-i\mathbf{G}\cdot\boldsymbol{\mu}^{(j)}}\)；由泊松求和 \(\sum_{\mathbf{R}\in\Lambda}\rho(\mathbf{r}+\mathbf{R})=\frac{1}{|\Omega|}\sum_{\mathbf{G}\in\Lambda^*}\hat\rho(\mathbf{G})e^{i\mathbf{G}\cdot\mathbf{r}}\)，截断 \(\Lambda^*\) 只产生低通滤波（DFT 密度本身也是低通），不像截断 \(\Lambda\) 那样在原胞边界留下不连续。
- 设计动机：朴素镜像求和即便只取 \(N=1\)（27 个相邻原胞）也要付 27 倍代价且仍不真正周期；直接预测 \(\hat\rho(\mathbf{G})\) 需要输出维度随 FFT 网格爆炸，已有工作（GPWNO 等）只能当作小补丁用；本文用"局部高斯参数 + 解析全局变换"把可学习量和分辨率解耦，网络规模与晶胞大小、平面波截断完全无关。
浮动各向异性高斯展开 + Gram 因式协方差:
- 功能：用比球谐基紧凑得多的表示捕捉无机晶体里又胖又弥散的价电子密度。
- 核心思路：每个高斯放在原子上但通过预测的位移 \(\mathbf{d}^{(j)}\) 自由漂移到键中点/空隙位置；协方差由 \(\boldsymbol{\Sigma}^{(j)}=\gamma^{(j)}\mathbf{A}^{(j)}\mathbf{A}^{(j)\top}\) 写成 Gram 形式，天然 SPD 且既能调向也能调形状；签名权重 \(w^{(j)}=\tanh(s^{(j)})\) 允许加减以重建价电子密度的精细结构。最后全局乘一个标量保证电子数等于 \(N_e\)。
- 设计动机：固定原子中心的 LCAO 在材料体系里要堆很多高 \(l\) 的弥散基函数，效率低；浮动高斯 + 各向异性协方差直接削掉对高阶球谐的依赖，参数量从 \(O(L_{\max}^2)\) 退化到几个 \(3\times3\) 张量，推理大幅加速且对元素种类不敏感。
EScAIP 注意力骨干 + dyadic 张量读出:
- 功能：在 PBC 图上既要 cheap 又要造出向量/张量量纲与几何对齐的特征，用来直接喂给 \((\boldsymbol{\mu},\boldsymbol{\Sigma},w)\) 的预测头。
- 核心思路：不走严格等变 SO(3) 张量积（MACE / Equiformer 类的高阶 irreps + Gaunt contractions 代价随 \(L_{\max}\) 暴涨），而是用多头注意力+前馈在标量流上堆容量；向量/张量头用"对单位边方向 \(\hat{\mathbf{e}}_{ij}\) 及其无迹 dyad \(\hat{\mathbf{e}}_{ij}\hat{\mathbf{e}}_{ij}^\top\) 做注意力加权求和"，复杂度只与边数线性相关、且不需要显式三元组枚举（避开 DimeNet/GemNet 的角度基）。等变性靠数据增强补，实验显示训练时随机旋转后旋转测试误差从 1.69% 回落到 1.33%，反而比未增强的 1.35% 还略好。
- 设计动机：probe/orbital 模型推理慢的另一半原因来自骨干，等变高阶张量在大晶胞上 cost-prohibitive；这一设计把"能不能用就行"的几何归纳偏置交给数据，把算力集中在注意力和 FFN 上，是 ELECTRAFI 端到端能跑进 0.1 秒级别的关键之一。

损失函数 / 训练策略¶

直接以 NMAE 作为训练目标（与既有密度模型评测约定一致）： \(\mathcal{L}=\mathrm{NMAE}(\rho_{\mathrm{pred}},\rho_{\mathrm{ref}})=\frac{\int_\Omega|\rho_{\mathrm{ref}}-\rho_{\mathrm{pred}}|\,dV}{\int_\Omega\rho_{\mathrm{ref}}\,dV}\)，在每个结构的完整实空间网格上算（不做子采样），最终汇报 \(\mathrm{NMAE}[\%]=100\times\mathrm{NMAE}\)。MP-Full 训练集 117,876 结构，最大密度文件超 4GB、单材料最多输出 \(1.57\times10^8\) 个浮点数；超参数 \(M\)（每价电子高斯数）按附录设置。

实验关键数据¶

主实验¶

所有测试在单张 A100-40GB；VASP DFT 跑在 24 核 Intel Xeon E5-2650（Broadwell）。

数据集	指标	ELECTRAFI	ChargE3Net	DeepDFT / GPWNO / InfGCN
MP-Full	NMAE [%]	0.58	0.54	0.80 (DeepDFT)
MP-Full	\(t_{\mathrm{inf}}\) [s]	0.17	78.73	—
MP-Full	Speedup	463×	—	—
GNoME	NMAE [%]	0.93	0.69	—
GNoME	\(t_{\mathrm{inf}}\) [s]	0.11	33.28	—
GNoME	Speedup	302×	—	—
ECD-HSE06	NMAE [%]	1.35	1.53	—
ECD-HSE06	\(t_{\mathrm{inf}}\) [s]	0.05	31.65	—
ECD-HSE06	Speedup	633×	—	—
MP-Mixed	NMAE [%]	1.24	—	11.50 / 4.83 / 5.11
Cubic	NMAE [%]	1.37	—	10.37 / 7.69 / 8.98 (SCDP 2.59)

NMAE 上 ELECTRAFI 在 4/5 个 benchmark 拿到 SOTA 或并列；在 MP-Full / GNoME 比 ChargE3Net 略差 0.04 ~ 0.24 个百分点，但推理快两到三个数量级。

端到端 DFT 加速实验（重点）¶

数据集	方法	NMAE	ML 推理时间	SCF 步数	DFT 时间	总时间	总时间节省
MP	SAD (默认)	—	0	16.80	266.34 s	266.34 s	0%
MP	SC (上限)	—	0	8.73	161.98 s	161.98 s	39.18%
MP	ELECTRAFI	0.55%	0.17 s	13.33	219.57 s	219.74 s	+17.50%
MP	ChargE3Net	0.50%	72.11 s	12.09	208.26 s	280.37 s	−5.27%
GNoME	SAD	—	0	13.49	119.98 s	119.98 s	0%
GNoME	SC (上限)	—	0	6.88	77.91 s	77.91 s	39.75%
GNoME	ELECTRAFI	0.88%	0.11 s	10.11	95.06 s	95.17 s	+20.68%
GNoME	ChargE3Net	0.59%	28.29 s	9.50	92.67 s	120.96 s	−0.82%

关键发现¶

推理时间不能忽略：ChargE3Net 在 DFT 墙钟上确实比 ELECTRAFI 略快（SCF 步数也少一点），但把推理加回来后净反慢 0.8%~5.3%；ELECTRAFI 因为 inference 几乎免费，把 SCF 减少直接变成 17%~21% 的端到端节省。
结构级稳定性差距巨大：ELECTRAFI 只在 13.4% (MP) / 7.2% (GNoME) 的结构上让总时间变长，ChargE3Net 在 84.3% / 74.6% 的结构上变长。
元素级互补：错误分析（附录 F）显示 ELECTRAFI 在碱金属/碱土/卤素/重元素（Ac 等）上更准（更平滑、更弥散的密度），ChargE3Net 在轻共价元素和开壳过渡金属上更准（局域定向键合）；与各自表示的归纳偏置一致。
旋转鲁棒性：架构不显式等变，未增强训练时旋转测试 NMAE 从 1.35% 升到 1.69%；加随机旋转增强后回到 1.33%（甚至略好于无旋转测试），说明软等变 + 注意力骨干可以靠数据学到等变性。
天花板分析：用收敛态密度 (SC) 做初猜的上限是约 40% 总时间节省；ELECTRAFI 拿到了其中约 50%，说明 NMAE 改进还有空间但已经触及"SCF 收敛由长波长慢模式决定"的瓶颈，单纯优化 grid-level NMAE 收益边际递减。

亮点与洞察¶

"局部预测 + 全局解析变换"是个非常优雅的解耦：把"网络该学什么"（局部、原子尺度的高斯参数）和"周期性 + 长程结构怎么处理"（解析傅里叶 + 泊松求和）彻底分开，网络输出维度与晶胞大小/FFT 截断完全无关，可以直接迁移到更大体系而不重训。
把"推理成本"提升为一等公民指标：作者明确指出当前 charge density benchmark 只看 NMAE 或 SCF 步数都不够，端到端 DFT 墙钟才是终极指标——这是 ML for DFT 子领域一个真正重要的方法论转向。
可迁移 trick：用浮动各向异性高斯 + Gram 因式 SPD 协方差代替球谐展开，这套表示在所有"用 3D 函数场近似物理量"的任务里都有用（电势、磁场、流场……）；解析傅里叶 + 泊松求和也可以直接搬到任何周期反演问题（MRI k-space、声学超表面、光子晶体）。
软等变 + 数据增强可行：在大规模周期体系里，硬塞 SO(3) 等变可能不再划算，本文展示"快骨干 + 旋转增强"是一条可行的实用路线。

局限与展望¶

作者承认：周期电荷密度的 NMAE 下限本身比分子体系高（MP/GNoME 上 0.6~0.9%，QM9 上能到 0.1% 量级），SCF 节省也只能拿到上限的 50% 左右；推测原因是 SCF 收敛由初猜在慢收敛长波长模式上的投影误差决定，单看全局 NMAE 找不准方向。
互补性未利用：ELECTRAFI 强在弥散/离子键，ChargE3Net 强在局域定向键合，两者天然互补但还没有 hybrid 模型；可以尝试"全局 Fourier 表示 + 局部 real-space 修正"。
不支持磁体系：当前模型只预测总价电子密度，没有 spin-resolved 输出，spin-polarized DFT（这类结构的初猜尤其昂贵）暂时享受不到加速；需要新的 spin-resolved benchmark。
数据集元数据不完整：MP/ECD 都缺一些跑端到端 DFT 实验所需的输入信息（augmentation occupancies、PAW charges、SCF 设置），评测公允性受限；呼吁未来 benchmark 提供完整 DFT 输入文件。
个人补充：评测全部在 PBE/HSE06 泛函上，对 meta-GGA 或者带强 Hubbard-U 的复杂磁性体系是否依然有 SCF 节省存疑；另外 \(M\)（每价电子高斯数）是个全局超参，对不同元素是否需要 element-aware 自适应也值得探究。

评分¶

新颖性: ⭐⭐⭐⭐⭐ "局部高斯 + 解析傅里叶 + 泊松求和"这套组合既物理直观又解决了 prior art 反复绕不开的实空间镜像求和痛点。
实验充分度: ⭐⭐⭐⭐⭐ 五个 benchmark + 真实 VASP 端到端测试 + 元素级错误分析 + 旋转鲁棒性消融，覆盖到位。
写作质量: ⭐⭐⭐⭐ Method 推导清晰，把"为什么不能直接镜像求和、为什么不能直接预测 \(\hat\rho\)"讲得很透；Related Work 部分对 orbital/probe/frequency-domain 三派的归类很有价值。
价值: ⭐⭐⭐⭐⭐ 第一个真正实现"ML 初猜 → 端到端 DFT 总时间下降"的周期密度模型，且代码开源；对 DFT 加速、材料发现、催化/电池模拟都有直接的落地意义。