Year: 2019

Abstract:
Recently various neural networks have been proposed for irregularly structured data such as graphs and manifolds. To our knowledge, all existing graph networks have discrete depth. Inspired by neural ordinary differential equation (NODE) for data in the Euclidean domain, we extend the idea of continuous-depth models to graphs, and propose graph ordinary differential equation (GODE). The derivative of hidden node states are parameterized with a graph neural network, and the output states are the solution to this ordinary differential equation. Noticing that NODEs are typically trained with the adjoint method, with the advantages of adaptive evaluation and a free-form continuous invertible model; however, their performance on benchmark image classification tasks is significantly inferior to discrete-layer models. We show the reason is adjoint method generates inaccurate gradient due to numerical error in reverse-mode integration. We then propose a memory-efficient framework for free-form ODEs with accurate gradient estimation, which is fundamental to deep-learning models. Furthermore, when applied to invertible blocks, our method achieves constant memory cost with GODE (NODE). We show our method for free-form ODEs generalizes to various model structures and achieves high accuracy for both NODE and GODE in benchmark tasks.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORKS

heading: 2.1 NEURAL NETWORKS AND DIFFERENTIAL EQUATIONS
text: 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.

heading: 2.2 GRAPH NEURAL NETWORKS
text: 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.

heading: 2.3 INVERTIBLE BLOCKS
text: 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.

heading: 3 ACCURATE GRADIENT ESTIMATION FOR TRAINING OF NODE

heading: 3.1 FROM DISCRETE MODELS TO CONTINUOUS MODELS
text: 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).

heading: 3.2 BACK-PROP WITH ADJOINT METHOD IS SENSITIVE TO NUMERICAL ERROR
text: 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.

heading: 3.3 MEMORY-EFFICIENT DIRECT BACK-PROPAGATION THROUGH ODE SOLVER
text: The adjoint method might be sensitive to numerical errors when solving the ODE in reverse-time. To resolve this, we propose to directly back-propagate through the ODE solver.
As in Fig. 1(a), the ODE solver uses discretization for numerical integration, evaluated at time points {t0, t1, ...tN}. Fig. 1(c) demonstrates the direct back-propagation with accurate hidden states h(ti), which can be achieved with two methods: (1) the activation z(ti) can be saved in cache for back-prop, but requires huge memory; or (2) we can accurately reconstruct z(ti) by re-building the computation graph directly at evaluated time points {ti}. Since the model is evaluated at the same time points ti in forward-time, it’s guaranteed that z(ti) = h(ti). Therefore direct back-prop is accurate, regardless of the stability of Eq. 2.
Similar to the continuous case, we can define the adjoint with discrete time. Then we have:
ai = ∂L
∂z(ti) , ai+1 = ai
∂z(ti+1)
∂z(ti) ,
dL dθ = N∑ i=1 ai ∂z(ti) ∂θ (7)
where ai is the adjoint for the ith step in discrete forward-time ODE solution. Eq. 7 can be viewed as a numerical discretization of Eq. 6. We show Eq. 6 can be derived from an optimization perspective. Detailed derivations of Eq. 6-7 are in appendix E and F.
Algorithm 1: Algorithm for accurate gradient estimation in ODE solver for free-form functions
Define model dz(t)dt = f(z(t), t), where f is a free-form function. Denote integration time as T .
Forward (f, T, z0, tolerance): t = 0, z = z0 state0 = f.state dict(), cache.save(state0) Select initial step size h = h0 (adaptively with adaptive step-size solver). time points = empty list() While t < T :
state = f.state dict(), accept step = False While Not accept step:
f.load state dict(state) with grad disabled:
z new, error estimate = step(f, z, t, h) If error estimate < tolerance:
accept step = True z = z new, t = t+ h, time points.append(t)
else: reduce stepsize h according to error estimate delete z new, error estimate and related local computation graph
cache.save(time points) return z, cache
Backward (f, T, z0, tolerance, cache): {t0, t1, t2, ...tN−1, tN} = cache.time points For ti in {t0, t1, t2, ...tN−1, tN} :
z(ti+1) = step(f, z(ti), ti, step = ti+1 − ti) For ti in {tN , tN−1, ..., t1, t0} :
Determine a(ti) = ∂L∂z(ti) and ∂z(ti) ∂θ
dL dθ = ∑N i=1 ai ∂z(ti) ∂θ return dLdθ
Details of our method are summarized in Algorithm 1. We discuss its properties below:
Summary of the algorithm (1) During forward pass, the solver performs a numerical integration, with the stepsize adaptively varying with error estimation. (2) During forward pass, the solver outputs the integrated value, and the evaluation time points {ti}. All middle activations are deleted to save memory. (3) During backward pass, the solver re-builds the computation graph, by directly evaluating at saved time points, without adaptive searching. (4) During backward pass, the solver performs a numerical version (Eq. 7) of reverse-time integration (Eq. 6).
Support for free-form continuous dynamics There’s no constraint on the form of f . Therefore, our algorithm is a generic method.
Memory consumption analysis (1) Suppose f has Nf layers, the number of forward evaluation step is Nt on average, and the evaluations to adaptively search for an optimal stepsize is K. A naive solver will take O(Nf × Nt × K), while our method consumes O(Nf × Nt) because all middle
activations are deleted during forward pass, and we don’t need to search for optimal stepsize in backward pass. (2) If we perform step-wise checkpoint method, where we only store z(ti) for all ti, and compute the gradient ∂z(ti+1)∂z(ti) for one ti at a time, then the memory consumption can be reduced toO(Nf+Nt). (3) Since the solver can handle free-form functions, it can also handle restricted form invertible block (see below). In this case, we don’t need to store z(ti), and the memory consumption can reduce to O(Nf ).
More memory-efficient solver with invertible blocks Restricting the form of f to invertible blocks (Gomez et al., 2017) allows for O(Nf ) memory consumption. For invertible blocks, input x is split into two parts (x1, x2) of the same size (e.g. x has shape N × C, where N is batch size, C is channel number; we can split x into x1 and x2 with shape N × C2 ). The forward and inverse of a bijective block can be denoted as: y2 = ψ ( x2, F (x1) ) y1 = ψ ( x1, G(y2) )  x1 = ψ −1 ( y1, G(y2) ) x2 = ψ −1 ( y2, F (x1)
) (8) where the output of a bijective block is denoted as (y1, y2) with the same size as (x1, x2). F and G are any differentiable neural networks, whose output has the same shape as the input. ψ(α, β) is a differentiable bijective function w.r.t α when β is given; ψ−1(α, β) is the inverse function of ψ.
Theorem 1 If ψ(α, β) is a bijective function w.r.t α when β is given, then the block defined by Eq. 8 is a bijective mapping.
Proof of Theorem 1 is given in appendix D. Based on this, we can apply different ψ functions for different tasks. Since x can be accurately reconstructed from y, there’s no need to store activations, hence is memory-efficient. Details for back-prop without storing activation are in appendix B.

heading: 4 NEURAL ORDINARY DIFFERENTIAL EQUATION ON GRAPH NETWORKS
text: We first introduce graph neural networks with discrete layers, then extend to the continuous case and introduce graph ordinary differential equations (GODE).

heading: 4.1 MESSAGE PASSING IN GNN
text: As shown in Fig. 2, a graph is represented with nodes (marked with circles) and edges (solid lines). We assign a unique color to each node for ease of visualization. Current GNNs can generally be represented in a message passing scheme (Fey & Lenssen, 2019):
message(v,u) = φ (k)(xuk−1, x v k−1, eu,v) (9)
aggregationu = ζv∈N (u)(message(v,u)) (10)
xuk = γ (k)(xuk−1, aggregationu) (11)
where xuk represents states of the uth node in the graph at kth layer and eu,v represents the edge between nodes u and v. N (u) represents the set of neighbor nodes for node u. ζ represents a differentiable, permutation invariant operation such as mean, max or sum. γ(k) and φ(k) are differentiable functions parameterized by neural networks.
For a specific node u, a GNN can be viewed as a 3-stage model, corresponding to Eq. 9-11: (1) Message passing, where neighbor nodes v ∈ N (u) send information to node u, denoted by message(v,u). The message is generated from function φ(·), parameterized by a neural network. (2) Message aggregation, where a node u aggregates all messages from its neighborsN (u), denoted as aggregationu. The aggregation function ζ is typically permutation invariant operations such as mean and sum, because graphs are invariant to permutation. (3) Update, where the states of a node are updated according to its original states xuk−1 and aggregation of messages aggregationu, denoted as γ(·).

heading: 4.2 CONTINUOUS-TIME MODELS ON GRAPHS
text: We can convert a discrete-time GNN to continuous-time GNN by replacing f in Eq. 2 with the message passing process defined in Eq. 9 to 11, which we call graph ordinary differential equation (GODE). A diagram of GODE is shown in Fig. 2. Because GODE is an ODE in nature, it can capture highly non-linear functions, thus has the potential to outperform its discrete-layer counterparts.
We demonstrate that the asymptotic stability of GODE could be related to the over-smoothing phenomena (Li et al., 2018). It’s demonstrated that graph convolution is a special case of Laplacian smoothing (Li et al., 2018), which can be written as Y = (I − γD̃−1/2L̃D̃−1/2)X where X and Y are the input and output of a graph-conv layer respectively, Ã = A + I where A is the adjacency matrix, and D̃ is the corresponding degree matrix of Ã, and γ is a positive scaling constant.
When modified from a discrete model to a continuous model, the continuous smoothing process is: dX
dt = −γD̃−1/2L̃D̃−1/2X (12)
Since all eigenvalues of the symmetrically normalized Laplacian are real and non-negative, then all eigenvalues of the above ODE are real and non-positive. Suppose all eigenvalues of the normalized Laplacian are non-zero. In this case, the ODE has only negative eigenvalues, hence the ODE above is asymptotically stable (Lyapunov, 1992). Hence as time t grows sufficiently large, all trajectories are close enough. In the experiments, this suggests if integration time T is large enough, all nodes (from different classes) will have very similar features, thus the classification accuracy will drop.

heading: 5 EXPERIMENTS

heading: 5.1 DATASETS
text: To evaluate our method on general NODE, we conducted experiments with a CNN-NODE on two benchmark image classification tasks (CIFAR10 and CIFAR100) (Krizhevsky et al., 2009).
We also evaluated our method on benchmark graph datasets, including 2 bioinformatic graph classification datasets (MUTAG and PROTEINS), 2 social network graph classification datasets (IMDBBINRAY, REDDIT-BINARY) (Yanardag & Vishwanathan, 2015), and 3 citation networks (Cora, CiteSeer and PubMed). For graph classification tasks, different from the experiment settings in Xu et al. (2018), we input the raw dataset into our models without pre-processing. For node classification tasks, we performed transductive inference and strictly followed the train-validation-test split by Kipf & Welling (2016), where less than 6% nodes are used as training examples. Details of datasets are summarized in appendix A.

heading: 5.2 MODEL STRUCTURES
text: For image classification tasks, we directly modify a ResNet18 into its corresponding NODE model. For each block, the function is dz(t)dt = f(z(t)) where f is a sequence of conv − bn − relu − conv − bn − relu layers. f is the same as residual branch in ResNet, and it can be replaced with any free-form functions.
For tasks on graph datasets, GODE can be applied to any graph neural network by simply replacing f in Eq. 2 with corresponding structures (free-form functions), or replacing F,G in Eq. 8 with other structures (invertible blocks). To demonstrate that GODE is easily generalized to existing structures, we used several different GNN architectures, including the graph convolutional network (GCN) (Kipf & Welling, 2016), graph attention network (GAT) (Veličković et al., 2017), graph network approximated with Chebyshev expansion (ChebNet) (Defferrard et al., 2016), and graph isomorphism network (GIN) (Xu et al., 2018). For a fair comparison, we trained GNNs with different depths of layers (1-3 middle layers, besides an initial layer to transform data into specified channels, and a final layer to generate prediction), and reported the best results among all depths for each model structure.
On the same task, different models use the same hyper-parameters on model structures, such as channel number. For graph classification tasks, we set the channel number of hidden layers as 32 for all models; for ChebNet, we set the number of hops as 16. For node classification tasks, we set the channel number as 16 for GCN and ChebNet, and set number of hops as 3 for ChebNet; for GAT, we used 8 heads, and set each head as 8 channels.
For every GNN structure, we experimented with different number of hidden layers (1,2,3), and calculated the mean and variance of accuracy of 10 runs.

heading: 5.3 COMPARISON OF BACK-PROPAGATION METHODS
text: We compared the adjoint method and direct back-propagation on the same network, and demonstrated direct back-prop generates higher accuracy. For CNN-NODE on classification tasks, we directly modify a ResNet18 into NODE18, and report resuls in Table. 1; for graph networks, we train a GODE model with a GCN to parameterize the derivative, and report results in Table. 2.
Empirical performance Direct back-propagation consistently outperformed the adjoint method for both tasks. This result validates our analysis on the instability of the adjoint method, which is intuitively caused by the instability of the reverse-time ODE. On image classification tasks, compared to adjoint method, our training method reduces error rate of NODE18 from 19% (37%) to 5%(23%) on CIFAR10 (CIFAR100). Furthermore, NODE18 has the same number of parameters as ResNet18, but outperforms deeper networks such as ResNet101 on both datasets. Our method also consistently outperforms the adjoint method on several benchmark graph datasets, as shown in Table. 2.
Robustness to ODE solvers We implemented adaptive ODE solvers of different orders, as shown in Table 1. HeunEuler, RK23, RK45 are of order 1, 2, 4 respectively, i.e., for each step forward in time f is evaluated 1, 2, 4 times respectively. During inference, using different solvers is equivalent to changing model depth (without re-training the network): for discrete-layer models, it generally causes huge error; for continuous models, we observe only around 1% increase in error rate. This suggests our method is robust to different orders of ODE solvers.
Support for free-form functions Our method supports NODE and GODE models with free-form functions; for example, f in NODE18 in Table. 1 is a free-form function.

heading: 5.4 GENERAL BIJECTIVE BLOCKS
text: We demonstrate that bijective blocks defined as Eq. 8 can be easily generalized: F andG are general neural networks, which can be adapted to different tasks; ψ(α, β) can be any differentiable bijective mapping w.r.t. α when β is given. We demonstrate two examples of ψ: (1) additive, forward is η = ψ(α, β) = α + β, inverse is α = ψ−1(η, β) = η − β; (2) linear sigmoid, forward is η = ψ(α, β) = α× sigmoid(β), inverse is α = ψ−1(η, β) = η/sigmoid(β). Results for different ψ are reported in Table 3. Note that we experimented with different depths and reported the best accuracy for each model, and performed a paired t-test on results from GODE and their discrete-layer counterparts. Most GODE models outperformed their corresponding discretelayer models significantly, validating the effectiveness of GODE; different ψ functions behaved similarly on our node classification tasks, indicating the continuous-time model is more important than coupling function ψ. We also validate the lower memory cost, with details in appendix B.

heading: 5.5 RESULTS ON GRAPH CLASSIFICATION TASK
text: Results for different models on graph classification tasks are summarized in Table 4. We experimented with different structures, including GCN, ChebNet and GIN; for corresponding GODE models (marked with ODE), we tested both free-form (marked with “free”) and invertible block (marked with “INV”). We performed paired t-test comparing GODE and its discrete-layer counterparts. For most experiments, GODE models performed significantly better. This indicates the continuous process model might be important for graph models.

heading: 5.6 IMPACT OF INTEGRATION TIME
text: For a NODE and GODE model, during inference, we test the influence of integration time. Results are summarized in Table. 5. When integration time is short, the network does not gather sufficient information from neighbors; when integration time is too long, the model is sensitive to over-smooth issue, as discussed in Sec. 4.2. We observe accuracy drop in both cases.

heading: 6 CONCLUSIONS
text: We propose GODE, which enables us to model continuous diffusion process on graphs. We propose a memory-efficient direct back-propagation method to accurately determine the gradient for general free-form NODEs, and validate its superior performance on both image classification tasks and graph data. Furthermore, we related the over-smoothing of GNN to asymptotic stability of ODE. Our paper tackles the fundamental problem of gradient estimation for NODE; to our knowledge, it’s the first paper to improve accuracy on benchmark tasks to comparable with state-of-the-art discrete layer models. It’s an important step to apply NODE from theory to practice.

heading: A DATASETS
text: We perform experiments on various datasets, including citation networks (Cora, CiteSeer, PubMed), social networks (COLLAB, IMDB-BINARY, REDDIT-BINARY), and bioinformatics datasets (MUTAG, PROTEINS). Details of each dataset are summarized in Table 1.
Table 1: Statistics of datasets Dataset Graphs Nodes Edges Features Classes Label rate Cora 1 2,708 5,278 1,433 7 0.052 CiteSeer 1 3,327 4,552 3,703 6 0.036 PubMEd 1 19,717 44,324 500 3 0.003 MUTAG 188 17.93 19.79 7 2 0.8 PROTEINS 1,113 39.06 72.82 3 2 0.8 IMDB-BINARY 1,000 19.77 96.53 - 2 0.8 REDDIT-BINARY 200 429.63 497.76 - 2 0.8

heading: B DETAILS ABOUT INVERTIBLE BLOCKS
text: We explain the structure and conduct experiments for the invertible block here.
Structure of invertible blocks Structure of invertible blocks are shown in Fig. 1. We follow the work of Gomez et al. (2017) with two important modifications: (1) We generalize to a family of bijective blocks with different ψ in Eq. 8 in the main paper, while Gomez et al. (2017) restrict the form of ψ to be sum. (2) We propose a parameter state checkpoint method, which enables bijective blocks to be called more than once, while still generating accurate inversion.
The algorithm is summarized in Algo. 2. We write the pseudo code for forward and backward function as in PyTorch. Note that we use “inversion” to represent reconstructing input from the output, and use “backward” to denote calculation of the gradient. To reduce memory consumption, in the forward function, we only keep the outputs y1, y2 and delete all other variables and computation graphs. In the backward function, we first “inverse” the block to calculate x1, x2 from y1, y2, then perform a local forward and calculate the gradient ∂[y1,y2]∂[x1,x2] .
Experiments In this section we demonstrate that our bijective block is memory efficient. We trained a GODE model with bijective blocks, and compared the memory consumption using our memory-efficient function as in Algo. 2 and a memory-inefficient method as in conventional backpropagation. Results were measured with a batchsize of 100 on MUTAG dataset.
Results are summarized in Table. 2. We measured the memory consumption with different depths, which is the number of ODE blocks. When depth increases from 10 to 20, the memory by conventional methods increases from 5.3G to 10.5G, while our memory-efficient version only increases from 2.2G to 2.6G. In theory, our bijective block takesO(1) memory, because we only need to store the outputs in cache, while deleting activations of middle layers. For memory-efficient network, the slightly increased memory consumption is because states of F,G need to be cached; but this step takes up minimal memory compared to input data.
Algorithm 2: Function for memory-efficient bijective blocks
Forward (cache, x1, x2, F , G, ψ) cache.save([F states, G states]) forward in Eq. 8
η1 = F (x1), y2 = ψ(x2, η1) η2 = G(y2), y1 = ψ(x1, η2)
delete η1, η2, x1, x2 delete computation graphs generated by F and G return cache, y1, y2
Backward(cache, y1, y2, F , G, ψ, ∂L∂y1 , ∂L ∂y2 ) Reset F and G states from cache Inverse from y1, y2 to x1, x2
η2 = G(y2), x1 = ψ −1(y1, η2) η1 = F (x1), x2 = ψ −1(y2, η1)
Local forward pass and gradient X1, X2 = x1.detach(), x2.detach() calculate Y1, Y2 from X1, X2 as Eq. 8 determine ∂[Y1, Y2]/∂[X1, X2, θF , θG]
∂L ∂[x1,x2] = ∂L ∂[y1,y2] ∂[Y1,Y2] ∂[X1,X2]
∂L ∂[θF ,θG] = ∂L ∂[y1,y2] ∂[Y1,Y2] ∂[θF ,θG]
delete Y1, Y2, X1, X2 return ∂L/∂[x1, x2], ∂L/∂[θF , θG]

heading: C PROOF FOR PROPOSITION 1
text: 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 λ.
Proof Denote s = T − t, where T is the end time. Notice that the reverse-time in t is equivalent to forward-time in s. Therefore, we have forward-time ODE:
dz(t)
dt = f(z(t), t) (1)
and reverse-time ODE: dz(s)
ds = −f(z(s), s) = g(z(s), s) (2)
Therefore, we have λ(Jf ) = −λ(Jg). For both forward-time and reverse-time ODE to be stable, the eigenvalues of J need to have non-positive real part. Therefore
Reλi(Jf ) ≤ 0, Reλi(Jg) = −Reλi(Jf ) ≤ 0, ∀i (3) The only solution is
Reλi(Jg) = −Reλi(Jf ) = 0, ∀i (4)

heading: D PROOF FOR THEOREM 1
text: Theorem 1 For bijective block whose forward and reverse mappings are defined as
Forward(x1, x2) =  y2 = ψ ( x2, F (x1) ) y1 = ψ ( x1, G(y2) ) Reverse(y1, y2) =  x1 = ψ −1 ( y1, G(y2) ) x2 = ψ −1 ( y2, F (x1)
) If ψ(α, β) is a bijective function w.r.t α when β is given, then the block is a bijective mapping.
Proof To prove the forward mapping is bijective, it is equivalent to prove the mapping is both injective and surjective.
Injective We need to prove, if Forward(x1, x2) = Forward(x3, x4), then x1 = x3, x2 = x4.
The assumption above is equivalent to
Forward(x1, x2) = Forward(x3, x4) ⇐⇒ y2 = ψ(x2, F (x1)) = ψ(x4, F (x3)) (5) ψ(x1, G(y2)) = ψ(x3, G(y2)) (6)
Since ψ(α, β) is bijective w.r.t α when β is given, from Eq.(6), we have x1 = x3. Similarly, condition on x1 = x3 and Eq.(5), using bijective property of ψ, we have x2 = x4. Therefore, the mapping is injective.
Surjective We need to prove ∀ [y1, y2], ∃ [x1, x2] s.t. Forward(x1, x2) = [y1, y2].
Given y1, y2, we construct x1 = ψ −1 ( y1, G(y2) ) , x2 = ψ −1 ( y2, F (x1) ) (7)
Then for the forward function, given bijective property of ψ, apply Forward and Reverse defined in the proposition statement,
z2 = ψ(x2, F (x1)) = ψ ( ψ−1 ( y2, F (x1) ) , F (x1) ) = y2 (8)
z1 = ψ(x1, G(y2)) = ψ ( ψ−1 ( y1, G(y2) ) , G(y2) ) = y1 (9)
Therefore we construct x1, x2 s.t. Forward(x1, x2) = [y1, y2]. Therefore the mapping is surjective. Therefore is bijective.

heading: E DERIVATION OF GRADIENT IN DISCRETE CASE
text: We use a figure to demonstrate the computation graph, and derive the gradient from the computation graph.
The loss is L, forward pass is denoted with black arrows, gradient back-propagation is shown with red arrows. We use p to denote each path from θ to L, corresponding to all paths in red that goes from L to θ.
dL dθ = ∑ p ∂Lp ∂θ = ∑ i ∂L ∂z(ti) ∂z(ti) ∂θ = ∑ i ai ∂z(ti) ∂θ (10)
ai = ∂L
∂z(ti) =
∂L
∂z(ti+1)
∂z(ti+1)
∂z(ti) = ai+1
∂z(ti+1)
∂z(ti) (11)

heading: F DERIVATION OF PARAMETER GRADIENTS IN CONTINUOUS CASE
text: In this section we derive the gradient of parameters in an neural-ODE model from an optimization perspective. Then we extend from continuous cases to discrete cases.
Notations With the same notations as in the main paper, we use z(t) to denote hidden states z at time t. Denote parameters as θ, and input as x, target as y, and predicted output as ŷ. Denote the loss as J(ŷ, y). Denote the integration time as 0 to T .
Problem setup The continuous model is defined to follow an ODE: dz(t)
dt = f(z(t), t, θ), s.t. z(0) = x (12)
We assume f is differentiable, since f is represented by a neural network in our case. The forward pass is defined as:
ŷ = z(T ) = z(0) + ∫ T 0 f(z(t), t, θ)dt (13)
The loss function is defined as: J(ŷ, y) = J(z(T ), y) (14)
We formulate the training process as an optimization problem:
argminθ 1
N N∑ i=1 J(ŷi, yi) s.t. dz(t) dt = f(z(t), t, θ), zi(0) = xi (15)
For simplicity, Eq. 15 only considers one ODE block. In the case of multiple blocks, z(T ) is the input to the next ODE block. As long as we can derive dLossdθ and dLoss dz(0) when dLoss dz(T ) is given, the same analysis here can be applied to the case with a chain of ODE blocks.
Lagrangian Multiplier Method We use the Lagrangian Multiplier Method to solve the problem defined in Eq. 15. For simplicity, only consider one example (can be easily extended to multiple examples cases), the Lagrangian is
L = J(z(T ), y) + ∫ T 0 λ(t)[ dz(t) dt − f(z(t), t, θ)]dt (16)
Karush-Kuhn-Tucker (KKT) conditions are necessary conditions for an solution to be optimal. In the following sections we start from the KKT condition and derive our results.
Derivative w.r.t. λ At optimal point, we have δLδλ = 0. Note that λ is a function of t, we derive the derivative from calculus of variation.
Consider a cotninuous and differentiable perturbation λ(t) on λ(t), and a scalar , L now becomes a function of ,
L( ) = J(z(0) + ∫ T 0 f(z(t), t, θ), y) + ∫ T 0 (λ(t) + λ(t))[ dz(t) dt − f(z(t), t, θ)]dt (17)
It’s easy to check the conditions for Leibniz integral rule, and we can switch integral and differentiation, thus:
dL d = ∫ T 0 λ(t)[ dz(t) dt − f(z(t), t, θ)]dt (18)
At optimal λ(t), dLd | =0 = 0 for all continuous differentiable λ(t). Therefore,
dz(t)
dt − f(z(t), t, θ) = 0, ∀t ∈ (0, T ) (19)
Derivative w.r.t z Consider perturbation z(t) on z(t), with scale . With similar analysis:
L( ) = J(z(T ) + z(T ), y) + ∫ T 0 λ(t)[ dz(t) + z(t) dt − f(z(t) + z(t), t, θ)]dt (20)
Take derivative w.r.t , it’s easy to check conditions for Leibniz integral rule are satisfied, when f and z(t) are Lipschitz continuous differentiable functions:
(1) f(z(t), t, θ) is a Lebesgue-integrable function of θ for each z(t) ∈ Rd, since we use a neural network to represent f , which is continuous and differentiable almost everywhere.
(2) for almost all θ, ∂f(z(t),t,θ)∂z(t) exists for almost all x ∈ R d.
(3) ∂f(z(t),t,θ)∂z(t) is bounded by g(θ) for all z(t) for almost all θ.
Then we calculate dL( )d , note that we can switch integral and derivative:
dL d | =0 = ∂J ∂z(T ) z(T ) + d d ∫ T 0 λ(t)[ dz(t) + z(t) dt − f(z(t) + z(t), t, θ)]dt (21)
= ∂J
∂z(T ) z(T ) + ∫ T 0 λ(t)[ dz(t) dt − ∂f(z(t), t, θ) ∂z(t) z(t)]dt (22)
= ∂J
∂z(T ) z(T ) + ∫ T 0 [λ(t) dz(t) dt + dλ(t) dt z(t)− dλ(t) dt z(t)− λ(t)∂f(z(t), t, θ) ∂z(t) z(t)]dt
(23)
= ∂J
∂z(T ) z(T ) + λ(t)z(t)|T0 − ∫ T 0 z(t)[ dλ(t) dt + λ(t) ∂f(z(t), t, θ) ∂z(t) ]dt (24)
= ∂J
∂z(T ) z(T ) + λ(T )z(T )− λ(0)z(0)− ∫ T 0 z(t)[ dλ(t) dt + λ(t) ∂f(z(t), t, θ) ∂z(t) ]dt (25)
= ( ∂J
∂z(T ) + λ(T ))z(T )− λ(0)z(0)− ∫ T 0 z(t)[ dλ(t) dt + λ(t) ∂f(z(t), t, θ) ∂z(t) ]dt (26)
Since the initial condition z(0) = x is given, perturbation z(0) at t = 0 is 0, then we have:
dL d | =0 = ( ∂J ∂z(T ) + λ(T ))z(T )− ∫ T 0 z(t)[ dλ(t) dt + λ(t) ∂f(z(t), t, θ) ∂z(t) ]dt = 0 (27)
for any z(t) s.t. z(0) = 0 and z(t) is differentiable.
The solution is: ∂J
∂z(T ) + λ(T ) = 0 (28)
dλ(t)
dt + λ(t)
∂f(z(t), t, θ)
∂z(t) = 0 ∀t ∈ (0, T ) (29)
Derivative w.r.t θ From Eq. 16,
dL dθ = − ∫ T 0 λ(t) ∂f(z(t), t, θ) ∂θ dt (30)
To sum up, first solve the ODE forward-in-time with Eq. 19, then determine the boundary condition by Eq. 28, then solve the ODE backward with Eq. 29, and finally calculate the gradient with Eq. 30. In fact λ corresponds to the negative adjoint.
From continuous to discrete case To derive corresponding results in discrete cases, we need to replace all integration with finite sum.
In discrete cases, the ODE condition turns into:
zi+1 − zi ti+1 − ti = f(zi, ti, θ) (31)
from Eq. 31, we can get:
∂L ∂zi = ∂L ∂zi+1 ∂zi+1 ∂zi = ∂L ∂zi+1 (I + ∂f(zi, ti, θ) ∂zi (ti+1 − ti)) (32)
Re-arranging terms we have:
[(− ∂L ∂zi+1 )− (− ∂L ∂zi )]/[ti+1 − ti] + (− ∂L ∂zi+1 ) ∂f(zi, ti, θ) ∂zi = 0 (33)
which is the discrete version of Eq. 29. Which also corresponds to our analysis in Eq. 10 and 11.

Review summaries:
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.