LARGE-SCALE GRAPH REGRESSION WITH VERY DEEP GNNS OPEN GRAPH BENCHMARK LARGE-SCALE CHALLENGE ENTRY (PCQM4M) Ravichandra Addanki∗ Peter W. Battaglia∗ David Budden∗ DeepMind DeepMind DeepMind [email protected] [email protected] [email protected] Andreea Deac∗ Jonathan Godwin∗ Wai Lok Sibon Li∗ DeepMind, Mila, Université de Montréal DeepMind DeepMind [email protected] [email protected] [email protected] Alvaro Sanchez-Gonzalez∗ Jacklynn Stott∗ Shantanu Thakoor∗ DeepMind DeepMind DeepMind [email protected] [email protected] [email protected] Petar Velickoviˇ c´∗ DeepMind [email protected] ABSTRACT Effectively and efficiently deploying graph neural networks (GNNs) at scale remains one of the most challenging aspects of graph representation learning. Many powerful solutions have only ever been validated on comparatively small datasets, often with counter-intuitive outcomes—a barrier which has recently been broken by the Open Graph Benchmark Large-Scale Challenge (OGB-LSC). Here we describe our OGB-LSC entry for the PCQM4M-LSC benchmark: a very deep (up to 50-layer) inductive graph regressor regularised by denoising objectives. Our model achieved an award-level (top-3) performance. In doing so, we demonstrate strong evidence of scalable self-supervised graph representation learning, and utility of very deep GNNs—both very important open issues. Our code is publicly available at: https://github.com/deepmind/deepmind-research/tree/master/ ogb_lsc/pcq. Keywords OGB-LSC · Graph Networks · Noisy Nodes 1 Introduction Effective high-dimensional representation learning necessitates properly exploiting the geometry of data [Bronstein et al., 2021]—otherwise, it is a cursed estimation problem. Indeed, early success stories of deep learning relied on imposing strong geometric assumptions, primarily that the data lives on a grid domain; either spatial or temporal. In these two respective settings, convolutional neural networks (CNNs) [LeCun et al., 1998] and recurrent neural networks (RNNs) [Hochreiter and Schmidhuber, 1997] have traditionally dominated. While both CNNs and RNNs are demonstrably powerful models, with many applications of high interest, it can be recognised that most data coming from nature cannot be natively represented on a grid. Recent years are marked with a gradual shift of attention towards models that admit a more generic class of geometric structures [Masci et al., 2015, Velickoviˇ c´ et al., 2017, Cohen et al., 2018, Battaglia et al., 2018, de Haan et al., 2020, Satorras et al., 2021]. ∗All authors contributed equally. Large-scale graph regression with very deep GNNs and denoising In many ways, the most generic and versatile of these models are graph neural networks (GNNs). This is due to the fact that most discrete-domain inputs can be observed as instances of a graph structure. The corresponding area of graph representaton learning [Hamilton, 2020] has already seen immense success across industrial and scientific disciplines. GNNs have successfully been applied for drug screening [Stokes et al., 2020], modelling the dynamics of glass [Bapst et al., 2020], web-scale social network recommendations [Ying et al., 2018] and chip design [Mirhoseini et al., 2020]. While the above results are certainly impressive, they likely only scratch the surface of what is possible with well-tuned GNN models. Many problems of real-world interest require graph representation learning at scale: either in terms of the amount of graphs to process, or their sizes (in terms of numbers of nodes and edges). Perhaps the clearest motivation for this comes from the Transformer family of models [Vaswani et al., 2017]. Transformers operate a self-attention mechanism over a complete graph, and can hence be observed as a specific instance of GNNs [Joshi, 2020]. At very large scales of natural language data, Transformers have demonstrated significant returns with the increase in capacity, as exemplified by models such as GPT-3 [Brown et al., 2020]. Transformers enjoy favourable scalability properties at the expense of their functional complexity: each node’s features are updated with weighted sums of neighbouring node features. In contrast, GNNs that rely on message passing [Gilmer et al., 2017]—passing vector signals across edges that are conditioned on both the sender and receiver nodes—are an empirically stronger class of models, especially on tasks requiring complex reasoning [Velickoviˇ c´ et al., 2019] or simulations [Sanchez-Gonzalez et al., 2020, Pfaff et al., 2020]. One reason why generic message-passing GNNs have not been scaled up as widely as Transformers is the lack of appropriate datasets. Only recently has the field advanced from simple transductive benchmarks of only few thousands of nodes [Sen et al., 2008, Shchur et al., 2018, Morris et al., 2020] towards larger-scale real-world and synthetic benchmarks [Dwivedi et al., 2020, Hu et al., 2020], but important issues still remain. For example, on many of these tasks, randomly-initialised GNNs [Velickoviˇ c´ et al., 2018], shallow GNNs [Wu et al., 2019] or simple label propagation-inspired GNNs [Huang et al., 2020] can perform near the state-of-the-art level at only a fraction of the parameters. When most bleeding-edge expressive methods are unable to improve on the above, this can often lead to controversial discussion in the community. One common example is: do we even need deep, expressive GNNs? Breakthroughs in deep learning research have typically been spearheaded by impactful large-scale competitions. For image recognition, the most famous example is the ImageNet Large Scale Visual Recognition Challenge (ILSVRC) [Russakovsky et al., 2015]. In fact, the very “deep learning revolution” has partly been kickstarted by the success of the AlexNet CNN model of Krizhevsky et al. [2012] at the ILSVRC 2012, firmly establishing deep CNNs as the workhorse of image recognition for the forthcoming decade. Accordingly, we have entered the recently proposed Open Graph Benchmark Large-Scale Challenge (OGB-LSC) [Hu et al., 2021]. OGB-LSC provides graph representation learning tasks at a previously unprecedented scale—millions of nodes, billions of edges, and/or millions of graphs. Further, the tasks have been designed with immediate practical relevance in mind, and it has been verified that expressive GNNs are likely to be necessary for strong performance. Here we detail our submitted model for the PCQM4M task, and our empirical observations while developing it. Namely, we find that the dataset’s immense scale provides a great platform for demonstrating clear outperformance of very deep GNNs, powered by denoising objectives [Godwin et al., 2021]. In doing so, we have provided meaningful evidence towards a positive resolution to the above discussion: deep and expressive GNNs are, indeed, necessary at the right level of task scale and/or complexity. Our final model has achieved award-level (top-3) ranking on PCQM4M. 2 Dataset description The PCQM4M-LSC dataset is an inductive graph regression dataset based on the PubChemQC project [Nakata and Shimazaki, 2017]. It consists of ∼4 million small molecules (described by their SMILES strings). The aim is to accelerate quantum-chemical computations: especially, to predict the HOMO-LUMO gap of each molecule. The HOMO-LUMO gap is one of the most important quantum-chemical properties, since it is related to the molecules’ reactivity, photoexcitation, and charge transport. The ground-truth labels for every molecule were obtained through expensive DFT (density functional theory) calculations, which may take several hours per molecule. It is believed that machine learning models, such as GNNs over the molecular graph, may obtain useful approximations to the DFT at only a fraction of the computational cost, if provided with sufficient training data [Gilmer et al., 2017]. The molecules are split with a 80:10:10 ratio into training, validation and test sets, based on their PubChem ID. 3 GNN Architectures We rely on a common encode-process-decode blueprint [Hamrick et al., 2018]. This implies that our input features are encoded into a latent space using node-, edge- and graph-wise encoder functions, and latent features are decoded to 2 Large-scale graph regression with very deep GNNs and denoising node-, edge- and graph- level predictions using appropriate decoder functions. The bulk of the computational processing is powered by a processor network, which performs multiple graph neural network layers over the encoded latents. n m To formalise this, assume that our input graph, G = (V; E), has node features xu 2 R , edge features xuv 2 R l n k and graph-level features xG 2 R , for nodes u; v 2 V and edges (u; v) 2 E. Our encoder functions fn : R ! R , m k l k fe : R ! R and fg : R ! R then transform these inputs into the latent space: (0) (0) (0) hu = fn(xu) huv = fe(xuv) hG = fg(xG) (1) Our processor network then transforms these latents over several rounds of message passing: (t+1) (t) H = Pt+1(H ) (2) (t) n (t)o n (t)o (t) where H = hu ; huv ; hG contains all of the latents at a particular processing step t ≥ 0. u2V (u;v)2E The processor network P is iterated for T steps, recovering final latents H(T ). These can then be decoded into node-, edge-, and graph-level predictions (as required), using analogous decoder functions gn, ge and gg: (T ) (T ) (T ) yu = gn(hu ) yuv = ge(huv ) yG = gg(hG ) (3) We will detail the specific design of f, P and g in the following sections. Generally, f and g are simple MLPs, whereas we use highly expressive GNNs for P in order to maximise the advantage of the large-scale datasets. Specifically, we use graph networks (GNs) [Battaglia et al., 2018]. All of our models have been implemented using the jraph library [Godwin et al., 2020]. 4 PCQM4M-LSC setup Input preprocessing For featurising our molecules within PCQM4M, we initially follow the baseline scripts provided by Hu et al. [2021] to convert SMILES strings into molecular graphs.