在物理建模和机器学习交叉的论文标题里等变学习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) 对称张量l0、l2 等不可约表示组合文章场景里最常见的是第一类以某种标量自由能为监督目标。最终输出只要做到旋转和平移不变即可。但要注意只做输出不变还不够中间特征最好也采用等变特征因为这样才能让网络通过方向性相互作用学习三维结构差异。否则即使输出是标量网络也会退化成只能使用距离信息的纯径向模型对角度排列不敏感。注意标量不变不等于等变。等变是一个包含输出变换规则的框架。预测自由能时模型对真实旋转应输出相同结果这只是等变中的“不变表示”当后续要输出矢量力或应力时必须使用真正的矢量等变操作。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 python3.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, dtypetorch.float32) for c in coords] self.atom_type [torch.as_tensor(a, dtypetorch.long) for a in atom_type] self.free_energy [torch.as_tensor([e], dtypetorch.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], dtypetorch.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表示 l2 的偶宇称张量。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, seed0): 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确认能跑通后再增加 l2 甚至 l3 特征。下表给出一组调整建议现象可能原因调整方向训练不收敛输出标签未归一化按体系粒子数归一化修正标签尺度旋转测试残差大模型使用绝对坐标或绝对平移拼接改为相对坐标和边向量迁移误差大训练状态单一加入不同外势和密度的训练数据高频振动不饱和截断半径过小分析邻居数增大半径网络过大反而差不可约表示阶数过高降低irreps_hidden先增大数据量5.5 数据口径错误经典密度泛函的数据有多种标签来源分子模拟自由能微扰、局部密度积分、参考泛函数值解。不同来源标签的单位和零点可能不同。建议数据加载后统一转换单位并在保存前把每个样本的粒子数、化学势、温度和标签一起写入文件头。训练以前先输出一张分布直方图确认没有异常零点或离群体系。6. 为真实实验准备的最佳实践清单进入正式研究和工程前可以在项目目录建立下面的检查清单把它当成生产级工作的最低门槛。确认对称性定义。论文体系只有旋转不变性还是也要求镜像不变性晶格取向是否涉及特殊欧拉角确认单位。坐标按原子长度还是约化长度能量按单个粒子还是按每个分子确认自由能分解。是否已经把理想气体部分从标签中扣除数据如何划分。是否包含按温度、化学势或粒子类型切分的留出集邻居下标是否可复现。坐标旋转后邻居序号是否重复损失函数是否包含导数。如果最终要迭代求平衡密度导数标签也必须有验证。等变测试是否作为 CI 流程的一部分。是否设置固定随机种子并在不同运行间对比模型正常波动。是否保存训练过程的原始预测残差而不是只保存平均指标。正式研究的最小实现顺序可以这样安排先做一个 50 到 100 个粒子的小体系用解析或非常简单的参考泛函生成数据训练一个仅包含 l0 和 l1 的等变模型验证训练循环和对称性检查代码通过。然后加入第二个体系类型和温度变化逐步扩大数据范围。只有当小体系中过拟合和旋转测试都通过后再引入更高阶特征和完整自由能标签。这样能避免大模型一上来就无法收敛分不清是数据问题还是代码问题。6.1 从总自由能扩展到更多可观测量的路线如果总自由能标量模型已经稳定下一阶段可以扩展的方向很多输出随密度变化的局域贡献、预测外势作用下的密度分布、利用自动微分构造化学势并做定点迭代、扩展到分子自由度。每扩展一步对称性要求也会同步变化。局域标量输出要求每个节点的特征随旋转发生相应变换压力或应力则要求张量输出而如果只预测各向同性自由能密度可能只需要标量特征与方向信息的间接作用。6.2 对论文复现最实用的三条建议第一不要把复杂的高阶等变卷积当作实验起点。先用最简单的等变图卷积跑通均匀流体或单一相态确认标签和验证流程正确。第二对自由能模型来说导出量比能量本身更值得关注。训练完成后一定要用自动微分计算化学势检查导数是否有物理意义和数值稳定性。第三公开代码时保留完整的随机数种子、训练配置 YAML、数据抽样逻辑和对称性检查脚本否则一年后很难判断某个误差变化是来自模型改进还是数据顺序差异。等变学习与三维经典密度泛函互相结合的关键点不是“换一个更花哨的神经网络”而是让模型的归纳偏置与真实物理保持一致。实现时只要抓住硬约束、标签分解、状态点切分和导数验证就可以将这类代理模型从 toy example 逐步推向真正可用于平衡密度求解和材料逆向设计的工具。