在物理建模和机器学习交叉的论文标题里,等变学习(equivariant learning)与三维经典密度泛函(three-dimensional classical density functional)是一组出现频率很高的关键词。它们放在一起,目标通常不是再做一个纯数据拟合的黑箱,而是训练一个尊重物理对称性、能在不同体系之间迁移的自由能泛函代理模型。经典密度泛函关心的是如何由三维局域密度分布计算自由能以及平衡密度;等变学习关心的是当坐标系发生旋转、平移或镜像变换时,模型输出的变化方式是否符合物理规律。两者结合后,模型可以避免靠数据增强硬凑对称性,也不会在训练集之外因为坐标系变化而产生不合理跳变。这篇文章不复制某篇具体论文的完整实验,而是围绕这一方向整理一条从物理概念、训练数据、模型骨架到验证与排错的实践主线,重点放在“一个可迁移的三维经典密度泛函代理模型需要怎样的等变结构”。
1. 等变学习在经典密度泛函建模里到底解决什么
1.1 经典密度泛函的基本物理对象
经典密度泛函理论处理的不是量子力学波函数,而是经典粒子系统的空间密度分布。设体系中粒子数密度为三维空间的函数 (\rho(\mathbf{r})),温度 (T)、化学势 (\mu) 和外势场 (V_{\rm ext}(\mathbf{r})) 已知,那么系统的巨势可以写成密度泛函:
[ \Omega[\rho] = F[\rho] + \int \rho(\mathbf{r}) \bigl(V_{\rm ext}(\mathbf{r}) - \mu\bigr) , d\mathbf{r} ]
其中 (F[\rho]) 是亥姆霍兹自由能泛函。平衡状态下密度分布满足一阶变分为零:
[ \frac{\delta F[\rho]}{\delta \rho(\mathbf{r})} = \mu - V_{\rm ext}(\mathbf{r}) ]
这个方程把自由能泛函和外势联系了起来。只要知道精确的 (F[\rho]),就能求解界面、受限流体、胶体、结晶前驱体等一系列复杂现象。但真实体系的自由能由粒子间相互作用决定,精确泛函几乎不可能解析写出。理想气体部分有明确形式,麻烦的是由粒子相互作用贡献的过量自由能 (F_{\rm ex}[\rho])。
机器学习要处理的问题自然浮现:用大量来自分子模拟、密度泛函微扰或实验反演的数据,学习一个 (F[\rho]) 的代理模型。由于真实系统里三维密度分布是空间连续对象,并且由大量粒子坐标或三维网格描述,模型必须高效处理 (N) 个局部环境,不能简单把全空间拍平成超大向量。
1.2 三维体系对机器学习代理模型的隐性约束
三维经典密度泛函的输入通常来自三种形式之一:粒子坐标集合、离散网格上的密度值、以某个原点为中心的局域密度场。无论采用哪种形式,物理量都满足一个基本要求:坐标系只是一个描述工具,不是物理本体。把整个体系刚性旋转、平移或镜像后,自由能这个标量不应改变。
如果模型直接学习
[ F_{\rm pred} = \mathrm{MLP}(\rho) ]
而输入使用未经处理的绝对坐标,那么模型必须靠自己从训练样本中学习“旋转后结果不变”。即使加了随机旋转数据增强,模型也只能逼近这种对称性,很难严格保证。真实测试样本如果旋转角度不在增强分布内,预测结果就可能漂移。
等变学习把这个问题从“希望模型学到”变成“从结构上强制满足”。模型不再消费绝对坐标,而是只使用相对位置、相对距离、方向信息,并通过具有旋转等变性的特征构建层逐级传递。这样,即使输入的全过程被旋转,模型内部特征也会以相同的旋转矩阵被变换,最终的自由能预测保持不变。
1.3 不变输出与等变输出的区别
需要先分清两个概念:标量自由能是“不变”目标,矢量或张量场是“等变”目标。
表格中列出常见情况和设计约束:
| 物理目标 | 数学要求 | 模型输出示例 | 网络设计方向 |
|---|---|---|---|
| 总自由能 (F) | 旋转平移后数值不变 | 单个实数 | 最终汇聚成标量,中间特征可使用等变表示 |
| 化学势或局域自由能密度 | 随体系刚性旋转随动 | 每个位置一个标量 | 局域预测后整体汇聚 |
| 密度梯度或力 (-\nabla F) | 输入旋转后输出按相同旋转矩阵变换 | 每点三维向量 | 等变向量特征或对能量场自动微分 |
| 应力、取向张量相关量 | 输入旋转后输出按张量规则变换 | (3\times3) 对称张量 | l=0、l=2 等不可约表示组合 |
文章场景里最常见的是第一类:以某种标量自由能为监督目标。最终输出只要做到旋转和平移不变即可。但要注意,只做输出不变还不够,中间特征最好也采用等变特征,因为这样才能让网络通过方向性相互作用学习三维结构差异。否则即使输出是标量,网络也会退化成只能使用距离信息的纯径向模型,对角度排列不敏感。
注意:标量不变不等于等变。等变是一个包含输出变换规则的框架。预测自由能时,模型对真实旋转应输出相同结果,这只是等变中的“不变表示”;当后续要输出矢量力或应力时,必须使用真正的矢量等变操作。
2. 三维可迁移密度泛函代理模型的建模路线
2.1 自由能泛函中学什么、不学什么
在实际建模中,最好不要让模型同时学习理想气体项和过量项。理想气体自由能已经有解析形式,把它也塞进神经网络会让模型浪费参数,而且会对密度为负或密度极低的区域产生不合理预测。常见做法是把已知解析部分从标签中扣除,只让模型学习过量自由能贡献。比如构造目标:
[ F_{\rm target} = F_{\rm reference} - F_{\rm id}[\rho] ]
之后预测出过量自由能,再在推理阶段加回解析项。这样做的好处是标签动态范围更小,跨状态迁移也更容易。训练数据可以来自不同温度、不同化学势下的模拟构型,而不是只依赖某个单一状态点的轨迹。
2.2 输入描述符的五种常见形式
三维经典密度泛函模型要处理高度几何化的输入。常见输入描述符如下:
- 原始粒子坐标:直接给原子三维坐标和原子类型,模型通过邻域图或消息传递建立局部环境。
- 吸附/受限区域网格密度:把密度场离散到三维网格,用卷积或连续卷积处理。此时网格分辨率决定精度。
- 中心粒子或探针密度图:以一个流体粒子为中心,描述周围粒子密度分布,适合构建局域单粒子泛函。
- 分子内坐标与方向向量:对非球形分子,除质心位置外还要提供朝向,例如分子键方向向量,这需要奇数 l 的旋转特征。
- 外势场分布:一部分模型不是直接输入密度,而是输入外部势 (V_{\rm ext}(\mathbf{r})),用模型隐式求解对应密度。要做到可迁移,通常仍需将外势场嵌入局部等变描述符。
无论选择哪种形式,模型的标定函数都会把多个粒子的空间关系映射成能量贡献。这时,“图”是非常自然的中间结构:节点是粒子,边是空间距离小于截断半径的粒子对。
2.3 输出头与自动微分扩展
如果只输出总自由能,网络可以在全局池化后输出一个标量。但很多实际应用需要一阶导信息,例如用
[ \frac{\partial F_{\rm pred}}{\partial \rho(\mathbf{r})} ]
近似化学势。如果模型输入是可微的密度场或粒子坐标,可以直接用自动微分从自由能预测得到这个导数。这样做有一个额外收益:即使训练标签里没有导数,只要输出自由能足够准,反向传播出来的导数也保留了物理约束。
对于粒子坐标输入,模型也可以输出每个粒子的力。此时网络最后需要返回每个原子的三维向量特征,而且这个向量必须在输入体系旋转时按旋转矩阵同步变换。严格实现要比输出标量复杂,因此建议先跑通自由能标量任务,再扩展到力。
2.4 可迁移性从哪里来
所谓可迁移,指的是模型在一个化学体系、一个热力学状态或一种外势形状上训练,却可以被应用到另一个相关体系或状态。在经典密度泛函任务中,可迁移性通常来自三方面:
- 局域化学环境共享:不同体系中粒子的第一近邻排布存在相似结构,网络只需要学会从局域构型映射到局域自由能贡献,而不是记住整段轨迹。
- 元素类型嵌入:用嵌入向量表示不同元素或粒种,而不是把每个元素当作独立 one-hot 分类标签的末端。
- 物理量的无量纲化和归一化:训练前把坐标、密度、能量分别除以特征长度、密度尺度和能量尺度,让模型更容易跨状态复用。
如果一个代理模型只能靠记忆全体系模式,那么它换一个容器形状或外势就会失效。等变图神经网络天然以局部环境为主,因此更适合作为可迁移代理模型的骨架。
3. 最小实验设计:数据、等变骨架与配置
这一节给出一个用于跑通流程的最小实验。正式研究需要替换成有物理意义的参考自由能数据,但下面代码足以验证等变管道和训练代码是否正确。
3.1 环境准备
示例使用 Python、PyTorch、e3nn 和常用分子工具库。版本变化快,建议安装时先查询对应文档。
conda create -n cdft-equiv python=3.9 conda activate cdft-equiv # PyTorch 需要根据自己的 CUDA 环境选择安装命令 pip install torch torchvision # 等变网络组件、图神经网络基础库、原子结构处理库 pip install e3nn torch-geometric ase pymatgen numpy scipy如果只做等变结构和训练流程验证,不实际运行分子模拟,也可以只安装torch和e3nn。ASe 和 pymatgen 用于读取真实分子模拟轨迹或晶体结构,不是最小模型必需品。
3.2 构造数据接口与合成验证样本
真实项目中数据往往来自 LAMMPS、GROMACS 或自研蒙特卡洛程序,保存为轨迹文件和每个快照的自由能标签。下面先定义一个数据集接口,字段包括坐标、原子类型和标量标签。
import torch from torch.utils.data import Dataset class DftFrameDataset(Dataset): """每个样本是一个三维体系快照。 Parameters ---------- coords: list of numpy arrays, shape (N, 3) atom_type: list of numpy arrays, shape (N,) free_energy: list of float """ def __init__(self, coords, atom_type, free_energy): self.coords = [torch.as_tensor(c, dtype=torch.float32) for c in coords] self.atom_type = [torch.as_tensor(a, dtype=torch.long) for a in atom_type] self.free_energy = [torch.as_tensor([e], dtype=torch.float32) for e in free_energy] def __len__(self): return len(self.coords) def __getitem__(self, idx): return { "coords": self.coords[idx], "atom_type": self.atom_type[idx], "energy": self.free_energy[idx], }为了在没有真实数据时先检查训练链路,可以用一个显式满足旋转不变性的合成函数生成标签。这里用任意粒子间的距离函数构造能量。合成标签没有任何物理意义,只用于验证模型能学习一个几何函数,真实实验必须用物理模拟参考值替换。
def synthetic_energy(coords): """只用于链路验证:由两两距离生成的旋转平移不变标量。""" d = coords[:, None, :] - coords[None, :, :] dist2 = (d * d).sum(dim=-1) mask = ~torch.eye(coords.shape[0], dtype=torch.bool) dist2 = dist2.masked_fill(~mask, float("inf")) # 距离越小,贡献越大 return (-torch.exp(-dist2 / 4.0)).sum()真正评估时,这段合成数据会被替换成经典 DFT 参考解或分子模拟数据的自由能标签。这样做的唯一价值是让数据接口、模型、损失和对称性检查先闭环。
3.3 等变网路骨架需要哪些模块
不强行要求在这个阶段手写完整论文架构。最小可理解骨架至少包含以下部分:
- 嵌入层:把原子类型整数映射为标量特征。这个特征初始时并不带方向信息。
- 边向量构造:由邻居坐标差生成相对向量 (\mathbf{r}_{ij})。
- 等变消息传递:用欧氏距离作为径向权重,用相对方向构造旋转方向特征,再和节点特征做张量积,得到更高阶的标量和向量特征。
- 门控非线性:对每个不可约表示使用合理的非线性激活方式,避免简单在向量分量上逐元素套 ReLU 破坏等变性。
- 汇聚输出:对节点特征做求和或注意力池化,合并成标量自由能。
下面是一个概念性骨架,目的是展示输入输出张量的形状变化,不是可直接复制的完整模块。
import torch from torch import nn from e3nn import o3 class EquivariantDFTModel(nn.Module): def __init__( self, num_types: int = 4, max_radius: float = 6.0, irreps_hidden: str = "32x0e + 16x1o + 8x2e", num_layers: int = 3, ): super().__init__() self.max_radius = max_radius self.irreps_hidden = o3.Irreps(irreps_hidden) self.embedding = nn.Embedding(num_types, 32) # 在实战中推荐调用 e3nn 的 MessagePassing 或 NequIP 的卷积层 # 而不是直接手写 TensorProduct。这里用列表表示多层结构。 self.layers = nn.ModuleList([ nn.Linear(32, self.irreps_hidden.dim) if i == 0 else nn.Linear(self.irreps_hidden.dim, self.irreps_hidden.dim) for i in range(num_layers) ]) self.head = nn.Linear(self.irreps_hidden.dim, 1) def build_graph(self, coords): # 示例只展示向量差计算,真正的邻居列表应使用 torch_geometric。 diff = coords[:, None, :] - coords[None, :, :] dist = torch.norm(diff, dim=-1) + 1e-6 return diff, dist def forward(self, data): _, dist = self.build_graph(data["coords"]) # 等变层稍后在这里循环。 # 关键点:必须保证网络不使用绝对坐标,只使用相对距离和方向。 h = data["coords"].shape[0] * 0.0 # 示意占位 return {"energy": h}上面代码的 forward 并没有真正完成等变卷积,只强调两点:输入侧要计算相对几何量,输出层最终得到一个标量。在实际项目中,应该把self.layers换成成熟的等变卷积层,例如 NequIP 中的InteractionLayer、MACE 中的 message passing,或 e3nn 社区常用卷积模块。
3.4 训练配置:把超参数显式化
经典密度泛函模型需要大量超参数,推荐用 YAML 保存,方便复现和做网格搜索。
model: name: equivariant_cdft_surrogate num_types: 6 max_radius: 5.0 irreps_hidden: "48x0e + 24x1o + 12x2e" num_layers: 3 radial_basis: "bessel_10" cutoff_embedding: 64 train: seed: 42 batch_size: 4 learning_rate: 0.001 epochs: 200 scheduler_gamma: 0.5 scheduler_step: 60 data: train_frames: 1000 valid_frames: 200 test_frames: 200其中max_radius控制每个粒子能看到多远的邻居。irreps_hidden中的0e表示标量特征,1o表示三维向量特征,2e表示 l=2 的偶宇称张量。0e、1o并不直接决定模型的表达能力,但它决定了网络能够携带多少方向信息。
3.5 损失函数与训练循环
如果只预测每个体系的总过量自由能,损失函数可以是最小均方误差:
def train_step(model, optimizer, batch): model.train() energy_pred = model(batch) loss = torch.nn.functional.mse_loss( energy_pred["energy"], batch["energy"], ) optimizer.zero_grad() loss.backward() optimizer.step() return loss.item()实际研究中经常还要配合化学势或力的标签,这时可以把损失写为:
[ \mathcal L = \lambda_E |F-F_{\rm ref}|^2 + \lambda_\mu |\mu-\mu_{\rm ref}|^2 ]
能量自由能和导数自由能同时监督,有助于提高拟合精度。下节将单独讨论验证问题,因为只把 train loss 压到很低并不能说明模型学到了正确的三维泛函。
4. 训练验证:不能只盯误差,要检查对称性与迁移性
4.1 数据划分要按状态而不是按帧随机切
在物理模拟数据里,相邻帧高度相关。如果随机划分训练集和测试集,两个集合可能来自同一条轨迹的连续状态,模型只需要“记住”轨迹就能得到很高的测试精度,迁移性评价会失真。更合理的方式是按状态点切分:
- 不同温度、不同压力或不同外势形状分别进入训练集或测试集。
- 训练集里出现的化学体系与测试集体系尽量有差别。
- 对可迁移性实验,至少保留一个体系做留出测试。
例如训练集使用简单球形流体,测试集使用同一类型但更强外势或更大密度梯度状态。这样才能知道模型是否学到了物理规律,而不是只拟合了训练分布。
4.2 旋转平移一致性要作为硬校验写进测试流程
即使网络结构声称是等变的,工程实现中的邻居排序、边界条件、坐标归一化都可能破坏这一性质。每次训练前都应该单独验证一组随机旋转、平移和镜像操作:
def check_rotation_invariance(model, sample, seed=0): torch.manual_seed(seed) rot_matrix = torch.linalg.svd(torch.randn(3, 3))[0] rot_matrix = rot_matrix.float() coords = sample["coords"].clone() coords_rot = coords @ rot_matrix.T # 如果模型内部使用了相对坐标,平移不会改变预测,此时也测试一下 coords_shifted = coords_rot + torch.tensor([10.0, -3.0, 5.0]) pred0 = model({**sample, "coords": coords})["energy"] pred1 = model({**sample, "coords": coords_rot})["energy"] pred2 = model({**sample, "coords": coords_shifted})["energy"] print("rot residual:", torch.abs(pred0 - pred1).item()) print("shift residual:", torch.abs(pred0 - pred2).item())如果残差大于 (10^{-4}) 量级,说明模型的某个环节仍在依赖绝对位置。常见来源是坐标直接进入线性层、特征里混入绝对坐标信息,或者按输入顺序拼接全局位置向量。
4.3 评估指标设计
只使用全局自由能误差不够。建议下表组成一组成套指标:
| 指标 | 计算对象 | 能判断什么 |
|---|---|---|
| 自由能 RMSE | 测试集总自由能 | 整体拟合精度 |
| 单粒子平均绝对误差 | 自由能除以粒子数 | 不同体系大小间的误差可比性 |
| 旋转残差 | 同一构型旋转前后预测差 | 模型是否严格满足旋转不变性 |
| 平移残差 | 同一构型平移前后预测差 | 模型是否误用绝对坐标 |
| 导数误差 | 化学势或力预测与参考值 | 一阶变分是否可靠,能否用于迭代求解密度 |
| 迁移误差 | 留出体系或状态 | 模型泛化能力 |
对经典密度泛函模型来说,导数比总自由能更难学。若目标是在平衡密度求解中使用,建议至少在一个测试集上计算密度更新后的收敛曲线,而不仅仅是查看能量误差。
4.4 与普通非等变基线的对比要控制变量
为了说明等变结构有价值,可以训练一个没有方向信息的纯距离基线模型作为对比。比如把每组坐标转换成所有原子对的径向距离矩阵,用 MLP 输出自由能。这个基线在简单相距作用的体系中可能足够;遇到角度相关性强的体系会变差。对比实验要注意控变量:训练数据、截断半径、特征维度、训练轮数都应一致。
比较后通常能看到两个现象:一是在同等样本量下,等变模型的测试误差更小;二是旋转扰动下等变模型不会出现明显误差上升,而普通基线即使训练时加入了数据增强,也可能残留偏差。
5. 常见问题与排查路径
5.1 模型声称等变,但旋转测试不过
旋转测试失败的顺序优先级:
- 检查是否还有坐标进入没有相对化的全连接层。
- 检查原子类型特征与边向量的拼接方式,不要把绝对坐标向量直接拼入特征。
- 检查邻居列表构建是否有随机哈希或字典序依赖,同样的原子集合被旋转后邻居顺序不应影响结果。
- 如果使用周期边界,检查最小镜像约定是否在旋转后仍成立。
如果经过以上检查仍失败,可以做小规模单样本测试,只保留两个原子,分别旋转 0 度和 90 度,打印中间层特征,定位第一步差异出现的位置。
5.2 训练 loss 不下降
可以从数据范围检查:自由能标签数量级差异很大时,先用每个体系的粒子数归一化。比如总自由能从几百到几千不等,而网络输出初始化靠近零且权重很小,梯度会被大量样本的大标签主导。
另一个常见原因是截断半径太小。经典密度泛函中的长程静电或分散相互作用如果无法被截断覆盖,模型就会丢失重要远距离信息。训练前统计每个粒子在给定max_radius内的平均邻居数,过低时模型无法感知密度长程变化。
5.3 迁移到新体系误差大
如果训练集内部拟合很好,但留出体系误差大,常见原因有:
- 标签没有扣除解析理想气体项,导致模型把理想项也带进了神经网络,跨状态拟合更难。
- 训练集状态范围太窄,模型没有见过不同密度量级的构型。
- 元素类型数太少或嵌入维度不足,导致新体系只能落在 embedding 以外。
- 坐标没有按特征尺度单位化,不同体系的粒子尺寸不同,直接使用纳米和埃混合数据。
改进顺序应是:先检查数据归一化和标签分解,再检查特征表示的物理单位,最后再增加模型复杂度。
5.4 高阶不可约表示带来的训练不稳定
直接使用很大的irreps_hidden,比如128x0e + 64x1o + 32x2e,在小数据集上容易产生过拟合且训练不稳定。不可约表示的阶数越高,参数越多,非线性门控越复杂。应从低阶开始尝试,例如32x0e + 16x1o,确认能跑通后再增加 l=2 甚至 l=3 特征。下表给出一组调整建议:
| 现象 | 可能原因 | 调整方向 |
|---|---|---|
| 训练不收敛 | 输出标签未归一化 | 按体系粒子数归一化,修正标签尺度 |
| 旋转测试残差大 | 模型使用绝对坐标或绝对平移拼接 | 改为相对坐标和边向量 |
| 迁移误差大 | 训练状态单一 | 加入不同外势和密度的训练数据 |
| 高频振动不饱和 | 截断半径过小 | 分析邻居数,增大半径 |
| 网络过大反而差 | 不可约表示阶数过高 | 降低irreps_hidden,先增大数据量 |
5.5 数据口径错误
经典密度泛函的数据有多种标签来源:分子模拟自由能微扰、局部密度积分、参考泛函数值解。不同来源标签的单位和零点可能不同。建议数据加载后统一转换单位,并在保存前把每个样本的粒子数、化学势、温度和标签一起写入文件头。训练以前先输出一张分布直方图,确认没有异常零点或离群体系。
6. 为真实实验准备的最佳实践清单
进入正式研究和工程前,可以在项目目录建立下面的检查清单,把它当成生产级工作的最低门槛。
- 确认对称性定义。论文体系只有旋转不变性,还是也要求镜像不变性?晶格取向是否涉及特殊欧拉角?
- 确认单位。坐标按原子长度还是约化长度?能量按单个粒子还是按每个分子?
- 确认自由能分解。是否已经把理想气体部分从标签中扣除?
- 数据如何划分。是否包含按温度、化学势或粒子类型切分的留出集?
- 邻居下标是否可复现。坐标旋转后邻居序号是否重复?
- 损失函数是否包含导数。如果最终要迭代求平衡密度,导数标签也必须有验证。
- 等变测试是否作为 CI 流程的一部分。
- 是否设置固定随机种子,并在不同运行间对比模型正常波动。
- 是否保存训练过程的原始预测残差,而不是只保存平均指标。
正式研究的最小实现顺序可以这样安排:
先做一个 50 到 100 个粒子的小体系,用解析或非常简单的参考泛函生成数据,训练一个仅包含 l=0 和 l=1 的等变模型,验证训练循环和对称性检查代码通过。然后加入第二个体系类型和温度变化,逐步扩大数据范围。只有当小体系中过拟合和旋转测试都通过后,再引入更高阶特征和完整自由能标签。这样能避免大模型一上来就无法收敛,分不清是数据问题还是代码问题。
6.1 从总自由能扩展到更多可观测量的路线
如果总自由能标量模型已经稳定,下一阶段可以扩展的方向很多:输出随密度变化的局域贡献、预测外势作用下的密度分布、利用自动微分构造化学势并做定点迭代、扩展到分子自由度。每扩展一步,对称性要求也会同步变化。局域标量输出要求每个节点的特征随旋转发生相应变换;压力或应力则要求张量输出;而如果只预测各向同性自由能密度,可能只需要标量特征与方向信息的间接作用。
6.2 对论文复现最实用的三条建议
第一,不要把复杂的高阶等变卷积当作实验起点。先用最简单的等变图卷积跑通均匀流体或单一相态,确认标签和验证流程正确。第二,对自由能模型来说,导出量比能量本身更值得关注。训练完成后一定要用自动微分计算化学势,检查导数是否有物理意义和数值稳定性。第三,公开代码时保留完整的随机数种子、训练配置 YAML、数据抽样逻辑和对称性检查脚本,否则一年后很难判断某个误差变化是来自模型改进还是数据顺序差异。
等变学习与三维经典密度泛函互相结合的关键点,不是“换一个更花哨的神经网络”,而是让模型的归纳偏置与真实物理保持一致。实现时只要抓住硬约束、标签分解、状态点切分和导数验证,就可以将这类代理模型从 toy example 逐步推向真正可用于平衡密度求解和材料逆向设计的工具。