NODE EMBEDDING FROM NEURAL HAMILTONIAN OR-BITS IN GRAPH NEURAL NETWORKS

Abstract

In the graph node embedding problem, embedding spaces can vary significantly for different data types, leading to the need for different GNN model types. In this paper, we model the embedding update of a node feature as a Hamiltonian orbit over time. Since the Hamiltonian orbits generalize the hyperbolic exponential maps, this approach allows us to learn the underlying manifold of the graph in training, in contrast to most of the existing literature that assumes a fixed graph embedding manifold. Our proposed node embedding strategy can automatically learn, without extensive tuning, the underlying geometry of any given graph dataset even if it has diverse geometries. We test Hamiltonian functions of different forms and verify the performance of our approach on two graph node embedding downstream tasks: node classification and link prediction. Numerical experiments demonstrate that our approach adapts better to different types of graph datasets than popular state-of-the-art graph node embedding GNNs.

1. INTRODUCTION

Graph neural networks (GNNs) (Yue et al., 2019; Ashoor et al., 2020; Kipf & Welling, 2017b; Zhang et al., 2022; Wu et al., 2021) have achieved good inference performance on graph-structured data such as social media networks, citation networks, and molecular graphs in chemistry. Most existing GNNs embed graph nodes in Euclidean spaces without further consideration of the dataset graph geometry. For some graph structures like the tree-like graphs (Liu et al., 2019) , the Euclidean space may not be a proper choice for the node embedding. Recently, hyperbolic GNNs (Chami et al., 2019; Liu et al., 2019) propose to embed nodes into a hyperbolic space instead of the conventional Euclidean space. It has been shown that tree-like graphs can be inferred more accurately by hyperbolic GNNs. Furthermore, works like Zhu et al. (2020b) have attempted to embed graph nodes in a mixture of the Euclidean and hyperbolic spaces, where the intrinsic graph local geometry is attained from the mixing weight. Embedding nodes in a hyperbolic space is achieved through the exponential map (Chami et al., 2019) , which is essentially a geodesic curve on the hyperbolic manifold as the projected curve of the cogeodesic orbits on the manifold's cotangent bundle (Lee, 2013; Klingenberg, 2011) . In our work, we propose to embed the nodes, via more general Hamiltonian orbits, into a general manifold, which generalizes the hyperbolic embedding space, i.e., a strongly constrained Riemannian manifold of constant sectional curvature equal to -1. From the physics perspective, the cotangent bundles are the natural phase spaces in classical mechanics (De León & Rodrigues, 2011) where the physical system evolves according to the basic laws of physics modeled as differential equations on the phase spaces. In this paper, we propose a new GNN paradigm based on Hamiltonian mechanics (Goldstein et al., 2001) with flexible Hamiltonian functions. Our objective is to design a new node embedding strategy that can automatically learn, without extensive tuning, the underlying geometry of any given graph dataset even if it has diverse geometries. We enable the node features to evolve on the manifold under the influence of neighbors. The learnable Hamiltonian function on the manifold guides the node embedding evolution to follow a learnable law analogous to basic physical laws. Main contributions. Our main contributions are summarized as follows: 1. We take the graph as a discretization of an underlying manifold and enable node embedding through a learnable Hamiltonian orbit associated with the Hamiltonian scalar function on its cotangent bundle. (Liu et al., 2019; Chami et al., 2019; Zhang et al., 2021b; Zhu et al., 2020b) propose to embed nodes into a hyperbolic space, thus yielding hyperbolic GNNs. Moreover, Zhu et al. (2020b) proposes a mixture of embeddings from Euclidean and hyperbolic spaces. This mixing operation relaxes the strong space assumption of using only one type of space for a dataset. In this paper, we embed nodes into a general learnable manifold via the Hamiltonian orbit on its symplectic cotangent bundle. This allows our model to flexibly adapt to the inherent geometry of the dataset. Graph Neural diffusion: Neural Partial Differential Equations (PDEs) have been applied to graphstructured data (Chamberlain et al., 2021b; a; Song et al., 2022) , where different diffusion schemes are assumed when performing message passing on graphs. To be more specific, the heat diffusion model is assumed in (Chamberlain et al., 2021b) and the Beltrami diffusion model is assumed in (Chamberlain et al., 2021a; Song et al., 2022) . (Rusch et al., 2022) models the nodes in the graph as coupled oscillators, i.e., a second-order ODE. While above mentioned graph neural diffusion schemes and our model all use ODEs, there is a fundamental difference between our model and graph neural flows. In the graph PDEs, they wrap the message passing function, e.g., constant aggregation function like the one in GCN, and attention-based aggregation function like the one in GAT, into an ODE function. In contrast, our model treats the node embedding process and node aggregation process as two independent processes, where we use the ODE function only to learn a suitable node embedding space which is then followed by a node aggregation step. To sum up, our ODE is actually a node embedding layer taking node features as the input whereas graph PDEs are node aggregation layers taking node features as well as graph adjacency matrix as the input. Notations: We use the Einstein summation convention (Lee, 2013) for expressions with tensor indices.

3. MOTIVATIONS AND PRELIMINARIES

In this section, we briefly review the concepts of the geodesic curve on a Riemannian manifold from the principle of stationary action in the form of Lagrange's equations. We then further generalize the geodesic curve to the Hamiltonian orbit associated with an energy function H, which is a conserved quantity along the orbit. We first summarize the motivation of our work as follows. Motivation I: from the hyperbolic exponential map to Riemannian geodesic. The geodesic curve gives rise to the exponential map that maps points from the tangent space to the manifold and has been



2. Our node embedding strategy can automatically learn, without extensive tuning, the underlying geometry of any given graph dataset even if it has diverse geometries. We empirically demonstrate its ability by testing on two graph node embedding downstream tasks: node classification and link prediction. 3. From empirical experiments, we observe that the oversmoothing problem of GNNs can be mitigated if the node features evolve through Hamiltonian orbits. By the conservative nature of the Hamiltonian equations, our model enables a stable training and inference process while updating the node features over time and layers. 2 RELATED WORK While our paper is related to Hamiltonian neural networks in the literature, we are the first, to our best knowledge, to model graph-structured data with Hamiltonian equations. In what follows, we briefly review Hamiltonian neural networks, Riemannian manifold GNNs, and physics-inspired GNNs. Hamiltonian neural networks. Among these physics-inspired deep learning approaches, Hamiltonian equations have been applied to conserve an energy-like quantity when training neural networks. The papers Greydanus et al. (2019); Zhong et al. (2020); Chen et al. (2021) train a neural network to infer the Hamiltonian dynamics of a physical system, where the Hamiltonian equations are solved using neural ODE solvers. The work Haber & Ruthotto (2017) builds a Hamiltonian-inspired neural ODE to stabilize the gradients so as to avoid vanishing and exploding gradients. The paper Huang et al. (2022) further studies the adversarial robustness of Hamiltonian ODE. In this paper, we focus on applying Hamiltonian equations to graph neural networks, which has not been investigated in the above-mentioned works.

