中科大提出首个可证明收敛的子图采样方法ICLR2023Spo

　　编辑：好困【新智元导读】中科大王杰教授团队提出局部消息补偿技术，解决采样子图边缘节点邻居缺失问题，弥补图神经网络（GNNs）子图采样方法缺少收敛性证明的空白，推动GNNs的可靠落地。
　　图神经网络（GraphNeuralNetworks，简称GNNs）是处理图结构数据的最有效的机器学习模型之一，也是顶会论文的香饽饽。
　　然而，GNNs的计算效率一直是个硬伤，在大规模图数据上训练GNNs常常会遇上邻居爆炸（neighborexplosion）问题节点表示和随机梯度的计算复杂度会随着图神经网络层数的增加而指数上升。
　　很多GNNs的学术研究都会倾向于选择小规模图数据集（千量级节点数）进行实验，避开GNNs的计算效率问题。但是，这一问题在工业界实际落地的场景中无法避免：在大规模图数据（十亿节点）〔3〕上，这些GNNs根本无法运行。
　　一个最简单粗暴的办法是：在每次模型训练或预测的时候，从全量图上切出一个子图，在子图上运行GNNs。这又会带新的问题：在子图上训练的GNNs能和全量图上训练的GNNs一样吗？子图边缘节点会不会丢失很多邻居信息？
　　为此，中科大MIRALab王杰教授团队提出了一种GNNs的子图采样训练方法本地消息补偿（LocalMessageCompensation，简称LMC）。
　　LMC具有极低的计算开销；并且，理论证明：LMC在子图上训练的GNNs的性能可媲美在全量图上训练的GNNs，同时LMC能加速GNNs收敛。相关成果论文已被ICLR2023接收为Spotlight。
　　作者列表：石志皓，梁锡泽，王杰
　　论文链接：https：openreview。netforum？id5VBBA91N6n
　　1。引言
　　基于消息传递机制的图神经网络（GNNs）在许多实际应用中取得了巨大成功。然而，在大规模图上训练GNNs会遇到众所周知的邻居爆炸（neighborexplosion）问题节点的依赖性随消息传递层的数量呈指数增长。
　　子图采样方法一类备受瞩目的小批量训练（minibatchtraining）技术在反向传播中丢弃小批量之外的消息，以此避免邻居爆炸问题，但同时以牺牲梯度估计的精度为代价。这对它们的收敛性分析和收敛速度都提出了重大挑战，严重限制了它们在现实场景中的进一步应用。
　　为了应对这些挑战，我们提出了一种具有收敛性保证的新型子图采样方法本地消息补偿（LocalMessagePassing，简称LMC）。据我们所知，LMC是首个具有可证明收敛性的子图采样方法。
　　LMC的关键思想是基于反向传播传递的消息传递建模来恢复在反向传播中被丢弃的消息。通过对正向和反向传播中丢弃的消息进行高效和有效的补偿，LMC计算出准确的小批量梯度，从而加速收敛。
　　进一步地，我们证明了LMC收敛到GNNs的一阶驻点（firstorderstationarypoints）。在大规模基准测试任务中的实验表明，LMC在效率方面明显优于最先进的子图采样方法。
　　2。背景与问题
　　2。1图神经网络
　　在实际问题中，图结构数据随处可见，例如知识图谱、分子、计算机网络、社交网络、神经元网络、文章引用网络等，如图1所示。
　　图1。图结构数据在实际问题中随处可见，图中展示了各式各样的图数据。
　　图神经网络（GraphNeuralNetworks，简称GNNs）通过消息传递范式〔1〕处理图数据，是当前处理图结构数据最有效的机器学习模型之一。在每个消息传递层中，GNNs迭代地聚合邻居节点的消息，以更新当前节点的表示。这种范式在许多实际应用中取得了巨大的成功，例如搜索引擎〔2〕、推荐系统〔3〕、材料工程〔4〕、分子性质预测〔5，6〕，以及组合优化〔7〕。
　　具体地，以半监督的结点分类任务为例，GNNs旨在通过最小化目标函数来学习结点嵌入以及参数，其中，是有标签结点的集合，是参数为的输出层与一个损失函数的组合，是结点的嵌入，是结点的标签，是结点特征，是图上所有边的集合。
　　一个层的GNN通过次有着不同参数的消息传递迭代来生成最终的结点嵌入：
　　其中，是第层的消息传递函数，参数为。
　　消息传递函数遵循聚合更新机制，即
　　其中是为结点的每个邻居生成消息的函数，是将邻居消息集合映射到最终消息的聚合函数，是组合从前的结点嵌入，消息，以及结点特征的更新函数。
　　2。2邻居爆炸
　　尽管GNNs在许多应用中取得了巨大的成功，这种消息迭代机制也给GNNs在大规模图数据上的训练带来了挑战。
　　使用有限的GPU内存将深度模型扩展到任意大规模数据的一种常见方法是通过小批量梯度近似全批次梯度。然而，对于图结构数据，由于众所周知的邻居爆炸问题，计算小批量节点的损失函数和相应的小批量梯度的成本是非常昂贵的。
　　具体地，对于GNNs而言，一个结点在第层消息传递中的嵌入递归地依赖于它邻居在第层的嵌入。因此，计算复杂度会随着消息传递层数的增加而指数级上涨，带来无法令人接受的计算开销。
　　2。3子图采样方法
　　为了解决邻居爆炸问题，最近的一些工作提出了各种各样的采样技术以减少消息传递所牵涉的节点个数。例如，结点采样方法〔8，9〕和层采样方法〔10，11，12〕会在消息传递中递归地采样邻居，从而估计结点嵌入以及对应的小批量梯度。
　　与这种递归的范式不同，子图采样方法〔13，14，15，16〕使用了一种更为简单、成本低廉的一次性采样范式（oneshotsamplingfashion）为不同的消息传递层采样同一个子图，该子图由同一小批量结点所构建。
　　通过丢弃小批量之外的消息，子图采样方法将消息传递过程限制在小批量中，使得复杂度随消息传递层数的增加而线性增长，极大降低了计算开销。此外，通过直接在子图上运行GNNs，子图采样方法适用于非常广泛的GNN结构。由于上述优势，子图采样方法近期收到了越来越多的关注。
　　然而，子图采样方法这种丢弃小批量外部消息的做法牺牲了梯度估计的精度，这给它们的收敛性分析和收敛速度带来了极大挑战：首先，近期工作〔9，17〕表明，不准确的小批量梯度会严重降低GNNs的收敛速度。其次，我们的实验表明，现有子图采样方法在批量大小较小时难以达到全梯度下训练的表现；而我们在实际应用中经常会将批量大小设置为一个较小的数字，以避免超出GPU的显存。
　　对此，我们发问：能否设计一个子图采样方法，它既有极低的计算开销，又有媲美全梯度训练的预测精度，同时还有严格的收敛性保证？
　　我们的回答是：LMC能做到！
　　3。方法：局部信息补偿LMC
　　我们的研究思路受到了VRGCN〔9〕的启发，其主要抓手是把节点或层级别的递归采样看成一个无偏的基线方法StandardSGD的近似，进而通过对于梯度的误差分析来证明收敛性。
　　然而，很难把子图采样方法看成StandardSGD的近似，因为子图采样在每一层都采样相同的子图，每一层的计算都会引入不可避免的偏差。因此，我们第一步是先提出一个BackwardSGD，它更便于我们的分析子图采样这种一次性采样范式。在BackwardSGD的基础上，我们分析如何给子图采样方法加入合适的补偿项，减少它的偏差，进而找到一个可证明收敛的子图采样算法。
　　3。1将反向传播建模为消息传递
　　梯度是容易计算的，所以我们主要介绍如何计算。令，为辅助变量，则有。由链式法则，我们能够基于迭代地计算：
　　以及
　　然后，我们可以使用处理向量Jacobian积的自动求导工具来计算梯度。
　　我们将反向传播（即迭代计算方程（3）的过程）建模为消息传递。为了看到这一点，我们只需注意到（3）等价于
　　其中是的第列。方程（5）分别使用、求和聚合，以及恒等映射作为生成函数、聚合函数，以及更新函数。
　　3。2LMC的基石：BackwardSGD
　　基于这一反向传播的消息传递建模，我们设计了一个SGD变体BackwardSGD，它能带来无偏的梯度估计。需要说明的是，BackwardSGD是我们主要方法LMC的基石。
　　给定一个被采样的小批量，假设我们已经获得了小批量中结点的准确嵌入和准确辅助变量。读者需要注意：这一假设是比较强的，事实上我们很难计算准确的嵌入和辅助变量。下一节介绍的主要方法LMC所做的就是利用子图采样来估计结点嵌入和辅助变量。
　　首先，BackwardSGD计算参数的小批量梯度：
　　然后，BackwardSGD计算参数的小批量梯度：
　　注意到：对于不同的层数，小批量梯度牵涉的小批量是同一个（即），这就给基于BackwardSGD设计子图采样方法提供了基础。
　　BackwardSGD的另一个吸引人的性质是：小批量梯度和是无偏的，如第4节中的定理1所示。详细的证明请参见原论文附录。
　　3。3本地消息补偿LMC
　　在上一节中，BackwardSGD所计算的小批量梯度依赖于小批量中节点的准确嵌入和准确辅助变量，而不是整张图。然而，BackwardSGD仍然不是可扩展的（scalable），因为邻居爆炸问题会使得准确结点嵌入和辅助变量的计算极其昂贵，所以事实上我们无法得到准确的和。
　　在这一节中，为了解决邻居爆炸问题，我们提出了新颖的子图采样方法本地消息补偿（LocalMessageCompensation，简称LMC）。
　　LMC首先通过不完全最新值（incompleteuptodatevalues）和历史值（historicalvalues）的凸组合来高效地估计和，然后利用方程（6）和（7）计算小批量梯度。在之后的理论分析中，我们证明了LMC收敛到GNNs的一阶驻点。在算法1和理论分析中，我们用表示一个第层、第次迭代时的量，而在其他地方我们省略上标，用来表示。
　　在每个训练迭代中，我们采样一个小批量结点，通过历史值和，以及不完全最新值和的凸组合来高效地估计和。
　　为便于读者理解方法的核心思想，我们将LMC与现有最先进方法GAS〔15〕的前向传播、反向传播计算图展示在图2。
　　图2。LMC与GAS前向传播与反向传播的计算图。可以看到，在前向传播和反向传播中，LMC均进行了小批量结点与一跳邻居之间的消息交互（即补偿），而GAS在反向传播中丢弃了小批量之外的消息。在前向传播中，我们将的临时嵌入设为，然后以的顺序更新中的历史嵌入。特别地，在第层，我们进行以下计算：
　　在反向传播中，我们将的临时辅助变量设为，然后以的顺序更新中的历史辅助变量。特别地，在第层，我们进行以下计算：
　　关于方法的具体细节、详细解释、计算复杂度分析等，请读者参见原论文。
　　我们分别称和为第层前向传播和反向传播的本地消息补偿。
　　4。理论分析
　　理论分析包含三个主要定理。从直观（说人话）的角度，它们分别在说：定理1：BackwardSGD的梯度是无偏的。这样，我们就基本可以保证BackwardSGD的收敛性。定理2：LMC所估计的梯度和BackwardSGD的梯度相差不大，能够被我们给出的上界所控制。定理3：LMC收敛到GNNs的一阶驻点，这也是我们的最终定理。
　　理论部分的核心思想是：LMC和BackwardSGD的收敛行为一致。
　　在本节中，我们做如下假设：在第个迭代中，小批量节点是从中均匀采样的，对应的有标签节点集是从中采样的。函数，，，，，是Lipschitz连续的，其中。范数，，，，，，，，，，，被常数所控制。
　　定理1。假设一个小批量是从结点集合中均匀采样的，并且对应的有标签结点集合也是从中均匀采样的，则方程（6）和（7）所计算的小批量梯度和是无偏的。
　　定理2。在上述假设下，令和，存在和使得
　　定理3。在上述假设下，再假定最优值被控制。令，，以及，LMC可保证在次迭代后找到一个驻点使得，其中是随机从中选取的，。
　　5。实验
　　在实验部分，我们做了4个大数据集的实验，图3列出了其中3个。LMC的训练集loss的收敛速度超过了所有其他方法，但是测试集上有个众所周知的泛化问题，尽管LMC在训练集上收敛快，但模型很快就过拟合了，所以LMC在测试集的准确率曲线提升看起来并不如训练集明显。
　　我们发现，LMC最终的预测准确率和SOTA方法GAS相差不大，这是因为earlystopping技术，即在训练过程中，测试集的曲线是震荡的，earlystopping汇报的大概率是测试集准确率曲线的最大值。为了突出训练过程中波动性的影响，我们在图3汇报的曲线用滑动窗口取了平均值构成实线，标准差构成阴影部分。可以看出，LMC在训练稳定性上明显超过GAS。GAS和LMC最终预测准确率的差距会在batchsize比较小的情况下有所体现（图6），这时METIS的作用会被削弱。
　　图3。收敛时间对比
　　我们进一步统计达到一个给定测试集准确率的时间。我们先跑一个全梯度GD，得到它的最高测试集准确率，然后再分别运行几种子图采样方法，计算达到这一准确率需要的时间。在REDDIT数据集上，LMC相对于GAS的加速比达到2倍。
　　图4。达到最高测试集准确率时间对比
　　更进一步，LMC的梯度估计和BackwardSGD差不多，所以梯度估计是更准确的。我们这里统计了计算过程中的相对误差，如图5所示，确实是LMC的估计误差最小。
　　图5。相对误差对比
　　进一步做了smallbatchsize下的实验，前面在子图采样算法中举了一个例子，子图规模很小的话，丢弃的节点就很多，很容易达到次优。如表三所示，我们的方法对batchsize更加鲁棒，因此在计算资源受限的情景下，LMC的优势会更加明显。
　　图6。不同批量大小的表现
　　最后是消融实验，相对于SOTA的GAS方法，我们对前向传播过程的补偿消息进行了改进，并且在反向传播也加入了一个补偿。如图7所示，我们发现，在batchsize很小的情况下，反向传播的补偿很重要，因为这一设定下，丢弃了很多消息，导致收敛到次优解。在batchsize较大的时候，采样子图一阶邻居是很大的，我们通过采样子图一阶邻居内部的消息传递，提高了历史信息的准确率，也能提高子图采样算法的性能。
　　图7。消融实验
　　参考资料：
　　〔1〕Hamilton，WilliamL。Graphrepresentationlearning。SynthesisLecturesonArtificalIntelligenceandMachineLearning14。3（2020）：1159。
　　〔2〕Brin，Sergey，andLawrencePage。Theanatomyofalargescalehypertextualwebsearchengine。ComputernetworksandISDNsystems30。17（1998）：107117。
　　〔3〕Fan，Wenqi，etal。Graphneuralnetworksforsocialrecommendation。Theworldwidewebconference。2019。
　　〔4〕Gostick，Jeff，etal。OpenPNM：aporenetworkmodelingpackage。ComputinginScienceEngineering
　　18。4（2016）：6074。
　　〔5〕Moloi，N。P。，andM。M。Ali。Aniterativeglobaloptimizationalgorithmforpotentialenergyminimization。ComputationalOptimizationandApplications30（2005）：119132。
　　〔6〕Kearnes，Steven，etal。Moleculargraphconvolutions：movingbeyondfingerprints。Journalofcomputeraidedmoleculardesign30（2016）：595608。
　　〔7〕Wang，Zhihai，etal。LearningCutSelectionforMixedIntegerLinearProgrammingviaHierarchicalSequenceModel。arXivpreprintarXiv：2302。00244（2023）。
　　〔8〕Hamilton，Will，ZhitaoYing，andJureLeskovec。Inductiverepresentationlearningonlargegraphs。Advancesinneuralinformationprocessingsystems30（2017）。〔9〕Chen，Jianfei，JunZhu，andLeSong。Stochastictrainingofgraphconvolutionalnetworkswithvariancereduction。arXivpreprintarXiv：1710。10568（2017）。
　　〔10〕Chen，Jie，TengfeiMa，andCaoXiao。Fastgcn：fastlearningwithgraphconvolutionalnetworksviaimportancesampling。arXivpreprintarXiv：1801。10247（2018）。
　　〔11〕Zou，Difan，etal。Layerdependentimportancesamplingfortrainingdeepandlargegraphconvolutionalnetworks。Advancesinneuralinformationprocessingsystems32（2019）。
　　〔12〕Huang，Wenbing，etal。Adaptivesamplingtowardsfastgraphrepresentationlearning。Advancesinneuralinformationprocessingsystems31（2018）。
　　〔13〕Chiang，WeiLin，etal。Clustergcn：Anefficientalgorithmfortrainingdeepandlargegraphconvolutionalnetworks。Proceedingsofthe25thACMSIGKDDinternationalconferenceonknowledgediscoverydatamining。2019。
　　〔14〕Zeng，Hanqing，etal。Graphsaint：Graphsamplingbasedinductivelearningmethod。arXivpreprintarXiv：1907。04931（2019）。
　　〔15〕Fey，Matthias，etal。Gnnautoscale：Scalableandexpressivegraphneuralnetworksviahistoricalembeddings。InternationalConferenceonMachineLearning。PMLR，2021。
　　〔16〕Zeng，Hanqing，etal。Decouplingthedepthandscopeofgraphneuralnetworks。AdvancesinNeuralInformationProcessingSystems34（2021）：1966519679。
　　〔17〕Cong，Weilin，etal。Minimalvariancesamplingwithprovableguaranteesforfasttrainingofgraphneuralnetworks。Proceedingsofthe26thACMSIGKDDInternationalConferenceonKnowledgeDiscoveryDataMining。2020。