做分子性质预测的朋友，大概都听过一个共识：量子化学计算精度高，但贵，贵到让人肉疼。一个稍微大点的分子，高精度方法可能得算几天。那能不能让神经网络直接学一个映射，输入分子结构、输出能量？能，但前提是网络得先“看懂”分子长什么样。这个问题，当年那篇聚焦量子化学的论文给出了一个影响深远的答案，也就是今天要聊的MPNN（Message 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而不是LSTM
GRU和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_samples=2000, max_atoms=12, atom_types=4, bond_types=3): 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]], dim=0) ) 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_steps=3): 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(dim=0) # 求和读出 return self.readout(h)训练循环不需要多少改造，就是常规的监督回归：
dataset = SyntheticMolDataset(num_samples=2000) loader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True) model = MPNNPredictor(node_dim=4, edge_dim=1, hidden_dim=32, num_steps=3) optimizer = torch.optim.Adam(model.parameters(), lr=1e-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(f"epoch {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的开山之作，更因为它提供了一种看待问题的方式：先定义消息怎么在图上流动，再定义状态怎么更新，最后定义怎么读出，这几乎是所有后来模型的共同骨架。我自己跑完这些实验的最大感受是，与其一口气追最新变体，不如先把这种“三段式”框架拆明白。理解了消息函数的自由度，你看到新论文时会自然地问：变了什么？为什么要变？这比死记模型结构有效得多。
第一讲到这里先把框架和最小实现收尾。下一讲我会沿这个思路展开，重点聊距离感知的消息设计是怎么把空间结构进一步编码进分子图里的，到时候还会给出一套能在真实开源分子数据集上跑出可靠误差的完整代码，我们下篇见。
