本⽂介绍⼀项近期的研究⼯作，试图建⽴能量约束扩散微分⽅程与神经⽹络架构的联系，从而原创性的提出了物理启发下的 Transformer，称作 DIFFormer。作为⼀种通⽤的可以灵活⾼效的学习样本间隐含依赖关系的编码器架构，DIFFormer 在各类任务上都展现了强大潜⼒。这项工作已被 ICLR 2023 接收，并在⾸轮评审就收到了四位审稿⼈给出的 10/8/8/6 评分（最终均分排名位于前 0.5%）。

  • 论⽂地址：https://arxiv.org/pdf/2301.09474.pdf
  • 项⽬地址：https://github.com/qitianwu/DIFFormer

如何得到有效的样本表征是机器学习领域的⼀⼤核⼼基础问题，也是深度学习范式在各类下游任务能发挥作用的重要前提。传统的表征学习⽅法通常假设每个输⼊样本是独⽴的，即分别将每个样本输⼊进 encoder ⽹络得到其在隐空间中的表征，每个样本的前向计算过程互不干扰。然⽽这⼀假设通常与现实物理世界中数据的⽣成过程是违背的：由于显式的物理连接或隐含的交互关系，每个观测样本之间可能存在相互的依赖。

这⼀观察也启发了我们去重新思考⽤于表征计算的 encoder ⽹络设计：是否能设计⼀种新型的 encoder ⽹络能够在前向计算中显式的利⽤样本间的依赖关系（尽管 这些依赖关系是未被观察到的）。在这个⼯作中，我们从两个物理学原理出发，将神经⽹络计算样本表征的前向过程看作给定初始状态的扩散过程，且随着时间的推移（层数加深）系统的整体能量不断下降（见下图）。

DIFFormer_依赖关系

DIFFormer 模型主要思想的示意图：将模型计算样本表征的前向过程看作⼀个扩散过程，随着时间的推移，节点之间存在信号传递，且任意节点对之间信号传递的速率会随着时间适应性的变化，使得系统整体的能量最⼩化。通过扩散过程和能量约束，最终的样本表征能够吸收个体和全局的信息，更有助于下游任务。

通过试图建⽴扩散微分⽅程与神经⽹络架构的联系，我们阐释了能量约束扩散过程与各类信息传递网络（如 MLP/GNN/Transformers）的联系，并为新的信息传递设 计提供了⼀种理论参考。基于此，我们提出了⼀种新型的可扩展 Transformer 模型，称为 DIFFormer（diffusionbased Transformers）。它可以作为⼀种通⽤的 encoder，在前向计算中利⽤样本间隐含的依赖关系。⼤量实验表明在⼩ / ⼤图节点分类、图⽚ / ⽂本分类、时空预测等多个领域的实验任务上 DIFFormer 都展现了强⼤的应⽤潜⼒。在计算效率上，DIFFormer 只需要 3GB 显存就可以实现⼗万级样本间全联接的信息传递。

动机与背景

我们⾸先回顾⼀个经典的热⼒学中的热传导过程：假设系统中有

个节点，每个节点有初始的温度，两两节点之间都存在信号流动，随着时间的推移节点的温度会不断更新。上述物理过程事实上可以类⽐的看作深度神经网络计算样本表征（embedding）的前向过程。

DIFFormer_数据集_02

将神经⽹络的前向计算过程看作⼀个扩散过程：每个样本视为流形上的固定位置节点，样本的表征为节点的信号，表征的更新视作节点信号的改变，样本间的信息传递看作节点之间的信号流动 

DIFFormer_人工智能_03

这⾥我们可以把每个样本看作⼀个离散空间中的节点，样本表征看作节点的信号。当模型结构考虑样本交互时（如信息传递），它可以被看作节点之间的信号流动，随着模型层数加深（即时间的推移），样本表征会不断被更新。

扩散过程的描述

⼀个经典的扩散过程可以由⼀个热传导⽅程（带初始条件的偏微分⽅程）来描述

DIFFormer_数据集_04

在离散空间中，梯度算⼦可以看作两两节点的信号差异，散度算子可以看作单个节点流出信号的总和，⽽扩散率（diffusivity）是⼀种对任意两两节点间信号流动速率的度量

由此我们可以写出描述 N 个节点每时每刻状态更新的扩散微分⽅程，它描述了每个状态下系统中每个节点信号的变化等于流向其他节点的信号总和：

DIFFormer_图结构_05

由扩散方程导出的信息传递

DIFFormer_数据集_06

下图概述了这三种信息传递模式： 

DIFFormer_依赖关系_07

我们研究最后⼀种信息传递⽅式，每层更新的样本表征会利⽤上⼀层所有其他样本的表征，在理论上模型的表达能⼒是最强的。但由此产⽣的⼀个问题是：要如何才能确定合适的每层任意两两节点之间的 diffusivity，使得模型能够产⽣理想的样本表征？

刻画⼀致性的能量函数

DIFFormer_依赖关系_08

能量约束的扩散过程

DIFFormer_人工智能_09

定理

DIFFormer_人工智能_10

  • DIFFormer-a：在计算相似度时引⼊⾮线性，从⽽提升模型学习复杂结构的表达能⼒：

DIFFormer_依赖关系_11

 

DIFFormer_图结构_12

两种模型 DIFFormer-s 和 DIFFormer-a 每层更新的运算过程（矩阵形式），红⾊标注的矩阵乘法操作是计算瓶颈。DIFFormer-s 的优势在于可以实现对样本数量 N 的线性复杂度，有利于模型扩展到⼤规模数据集

模型扩展
更进⼀步的，我们可以引⼊更多设计来提升模型的适⽤性和灵活度。上述的模型主要考虑了样本间的 all-pair attention。对于输⼊数据本身就含有样本间图结构的情况，我们可以加⼊现有图神经⽹络（GNN）中常⽤的传播矩阵（propagation matrix）来融合已知的图结构信息，从⽽定义每层的样本表征更新如下

DIFFormer_依赖关系_13

DIFFormer 的全局输⼊包含样本输⼊特征 X 以及可能存在的图结构 A（可以省略），通过堆叠 DIFFormer layer 更新计算样本表征。在每层更新时，需要计算⼀个全局 attention（具体的可以使⽤ DIFFormer-s 和 DIFFormer-a 两种实现），如果考虑输⼊图结构则加⼊ GCN Conv

另⼀个值得探讨的问题，是如何处理⼤规模数据集（尤其是包含⼤量样本的数据集，此时考虑全局 all-pair attention ⾮常耗费资源）。在这种情况下我们默认使⽤线性复杂度的 DIFFormer-s 的架构，并且可以在每个训练 epoch 对数据集进⾏ random mini-batch 划分。由于线性复杂度，我们可以使⽤较⼤的 batch size 也能使得模型在单卡上进⾏训练（详⻅实验部分）。

DIFFormer_图结构_14

对于包含⼤量样本的数据集，我们可以对样本进⾏随机 minibatch 划分，每次只输⼊⼀个 batch 的样本。当输⼊包含图结构时，我们可以只提取 batch 内部样本所组成的⼦图输⼊进⽹络。由于 DIFFormer-s 只需要对 batch size 的线性复杂度，在实际中就可以使⽤较⼤的 batch size，保证充⾜的全局信息

实验结果

为了验证 DIFFormer 的有效性和在不同场景下的适⽤性，我们考虑了多个实验场景，包括不同规模图上的节点分类、半监督图⽚ / ⽂本分类和时空预测任务。

图节点分类实验
此时输⼊数据是⼀张图，图中的每个节点是⼀个样本（包含特征和标签），⽬标是利⽤节点特征和图结构来预测节点的标签。我们⾸先考虑⼩规模图 的实验，此时可以将⼀整图输⼊ DIFFormer。相⽐于同类模型例如 GNN，DIFFormer 的优势在于可以不受限于输⼊图，学习未被观测到的连边关系，从⽽更好的捕捉⻓距离依赖和潜在关系。下图展示了与 SOTA ⽅法的对⽐结果。

DIFFormer_数据集_15

进⼀步的我们考虑在⼤规模图上的实验。此时由于图的规模过⼤，⽆法将⼀整图直接输⼊模型（否则将造成 GPU 过载），我们使⽤ mini-batch 训练。具体的，在每个 epoch，随机的将所有节点分为相同⼤⼩的 mini-batch。每次只将⼀个 mini-batch 的节点输⼊进⽹络；⽽对于输⼊图，只使⽤包含在这个 mini-batch 内部的节点所组成的⼦图输⼊进⽹络；每次迭代过程中，DIFFormer 也只会在 mini-batch 内部的节点之间学习 all-pair attention。这样做就能⼤⼤减⼩空间消耗。⼜因为 DIFFormer-s 的计算复杂度关于 batch size 是线性的，这就允许我们使⽤很⼤的 batch size 进⾏训练。下图显示了在 ogbn-proteins 和 pokec 两个⼤图数据集上的测试性能，其中对于 proteins/pokec 我们分别使⽤了 10K/100K 的 batch size。此外，下图的表格也展示了 batch size 对模型性能的影响，可以看到，当使⽤较⼤ batch size 时，模型性能是⾮常稳定的。

DIFFormer_人工智能_16

图⽚ / ⽂本分类实验

第⼆个场景我们考虑⼀般的分类问题，输⼊是⼀些独⽴的样本（如图⽚、⽂本），样本间没有已观测到的依赖关系。此时尽管没有输⼊图结构， DIFFormer 仍然可以学习隐含在数据中的样本依赖关系。对于对⽐⽅法 GCN/GAT，由于依赖于输⼊图，我们这⾥使⽤ K 近邻⼈⼯构造⼀个样本间的图结构。

DIFFormer_依赖关系_17

时空预测

进⼀步的，我们考虑时空预测任务，此时模型需要根据历史的观测图⽚段（包含上⼀时刻节点标签和图结构）来预测下⼀时刻的节点标签。这⾥我们横向对⽐ 了 DIFFormer-s/DIFFormer-a 在使⽤输⼊图和不使⽤输⼊图（w/o g）时的性能，发现在不少情况下不使⽤输⼊图模型反⽽能给出的较⾼预测精度。这也说明了在这类任务中，给定的观测图结构可能是不可靠的，⽽ DIFFormer 则可以通过从数据中学习依赖关系得到更有⽤的结构信息。

DIFFormer_依赖关系_18

扩散过程下的统⼀视⻆

从能量约束的扩散过程出发，我们也可以将其他信息传递模型如 MLP/GCN/GAT 看作 DIFFormer 的特殊形式，从⽽给出统⼀的形式化定义。下图概括了⼏种⽅法对应的能量函数和扩散率。相⽐之下，从扩散过程来看， DIFFormer 会考虑任意两两节点之间的信号流动且流动的速率会随着时间适应性的变化，⽽ GNN 则是将信号流动 限制在⼀部分节点对之间。从能量约束来看，DIFFormer 会同时考虑局部（与⾃身状态）和全局（与其他节点）的⼀致性约束，⽽ MLP/GNN 则是分别侧重于⼆者之⼀， 且 GNN 通常只考虑输⼊图中相邻的节点对约束。

DIFFormer_图结构_19

总结与讨论

在这个⼯作中，我们讨论了如何从扩散⽅程出发得到 MLP/GNN/Transformer 的模型更新公式，⽽后提出了⼀个能量约束下的扩散过程，并通过理论分析得到了最优 扩散率的闭式解。基于理论结果，我们提出了 DIFFormer。总的来说，DIFFormer 主要具有以下两点优势：

DIFFormer_图结构_20

DIFFormer 作为⼀个通⽤的 encoder，可以被主要应⽤于以下⼏种场景：

DIFFormer_数据集_21