第4章 消息函数设计：线性变换、MLP与共享参数

消息函数是MPNN的核心。说白了，它决定了邻居信息怎么编码、怎么流进当前节点。我刚开始做图网络的时候，总觉得消息函数随便选个线性层就行。后来在分子性质预测项目里吃了几次亏，才明白这里头的门道不少。

这一章我们聚焦消息函数的三种主流设计：线性变换、MLP和共享参数。每种方案各有优劣，适用场景也不同。你想想看，同样是做消息传递，换个函数设计，模型效果可能天差地别。

4.1 消息函数在MPNN中的角色

在MPNN的每一层，每个节点会聚合邻居的信息来更新自己的表示。消息函数负责把源节点的特征、目标节点的特征、边特征这三者组合起来，生成一条"消息"。

用公式表达就是：

m_{vw} = M(h_v, h_w, e_{vw})

其中 h_v 是源节点特征，h_w 是目标节点特征，e_{vw} 是边特征。M 就是消息函数。你设计什么样的 M，直接决定了模型能学到什么级别的结构信息。

4.2 线性变换：最朴素的方案

线性变换是最基础的消息函数。直接对源节点特征做一次线性映射，加上偏置，完事。

m_{vw} = W * h_v + b

这种做法简单、速度快、参数量少。我在早期的一个小规模社交网络项目里用过，训练效率确实高。但问题是——它没有考虑目标节点特征，也没有非线性能力。表达能力有限。

什么时候用线性变换？我觉得有两种场景比较合适：

• 图规模超大（百万级节点以上），参数量必须压得很低
• 作为 baseline，快速验证整体流程是否跑通

4.3 MLP：表达能力的核心武器

MLP（多层感知机）是目前最主流的消息函数实现。它一般会把源节点特征、目标节点特征、边特征拼接起来，再送入一个多层全连接网络。

m_{vw} = MLP( concat( h_v, h_w, e_{vw} ) )

这个方案的优势很明显：非线性激活让模型能学到更复杂的模式。我在分子性质预测项目里试过，用两层MLP做消息函数，效果比线性变换提升了将近12%。

但MLP也有坑。我曾经在一个图分类任务里把MLP堆到5层，结果训练loss死活降不下去。后来查了半天，发现是梯度消失了。加了残差连接和LayerNorm才救回来。

MLP的典型代码结构：

class MessageMLP(nn.Module):
    def __init__(self, in_dim, hidden_dim, out_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, out_dim)
        )

    def forward(self, h_src, h_dst, e):
        x = torch.cat([h_src, h_dst, e], dim=-1)
        return self.net(x)

这段代码里我把拼接放在forward里做，这样更灵活。你可以在外部控制输入特征的组合方式。

4.4 共享参数：效率与泛化的平衡

共享参数的核心思路是：不同类型的边或者不同位置的层，共用同一套消息函数参数。

常见的共享方式有两种：

共享方式	做法	适用场景
边类型共享	同一种边类型共用一套参数	异构图、关系图
跨层共享	不同MPNN层共用同一套MLP	深层网络、参数敏感场景

我记得在做一个知识图谱推理项目时，图里有30多种关系。如果每种关系都单独学一套MLP，参数量直接爆炸。后来改用边类型共享，把相似关系归到同一组，参数量降了60%，效果只掉了不到3%。值。

跨层共享的实现也很直接：

shared_mlp = MessageMLP(in_dim=128, hidden_dim=64, out_dim=64)

# 第1层消息
msg_1 = shared_mlp(h_src_1, h_dst_1, e)
# 第2层消息
msg_2 = shared_mlp(h_src_2, h_dst_2, e)

你看，同一个MLP对象在不同层重复使用。这样参数总量不变，层数可以加深。

4.5 三种方案对比与选型

下面这张表是我自己的实践经验总结，供你参考：

方案	表达能力	参数量	训练速度	推荐场景
线性变换	低	少	快	超大图、baseline
MLP	高	中	中	通用、中小规模图
共享参数	中高	少	快	异构图、深层网络

你选哪种？我的建议是：默认用MLP（2层），如果参数量扛不住再切共享参数。线性变换只做baseline或极端大图场景。

4.6 消息函数设计路径总览

下面这张图把三种方案放在一起对比，你可以直观地看到它们的输入输出关系和核心差异：

4.7 避坑总结

最后，我把这些年踩过的坑集中说一下：

• MLP层数别贪多——3层以上收益递减，还容易梯度问题。一定要加深的话，加上残差和Norm。
• 共享参数别忘了做分布适配——不同层/不同边的输入分布差异可能很大，Norm是必需品。
• 线性变换不是万能的——它做不了特征交互。如果你发现模型效果怎么调都上不去，先检查消息函数是不是太简单了。
• 输入拼接顺序要统一——concat(h_v, h_w, e) 和 concat(h_w, h_v, e) 是不同的。定好规矩，别混用。