Machine-learned, finite temperature Fermi-operator expansions suitable for GPUs and AI-hardware

本文提出了一种基于 SP2 方案、将电子结构计算映射至深度神经网络架构的机器学习增强型有限温度递归费米算符展开方法，该方法通过利用优化的矩阵 - 矩阵乘法与仿射重缩放来避免显式对角化与模型重新训练，从而在 GPU 和人工智能硬件上实现数量级的加速。

要点

以下是用通俗语言和日常类比对该论文的解读。

宏观图景：一种更快的原子模拟方法

想象一下，你试图预测一群人在一个房间（材料）里会如何移动和互动（电子）。在量子物理的世界里，这极其困难。要得到精确答案，你通常必须解决一个庞大而复杂的谜题，称为“对角化”。

把对角化想象成试图通过阅读每一本书的每一页来对一百万本书进行排序，以找到正确的顺序。这很准确，但非常耗时，尤其是当房间变大时。

这篇论文的作者构建了一个捷径。他们不再阅读每一页，而是创建了一个“智能猜测”机器，能够几乎瞬间学会如何排序书籍。他们称之为机器学习费米算符展开。

问题：热人群与冷人群

过去，这些捷径只有在“人群”非常冷（零温度）时才有效。在冷人群中，每个人都站在一条非常可预测的队列中静止不动。数学很简单：你要么在队列里，要么不在。

然而，在现实世界中，事物往往是“热”的。当电子变热时，它们会变得焦躁不安。一些原本站在队列里的人可能会走出来，而一些正在等待的人可能会走进去。这就产生了一个“模糊”的边界，人们处于半进半出的状态。

以前的捷径在这里失败了，因为它们太僵化了。它们无法处理热人群的这种“模糊性”。

解决方案：教神经网络进行“压缩”

作者意识到，用于对冷人群排序的数学结构，与深度神经网络（用于识别人脸或写诗的那种人工智能）的结构完全相同。

• 旧方法（SP2）： 想象一台机器，它接收一个数字，要么将其平方（$x^2$），要么执行特定的减法（$2x - x^2$）。它一遍又一遍地重复这个过程，将数字“压缩”，直到它们变成 0 或 1。这对冷人群非常有效。
• 新方法（MLSP2）： 作者给这台机器装上了一个“大脑”。他们不再使用固定规则，而是利用机器学习来训练这台机器。他们教会它调整自己的内部旋钮（系数），以便能够完美地处理“模糊”的热人群。

可以这样理解：

• 旧机器： 一个僵硬的印章，只能打印“是”或“否”。
• 新机器： 一个灵活的 3D 打印机，它学习如何精确地塑造“是”和“否”，根据人群的热度在两者之间创造出平滑、完美的曲线。

魔法技巧：一个模型适应多种温度

通常，如果你改变模拟的温度，就必须从头重新训练你的 AI 模型。那需要耗费永恒的时间。

作者发现了一个巧妙的技巧，称为仿射重缩放。
想象你有一张城市地图。如果你想放大或缩小，你不需要重画整个城市；你只需要拉伸或缩小地图即可。

作者发现，他们只需要针对特定的“缩放级别”（特定的温度和化学势）训练一次他们的 AI 模型。然后，对于该范围内的任何其他温度，他们只需在将输入数据（哈密顿矩阵）输入模型之前将其“拉伸”。模型不需要重新学习任何东西；它只是以略微不同的比例看到数据，并给出正确的答案。

这意味着他们可以运行温度不断变化的模拟（如化学反应），而无需停下来重新训练 AI。

硬件：利用 AI 芯片进行科学计算

该论文强调，这种方法专为现代计算机芯片构建，特别是GPU（图形处理单元）和张量核心（专为 AI 设计的芯片）。

• 类比： 传统的对角化就像一位大师级木匠手工雕刻每一件家具。它精确但缓慢。
• 新方法： 这就像使用高速 3D 打印机。它利用 AI 芯片的特定架构，以极快的速度执行大规模计算（矩阵乘法）。

作者在 Nvidia RTX 6000 Ada GPU 上测试了这种方法。他们发现，与当今科学家使用的标准高度优化方法相比，他们的方法快 9 到 16 倍，同时仍保持高精度。

结果总结

1. 速度： 在计算材料中电子的行为方面，他们实现了巨大的加速（高达 16 倍），特别是在现代 AI 硬件上。
2. 精度： 他们能够以极高的精度模拟“热”电子（分数占据），这是以前的捷径无法很好做到的。
3. 效率： 通过训练模型一次并利用数学技巧对输入进行重缩放，他们避免了在模拟中每次温度变化时重新训练模型的需求。
4. 无“魔法”对角化： 他们完全避免了对角化缓慢、繁重的数学运算，转而依赖 AI 芯片喜欢执行的快速、重复的乘法步骤。

简而言之，作者将一种缓慢、僵化的数学过程转变为一个快速、灵活、由 AI 驱动的工具，该工具在现代计算机芯片上运行效率极高，使科学家能够比以前更快地模拟复杂材料。

技术摘要

技术摘要：机器学习有限温度费米算符展开

问题陈述
电子结构计算，特别是基于 Kohn-Sham 密度泛函理论（KS-DFT）的计算，受限于对角化哈密顿矩阵以解决特征值问题所带来的立方级缩放成本。尽管递归费米算符展开方案（如二阶谱投影 SP2 方法）提供了一种无需对角化即可直接计算密度矩阵的途径，但现有高效实现仅限于零电子温度。在零温度下，密度矩阵是幂等的（占据数严格为 0 或 1）。然而，许多物理系统——例如金属或处于升高电子温度下的系统——需要分数轨道占据数以准确模拟简并特征态或热展宽。

此前将 SP2 推广至有限温度的尝试涉及截断递归以引入热展宽。然而，这些截断展开本质上是近似的，无法复现精确的费米函数，特别是在精度至关重要的化学势附近。其他替代方法（如切比雪夫展开或帕德近似）要么需要高得多的多项式阶数以抑制吉布斯振荡，要么因反复求解线性系统而产生巨大的计算开销。

方法论
作者提出了一种框架，通过将递归 SP2 方法的代数结构映射到深度神经网络（DNN）架构，将其推广至有限温度。核心见解在于，递归 SP2 更新步骤类似于神经网络的层。通过将展开系数视为可训练的权重和偏置，作者构建了能够以任意温度近似具有分数占据数的费米分布函数的机器学习模型。

关键方法论组件包括：

1. 神经网络架构：

   • MLSP2（机器学习 SP2）： SP2 的推广形式，其中二次更新规则（$X^2$ 或 $2X-X^2$）被替换为带有累加项的可学习二次多项式（$ax^2 + bx + c$）。这使得模型能够近似费米函数的精确热展宽，而非截断的阶跃函数。
   • Max-SP2： 一种更具表达力的架构，包含“跳跃连接”，其中每一层是所有先前层线性组合的平方。
   • Skip-SP2： Max-SP2 的压缩版本，利用近期层和累加器的有限内存，在表达能力和内存使用之间取得平衡。

2. 熵近似：
   作者还开发了一种递归方案来近似电子熵函数 $s(x)$，这对于计算电子自由能是必要的。他们利用费米函数与其补函数的缩放乘积 $f(x)(1-f(x))$ 作为初始猜测，随后通过递归二次展开进行细化，该展开经过训练以匹配化学势处真实熵的二阶导数。

3. 训练与优化：
   模型在单位区间 $[0, 1]$ 内的标量输入上进行训练，而非完整矩阵，使用带有测地加速的 Levenberg–Marquardt 算法。训练数据的采样权重与费米函数的导数成正比，以最小化化学势附近的最大误差。

4. 仿射重缩放与可迁移性：
   一项关键创新是利用仿射重缩放，从而在模拟参数变化时无需重新训练。通过归一化哈密顿量（$H'$）、化学势（$\mu'$）和逆温度（$\beta'$），在特定参数 $(\beta_0, \mu_0)$ 下训练的单一模型可应用于其他参数的广泛“有效区域”。这是通过将输入哈密顿量重缩放以匹配训练条件来实现的，从而允许同一组权重在模拟过程中跨越不同的温度和化学势使用。

5. 硬件实现：
   算法针对现代 GPU 和 AI 硬件（特别是 NVIDIA Tensor Cores）进行了优化。作者利用混合精度算术（FP16/FP32）高效执行矩阵平方运算，并利用哈密顿量的对称性来减少所需的乘法次数和数据传输量。

关键结果

• 精度： MLSP2 模型在费米函数近似方面的误差量级为 $10^{-7}$，显著优于截断的 SP2 方法（误差约为 $10^{-2}$），并在许多区间内与双精度对角化的精度相匹配。
• 性能： 在 NVIDIA RTX 6000 Ada GPU 上，对于中等规模的矩阵，MLSP2 方法相比双精度对角化（使用 cuSOLVER）实现了 16 倍的加速，对于更大规模的矩阵则实现了 9 倍的加速。即使与单精度对角化相比，MLSP2 在保持更高稳定性和精度的同时，也提供了 2 倍至 5 倍的加速。
• 可扩展性： 该方法仅依赖于高度优化的矩阵 - 矩阵乘法内核，避免了对角化。达到目标精度所需的层数随逆温度（$\beta$）呈对数级缩放，从而即使在低温下也能实现高效计算。

意义与主张
本文主张，该方法为有限温度电子结构计算提供了一种稳健且可推广的解决方案，避免了传统对角化带来的计算瓶颈。通过机器学习推广 SP2，作者使得以传统方法一小部分成本计算具有分数占据数系统的密度矩阵成为可能。

其意义在于能够执行动力学有限温度模拟（如量子分子动力学），其中化学势和电子温度在时间步长之间波动。仿射重缩放策略确保了单一预训练模型可在整个模拟过程中重复使用而无需重新训练，从而使该方法适用于大规模应用。此外，该方法专门针对利用现代 AI 硬件（Tensor Cores）的性能特征进行了定制，在保持高数值精度的同时，提供了超越厂商优化的对角化例程的显著加速。