1. 项目概述：当AI遇见量子化学
“AI赋能密度泛函理论：量子张量学习与机器学习泛函”，这个标题听起来很学术，但它的内核其实非常“工程化”。简单来说，我们正在尝试用当下最火的机器学习方法，去解决一个困扰了计算化学和材料科学领域几十年的老问题：如何又快又准地计算材料的电子结构。
我从事计算模拟工作十几年，从早期的半经验方法，到后来的第一性原理计算，再到如今探索AI辅助的路径，可以说每一步都踩在算力需求和精度瓶颈的钢丝上。密度泛函理论（DFT）是当前材料模拟的绝对主力，它用一个巧妙的“电子密度”概念，把复杂的多电子问题简化成了可解的方程。但DFT的精度，完全取决于一个叫做“交换关联泛函”的近似公式。现有的泛函，比如经典的PBE、B3LYP，要么计算快但精度一般，要么精度高但计算量巨大，像HSE06这种，算个稍大的体系就能让超算集群跑上好几天。
这个项目的核心目标，就是让AI去学习并构建一个更好的“交换关联泛函”。它不是简单地用AI去拟合DFT的计算结果，那样只是换了个黑箱。我们想做的是，让AI深入理解电子密度分布与体系能量之间的内在物理规律，学习量子力学中那些复杂的、高维的张量关系（这就是“量子张量学习”的含义），从而生成一个既保持物理可解释性，又具备超高精度和泛化能力的“机器学习泛函”。
这能解决什么问题？想象一下，新材料研发中，我们需要预测一种新型电池材料的能量密度、催化剂的反应活性、或者半导体的带隙。传统DFT计算要么太慢，要么结果不可靠。如果有一个AI“训练”出来的泛函，能在保持DFT计算框架的前提下，用接近普通泛函的速度，达到甚至超越高精度泛函的准确度，那无疑将是一场革命。它意味着我们可以以前所未有的速度和规模，在计算机里“筛”材料，极大地加速从实验室发现到实际应用的进程。
2. 核心思路与技术路线拆解
2.1 为什么是“量子张量学习”？
要理解这个方向，得先拆解“量子张量”这个概念。在量子化学里，描述一个多电子体系的波函数，其复杂程度随电子数指数增长。DFT的伟大之处在于，它用三维的电子密度函数替代了高维的波函数。但是，电子密度本身包含的信息对于精确描述电子关联效应仍然是不够的。更底层的物理信息，比如电子对密度、动能密度、以及它们的梯度，通常以张量的形式存在。
“张量”在这里可以通俗地理解为一种“多维数组”。电子密度是一个标量场（0阶张量），它的梯度是矢量场（1阶张量），更高阶的导数或相关物理量则构成更高阶的张量。传统的泛函设计，是物理学家基于模型体系（如均匀电子气）和物理直觉，手工构造这些标量、矢量输入与交换关联能量之间的函数关系。这个过程充满了尝试和妥协。
“量子张量学习”的思路，则是将这些不同阶的张量（电子密度、密度梯度、动能密度等）作为机器学习模型的输入特征。机器学习模型，特别是深度神经网络，天生擅长处理高维、非线性的映射关系。我们不再需要手动设计那个复杂的函数形式，而是让神经网络自己去从海量的、精确的量子化学计算数据（如耦合簇CCSD(T)的结果，被视为“金标准”）中，学习出从输入张量到交换关联能量的最优映射。
这里的挑战在于，如何让学习过程尊重物理规律。例如，交换关联能量必须满足尺度变换不变性、旋转不变性等基本对称性。我们不能直接把原始张量扔进一个全连接网络，那样学出来的模型很可能不满足这些物理约束，导致泛化能力极差。因此，我们需要设计具有物理意识（Physics-Informed）的神经网络架构，确保其输入输出关系天然满足这些对称性，这就是“等变神经网络”或“张量网络”等先进模型发挥作用的地方。
2.2 机器学习泛函的两种范式
目前，构建机器学习泛函主要有两种技术路线，各有优劣。
第一种是“直接能量学习”范式。这种范式最为直观：构建一个机器学习模型，其输入是描述整个体系电子结构的一组特征（通常是基于电子密度及其导数的局部描述符），输出直接是整个体系的交换关联能。训练时，使用高精度量子化学方法计算得到的总能量，减去其他能量项（如动能、外势能等），得到“目标”交换关联能作为标签。
注意：这里的关键是特征工程。特征必须足够丰富以捕获电子关联的物理本质，同时又要是局域的，以保证计算的可扩展性。常用的特征包括电子密度、密度梯度模方、动能密度，以及由它们构造出的更复杂的标量，如约化密度梯度。这些特征在每个空间格点上计算，模型则在每个格点上独立地预测交换关联能密度，最后积分得到总能量。
这种方法的优点是端到端，理论上可以逼近任何函数形式。但缺点也很明显：模型学习的是总能量，它可能通过复杂的补偿效应来拟合数据，而未必真正学到了正确的物理规律。这可能导致在训练集之外的材料或化学环境中表现不佳。
第二种是“泛函导数学习”范式，也被称为“学习Kohn-Sham势”。在DFT的实际计算中，我们不仅需要交换关联能，更需要它的函数导数——交换关联势，因为这个势直接进入Kohn-Sham方程，决定了电子的轨道和最终的电子密度。这种范式让机器学习模型直接学习从电子密度到交换关联势的映射。
训练这种模型更为复杂，通常需要一种自洽循环：给定一个初始密度，模型预测一个势，解Kohn-Sham方程得到新密度，比较新密度与初始密度（或与高精度方法得到的参考密度），用差异来更新模型参数。这个过程模拟了真实的DFT自洽计算。
这种范式的优势在于，它直接建模了DFT计算的核心环节，强迫模型学习更本质的物理关系，因此通常具有更好的可移植性和泛化能力。但它的实现难度和计算成本也高得多。
我们的项目更倾向于探索第二种范式，或者将两种范式结合。我们认为，让AI学习产生正确的Kohn-Sham势，是确保机器学习泛函能够无缝嵌入现有DFT软件生态、并稳健工作的关键。
3. 实操构建：从数据到可用的泛函
3.1 数据集的准备与挑战
任何机器学习项目的基石都是数据。对于训练一个通用的机器学习泛函，我们需要一个覆盖元素周期表大部分区域、包含多种化学键类型（金属键、共价键、离子键、氢键、范德华作用）、以及不同电子关联强度体系的数据集。
目前学术界有一些公开基准数据集被广泛使用，例如：
- QM9：包含13.4万个有机小分子，提供了在B3LYP/6-31G(2df,p)级别计算的结构、能量和多种性质。虽然精度不算最高，但数据量庞大，适合初步训练和验证模型架构。
- MD17：包含几个小分子的分子动力学轨迹，提供了原子位置和对应的力，对于训练能够预测原子受力的模型至关重要，因为力是能量的负梯度。
- 3BPA：一个包含3万个二肽构象的数据集，专注于生物相关分子的构象能。
- 更重要的是高精度数据集：如使用CCSD(T)/CBS（耦合簇方法，基组极限）级别计算的小分子数据集（如GMTKN55数据库的子集），或者使用量子蒙特卡洛方法计算的固态数据。这些数据的获取成本极高，但它们是训练高精度泛函的“黄金标准”。
在准备数据时，有几个必须注意的坑：
- 数据一致性：所有数据必须基于同一套物理定义和计算设置生成。例如，能量零点必须统一，是否包含相对论效应需要明确。
- 数值噪声：量子化学计算本身存在数值误差，特别是对于力（能量的梯度）的计算，噪声更大。在训练对力敏感的模型时，需要对数据进行平滑处理或采用对噪声不敏感的损失函数。
- 数据分布：要警惕数据集的偏见。如果数据集中全是有机分子，训练出的泛函在金属或离子固体上很可能失效。因此，构建一个均衡的、具有代表性的数据集是首要任务，通常需要计算化学家和数据科学家的紧密合作。
3.2 模型架构的选择与实现
确定了学习范式（比如学习Kohn-Sham势）后，就要选择具体的神经网络模型。全连接网络（MLP）在处理这类问题时已经力不从心，因为它无法有效处理空间的连续性和对称性。
目前看来最有前景的架构是等变图神经网络。它的设计思想非常贴合我们的问题：
- 将体系视为图：原子是节点，原子间的相互作用（或距离）构成边。这自然描述了分子或晶体结构。
- 等变性：模型输出（如每个原子贡献的势能）会随着输入（如原子坐标）的旋转、平移、镜像等空间变换而进行相应的协变或不变变换。这保证了模型的预测遵守基本的物理规律。
一个典型的流程是：
- 输入原子种类（原子序数）和坐标。
- 模型通过多层消息传递，让每个原子“感知”其化学环境的信息。
- 最终，模型输出每个原子位置处的电子密度值（用于初始化）或直接输出一个与密度相关的势场。
在代码层面，现在有一些优秀的框架可以大大降低实现难度，比如DeePMD-kit、SchNetPack、以及e3nn库。以e3nn为例，它专门用于构建等变神经网络，我们可以这样定义一个有物理意识的层：
import torch import e3nn # 定义输入和输出的不可约表示（irreps） # 标量（l=0）：类型1，宇称偶（+） # 矢量（l=1）：类型1，宇称奇（-） irreps_in = ‘1x0e’ # 输入特征：1个标量（如原子类型编码） irreps_out = ‘1x0e + 1x1o’ # 输出：1个标量（能量）和1个矢量（力） # 创建一个等变的TensorProduct层 tp = e3nn.o3.FullyConnectedTensorProduct(irreps_in, irreps_in, irreps_out) # 假设有一些节点特征和几何信息 node_feats = torch.randn(10, 1) # 10个原子，每个原子1个标量特征 edge_vectors = torch.randn(45, 3) # 45条边，每条边有3维方向矢量 edge_lengths = torch.norm(edge_vectors, dim=1, keepdim=True) # 边长为标量 # 将边长作为标量特征，方向作为几何信息输入层 output = tp(node_feats, node_feats, edge_lengths, edge_vectors)这个简单的例子展示了如何构建一个能同时输出标量（能量）和矢量（力）的等变层。在实际模型中，我们会堆叠很多这样的层，并引入注意力机制等，让模型能够学习长程相互作用。
3.3 损失函数设计与训练技巧
训练一个机器学习泛函，损失函数的设计是灵魂。我们不能只简单地最小化预测能量和参考能量之间的均方误差（MSE）。
一个更有效的复合损失函数通常包含以下几项：
- 能量损失：预测的总能量与参考总能量之差的MSE。这是基础。
- 力损失：预测的原子受力（能量的负梯度）与参考力之间的MSE。引入力损失至关重要，因为它为模型提供了远比能量更丰富的梯度信息，能极大地提升模型的精度和收敛速度。通常会给力损失一个较大的权重。
- 势损失：如果我们采用学习Kohn-Sham势的范式，还需要加入预测的交换关联势与参考势之间的差异。参考势可以通过高精度计算得到的电子密度，用数值方法反推出来，这个过程比较繁琐。
- 正则化损失：为了防止过拟合，需要加入L1或L2正则化项。对于泛函模型，物理驱动的正则化可能更有效，例如惩罚那些在电子密度极稀或极密区域行为不合理的预测。
训练过程本身也是一个挑战。由于DFT自洽计算的存在，我们常常需要在一个“循环”中训练模型：
- 外层循环：用当前模型参数进行DFT自洽计算，得到预测的电子密度和能量。
- 内层循环：比较预测结果与参考数据，计算损失，反向传播更新模型参数。
这本质上是一个双层优化问题，计算开销巨大。实践中，我们常采用“离线学习”策略：先使用高精度方法计算一大批不同电子密度分布对应的“理想”Kohn-Sham势，构建一个“（密度，势）”配对的数据集，然后直接在这个数据集上训练模型学习映射关系，暂时绕过自洽循环。模型初步训练好后，再放入真实的DFT代码中进行微调和测试。
4. 集成与应用：让AI泛函“跑起来”
4.1 与现有DFT代码的对接
训练出一个神经网络模型只是第一步，如何将它变成一个能被主流DFT软件（如VASP、Quantum ESPRESSO、ABINIT、GPAW）调用的“泛函”，才是工程上的关键。
目前主流的有两种集成模式：
- Libxc模式：大多数DFT程序通过调用Libxc这个库来获取交换关联泛函的能量和势。我们可以仿照Libxc的接口，开发一个“机器学习插件”。这个插件以共享库（.so或.dll）的形式存在，它内部封装了我们的神经网络模型。当DFT程序计算到某个空间格点时，它会将当地的电子密度、梯度等参数传递给这个插件，插件调用模型进行计算，并返回该点的交换关联能密度和势。这种方式侵入性小，理论上任何支持Libxc的DFT程序都能使用。
- 内置模式：直接将机器学习模型的调用代码写入DFT程序的源代码中。这种方式性能优化潜力更大，可以与程序本身的并行化、网格积分等流程深度结合，但可移植性差，维护困难。
以相对简单的Libxc插件模式为例，其核心是实现Libxc定义的几个标准函数接口：
xc_func_init: 初始化函数，在这里加载我们训练好的模型权重。xc_lda_exc_vxc: 对于LDA（局域密度近似）类型的泛函，这个函数被调用。输入是当前格点的密度，输出是该点的交换关联能密度和势。xc_gga_exc_vxc: 对于GGA（广义梯度近似）类型的泛函，输入还包括密度梯度。
我们的任务就是在这些函数内部，将DFT程序传递过来的标量或矢量参数，组织成我们的神经网络模型所需的输入张量格式，进行前向推理，然后将结果返回。这里需要注意数据精度（float/double）、内存布局以及计算效率，因为每个自洽迭代中，这个函数会被调用数百万甚至数十亿次（取决于体系大小和网格精度）。
4.2 实际计算流程与参数设置
当机器学习泛函成功集成后，使用它进行计算与使用传统泛函在流程上几乎没有区别，这降低了使用门槛。以下是一个概念性的输入文件示例（以某种DFT程序格式为例）：
&SYSTEM calculation = ‘scf’ functional = ‘ML-XC’ # 指定使用我们的机器学习泛函‘ML-XC’ ... / &ELECTRONS conv_thr = 1.0e-8 # 由于ML泛函可能更“敏感”，建议使用更严格的自洽收敛阈值 ... / &ML_XC # 机器学习泛函的专属参数块 model_path = ‘./best_model.pth’ # 训练好的模型权重文件路径 descriptor_type = ‘el_density+gradient+kinetic’ # 模型使用的输入描述符类型 cutoff_radius = 6.0 # 原子局域环境截断半径（单位：玻尔） gpu_id = 0 # 指定使用哪块GPU进行模型推理加速 /在实际计算中，有几个参数需要特别关注：
- 截断半径：对于基于原子环境的模型，需要设定一个截断半径来定义每个原子的“邻居”。半径太小会丢失长程作用信息，太大会急剧增加计算量。对于分子，6-8玻尔通常足够；对于具有长程库仑作用的体系，可能需要特殊处理。
- 网格精度：DFT需要在实空间网格上进行积分。机器学习泛函的预测在网格点间应该是平滑的。如果网格太粗糙，可能会错过密度变化的细节，影响精度；太精细则计算量暴增。建议开始时使用与传统泛函相同或稍密的网格设置。
- 自洽场收敛：神经网络泛函的势能面可能与传统泛函不同，可能导致自洽迭代收敛行为发生变化。有时需要调整电子弛豫的算法（如使用更稳健的Kerker混合）或降低混合因子来保证收敛。
4.3 性能优化与推理加速
推理速度是机器学习泛函能否实用的生命线。在DFT的每个自洽迭代中，需要对成千上万个网格点进行模型评估。优化至关重要：
- 模型轻量化：在保证精度的前提下，尽可能减少模型的层数和参数量。可以使用知识蒸馏、剪枝、量化等技术来压缩模型。
- 批处理推理：不要对每个网格点单独调用模型。将空间相邻的多个格点的特征数据（密度、梯度等）打包成一个批次（Batch），一次性送入模型进行并行计算，能极大利用GPU的并行计算能力。
- 缓存与预计算：在自洽迭代过程中，原子位置不变，因此每个原子的邻居列表、距离、方向矢量等几何信息是固定的。这些信息可以预先计算并缓存，避免在每次模型调用时重复计算。
- 使用专用算子库：利用像TensorRT或OpenVINO这样的推理优化框架，将PyTorch或TensorFlow模型转换为高度优化的计算图，可以在特定硬件（如NVIDIA GPU或Intel CPU）上获得显著的加速比。
一个经过良好优化的机器学习泛函插件，其单点推理时间应该控制在与一次中等复杂度GGA泛函计算相当的水平。虽然初期开销会比传统泛函大，但考虑到它可能一步达到需要昂贵杂化泛函才能获得的精度，其综合性价比优势会非常明显。
5. 效果验证、常见问题与未来展望
5.1 如何评估一个机器学习泛函的好坏？
不能只看它在训练集上的表现，必须进行严格、全面的基准测试。一个标准的测试流程应包括：
- 分子基准测试：使用像GMTKN55这样的数据库，它包含了55个子集，涵盖反应能、势垒、非共价相互作用等各类化学性质。计算每个子集的平均绝对误差（MAE）和均方根误差（RMSE），并与主流泛函（PBE, B3LYP, ωB97X-D, M06-2X等）以及高精度波函数方法（如CCSD(T)）进行对比。一个好的机器学习泛函应该在绝大多数子集上接近或超越最好的传统泛函。
- 固态性质测试：计算晶格常数、结合能、体模量、电子带隙等。这对于材料应用至关重要。需要与实验值以及使用HSE06等昂贵泛函的计算结果对比。特别注意带隙的预测，这是传统DFT的软肋。
- 势能面扫描：选取一个简单的双原子分子（如N2），计算其键长-能量曲线（解离曲线）。观察机器学习泛函预测的平衡键长、振动频率和解离极限是否正确。错误的泛函可能会产生不合理的势能面，比如出现虚假的极小值点。
- 泛化能力测试：用完全不含在训练集中的新元素、新成键类型或新体系（如表面吸附、缺陷体系）来测试模型。这是检验其是否真正学习到物理规律，而非简单记忆数据的关键。
5.2 实操中遇到的典型问题与解决思路
在开发和测试我们自己的原型时，踩过不少坑：
模型在训练集上完美，在新体系上崩溃
- 现象：对一个有机分子数据集训练的模型，在预测金属团簇时能量严重错误。
- 诊断：根本原因是数据偏差和特征局限性。模型从未见过金属体系中高度离域的电子密度分布。
- 解决：扩充训练集，必须包含金属、离子化合物、半导体等多样化的体系。同时，审视输入特征是否足以描述金属键的特性，可能需要引入像电子局域函数（ELF）这类更能刻画电子离域的特征。
自洽计算不收敛或振荡
- 现象：使用ML泛函时，电子密度在迭代中剧烈振荡，无法达到收敛阈值。
- 诊断：机器学习模型预测的交换关联势可能存在数值噪声或不光滑的区域，导致Kohn-Sham方程难以求解。
- 解决：首先，检查模型预测的势场在空间上是否平滑。可以在损失函数中加入对势场二阶导数的正则化项，强制其光滑。其次，调整DFT计算中的电子弛豫参数，如使用更小的混合因子、启用Kerker预处理、或换用更稳定的对角化算法。
计算速度慢得无法接受
- 现象：集成后的ML泛函计算一个中等体系的时间是传统GGA的100倍。
- 诊断：可能是网格点推理没有批处理，或者模型本身过于庞大，又或者频繁在CPU和GPU之间拷贝数据。
- 解决：实现高效的批处理推理；对模型进行剪枝和量化；确保特征计算和模型推理都在同一设备（GPU）上完成，避免数据传输瓶颈；使用TensorRT等工具进行图优化和内核融合。
力预测不准导致结构优化失败
- 现象：用ML泛函进行几何结构优化时，原子受力方向怪异，优化出的结构不合理。
- 诊断：能量和力的预测不一致。在机器学习中，力是能量的负梯度，通常通过自动微分获得。如果模型本身关于坐标的微分性不好，或者训练时力损失的权重不够，就会导致此问题。
- 解决：确保模型架构是光滑且可微的。在训练时，显著提高力损失项的权重（有时能量和力损失的权重比可以达到1:100甚至更高）。在发布泛函前，务必进行严格的力测试。
5.3 当前的局限与未来的方向
尽管前景广阔，但AI赋能的密度泛函理论仍处于早期阶段，面临诸多挑战：
- 数据饥渴：构建一个覆盖化学空间全局的高精度数据集成本极高。如何用更少的数据训练出更好的模型，比如通过主动学习、迁移学习、或利用物理规律生成合成数据，是重要方向。
- 外推能力：模型在训练数据分布之外的区域（如极高的压力、极端的电子密度）表现不可靠。将已知的渐进行为（如高密度极限、低密度极限）作为硬约束嵌入模型，是提升外推能力的有效手段。
- 可解释性：神经网络是黑箱。我们如何理解它做出的预测？发展可解释AI技术，例如通过归因分析找出是哪些输入特征对最终预测贡献最大，有助于我们反过来加深对交换关联物理的理解，甚至可能发现新的物理规律。
- 动力学与激发态：目前的努力主要集中在基态性质。如何将机器学习泛函扩展到含时DFT以处理激发态，或用于分子动力学模拟，是下一个前沿。这需要模型能正确描述势能面随时间和电子态的变化。
从我个人的实践来看，这条路虽然艰难，但每一步都令人兴奋。我们不再仅仅是泛函的使用者，某种程度上成为了泛函的“设计者”。每一次调试模型、分析错误，都像是在与量子力学的基本规律进行一场对话。最大的体会是，成功的机器学习泛函项目必须是一个深度交叉的团队作业：计算化学家提供物理洞见和高质量数据，机器学习专家设计稳健的模型和训练策略，软件工程师实现高性能的集成和部署。当看到自己训练的模型第一次成功预测出一种新材料稳定存在的相，并且与后来的实验结果吻合时，那种成就感是无可替代的。这个领域没有银弹，每一个百分点的精度提升，都需要对数据、模型和物理的深刻理解与反复打磨。
