Bandit Samplers for Training Graph Neural Networks
Ziqi Liu Ant Financial Services Group、Zhengwei Wu††Ant Financial Services Group Zhiqiang Zhang Ant Financial Services Group、Jun Zhou Ant Financial Services Group、Shuang Yang Ant Financial Services Group、Le Song Ant Financial Services Group Georgia Institute of Technology、Yuan Qi Ant Financial Services Group Equal Contribution.
Abstract
Several sampling algorithms with variance reduction have been proposed for accelerating the training of Graph Convolution Networks (GCNs). However, due to the intractable computation of optimal sampling distribution, these sampling algorithms are suboptimal for GCNs and are not applicable to more general graph neural networks (GNNs) where the message aggregator contains learned weights rather than fixed weights, such as Graph Attention Networks (GAT). The fundamental reason is that the embeddings of the neighbors or learned weights involved in the optimal sampling distribution are changing during the training and not known a priori, but only partially observed when sampled, thus making the derivation of an optimal variance reduced samplers non-trivial. In this paper, we formulate the optimization of the sampling variance as an adversary bandit problem, where the rewards are related to the node embeddings and learned weights, and can vary constantly. Thus a good sampler needs to acquire variance information about more neighbors (exploration) while at the same time optimizing the immediate sampling variance (exploit). We theoretically show that our algorithm asymptotically approaches
中文速览
训练图神经网络（Graph Neural Networks, GNNs）时，"邻居爆炸"问题会导致计算量和内存占用急剧膨胀，现有的重要性采样方法虽能缓解这一问题，但要计算最优采样分布，需要提前知道所有邻居的隐层嵌入和注意力权重，而这些值在训练过程中持续变化、无法预先获得，导致现有方法只能用近似值凑合，且对图注意力网络（Graph Attention Networks, GAT）这类权重可学习的模型根本不适用。本文将采样分布的优化重新建模为一个对抗性多臂老虎机（adversarial bandit）问题——采样器既要探索更多邻居以积累方差信息，又要即时利用已有信息降低采样方差——并提出了基于EXP3和EXP3.M两种在线学习算法的采样策略，无需预先知道最优分布，而是随训练迭代不断逼近它。理论上证明了所提算法的估计方差渐近收敛到最优方差的3倍以内，实验上也在多个公开数据集上验证了其在收敛速度和采样方差方面优于现有方法。这项工作的重要性在于，它是首个同时适用于GCN和GAT等更广泛GNN架构的有理论保证的方差缩减采样框架。
原文 arXiv:2006.05806;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2006.05806v2