图池化：可微分池化

可视化演示

展示DiffPool层次化图池化过程：节点通过可学习软分配矩阵聚类为超节点，逐层压缩图结构，以及分配矩阵的热力可视化。

池化动机

一家创业公司从 20 人扩张到 2 万人后，老板再也不能靠全员大会做决策了——两万人轮流发言，会议还没开完公司就破产了。现实中的解法是引入管理层：先由小组长汇总意见，再交给部门经理提炼，最后上报 CEO。这个“层层收拢”的结构，恰恰是图神经网络在大规模图上遇到的核心问题。

图神经网络（GNN）堆叠多层后，每个节点的感受野越来越大，但节点数量始终不变。面对蛋白质交互图这种动辄数万节点的大图，直接输出整图表示既昂贵又低效。再看 CNN，它在图像上通过池化逐层缩小空间尺寸，提炼出越来越抽象的特征。图的“组织架构”也需要类似的机制：图池化。

DiffPool 的动机正是让图在层层传递中不断“缩编”：每经过一层池化，节点数按比例缩小，同时把局部结构信息凝聚到更高层的表示里。定义压缩率为 $r$，则层与层之间的节点规模关系如下：

回到公司的比喻：$r$ 就像“管理层级幅度”——每个部门平均管理 $1/r$ 名员工。组织架构设计的本质，就是在保留足够细节和提炼高层语义之间寻找平衡。

全局池化

最粗暴的“组织架构”是全员大会：让每个员工发言，然后取大家的平均意见，或者只听嗓门最大的人。前者对应 mean pooling，后者对应 max pooling。这两种操作在简单场景下很快，但问题也很明显——平均会淹没少数关键信息，取最大又会忽略整体结构。

在 GNN 中，全局池化（Global Pooling）就是直接把所有节点表示压成一个图级向量。给定节点特征 $h_v$，图表示 $h_G$ 可以这样计算：

全局池化没有任何可学习参数，也不感知图的结构——好比把整个公司两万人压成一位员工去参加董事会，这位员工只能说出“平均工资大概 1.5 万”，却说不清研发和销售之间的协作关系。面对蛋白质交互图这类层次丰富的结构，这样的“全员大会”式池化就成了瓶颈：信息被严重压扁，层次结构荡然无存。DiffPool 要做的，正是抛弃这场低效的大会，让组织架构自己“长”出来。

聚类重组

聪明的 HR 从不搞全员大会。她翻看员工的协作记录——谁和谁天天开会、谁和谁在同一个代码仓库提交——发现公司里天然存在好几个默契的小团队。于是她把大家重组为部门，部门在更高管理层眼中就是一个“超级员工”。这就是 DiffPool 的聚类重组思想：把一组节点软聚类成一个超节点（supernode）。

聚类不是硬性选一个代表，而是让每个节点按比例融入多个超节点，这就是软分配。形式化地，用一个分配矩阵 $S \in \mathbb{R}^{N \times k}$ 记录每个节点属于每个超节点的概率。然后，超节点的特征就是所有节点特征的加权和，超节点之间的边权则通过原图的边按两端节点的归属概率聚合得到：

部门成立后，公司把部门当作“新员工”继续运作。节点数从 $N$ 降到 $k$，复杂度下来了，但语义却升上去了——这正是 DiffPool 分层抽象的起点。

分配矩阵

HR 不是凭空拍脑袋，她身边有一支数据分析团队（GNN）盯着协作网络，每分钟产出一张“员工去向建议表”：员工 A 60% 去技术部、40% 去产品部。这张表就是分配矩阵 $S$。$S$ 不是人为指定的，而是由 GNN 根据当前图的节点特征和拓扑结构动态生成。

具体地，DiffPool 每层并行跑两个 GNN：一个用于生成节点嵌入 $Z$，另一个用于生成分配打分 $\tilde S$，再对 $\tilde S$ 的每一行做 softmax，使每行成为一个概率分布。整个过程可写成：

“可微分”三个字是 DiffPool 的灵魂。HR 的建议表不是一次性文件，而是随公司战略动态更新的活文档：如果这季度主攻技术产品，分配就偏向研发部门；如果下季度转向海外市场，分配又随之调整。固定池化（比如按度排序取 Top-k 节点）则像刻在石头上的架构图，永远无法根据任务反馈自我修正。

训练目标

公司重组有两条铁律：别拆散配合默契的团队；每位员工要有明确的归属岗位。如果 HR 乱分组，把天天结对写代码的同事硬拆到不同部门，项目立刻停摆；如果让每个人同时挂名 8 个部门，绩效也没法考核。DiffPool 的训练目标，正是这两条铁律的数学化。

光靠下游任务损失还不够，因为模型可能学到“拆散强连接”或“模糊分配”这两种病态解。于是 DiffPool 增加两个辅助损失。连接池化损失（Link Prediction Loss）惩罚“把有强边的节点分到不同超节点”：如果原图 $A$ 中的两个节点有边，那么它们在 $S$ 中的归属向量应该相近。熵正则化（Entropy Regularization）惩罚模糊的分配向量，鼓励每个节点的归属接近 one-hot，避免信息混叠。总损失为：

两条铁律保证重组后的公司执行力不泄漏、职责不扯皮。在训练过程中，$\alpha$ 和 $\beta$ 就像 HR 手中的两杆秤：既要让部门协作顺畅（结构保真），又要让每个员工定位清晰（分配确定），最终端到端地服务公司战略（任务损失）。

层次堆叠

公司规模还在涨：两万人变成十万人，部门之上还需要事业部，事业部再汇报给 CEO。每一层都是一次“超节点化”，组织的层级天然对齐信息的抽象层次。DiffPool 同样可以递归堆叠：第 $\ell$ 层输出的超节点图，直接作为第 $\ell+1$ 层的输入。

在最底层，节点是原子、用户或氨基酸；往上一层是官能团、兴趣小组或二级结构；再往上是分子骨架、社区或结构域。聚类数逐层递减，构成一座信息金字塔。递归关系如下：

实际使用中，每层聚类数 $k$ 是超参数（如 16 → 8 → 4 → 2），通过限制 $k$ 可以精确控制计算预算。对蛋白质交互图这类天然具有多级结构的图数据，堆叠 3~4 层往往比单层效果好得多。多级组织架构让十万人的公司仍然能快速决策——DiffPool 让大规模图在深度网络上也能高效训练与推理，这正是它成为层次图表示学习基石的原因。

---

AI 应用场景

Python 演示

import torchimport torch.nn as nnimport torch.nn.functional as F  class GraphConv(nn.Module):    def __init__(self, in_features, out_features):        super().__init__()        self.linear = nn.Linear(in_features, out_features)        self.bias = nn.Parameter(torch.zeros(out_features))     def forward(self, A, X):        X = self.linear(X)        X = torch.mm(A, X) + self.bias        return F.relu(X)  class DiffPoolLayer(nn.Module):    def __init__(self, in_feat, assign_dim):        super().__init__()        self.embed_net = GraphConv(in_feat, in_feat)        self.assign_net = GraphConv(in_feat, assign_dim)     def forward(self, A, X):        Z = self.embed_net(A, X)        S_logits = self.assign_net(A, X)        S = torch.softmax(S_logits, dim=1)        self.last_S = S        A_new = torch.mm(torch.mm(S.t(), A), S)        X_new = torch.mm(S.t(), Z)        link_loss = torch.norm(A - torch.mm(S, S.t()), p='fro')        ent_loss = (-S * torch.log(S + 1e-9)).sum() / S.size(0)        return A_new, X_new, link_loss, ent_loss  def build_toy_graph():    A = torch.ones(4, 4) - torch.eye(4)    A = A / A.sum(dim=1, keepdim=True)    X = torch.tensor([[1.0, 0.0], [0.5, 1.0], [1.0, 1.0], [0.0, 0.5]])    return A, X  def train():    A, X = build_toy_graph()    pool = DiffPoolLayer(in_feat=2, assign_dim=2)    opt = torch.optim.Adam(pool.parameters(), lr=0.05)    for epoch in range(100):        A2, X2, link, ent = pool(A, X)        loss = -torch.sum(X2[0] * X2[1]) + 0.1 * link + 0.1 * ent        opt.zero_grad()        loss.backward()        opt.step()        if epoch % 20 == 0:            print('epoch', epoch, 'loss', round(loss.item(), 3))    print('S:', pool.last_S.tolist())  if __name__ == '__main__':    train()

1. S = torch.softmax(S_logits, dim=1)对分配打分按行做 softmax，得到每行和为 1 的软分配矩阵，这是“可微分”的关键。
2. self.last_S = S保存本次前向的分配矩阵，便于训练后直接查看学到的分组。
3. A_new = torch.mm(torch.mm(S.t(), A), S)用 $S^T A S$ 把原图边权聚合到超节点之间，完成拓扑压缩。
4. X_new = torch.mm(S.t(), Z)用 $S^T Z$ 生成超节点的特征矩阵。
5. link_loss = torch.norm(A - torch.mm(S, S.t()), p='fro')计算连接池化损失，惩罚拆散有边相连的节点。
6. ent_loss = (-S * torch.log(S + 1e-9)).sum() / S.size(0)计算分配向量的平均信息熵，鼓励归属更明确。
7. loss = -torch.sum(X2[0] * X2[1]) + 0.1 * link + 0.1 * ent用任务损失加两个辅助损失的加权和，实现端到端优化。

常见误区

面试题