基于扩散的 Transformer：神经消息传递统一框架

为具有特定几何（例如观察到的或未观察到的）结构的数据学习表示是一项基本挑战，其中消息传递神经网络（MPNN）已成为事实上的模型解决方案。在本文中，受物理系统启发，我们提出了一种受能量约束的扩散模型，将流形上的扩散偏置与能量最小化的层级约束相结合。我们指出，扩散算子与能量函数在扩散过程中隐式下降的能量函数之间存在一一对应关系，而求解能量约束扩散系统的有限差分迭代可诱导各种类型 MPNN 在观察到的或潜在的结构上的传播层。这导致一种统一的数学框架，可概括其计算流程可表述为消息传递（或其特殊情况）的常见神经网络架构，包括 MLP、GNN 和 Transformer。基于这些洞见，我们构思了一类新的神经消息传递模型，称为扩散激发的 Transformer（DIFFormer），其全局注意力层源自以能量约束扩散框架为原则的框架。我们在从真实网络到图像、文本和物理粒子的多样化数据集上进行评估，展示了该新模型在数据结构为观察（作为图）、部分观察或完全未观察的情况下取得令人满意的性能。