瓶颈在固定基组而非泛函
Kohn-Sham DFT 把多电子问题转化为对单电子轨道的优化,而在分子计算中,这些轨道几乎总是展开在一组固定在原子上的高斯型函数里。被优化的只有展开系数;高斯的宽度和形状按元素一次性设定,存于静态表格。这一选择换来了闭式积分和廉价的自洽场迭代,但也固定了精度和成本随体系规模增长的方式。
论文直接针对这一约束,而不是交换关联泛函。它的主张是:基组本身可以成为可学习的对象,从而让表示去适应分子,而不是事先从 cc-pVXZ 这类标准族中挑选。
浮动高斯取代原子中心函数
GS-DFT 把占据的 Kohn-Sham 轨道表示为 M 个归一化各向异性高斯原语(称为 splat)的线性混合。每个 splat 带九个自由参数:三维空间中的一个中心、三个沿主轴设定宽度的对数特征值,以及一个确定朝向的单位四元数。占据轨道再由这些共享原语与一个系数矩阵混合而成。
由于每个 splat 都是正的,轨道的节点结构必须来自系数符号相反的相邻 splat。没有任何东西把 splat 绑定在原子核上,因此这团高斯云可以在优化过程中移动、拉伸和重新取向。作者把这一构造描述为把 3D 高斯泼溅的渲染器换成量子力学,并证明该参数化在整个正交群下等变,因此分子的刚体运动可以通过变换 splat 参数来吸收,而不必重新优化。
两个求解器组件让高斯云可训练
从这一表示到可用的求解器之间有两道障碍。第一,这些原语彼此并不正交,而轨道必须正交。作者以闭式求出 splat 的重叠矩阵,并在前向传播中通过 Löwdin 正交化强制正交归一,该操作对系数和几何都可微。正则化的特征分解在特征值简并时保持梯度有限,低于下限的特征值对不接收梯度;远离简并时,正则化在收敛处没有影响。
第二,Hartree 能量耦合了每一对 splat 乘积,精确计算需要 O(M^4)。作者利用两个高斯的乘积仍是高斯这一事实,只保留重叠超过阈值的乘积对,作为密度拟合的辅助基组。阈值为零时保留全部乘积对,O(M^4) 的成本随之回来;实际计算中每个 splat 只有常数个重叠邻居,因此辅助函数为 O(M)。辅助基组在 K 个优化步内冻结,然后刷新。在 M=468 的丙氨酸二肽上做扫描时,10^-8 的筛选已经去掉一半乘积对和 46% 的实际耗时,再降低四个数量级只多省 12%;作者采用 10^-7,位于乘积对数量曲线的拐点。
每参数精度是核心比较
核心实证主张是在相同参数数量下做的比较。作者以自由参数数量为横轴,绘制了每原子能量残差、每电子密度总变差和核力最大偏差,并分别以完全基组极限和无限分离极限为基准,覆盖水、乙醇、丙氨酸二肽和更长的链。优化后的 splat 基组在这三个观测量上都系统性收敛,论文称其以远更少的参数达到了最大常规基组的精度。
第二项比较是定性的而非数值的:在相同参数数量下,优化后的基组能捕捉拉伸键和阴离子物理,而固定基组只有借助专门的基组增广才能恢复这些性质。附录报告了裸阴离子和解离能阶梯实验来支持这一点。这些是与能量残差不同的度量,论文没有把它们合并成单一分数。
峰值内存随高斯云规模二次增长
让大体系变得可行的扩展性结果针对的是峰值内存,而不是计算时间。作者报告峰值内存随高斯云规模二次增长,并借此在单个四 GPU 节点上、不做任何修改地以 triple-zeta 尺度模拟了最多 2,742 个原子、10,406 个电子的体系。论文还展示了在含 4,299 个核的 45 个丙氨酸残基链上的分片运行,以相对最佳能量的实际耗时作图。
这是一个关于内存的论证,不应被读成 GS-DFT 比常规 DFT 更快。讨论部分在优化方面说的恰恰相反:GS-DFT 需要更多优化步数,因而收敛更耗时,作者把这视为物理启发的自然梯度方法和摊销方法的切入点。
核力无需 Pulay 修正
核力是基态能量对核坐标的负梯度,驱动几何优化和分子动力学。在原子中心基组中,每个函数随其原子核移动,因此对能量求导需要每个单电子和双电子积分的导数;这一额外项就是 Pulay 力,它在收敛时并不消失。
由于 splat 中心本身也被优化,基组的运动属于内层最小化问题,包络定理把它与系数一起消去,剩下的就是静电力。作者在非简并假设下证明了这一点,并给出了优化提前停止时力误差随残余梯度增长的界,因此未收敛的运行产生的力误差与该残差成正比,而不是不受控制。
证据尚未回答的问题
报告的大体系运行说明求解器能在该硬件上以该规模运行,而不是与成熟代码在相同体系上的基准对比。精度比较是针对作者所选分子、以基组极限为基准做出的,参数效率的主张是在相同参数数量下陈述的,而不是在相同实际耗时下。
优化成本是论文承认的弱点:更多步数、更多收敛时间。论文自己的定位是,这为更好的优化器和摊销方法提供了目标,这意味着 GS-DFT 目前的实用理由建立在内存和参数效率上,而不是速度上。
权衡最终落在哪里
论文中最具决策意义的分界在于资源类型。如果约束条件是内存,或达到给定精度所需的基组参数数量,那么 GS-DFT 正是为此立论:峰值内存随高斯云规模二次增长,以远更少的参数获得大基组精度,并在一个四 GPU 节点上完成 2,742 个原子、10,406 个电子的 triple-zeta 尺度运行。如果约束条件是求解时间,论文并未主张优势,并明确把收敛成本列为限制。
第二个分界是相同参数数量下固定基组与自适应基组的对比。在那里,论文的主张不只是能量更低,而是能触及拉伸键和阴离子物理,而固定基组只有通过专门增广才能达到。这是值得继续关注的比较:GS-DFT 的价值在于一种自适应表示,把参数花在分子需要的地方;悬而未决的问题是,更快的优化能否弥合时间上的差距。
本文引用的文献
Scaling Density Functional Theory with Gaussian Splatting
GS-DFT 把分子轨道表示为可学习的浮动高斯云,并以变分能量为目标通过梯度下降优化,以远更少的参数达到最大常规基组的精度,在单个四 GPU 节点上以 triple-zeta 尺度模拟最多 2,742 个原子(10,406 个电子),代价是需要更多优化步数。
