NequIP与Allegro：机器学习势能函数
1. 引言
传统分子动力学（MD）模拟依赖经验力场（如CHARMM、AMBER）或量子力学（QM）计算。前者速度快但不准确，后者准确但计算量巨大（只能模拟数百个原子、皮秒级别）。**机器学习势能函数（ML Potentials）**通过学习量子化学精度的能量和力，实现”又快又准”的分子模拟。1
ML势能函数的核心挑战
- 数据效率：训练数据（DFT计算）昂贵，需要在有限数据下达到高精度
- 可扩展性：从几十个原子到数百万原子
- 对称性：必须精确满足物理对称性（平移、旋转不变性）
- 能量守恒：力是能量的负梯度，必须保证
2. NequIP架构详解
2.1 核心设计思想
NequIP（Neural Equivariant Interatomic Potentials）的核心创新是：
使用E(3)等变消息传递，在消息中编码方向信息
2.2 架构概览
输入：原子类型 + 3D坐标
↓
Embedding层：原子类型 → 初始特征
↓
等变消息传递层（重复L次）
↓
Readout层：聚合 → 原子能量贡献
↓
输出：总能量 + 原子受力
2.3 等变消息传递
NequIP的消息传递在球谐函数域进行：
class NequIPConvolution(nn.Module):
"""
NequIP等变卷积层
核心：将方向信息编码为球谐函数
"""
def __init__(self, l_max, num_species, hidden_channels):
super().__init__()
self.l_max = l_max
# 边方向编码：计算 r_ij 的球谐函数
self.spherical_harmonics = o3.SphericalHarmonics(
o3.Irreps([(l, 1 if l % 2 == 0 else -1) for l in range(l_max + 1)])
)
# 径向MLP（只依赖于距离，不破坏等变性）
self.radial_mlp = nn.Sequential(
nn.Linear(1, 32),
nn.SiLU(),
nn.Linear(32, 64),
nn.SiLU(),
nn.Linear(64, (l_max + 1) ** 2 * hidden_channels)
)
# 等变张量积
self.tp = o3.TensorProduct(
o3.Irreps([(l, 1 if l % 2 == 0 else -1) for l in range(l_max + 1)]),
o3.Irreps([(l, 1) for l in range(l_max + 1)]),
# ...
)
def forward(self, node_features, edge_vectors, edge_lengths):
# 1. 编码边方向
Y = self.spherical_harmonics(edge_vectors) # [B, N, N, (l_max+1)^2]
# 2. 径向特征（距离的函数）
radial = self.radial_mlp(edge_lengths.unsqueeze(-1)) # [B, N, N, H]
# 3. 消息 = 径向特征 × 球谐函数
messages = Y.unsqueeze(-1) * radial.unsqueeze(-2) # [B, N, N, (l_max+1)^2, H]
# 4. 张量积聚合
# 聚合来自邻居的消息
output = self.aggregate_messages(node_features, messages)
return output2.4 关键创新点
| 创新 | 说明 | 优势 |
|---|---|---|
| 连续滤波 | 卷积核是连续的径向函数，而非离散 | 更好的泛化能力 |
| 等变消息 | 消息包含方向信息 | 更高的数据效率 |
| 自相互作用 | 支持特征自旋转变 | 增加表达能力 |
| 门控机制 | 标量门控非线性 | 保持等变性 |
2.5 训练流程
# NequIP训练示例（简化）
from nequip.train import Trainer
from nequip.model import model_builder
# 1. 准备数据集
dataset = nequip.dataset.AtomicDataModule(
train_url="train.xyz",
validation_url="val.xyz",
batch_size=16,
num_workers=4
)
# 2. 构建模型
model = model_builder(
r_max=6.0, # 截断半径
num_species=3, # 原子种类数
l_max=3, # 球谐函数最高阶
hidden_channels=64 # 隐藏维度
)
# 3. 训练配置
trainer = Trainer(
model=model,
optimizers={
"optimizer": torch.optim.AdamW,
"lr": 1e-2,
},
loss={
"energy": 1.0, # 能量损失权重
"forces": 100.0 # 力损失权重（通常更大）
},
metrics={
"energy": ["mae", "rmse"],
"forces": ["mae"]
}
)
# 4. 训练
trainer.fit()3. Allegro架构详解
3.1 核心设计思想
Allegro的核心创新是逐对（Pairwise）分解：
不是对每个原子做消息传递，而是直接学习原子对之间的等变交互
3.2 架构对比
| 特性 | NequIP | Allegro |
|---|---|---|
| 消息传递 | 原子中心 | 原子对中心 |
| 邻居聚合 | 先消息后聚合 | 先分解后组合 |
| 复杂度 | ||
| 等变性 | 完全等变 | 完全等变 |
3.3 Allegro的核心公式
Allegro学习一个逐对等变函数 ：
然后通过张量积组合：
class AllegroLayer(nn.Module):
"""
Allegro等变层
关键：分解为逐对项的组合
"""
def __init__(self, l_max, hidden_dim, num_basis):
super().__init__()
# 逐对编码器：r_ij -> 逐对特征
self.pair_encoder = nn.ModuleList([
nn.Sequential(
nn.Linear(num_basis, hidden_dim),
nn.SiLU(),
nn.Linear(hidden_dim, hidden_dim)
) for _ in range(l_max + 1)
])
# 张量积组合
self.tp_rules = self._build_tensor_product_rules(l_max)
def forward(self, edge_attrs, edge_lengths):
"""
Args:
edge_attrs: [B, num_edges, num_basis] 边特征
edge_lengths: [B, num_edges] 边长度
"""
# 1. 逐对特征编码
pair_features = []
for l in range(self.l_max + 1):
pf = self.pair_encoder[l](edge_attrs) # [B, E, H]
pair_features.append(pf)
# 2. 张量积组合
combined = self.combine_pairs(pair_features, self.tp_rules)
return combined
def combine_pairs(self, features, rules):
"""
根据CG规则组合逐对特征
"""
output = []
for (l1, l2, l3), mode in rules:
# 组合特征
combined = torch.einsum('bij,bjk->bik', features[l1], features[l2])
output.append(combined)
return output3.4 高阶等变性
Allegro支持任意阶的等变特征，通过Wigner-D矩阵实现：
# Allegro支持高阶张量特征
class HighOrderAllegro(nn.Module):
def __init__(self, max_l, hidden_dim):
super().__init__()
# 阶数从0（标量）到max_l
self.irreps = o3.Irreps([(l, 1 if l % 2 == 0 else -1)
for l in range(max_l + 1)])4. 与LAMMPS集成
4.1 NequIP-LAMMPS接口
NequIP可以直接导出为LAMMPS势能文件：
# 1. 训练并导出模型
nequip-deploy deploy \
--train-dir /path/to/training \
--output /path/to/deployed.pth \
--final
# 2. 转换为LAMMPS格式
nequip2lammps \
--model deployed.pth \
--species O H C \
--output nequip.eam.fs4.2 MD模拟示例
# 使用OpenMM + NequIP进行MD模拟
import torch
from nequip.data import AtomicData
from nequip.model import model_builder
# 加载模型
model = model_builder(
file_name="deployed.pth",
device="cuda"
)
model.eval()
# 初始化OpenMM系统
from openmm import *
from openmm.app import *
# 创建分子系统
pdb = PDBFile("molecule.pdb")
forcefield = ForceField("tip3p.xml")
# 替换力场为NequIP
system = forcefield.createSystem(
pdb.topology,
nonbondedMethod=NoCutoff,
constraints=None
)
# 添加自定义力（使用NequIP势能）
class NequIPForce(CustomExternalForce):
def __init__(self, model, positions):
super().__init__("0")
self.model = model
# 设置原子索引
for i in range(len(positions)):
self.addParticle(i)
def calcForceAndEnergy(self, positions, velocities, forces, energy):
# 转换为NequIP格式
data = AtomicData.from_points(
positions=positions,
species=["H", "O", "C"], # 根据实际分子
r_max=6.0
)
# 计算能量和力
with torch.no_grad():
output = self.model(data.to("cuda"))
energy[0] = output["energy"].item()
forces[:, :] = output["forces"].cpu().numpy()5. 应用案例
5.1 液态水模拟
| 方法 | 能量RMSE (meV/atom) | 力RMSE (meV/Å) |
|---|---|---|
| 经典力场 (TIP4P) | ~50 | - |
| NequIP | ~8 | ~50 |
| Allegro | ~6 | ~40 |
5.2 化学反应
NequIP已成功应用于：
- 光合作用：光捕获复合物的能量转移
- 催化：甲醇合成的表面反应
- 电化学：电极-电解质界面
6. 训练数据生成
6.1 DFT计算流程
# 使用ASE进行DFT计算
from ase import Atoms
from ase.calculators import Vasp
# 1. 准备分子结构
atoms = Atoms("H2O", positions=[[0, 0, 0], [0.96, 0, 0], [-0.24, 0.96, 0]])
# 2. DFT计算
calc = Vasp(
xc="PBE",
kpts=(4, 4, 4),
encut=500
)
atoms.calc = calc
# 3. 获取能量和力
energy = atoms.get_potential_energy()
forces = atoms.get_forces()
# 4. 保存为训练数据
with open("h2o_train.xyz", "a") as f:
f.write(f"{len(atoms)}\n")
f.write(f"energy={energy} pbc=F\n")
for sym, pos, f in zip(atoms.symbols, atoms.positions, forces):
f.write(f"{sym} {pos[0]} {pos[1]} {pos[2]} {f[0]} {f[1]} {f[2]}\n")6.2 主动学习
# 不确定性驱动的主动学习
class ActiveLearningLoop:
def __init__(self, model, calculator, uncertainty_threshold=0.1):
self.model = model
self.calculator = calculator
self.threshold = uncertainty_threshold
def select_new_configurations(self, trajectory, n_samples=10):
"""
从MD轨迹中选择不确定性高的构型
"""
uncertainties = []
for config in trajectory:
# 使用多个随机dropout前向传播估计不确定性
preds = [self.model(config + noise) for _ in range(10)]
uncertainty = np.std([p["energy"] for p in preds])
uncertainties.append((config, uncertainty))
# 选择不确定性最高的构型
uncertainties.sort(key=lambda x: x[1], reverse=True)
return [c for c, u in uncertainties[:n_samples]
if u > self.threshold]
def run(self, initial_configs, n_iterations=5):
"""
运行主动学习循环
"""
training_data = list(initial_configs)
for i in range(n_iterations):
# 1. 训练模型
self.train(training_data)
# 2. 运行短MD模拟
trajectory = self.run_md_simulation()
# 3. 选择新构型
new_configs = self.select_new_configurations(trajectory)
# 4. 计算DFT能量
for config in new_configs:
dft_result = self.calculator.calculate(config)
training_data.append(dft_result)7. 性能对比
| 模型 | 参数数量 | MDPS@H(O) | 可扩展性 |
|---|---|---|---|
| NequIP-3L | ~300K | 98.5% | 好 |
| Allegro | ~500K | 98.8% | 优秀 |
| SchNet | ~500K | 95.2% | 优秀 |
| PhysNet | ~1M | 97.0% | 一般 |
参考文献
Footnotes
-
Batzner et al. “E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Molecular Dynamics.” arXiv:2201.01288, 2022. ↩