Ordinary differential equations on graph networks

1 INTRODUCTION . Convolutional neural networks ( CNN ) have achieved great success in various tasks , such as image classification ( He et al. , 2016 ) and segmentation ( Long et al. , 2015 ) , video processing ( Deng et al. , 2014 ) and machine translation ( Sutskever et al. , 2014 ) . However , CNNs are limited to data that can be represented by a grid in the Euclidean domain , such as images ( 2D grid ) and text ( 1D grid ) , which hinders their application in irregularly structured datasets . A graph data structure represents objects as nodes and relations between objects as edges . Graphs are widely used to model irregularly structured data , such as social networks ( Kipf & Welling , 2016 ) , protein interaction networks ( Fout et al. , 2017 ) , citation and knowledge graphs ( Hamaguchi et al. , 2017 ) . Early works use traditional methods such as random walk ( Lovász et al. , 1993 ) , independent component analysis ( ICA ) ( Hyvärinen & Oja , 2000 ) and graph embedding ( Yan et al. , 2006 ) to model graphs , however their performance is inferior due to the low expressive capacity . Recently a new class of models called graph neural networks ( GNN ) ( Scarselli et al. , 2008 ) were proposed . Inspired by the success of CNNs , researchers generalize convolution operations to graphs to capture the local information . There are mainly two types of methods to perform convolution on a graph : spectral methods and non-spectral methods . Spectral methods typically first compute the graph Laplacian , then perform filtering in the spectral domain ( Bruna et al. , 2013 ) . Other methods aim to approximate the filters without computing the graph Laplacian for faster speed ( Defferrard et al. , 2016 ) . For non-spectral methods , the convolution operation is directly performed in the graph domain , aggregating information only from the neighbors of a node ( Duvenaud et al. , 2015 ; Atwood & Towsley , 2016 ) . The recently proposed GraphSAGE ( Hamilton et al. , 2017 ) learns a convolution kernel in an inductive manner . To our knowledge , all existing GNN models mentioned above have a structure of discrete layers . The discrete structure makes it hard for the GNN to model continuous diffusion processes ( Freidlin & Wentzell , 1993 ; Kondor & Lafferty , 2002 ) in graphs . The recently proposed neural ordinary differential equation ( NODE ) ( Chen et al. , 2018 ) views a neural network as an ordinary differential equation ( ODE ) , whose derivative is parameterized by the network , and the output is the solution to this ODE . We extend NODE from the Euclidean domain to graphs and propose graph ordinary differential equations ( GODE ) , where the message propagation on a graph is modeled as an ODE . NODEs are typically trained with adjoint method . NODEs have the advantages of adaptive evaluation , accuracy-speed control by changing error tolerance , and are free-form continuous invertible models ( Chen et al. , 2018 ; Grathwohl et al. , 2018 ) . However , to our knowledge , in benchmark image classification tasks , NODEs are significantly inferior to state-of-the-art discrete-layer models ( error rate : 19 % for NODE vs 7 % for ResNet18 on CIFAR10 ) ( Dupont et al. , 2019 ; Gholami et al. , 2019 ) . In this work , we show this is caused by error in gradient estimation during training of NODE , and propose a memory-efficient framework for accurate gradient estimation . We demonstrate our framework for free-form ODEs generalizes to various model structures , and achieves high accuracy for both NODE and GODE in benchmark tasks . Our contribution can be summarized as follows : 1 . We propose a framework for free-form NODEs to accurately estimate the gradient , which is fundamental to deep-learning models . Our method significantly improves the performance on benchmark classification ( reduces test error from 19 % to 5 % on CIFAR10 ) . 2 . Our framework is memory-efficient for free-form ODEs . When applied to restricted-form invertible blocks , the model achieves constant memory usage . 3 . We generalize ODE to graph data and propose GODE models . 4 . We demonstrate improved performance on different graph models and various datasets . 2 RELATED WORKS . 2.1 NEURAL NETWORKS AND DIFFERENTIAL EQUATIONS . There have been efforts to view neural networks as differential equations . Lu ( 2017 ) viewed a residual network as a discretization of a differential equation and proposed several new architectures based on numerical methods in ODE solver . Haber & Ruthotto ( 2017 ) proposed a stable architecture based on analysis of the ODE . Chen et al . ( 2018 ) proposed neural ordinary differential equation ( NODE ) , which treats the neural network as a continuous ODE . NODE was later used in a continuous normalizing flow for generative models ( Grathwohl et al. , 2018 ) . There have been many studies on the training of NODE . The adjoint method has long been widely used in optimal control ( Stapor et al. , 2018 ) and geophysical problems ( Plessix , 2006 ) , and recently applied to ODE ( Chen et al. , 2018 ) . Dupont et al . ( 2019 ) proposed augmented neural ODEs to improve the expressive capacity of NODEs . However , to our knowledge , none of the methods above discusses the inaccurate gradient estimation issue ; empirical performances of NODE in benchmark classification tasks are significantly inferior to state-of-the-art discrete-layer models . 2.2 GRAPH NEURAL NETWORKS . GNNs can be divided into two categories : spectral methods and non-spectral methods . Spectral GNNs perform filtering in the Fourier domain of a graph , thus need information of the whole graph to determine the graph Laplacian . In contrast , non-spectral GNNs only consider message aggregation around neighbor nodes , therefore are localized and generally require less computation ( Zhou et al. , 2018 ) . We first briefly introduce several spectral methods . Bruna et al . ( 2013 ) first introduced graph convolution in the Fourier domain based on the graph Laplacian , however the computation burden is heavy because of non-localized filters . Henaff et al . ( 2015 ) incorporated a graph estimation procedure in spectral networks and parameterized spectral filters into a localized version with smooth coefficients . Defferrard et al . ( 2016 ) used Chebyshev expansion to approximate the filters without the need to compute the graph Laplacian and its eigenvectors , therefore significantly accelerated the running speed . Kipf & Welling ( 2016 ) proposed to use a localized first-order approximation of graph convolution on graph data and achieved superior performance in semi-supervised tasks for node classification . Defferrard et al . ( 2016 ) proposed fast localized spectral filtering on graphs . Non-spectral methods typically define convolution operations on a graph , only considering neighbors of a certain node . MoNet ( Monti , 2017 ) uses a mixture of CNNs to generalize convolution to graphs . GraphSAGE ( Hamilton et al. , 2017 ) samples a fixed size of neighbors for each node for fast localized inference . Graph attention networks ( Veličković et al. , 2017 ) learn different weights for different neighbors of a node . The graph isomorphism network ( GIN ) ( Xu et al. , 2018 ) has a structure as expressive as the Weisfeiler-Lehman graph isomorphism test . 2.3 INVERTIBLE BLOCKS . Invertible blocks are a family of neural network blocks whose forward function is a bijective mapping . Therefore , the input to a bijective block can be accurately reconstructed from its outputs . Invertible blocks have been used in normalizing flow ( Rezende & Mohamed , 2015 ; Dinh , 2016 ; Kingma & Dhariwal , 2018 ; Dinh et al. , 2014 ; Kingma et al. , 2016 ) , where the model is required to be invertible in order to calculate the log-density of data distribution . Later on , Jacobsen et al . ( 2018 ) used bijective blocks to build invertible networks . Gomez et al . ( 2017 ) proposed to use invertible blocks to perform back propagation without storing activation , which achieves a memory-efficient network structure . They were able to discard activation of middle layers , because each layer ’ s activation can be reconstructed from the next layer with invertible blocks . 3 ACCURATE GRADIENT ESTIMATION FOR TRAINING OF NODE . 3.1 FROM DISCRETE MODELS TO CONTINUOUS MODELS . We first consider discrete-layer models with residual connection ( He et al. , 2016 ) , which can be represented as : xk+1 = xk + fk ( xk ) ( 1 ) where xk is the states in the kth layer ; fk ( · ) is any differentiable function whose output has the same shape as its input . When we add more layers with shared weights , and let the stepsize in Eq . 1 go to infinitesimal , the difference equation turns into a neural ordinary differential equation ( NODE ) ( Chen et al. , 2018 ) : dz ( t ) dt = f ( z ( t ) , t ) ( 2 ) We use z ( t ) in the continuous case and xk in the discrete case to represent hidden states . f ( · ) is the derivative parameterized by a network . Note that a key difference between Eq . 1 and 2 is the form of f : in the discrete case , different layers ( different k values ) have their own function fk ; while in the continuous case , f is shared across all time t. The forward pass of model with discrete layers can be written as : x0 = input , x1 = x0 + f0 ( x0 ) , ... , xK = xK−1 + fK−1 ( xK−1 ) ( 3 ) where K is the total number of layers . Then an output layer ( e.g . fully-connected layer for classification ) is applied on xK . The forward pass of a NODE is : z ( T ) = z ( 0 ) + ∫ T t=0 dz ( t ) dt dt = input + ∫ T t=0 f ( z ( t ) , t ) dt ( 4 ) where z ( 0 ) = input and T is the integration time , corresponding to number of layers K in the discrete case . The transformation of states z is modeled as the solution to the NODE . Then an output layer is applied on z ( T ) . Integration in the forward pass can be performed with any ODE solver , such as the Euler Method , Runge-Kutta Method , VODE solver and Dopris Solver ( Milne & Milne , 1953 ; Brown et al. , 1989 ; Ascher et al. , 1997 ) . 3.2 BACK-PROP WITH ADJOINT METHOD IS SENSITIVE TO NUMERICAL ERROR . The adjoint method is widely used in optimal process control and functional analysis ( Stapor et al. , 2018 ; Pontryagin , 2018 ) . We follow the method by ( Chen et al. , 2018 ) . Denote model parameters as θ , which is independent of time . Define the adjoint as : a ( t ) = ∂L ∂z ( t ) ( 5 ) where L is the loss function . Then we have da ( t ) dt = −a ( t ) T ∂f ( z ( t ) , t , θ ) z ( t ) , dL dθ = − ∫ 0 T a ( t ) T ∂f ( z ( t ) , t , θ ) ∂θ dt ( 6 ) with detailed proof from optimization perspective in appendix F. Then we can perform gradient descent to optimize θ to minimize L. Eq . 6 is a reverse-time integration , which can be solved with any ODE solver ( Chen et al. , 2018 ) . To evaluate ∂f ( z ( t ) , t , θ ) ∂θ , we need to determine z ( t ) by solving Eq . 2 reverse-time ( Directly storing z ( t ) during forward pass requires a large memory consumption , because the continuous model is equivalent to an infinite-layer model ) . To summarize , in the forward pass we solve Eq . 2 forward in time ; in the backward pass , we solve Eq . 2 and 6 reverse in time , with initial condition determined from Eq . 5 at time T . We give an intuition why the reverse-time ODE solver causes inaccurate gradient in adjoint methods . The backward pass ( Eq . 6 ) requires determining f ( z ( t ) , t , θ ) and ∂f ( z ( t ) , t , θ ) ∂θ , which requires determining z ( t ) by solving Eq . 2 reverse-time . As shown in Fig . 1 ( a , b ) , the hidden state solved forward-time ( z ( ti ) ) and the hidden state solved reverse-time ( h ( ti ) ) may not be equal ; this could be caused by the instability of reverse-time ODE , and is represented by the mismatch between z ( t ) ( dashed curve ) and h ( t ) ( solid curve ) . Error h ( t ) − z ( t ) will cause error in gradient dLdθ . Proposition 1 For an ODE in the form dz ( t ) dt = f ( z ( t ) , t ) , denote the Jacobian of f as Jf . If this ODE is stable both in forward-time and reverse-time , then Re ( λi ( Jf ) ) = 0 ∀i , where λi ( Jf ) is the ith eigenvalue of Jf , and Re ( λ ) is the real part of λ . Detailed proof is in appendix C. Proposition 1 indicates that if the Jacobian of the original system Eq . 2 has eigenvalues whose real-part are not 0 , then either the reverse-time or forward-time ODE is unstable . When |Re ( λ ) | is large , either forward-time or reverse-time ODE is sensitive to numerical errors . This phenomenon is also addressed in Chang et al . ( 2018 ) . This instability affects the accuracy of solution to Eq . 2 and 6 , thus affects the accuracy of the computed gradient .

Summary: The authors discussed that most graph NNs to date considered discrete layers and hence are difficult to model diffusion processes on graphs. This paper proposed the Neural ODE on a graph, termed with a graph ODE, to tackle this problem. The authors gave a sufficient condition under which the adjoint method gets instable and pointed out potential issues of training Neural ODEs using it. To overcome the instability issue, the author proposed to backpropagate errors directly at discretized points. Since the naive implementation of the direct method is memory-consuming, the authors used invertible blocks as building blocks of graph ODEs, which do not store the intermediate activations for backward propagation. Finally, the authors conducted empirical studies to see the effectiveness of the proposed method.