设计针对图结构数据的 Transformer 模型（通常简称 graph Transformer）目前已成为了一个备受关注的研究方向。尽管现有方法在表达能力上取得了显著进展，但它们在面对拓扑结构变化时的泛化能力仍是一个开放性问题。
具体而言，大多数现有工作主要关注特征和标签的分布偏移，却忽略了训练和测试图拓扑可能来自不同分布的“拓扑分布偏移”（topological distribution shifts）现象。
这种偏移在许多现实世界的关键场景中普遍存在，例如分子结构具有不同的药物相似性，在药物发现中，即使分子结构只有微小差异（拓扑变化），其药理活性也可能截然不同。
为应对这一挑战，我们提出了 Advective Diffusion Transformer（简称AdvDIFFormer），一个受物理学启发的图 Transformer 模型。
该模型源于对流扩散方程（Advective Diffusion Equation），这类方程描述了一个定义在同时包含观测和潜在拓扑结构上的连续消息传递过程，能够巧妙地结合图的局部结构信息与全局依赖关系，并提供在拓扑结构变化下可证明的泛化能力。
论文题目：
Supercharging Graph Transformers with Advective Diffusion
论文链接：
https://proceedings.mlr.press/v267/wu25a/wu25a.pdf
代码链接：
https://github.com/qitianwu/AdvDIFFormer
模型介绍：物理启发的图表示学习
近期相关工作 [1,3] 阐释了图学习模型（即图神经网络）与扩散方程的本质联系，而工作 [2,3] 则进一步把 Transformer 和扩散方程联系了起来。AdvDIFFormer 的核心思想是利用对流扩散方程来启发设计新的 Transformer。
对流扩散方程常用于描述复杂系统中物理量的时空变化，其中“扩散”项由浓度梯度驱动（物质从高浓度区域向低浓度区域传播），而“对流”项则由扩散物质的宏观运动引起。我们将这种物理直觉映射到图上的消息传递过程。
▲ 对流扩散方程描述了系统中驱动物理量变化的两种效应：扩散效应是一种内在的由浓度差导致的运动，其不会随环境而改变；对流效应则是由宏观运动导致的，其会依赖于特定的环境。受此启发，一个理想的可泛化的模型，应该能学习到“扩散效应”所对应的映射关系。
在图 上，我们将每个节点 的特征 视为在时间 的“浓度”。图上的对流扩散方程可以形式化为：
其中 是初始条件，通过一个编码器 从原始节点特征 得到； 是一个超参数，用于平衡对流（局部信息）和扩散（全局信息）的影响。该方程的闭合形式解为 。
AdvDIFFormer 的两大核心机制：
非局部扩散作为全局注意力（Non-local diffusion as global attention）：扩散过程由浓度梯度驱动，其扩散系数在不同环境中保持不变，反映了数据固有的、不受外部结构变化影响的潜在关联。
这恰好对应了图上的全局注意力机制，它允许任意节点对之间进行即时信息流，并捕获与底层数据流形相关的潜在交互。我们将耦合矩阵 C 实例化为一个全局注意力矩阵，用于计算任意节点对之间的相似性：
其中 是一个可学习的成对相似性函数（例如，可以是一个多层感知机或点积注意力）。这种设计使得模型能够捕获到图中节点之间的长距离依赖关系，而不受限于局部邻域。
对流作为局部消息传递（Advection as local message passing）：对流过程由定向运动驱动，这是一种外部力，其速度取决于具体环境。这与环境敏感的图拓扑结构类似，后者在特定环境中对预测标签具有有用的信息。
我们将速度 实例化为归一化图邻接矩阵，即 ，这反映了观测到的局部结构信息。其中 是邻接矩阵， 是度矩阵。这一部分实现了传统的图神经网络（GNN）中的局部消息传递，确保模型能够有效利用图中显式的局部连接信息。
通过将非局部扩散（通过注意力实现）和对流（通过 MPNN 实现）巧妙地融合在一个统一的模型框架中，AdvDIFFormer 能够同时处理图的全局和局部信息，从而在面对拓扑结构变化时展现出更强的可泛化性。
▲ AdvDIFFormer 模型架构概览。模型通过对流-扩散方程融合了全局注意力（扩散）和局部消息传递（对流）机制
理论分析：泛化能力与拓扑鲁棒性
下面我们深入分析 AdvDIFFormer 在拓扑结构变化下的泛化能力，这也是以往图学习模型常常忽视的关键问题。
我们感兴趣的是当模型从由环境 生成的训练数据迁移到由 生成的测试数据时，其泛化误差如何变化。这种数据分布偏移导致了图拓扑 从 分布变为 分布。
我们首先将模型在拓扑结构变化下的泛化误差分解为三个误差项：
其中：
是分布内泛化误差，与测试数据无关，主要取决于模型复杂度、训练样本量等。
这一项 反映了模型输出的节点表示 随拓扑变化 而变化的幅度。这是我们关注的核心项，因为它直接衡量了模型对拓扑变化的敏感度。
这一项 衡量了不同环境下标签 的固有差异。
我们的主要贡献在于对 的分析，并证明 AdvDIFFormer 能够显著控制这一项。
定理 3.2：
对于在定义的生成机制下产生的任何图数据，如果 是单射的（保证了节点身份的唯一性），那么 AdvDIFFormer 模型可以将节点表示的变化幅度 降低到任意阶 ，其中 是任意多项式函数， 且 。
这一结论的含义是，即使图的拓扑结构发生了变化，AdvDIFFormer 产生的节点表示 也只会以拓扑变化幅度 的多项式形式发生变化，而不是指数级变化。这意味着模型在处理拓扑结构变化时，能够以任意速度控制节点表示的变化率，从而保证了所需的泛化水平。
推论 3.3：
在定理 3.1 和 3.2 的相同条件下，AdvDIFFormer 模型的泛化误差界 可以降低到任意多项式阶，即 。
这一结果表明与常见的图扩散模型（其泛化误差可能随拓扑结构变化呈指数增长）不同，AdvDIFFormer 的泛化误差可以在任意速度下被控制，从而在拓扑结构变化下具有可证明的泛化能力，为模型在现实世界复杂场景中的应用提供了理论基础。
▲ 在合成数据上 AdvDIFFormer 与其他图扩散模型的对比。随着拓扑分布偏移的加重，其他方法的泛化误差显著增加，而两种 AdvDIFFormer 实现的泛化误差几乎不变
AdvDIFFormer 的两种实现方式
由于对流扩散方程的闭合形式解中涉及矩阵指数 的计算，这在实际操作中是难以处理的。因此，论文提出了两种基于数值近似的实现版本。
AdvDIFFormer-i（基于线性系统近似）：该版本利用 Padé-Chebyshev 理论的扩展 [4,5]，通过多个矩阵逆的组合来近似矩阵指数。具体而言，形如 的矩阵指数可以近似表示为：
其中 和 是预定义的参数， 是近似的阶数。为了增强神经网络的能力，AdvDIFFormer-i 将此方案扩展到多头网络，每个头都通过独立参数化的注意力网络进行传播。矩阵逆的计算则通过深度学习工具（如 PyTorch）中可用的线性系统求解器实现，并支持自动微分。具体模型的数学表达为：
AdvDIFFormer-s（基于几何级数近似的线性复杂度实现）：当图的规模较大时，矩阵逆的计算成本会非常高。为了提高可扩展性，AdvDIFFormer-s 采用几何级数近似来计算矩阵指数。这种方法可以将矩阵指数近似为有限的级数和，即：
其中 是级数的截断长度。在 AdvDIFFormer 中，闭合形式的 PDE 解对应于 。因此，几何级数求和可以表示为 在 上的加权和。AdvDIFFormer-s 将这种加权和推广到一个单层神经网络中实现，以进一步优化计算效率。具体模型的数学表达为：
在具体实现中，我们沿用了 [2] 中提出的线性注意力实现方式，最终模型的计算复杂度控制在相对节点数目的线性级别
实验结果
我们在多个领域（包括信息网络、分子图和蛋白质相互作用网络）的不同预测任务中验证了AdvDIFFormer的泛化性能。实验设计旨在模拟实际中可能遇到的拓扑结构变化，以全面评估模型的泛化能力。
信息网络：我们考虑在两个广泛使用的信息网络数据集上的节点分类任务。
Arxiv数据集：包含来自 arXiv 论文引用网络的节点（论文）和边（引用关系）。
Twitch 数据集：包含 Twitch 流媒体平台的用户和关注关系。实验结果表明，在这些拓扑结构会发生变化的数据集上，AdvDIFFormer 表现出了显著优越的泛化能力。
▲ 在 Arxiv 和 Twitch 数据集上的准确率对比。AdvDIFFormer 在不同的测试设置下均表现出更强的鲁棒性
蛋白质相互作用：我们使用了 DPPIN (Dynamic Protein-Protein Interaction Network) 数据集，该数据集包含 12 个动态蛋白质-蛋白质相互作用网络，涵盖了酵母在不同实验条件下的基因表达和蛋白质相互作用信息。
我们考虑了两种预测任务来全面评估模型：
节点回归：预测当前时间的基因表达值（RMSE 衡量）；
边回归：预测蛋白质之间的共表达相关系数（RMSE 衡量）。
AdvDIFFormer 在这些任务中均表现出色，尤其在最差情况下的测试性能中位居前列，这进一步验证了其在复杂生物网络中处理动态拓扑的有效性。
▲ 在 DPPIN 数据上的准确率对比
分子映射算子生成：在分子映射（HAM）数据集上，预测任务可以被建模为图分割问题，目标是预测分子中哪些原子属于同一个粗粒度映射操作。
实验设置通过分子质量进行数据划分（训练集为小分子，测试集为大分子），这要求模型从训练集中的小分子泛化到测试集中的大分子，本质上是应对分子拓扑结构的变化。AdvDIFFormer 能够更准确地估计分子结构，相比其他模型展现出更理想的泛化能力。
▲ 在 HAM 映射算子生成的结果对比
这些广泛的实验结果一致表明，AdvDIFFormer 在处理各种图学习任务中的拓扑结构变化方面，均能提供卓越的性能和强大的泛化能力。
结语
本文提出了 AdvDIFFormer，一个基于对流扩散方程的图 Transformer 模型，旨在解决图学习模型在拓扑结构变化下的泛化难题。
我们不仅从理论上证明了其控制泛化误差的能力，使其能够适应不同拓扑结构的图数据，更在信息网络、分子筛选和蛋白质相互作用等多个领域的经验任务中验证了其卓越的性能。
这项工作为构建更具鲁棒性和泛化能力的图学习模型提供了一条新颖且富有前景的技术路径，有望推动图表示学习在现实世界复杂动态系统中的广泛应用。
（文：PaperWeekly）