A N E NERGY -B ASED V IEW OF G RAPH N EURAL N ETWORKS
A BSTRACT
Graph neural networks are a popular variant of neural networks that work with graph-structured data.
In this work, we consider combining graph neural networks with the energy-based view of Grathwohl et al.
with the aim of obtaining a more robust classiﬁer.
We successfully implement this framework by proposing a novel method to ensure generation over features as well as the adjacency matrix and evaluate our method against the standard graph convolutional network (GCN) architecture (Kipf & Welling (2016)).
Our approach obtains comparable discriminative performance while improving robustness, opening promising new directions for future research for energy-based graph neural networks.
I NTRODUCTION
Graph neural networks (GNNs) are a generalization of neural networks that operate on graphstructured data, typically in the form of an adjacency matrix or graph laplacian, and feature vectors deﬁned for the nodes.
They have found success in tasks such as link prediction, node classiﬁcation, and graph classiﬁcation (e.g., Izadi et al.
(2020), Wang et al.
(2019b), Zhang et al.
Alongside the empirical success, recent theoretical work has elucidated properties on their depth (Oono & Suzuki (2019)), architectural alignment with algorithms (Xu et al.
(2019)), and their discriminative power (Xu et al.
Recently, the work of Grathwohl et al.
(2019) proposed viewing traditional classiﬁers as energy-based models, adjusting the softmax transfer function to the Boltzmann distribution, and adding an inner stochastic Langevin gradient descent (Welling & Teh (2011)) loop over new samples.
We extend this framework for GNNs, and introduce a novel method to ensure generation over the adjacency matrix itself along with node features.
We evaluate our approach for the node classiﬁcation task on the Cora, Pubmed, and Citeseer datasets and explore its generative capabilities.
B ACKGROUND AND P RIOR W ORK
One of the most popular GNN architectures is the GCN architecture of Kipf & Welling (2016), which utilizes a normalized adjacency matrix with self-connections.
Consider the adjacency matrix of a graph, A ∈ R n × n .
We deﬁne ˆ A
=
A + I n , where I n is the n × n identity matrix.
Let ˆ D ii ≡ (cid:80)
j ˆ A ij .
Layer-wise propagation is given by:
[Equation]
Equation 1 follows from Kipf & Welling (2016), and multiplication by ˆ D − 1 / 2 ensures symmetric normalization of the graph adjacency matrix ˆ A .
H 0 = X ∈ R n × f is the data matrix with row-wise examples, where f is the number of features.
σ ( · ) is a non-linear activation function, and W l is a linear transformation which typically reduces the dimensionality of the data.
The output of the ﬁnal layer is k -dimensional, where k is the number of classes. This is passed through a softmax function, which is then optimized through the cross-entropy loss.
In the work of Grathwohl et al.
(2019), it was noted that the softmax function can be changed in the following way: with typical machine learning classiﬁcation, we are interested in a function f θ ( x ) , which maps the input data x ∈ R d
→ R k to k real valued numbers, where k is the number of classes, and θ are trainable parameters.
When this is passed through a softmax function, this can be viewed as a conditional distribution of y conditioned on x .
In viewing this as a conditional distribution, we can also see that the joint distribution is given by normalization by the partition function.
[Equation]
By marginalization of y , we have the probability density function of x :
[Equation]
By viewing this as a Boltzmann distribution, the energy function is given as:
[Equation]
Combining energy-based methods and graph neural networks has been explored in the works of Cao & Shen (2020) and Ma et al.
In Ma et al.
(2019), the energy is given as the sum over the node embeddings produced by a graph neural network normalized by the number of nodes, which is input into a multilayer perceptron and passed through a ReLU.
In Cao & Shen (2020), an analogy to the coulomb potential is used in designing a feature pooling mechanism.
Our work focuses on extending the framework of Grathwohl et al.
for GNNs and the exact approach is detailed in the following section.
M ETHODOLOGY
In Algorithm 1, we outline the GCN-JEM algorithm. In essence, the algorithm is a stochastic Langevin gradient descent (SLGD) loop nested inside of an stochastic gradient descent (SGD) loop.
The inner SLGD loop samples the generative features over some simple distribution, and minimizes the energy of the new features using SGLD. The outer SGD loop computes the cross-entropy of the classiﬁer, and combines the loss for the classiﬁer as well as the generative loss, and computes the gradient. Finally, within the SGD loop, we generate new links in the adjacency matrix for nodes that are close in energy . The key factor for convergence when using stochastic Langevin gradient descent training with graph structured data is the renormalization by the total graph energy.
The intuitive justiﬁcation is that the gradient correction for the features of a node takes into account the total energy of the graph, since the messages are passed over the entire graph in the forward pass. The generative loss, L gen , ensures that the energy of the original classiﬁer over the selected indices is close to the energy of the generative features. This encourages their probabilities to be close to each other. In addition to the generation over features, we have incorporated generation over the adjacency matrix itself.
The motivation for this stems from the fact that for traditional machine learning models, samples are thought to be independent but that need not be the case for graph structured data as nodes are connected by edges. We use the energy of a sample as a quantitative measure to incorporate edges between the sampled data and existing graph. We compare the energy of two nodes, and add a link to the graph if the difference between the energy of two nodes is within the threshold τ . For stability, we do this once every 50 epochs.
In the complex networks literature (e.g., Krioukov et al. (2010)), which has connections to the statistical mechanics literature, the connection probability of two nodes can be given by the Fermi-Dirac distribution. Under this assumption, two nodes that are close in energy have a high probability of being connected.  E XPERIMENTAL RESULTS To evaluate our approach for training an energy-based GNN, three citation network datasets are considered – Cora, Pubmed and Citeseer.
In Table 1, the reported accuracies are the average performance on the test set for 5 runs with random initialization, with a train-test split following that of θ in matrix with samples,
X . while not converged do
L clf ( θ ) =
xent
( f θ ( A , X ) , y ) Batch sample x i and y i from dataset X , with number of samples ζ . Let this set be S . for all i ∈ S do
Sample ˆ x i 0 ∼ B with probability 1 − ρ , else ˆ x i 0 ∼ U ( − 1 , 1) .
Replace x i with ˆ x i 0 in X , creating ˆ X 0 .
Also replace x i with ˆ x i 0 in S .
for t ∈ [0 , . . .
, η ] do
Let Z t =
LogSumExp y (cid:48)
f θ ( A , ˆ X t )
[ y (cid:48) ] (energy of the graph with the new features) ˆ x {+it +
+} =
ˆ x it + α ∂ LogSumExp y (cid:48)
( f θ ( A , ˆ x it )
[ y (cid:48) ]) /Z
t
∂ ˆ x it + σ N (0 , I ) end end
Let
Z gen =
LogSumExp
y (cid:48)
f θ ( A , ˆ X )[ y (cid:48) ] (energy of the graph with the new features)
Let Z clf =
LogSumExp y (cid:48)
f θ ( A , X ) [ y (cid:48) ] (energy of the graph with the original features)
L gen ( θ ) = | (cid:80) i ∈ S
LogSumExp y (cid:48)
( f θ ( A , x i ) [ y (cid:48) ]) /Z
clf − LogSumExp y (cid:48)
( f θ ( A , ˆ x it ) [ y (cid:48) ]) /Z
gen |
L ( θ ) =
L clf ( θ ) + L gen ( θ )
Compute gradients ∂L
( θ ) ∂θ and propagate.
Add S to B .
for All the new features , ˆ x it in S do if | LogSumExp y (cid:48)
( f θ ( A , ˆ x it )
[ y (cid:48) ]) − LogSumExp y (cid:48)
( f θ ( A , ˆ x jt )
[ y (cid:48) ]) | ≤ τ then
Add edge between nodes i and j in adjacency matrix, ˜ A . end end
Set A =
˜
A end
Kipf & Welling (2016).
The deep graph library (DGL) framework (Wang et al. (2019a)) is utilized for implementing the GCN upon which the energy based framework was implemented.
The hyperparameters used in the experiments are given in Table 2.
GCN-JEMO corresponds to the model with an additional term in the loss function which encourages the weight matrix W l for each layer l to be orthogonal.
We observe better performance with this additional constraint that promotes weight orthogonality:
[Equation]
(5) where (cid:107) .
(cid:107) F denotes the standard Frobenius norm ( (cid:107) A (cid:107) 2 F = (cid:80) ij A 2 ij ) . [Table] GCN-JEM (ours)
[Table]
GCN-JEMO (ours) [Table] [Table] It should be noted that some papers for the Cora and PubMed datasets appear to use different training/test splits than the one used in the original GCN paper.
To be consistent with the literature, we use the original train/test splits in the GCN paper (Kipf & Welling (2016)).
Our technique is in a sense architecture agnostic and can be used with any graph neural network with a similar message-passing scheme.
In our study, we use the vanilla GCN (Kipf & Welling (2016)) as a baseline comparison.
In theory, we could adapt our technique with any of the leading architectures. [Table]
D ISCUSSION
Our experimental results do not consistently improve the accuracy across the three datasets (an improvement for CORA is obtained), but it should be noted that the main focus of Grathwohl et al.
(2019) was to improve the robustness of classiﬁers while having discriminative performance comparable to other top performing models.
In Grathwohl et al.
(2019), the authors compare their method to other hybrid models, and in fact lose accuracy compared to a purely discriminative model.
As such, the generative capabilities of the proposed model is explored.
[Figure]
Spectrum of the adjacency matrix:
In Figure 1, the spectrum of the original adjacency matrix and the spectrum of the generative adjacency matrix for the Cora dataset is shown.
The x -axis is the eigenvalue, and the y -axis is the count of the eigenvalue with log scaling.
The number of bins is set at 100 .
It is seen that the spectrum of the generative adjacency matrix resembles that of the spectrum of the original adjacency matrix.
By only adding links to nodes that are close in energy, it is thought that the generative ˜ A is only adding short cycles.
The number of closed paths of length n can be computed with the following formula:
[Figure]
closed paths of length n } = (cid:88)
i λ ni (6)
Where λ i , i ∈ { 1 , · · · , n } is the spectrum of the adjacency matrix.
The cycles of length n = 3 are computed at 9780 for the original matrix, and 10674 for the generative matrix. [Figure]
Spectrum of the feature covariance matrix:
In Figure 2, the spectrum of the covariance of the original features as well as the covariance of the generative features for the Cora dataset are shown.
The x -axis is the eigenvalue, and the y -axis is the count with log scaling.
The number of bins is set at 100 .
The generative features add high covariance eigenvalues, which may add robustness against adversarial examples.
It is hypothesized that if the features have more variance, then the classiﬁer is less likely to be overtuned to the dataset. [Figure] Conﬁdence in predictions: In Figure is the max of the The is the prediction and the scaling
One metric to measure the calibration of a classiﬁer is the Expected Calibration Error (ECE).
First, one computes the conﬁdence in the predictions of the dataset.
Then, one bins these into equally spaced buckets.
The ECE measures the absolute difference between the average accuracy over that bucket and the average conﬁdence, weighed by the cardinality of the bucket over the total number of samples.
For a perfectly calibrated classiﬁer, this value will be 0 for any choice of M , the number of buckets.
[Equation]
We compute an ECE of 0 .
52 for the generative model and 0 .
67 for the discriminative model over the test set of Cora (lower is better).
5. F UTURE DIRECTIONS
Combining the energy-based framework with graph neural networks presents promising new directions for further research, both from a graph generation and theoretical point-of-view.
One point of further exploration is exploring the limitations of depth as set out in Oono & Suzuki (2019).
The paper posits that GNNs are limited in depth due to the action of repeated application of the graph operator.
It is thought that this can be ameliorated by perturbing A at every layer within the generative framework proposed in this work.
Another direction for future work would be in utilizing our generative model for tasks in computational chemistry and bioinformatics tasks, where graph structure is abound, and where many tasks are inverse problems–given some chemical property that one requires, what is the structure that gives that property?
Such tasks are thought to be apt for a generative graph model.