做分子性质预测的朋友大概都听过一个共识量子化学计算精度高但贵贵到让人肉疼。一个稍微大点的分子高精度方法可能得算几天。那能不能让神经网络直接学一个映射输入分子结构、输出能量能但前提是网络得先“看懂”分子长什么样。这个问题当年那篇聚焦量子化学的论文给出了一个影响深远的答案也就是今天要聊的MPNNMessage Passing Neural Network消息传递神经网络。MPNN不是某一个具体的网络结构而是一整套通用框架把分子看作一张图节点是原子边是化学键然后通过“消息传递”在图上迭代更新原子表示最后汇总成整个分子的性质。这篇是系列第一篇我会把MPNN从动机、数学定义到最小实现完整过一遍不堆公式但该讲清楚的地方一个都不少。适合两类人一是还没接触过GNN、想从分子场景切入的初学者二是已经跑过分子描述符模型、想换成基于图的端到端方法的实践者。1. 为什么量子化学预测需要图神经网络1.1 分子的“语言”不是向量很多传统方法处理分子时第一件事是把结构转成某个固定长度的向量比如各种分子指纹Morgan指纹、MACCS keys或者从SMILES字符串统计的拼图特征。它们确实有效尤其在数据量不大的经典机器学习模型里我也用这些东西跑过不少实验速度很快效果也够用。但这里有个绕不过去的坎指纹本质上是把二维拓扑信息压扁后做的稀疏编码它保留了什么、丢掉了什么都由人工规则决定模型自己学不到“哪些局部结构真正重要”。打个比方指纹相当于给分子拍了一张模糊的“全景照”你能看出大概轮廓但具体哪个原子连着哪个原子、成键的电子云分布如何细节全丢了。量子化学性质恰恰对细节极敏感同一个官能团出现在不同位置对整体电子结构的扰动完全不同。你让一个只看全景照的网络去预测这种差值它只能靠记忆样本分布去凑泛化能力可想而知。图表示就不一样了。分子天然是图原子是节点化学键是边键级、键长、电荷都可以作为节点或边上的属性。网络看着这张完整的连接关系去学特征信息传播路径和化学结构的空间邻域高度一致不需要人为设计指纹规则。1.2 量子化学预测问题的输入输出形态从机器学习角度看量子化学预测是一个典型的“图到标量”或“图到向量”的回归任务。输入是分子的原子类型、坐标或连接关系输出可能是总能量、某个轨道的能级、偶极矩、热力学性质等。MPNN这篇论文里重点做了原子化能和内能的预测这些值本质上由电子结构决定而电子结构又高度依赖于原子间的连接和临近关系。这里还有两个绕不开的性质值得提前说清楚。第一分子没有方向你旋转它、平移它性质不变。第二你给原子重新排个序同样的分子还是同一个分子。这意味着模型内部对节点的处理必须是置换等变的对分子级的输出必须是置换不变的。MPNN通过固定的消息聚合函数做到了这两点聚合操作对所有邻居的次序没有依赖于是天然满足物理上的对称性。1.3 传统指纹方案在哪一环开始失效我实际测试过指纹加随机森林和GNN在同一个分子数据集上的差异。数据量小、分布集中时两者差距不大指纹甚至训练得更快因为特征维数低、模型简单。一旦增大分子多样性指纹模型的误差就开始明显上涨。原因很直接指纹特征空间是离散的、稀疏的分子稍一变化特征向量就发生剧烈且非结构化的改变模型很难学到平滑的映射。而消息传递模型每一层都在“看图说话”原子的表示随着周围环境的展开逐步丰富等于是把“哪些结构模式重要”这件事也交给网络自己学了。另外一个非常现实的问题是可扩展性。指纹只能编码预先定义好的化学副词库如果分子含硼、含硅这类相对少见的元素标准指纹库往往无法完整覆盖。图的表示方式就没有这个问题节点特征直接由元素序号、杂化状态等连续或离散属性构成新元素来了也能处理不需要靠模式库背单词。2. MPNN的核心设计思路拆解2.1 消息阶段邻居之间如何“说话”MPNN最核心的概念是节点之间通过边上传递的“消息”来交换信息。在每一轮迭代里每个原子会向它所有的邻居发一条消息消息的内容由发送方原子当前的状态、接收方原子当前的状态以及它们之间化学键的属性共同决定。这个设计意图很清晰模型中边的信息不是装饰品。化学键的特征比如单键、双键、芳香键直接决定了电子共享的程度和空间取向如果消息公式里不含键类型网络就只能靠原子身份硬猜键对性质的影响效果差很多。所以最常见的消息函数是![edge_index表示法下消息 拼接[源节点特征, 目标节点特征, 边特征]过一个小多层感知机]把这三个部件拼在一起过神经网络是初代MPNN给出的一个很稳妥的默认做法。它让网络既知道“谁在给我发消息”又知道“我和它之间是什么关系”信息完整性很高。2.2 更新阶段如何“消化”接收到的消息每个节点收齐所有邻居发来的消息之后通常先把消息按位求和或取平均得到一个聚合消息向量再和自己的旧状态一起送入一个更新函数得到新一轮的状态。论文里用的是门控循环单元GRU也就是说每个原子的状态更新被建模成“读入一条聚合消息以旧状态为记忆得到一个融合的新状态”。这种带门的更新方式极大缓解了多层消息传递时的梯度衰减问题实践上非常实用。我做实验的时候对比过“GRU更新”和“简单加一层ReLU全连接更新”两种方案结论很明确在没有跳跃连接的前提下GRU在深网络上的稳定性和收敛速度明显好于简单的MLP更新。原因不难理解GRU本身就像一个带遗忘机制的状态机不重要的历史信息会被门控主动屏蔽重要的会被保留而MLP更新每次都硬灌一个变换进去不区分信息主次。2.3 读出阶段从节点表示到分子性质消息传递迭代固定步数之后每个原子节点都攒了一个包含局部甚至全局环境信息的向量。接下来要做的是把这些向量汇总成整个分子的性质这个汇总过程就是读出Readout阶段。读出函数必须保证置换不变性即打乱节点顺序不影响输出。最简单的做法是求和、取平均、取最大这类对称聚合。我最初跑MPNN实验时用的是求和效果稳定代价是分子尺寸差异会被完整保留对大分子和小分子的预测偏差不同换成平均后分子大小的影响被抹平了但小分子上细微的电子差异也容易被冲淡。后来我比较了几个方案发现更稳妥的路线是先求和再加一层MLP并且在分子尺度差异较大的数据集上考虑用带注意力权重的Set2Set读出它能把节点特征压缩成一个固定维度的上下文向量兼顾了置换不变性和表达力。2.4 用“邻居串门”来理解整个过程如果把一个分子想象成一个小区每个原子是一户人家那消息传递就是一次社区大会第一轮每户人家只说一说自己和最亲近邻居的关系第二轮信息已经传到两跳之外的住户那里因为你的邻居会把他在第一轮听到的内容转述给你多轮之后每户人家掌握的都不再是自家的孤立信息而是附近整个社区的结构画像。分子性质往往是由局部电子结构决定的比如某个官能团的电子得失倾向而这种“邻居串门”式传播恰好能让模型在固定的层数内捕获到足够大感受野的结构信息。3. 初代MPNN的数学表述与算法流程3.1 符号约定与核心公式严格看论文里的表述一张图被定义为三元组节点集合、节点特征、边特征。每个节点有一个初始隐藏状态每轮迭代分成两个阶段。第一阶段是消息阶段对所有边消息由发送方、接收方的当前状态加上边属性共同决定形式上写成消息函数作用于三元组消息向量 消息函数(源节点状态, 目标节点状态, 边特征)随后目标节点把来自所有邻居的消息累加起来形成聚合消息聚合函数在初代实现里就是求和。第二阶段是更新阶段每个节点把自身旧状态和聚合消息放进更新函数得到新状态。这里的更新函数对应的是GRU运算可以理解为一次带记忆的序列更新。整个流程循环固定轮数最后一轮结束后进入读出阶段。读出函数把全部节点状态作为输入输出图级别的预测值。论文里给出的实现是“逐模块的和再进回归网络”具体为分子预测值 回归网络(所有节点状态的加和)这个方法在实践中超级省参数也相当稳健。3.2 为什么活泼的更新函数选GRU而不是LSTMGRU和LSTM的区别在于门控设计更精简参数更少。在分子图上每个时间步处理的是不同节点的状态序列不像经典NLP那样有严格的时间前后关系复杂的遗忘门不一定有用。GRU的两个门足以控制“留下什么、忘掉什么”训练速度更快在小规模分子数据集上也不容易过拟合。我做对比实验时LSTM版本收敛稍慢精度提升几乎为零所以这类场景下优先选GRU是合理决策。3.3 MPNN是一整套通用框架很多时候我们把MPNN和GCN、GAT这些放一起比较其实它们不在一个层面上。MPNN描述的是“消息传递 更新 读出”这三段式流程GCN、GAT、SchNet这些具体模型都可以看作在MPNN框架里选择了不同的消息函数、更新函数和读出函数。比如GCN是“把邻居特征平均并与自身特征加权求和”的特殊实例GAT是在消息聚合时用注意力权重替换普通加权的实例处理分子三维构型的SchNet则是把“原子之间的欧氏距离”嵌入进消息函数的距离滤波卷积。理解了MPNN这个母框架后续你再去看具体变体只需要关注三个问题它的消息公式是什么更新公式是什么读出公式是什么这算是我阅读GNN论文时最实用的一个技巧。3.4 符号变体与现代实现的共通点初代论文里给出的数学公式如今看起来有一点繁琐因为它把节点状态、边隐藏状态分开写了还在边上也维护了一套隐藏状态。我一开始照搬这个设定发现边上的额外隐藏状态在普通分子数据上并不总是带来收益倒是让代码复杂了不少。现代库的实现往往简化掉“边状态随时间更新”这一步只保留固定的边属性特征作为消息输入。两者在多数任务上性能相当简化的写法更易维护。真正的争议不在边状态是否需要更新而在于消息聚合时是否要把接收方自身的特征一起拼进消息里。初代公式包含了接收方状态这给了网络更强的目标感知能力但也会让消息计算量翻倍。很多后续工作选择把接收方状态交给更新阶段处理消息里只放源节点状态加边特征效果也不错。这类取舍我在第五部分还会具体谈。4. 跑通一个最小MPNN实现4.1 准备一个最小分子图数据集自己动手实现之前我不建议一上来就冲开源大分子数据集因为数据解析、特征归一化、训练周期这些问题会把理解框架本身这件事淹没。更好的路径是先用一个可复现的合成分子集验证模型能拟合再切到真实数据。下面这个合成数据集非常简单每个“分子”由若干原子组成每个原子有一个整数类型编号每条边有一个整数键类型。目标性质则人为设定为一阶线性函数比如“原子类型权重之和加上键数量的某个倍数”。这样生成的目标函数模型很容易学习如果你实现的MPNN是对的训练几轮后loss就能降到很低如果实现有bug这个简单的拟合任务会立刻暴露问题。import random import torch from torch.utils.data import Dataset class SyntheticMolDataset(Dataset): def __init__(self, num_samples2000, max_atoms12, atom_types4, bond_types3): super().__init__() self.samples [] for _ in range(num_samples): n_atoms random.randint(3, max_atoms) # 用一个简单的树状结构作为分子骨架 edges [] for i in range(1, n_atoms): parent random.randint(0, i - 1) btype random.randint(0, bond_types - 1) edges.append((parent, i, btype)) # 节点特征one-hot 原子类型 node_types [random.randint(0, atom_types - 1) for _ in range(n_atoms)] x torch.zeros(n_atoms, atom_types) for i, t in enumerate(node_types): x[i, t] 1.0 # 边特征one-hot 键类型 edge_index [] edge_attr [] for u, v, bt in edges: edge_index.append([u, v]) edge_index.append([v, u]) edge_attr.append([bt]) edge_attr.append([bt]) # 合成目标函数原子权重加键数量 atom_weights torch.tensor([0.3, 0.7, 1.2, 0.5]) y x atom_weights y y.sum() 0.05 * len(edge_index) // 2 y y.unsqueeze(0) self.samples.append((x, torch.tensor(edge_index).t().contiguous(), torch.tensor(edge_attr).float(), y)) def __len__(self): return len(self.samples) def __getitem__(self, idx): x, edge_index, edge_attr, y self.samples[idx] return x, edge_index, edge_attr, y这里我故意把图做成了树状结构并且把每条边正反两个方向都加进去了。原因有两个一是保证消息传递是双向的每个原子都能看到邻居的信息二是让邻接矩阵不必为对称矩阵做额外处理代码更统一。4.2 用PyTorch定义MPNN核心层下面这段代码是一个极简MPNN实现消息函数是“拼接源节点、目标节点、边属性”后过两层全连接更新函数是GRU。为了可读性我用了循环遍历边上消息的方式实际大批量训练时建议用稀疏矩阵的scatter操作替代。import torch import torch.nn as nn import torch.nn.functional as F class MPNNLayer(nn.Module): def __init__(self, node_dim, edge_dim, hidden_dim): super().__init__() # 消息函数 self.message_net nn.Sequential( nn.Linear(node_dim * 2 edge_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, node_dim) ) # 更新函数 self.gru nn.GRUCell(node_dim, node_dim) def forward(self, x, edge_index, edge_attr): # x: [N, node_dim] # edge_index: [2, E] # edge_attr: [E, edge_dim] src edge_index[0] dst edge_index[1] messages [] for i in range(edge_index.size(1)): s, d src[i].item(), dst[i].item() message self.message_net( torch.cat([x[s], x[d], edge_attr[i]], dim0) ) messages.append(message) messages torch.stack(messages) # [E, node_dim] # 聚合消息目标节点维度上做 scatter add aggregated torch.zeros_like(x) dst_expanded dst.unsqueeze(-1).expand_as(messages) aggregated aggregated.scatter_add_(0, dst_expanded, messages) # 更新节点状态 out self.gru(aggregated, x) return out这个实现把消息聚合、更新两步完整包含在一个层里循环num_steps次就完成了整个消息传递阶段。读者可以观察到scatter_add操作天然满足置换不变性因为聚合结果和输入节点的顺序无关。这正是MPNN能保持物理对称性的关键所在。4.3 完整的模型封装与训练循环有了数据、有了层接下来就是把它们拼起来。我用一个简单的求和读出作为图级预测头所有节点状态加和后过两层MLP。为了让结果可复现固定随机种子学习率调为1e-3。class MPNNPredictor(nn.Module): def __init__(self, node_dim, edge_dim, hidden_dim, num_steps3): super().__init__() self.num_steps num_steps self.layer MPNNLayer(node_dim, edge_dim, hidden_dim) self.readout nn.Sequential( nn.Linear(node_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 1) ) def forward(self, x, edge_index, edge_attr): h x for _ in range(self.num_steps): h self.layer(h, edge_index, edge_attr) h h.sum(dim0) # 求和读出 return self.readout(h)训练循环不需要多少改造就是常规的监督回归dataset SyntheticMolDataset(num_samples2000) loader torch.utils.data.DataLoader(dataset, batch_size32, shuffleTrue) model MPNNPredictor(node_dim4, edge_dim1, hidden_dim32, num_steps3) optimizer torch.optim.Adam(model.parameters(), lr1e-3) loss_fn nn.MSELoss() for epoch in range(30): total_loss 0.0 for x, edge_index, edge_attr, y in loader: optimizer.zero_grad() pred model(x, edge_index, edge_attr) loss loss_fn(pred, y) loss.backward() optimizer.step() total_loss loss.item() * len(x) print(fepoch {epoch:02d}, loss {total_loss / len(dataset):.6f})在我的机器上这个模型大概跑到第8个epoch时loss降到0.01以下。由于数据是人为生成的线性目标模型可解释性非常强读出层几乎就是学到一个近似“节点特征求和”的算子把目标函数里的原子权重线性组合再算出来。这验证了消息传递和读出流程搭建是正确的没有梯度传播黑洞。4.4 从合成数据切到真实分子数据的关键差异合成数据证明模型实现正确之后再切到真实量子化学数据就有底气了。真实场景下需要注意这么几点第一节点特征不再只是一阶one-hot往往还包括原子序号、杂化类型、部分电荷、手性标记等第二边特征也不再是键类型这么简单还可能包含键长、是否成环、是否共轭等第三标签的数值范围差异很大能量值通常是负几十到负几千需要做均值方差归一化或者直接预测“能量与参考值的差值”训练会更稳定。我最开始完全按照MPNN论文原始设定去做真实数据拟合输入特征只给原子序号和边类型结果在几个中等分子集上误差明显偏大。后来仔细排查发现瓶颈不在模型而在特征信息不足仅靠键类型无法区分芳香键在不同化学环境里的差异。给边特征增加“是否在环内”的布尔向量后误差立刻下降了十几个百分比。这类工程经验比模型结构本身更影响最终效果。5. 实操心得与常见问题排查5.1 训练不收敛先查这四件事如果复现MPNN时loss要么不动、要么直接变成NaN我建议按顺序排查。第一检查标签是否做了归一化量子化学能量标签动辄几百上千不归一化的话梯度尺度容易失控第二检查学习率我一般从5e-4配Adam起步太大数据量的场景降到1e-4第三检查消息聚合是否出错尤其要注意scatter_add的目标维度和返回张量的形状是否对齐这一步最容易悄悄出bug第四检查GRU的输入维度GRUCell要求当前时刻输入维度等于隐藏层维度如果你的消息输出维度和节点状态维度不一致运行时会直接报错但如果你把node_dim和hidden_dim误设成不同值会在训练很久后才发现性能差得离谱。5.2 消息传递层数越多越好吗不是。MPNN的每一层都在扩大感受野但层数过多会造成过度平滑所有原子的状态逐渐趋同节点表示失去区分度。在分子图上3到5层通常就够用了因为分子的直径本身不大几乎所有的原子间关系都在6跳以内。设计时还有一个容易被忽略的问题层数增多参数也增多分子数据集普遍不大几万到几十万级过拟合风险比一般图任务更突出。我的习惯是先用3层训练出一个baseline再尝试5层比较验证集误差后再决定要不要加深。5.3 聚合操作选求和、平均还是注意力这个话题几乎每次都要争论。求和的优点是保留分子大小信息在很多热力学性质上正好需要这种信息缺点是支撑离散分子尺寸的波动会直接干扰模型预测尤其是垂直方向的数据分布不均衡时。求平均的好处是对尺度的鲁棒性更好但会丢失总强度信息比如偶极矩这类本身就累积的性质平均聚合往往会出大问题。注意力机制的引入让聚合权重可根据内容自适应我在处理含多个官能团的大分子时经常用它只是训练成本会高一些。如果拿不准我建议从求和开始因为那个它最简单、最稳健然后再用Bias-Variance的思路判断模型误差是来自欠拟合还是过拟合决定要不要切换聚合策略。5.4 边状态到底要不要随时间更新这个问题是重读初代MPNN论文时很难绕过的一个坑。论文原始设定里边隐藏状态本身也会随迭代更新相当于信息既存节点上也存边上这会显著增加存储和计算开销。实际实现时很多库把边属性当作固定输入处理不再每层更新。我做过的对比实验显示当数据集不大、特征已经包含了足够化学信息时是否更新边状态对最终精度影响很小但当任务特别依赖键属性之间的语义交互时比如手性识别边状态更新是能带来收益的。考虑到代码复杂度和显存占用第一版实现通常不值得上边状态更新先跑通再优化是正确的顺序。5.5 数值稳定性和特征尺度分子图的特征维度不高的前提下计算稳定性主要受两个因素影响。第一是边特征多为小整数比如键类型的0、1、2直接送入线性层时没问题如果键长这类连续值也加进来最好先做标准化否则消息网络中较大数值的维度会在训练初期主导梯度。第二是读出阶段的求和操作节点数量从几十到几百不等非归一化的总和值尺度方差很大配合归一化标签后模型在分子量差异大的数据集中往往不稳定。一个折中方案是“求和再过LayerNorm”我试过之后发现收敛速度和稳定性双双改善代价只是多几行代码。5.6 对称性和性能的工程取舍最后一个实操心得是关于批量训练的。我上面给的代码是一次处理一个分子在真实数据集上效率不足。批量训练时多个分子的图会被合并成一个大图节点编号做了偏移边索引也相应平移。工程上这很好处理唯一要注意的是聚合消息时不要把不同分子之间的边算到同一个节点上。PyTorch Geometric等库的Batch机制就是这个原理。如果你要自己写建议在聚合后按分子索引做分组抠出每个分子的节点表示再做单独的读出否则预测结果会串图。最后的一点个人体会MPNN之所以值得认真理解不仅仅因为它是量子化学GNN的开山之作更因为它提供了一种看待问题的方式先定义消息怎么在图上流动再定义状态怎么更新最后定义怎么读出这几乎是所有后来模型的共同骨架。我自己跑完这些实验的最大感受是与其一口气追最新变体不如先把这种“三段式”框架拆明白。理解了消息函数的自由度你看到新论文时会自然地问变了什么为什么要变这比死记模型结构有效得多。第一讲到这里先把框架和最小实现收尾。下一讲我会沿这个思路展开重点聊距离感知的消息设计是怎么把空间结构进一步编码进分子图里的到时候还会给出一套能在真实开源分子数据集上跑出可靠误差的完整代码我们下篇见。