Is Distance Matrix Enough for Geometric Deep Learning?

TL;DR

A connection between geometric deep learning (GDL) and traditional graph representation learning ( GRL) is established, showing that those highly expressive GNN models originally designed for GRL can also be applied to GDL with impressive performance, and that existing complicated, equivariant models are not the only solution.

Abstract

Graph Neural Networks (GNNs) are often used for tasks involving the 3D geometry of a given graph, such as molecular dynamics simulation. While incorporating Euclidean distance into Message Passing Neural Networks (referred to as Vanilla DisGNN) is a straightforward way to learn the geometry, it has been demonstrated that Vanilla DisGNN is geometrically incomplete. In this work, we first construct families of novel and symmetric geometric graphs that Vanilla DisGNN cannot distinguish even when considering all-pair distances, which greatly expands the existing counterexample families. Our counterexamples show the inherent limitation of Vanilla DisGNN to capture symmetric geometric structures. We then propose $k$-DisGNNs, which can effectively exploit the rich geometry contained in the distance matrix. We demonstrate the high expressive power of $k$-DisGNNs from three perspectives: 1. They can learn high-order geometric information that cannot be captured by Vanilla DisGNN. 2. They can unify some existing well-designed geometric models. 3. They are universal function approximators from geometric graphs to scalars (when $k\geq 2$) and vectors (when $k\geq 3$). Most importantly, we establish a connection between geometric deep learning (GDL) and traditional graph representation learning (GRL), showing that those highly expressive GNN models originally designed for GRL can also be applied to GDL with impressive performance, and that existing complicated, equivariant models are not the only solution. Experiments verify our theory. Our $k$-DisGNNs achieve many new state-of-the-art results on MD17.

Figures (10)

• Figure 1: A pair of geometric graphs that are non-congruent but cannot be distinguished by Vanilla DisGNN. The nodes of the geometric graphs are taken from regular icosahedrons and have identical node features. Only the red nodes are part of the two geometric graphs; the grey nodes and the “edges” are included solely for visualization purposes.
• Figure 2: A: High-order geometric information contained in the distance matrix of $k$-tuples. We mark different orders with different colors, with brown, green, blue for 2-,3-,4-order respectively. (A1) High-order geometric information contained in 4-tuples, including distances, angles and dihedral angles. (A2) More 3-order geometric features, such as vertical lines, middle lines and the area of triangles. B: Examples explaining that neighboring 3-tuples can form a 4-tuple. Blue represents the center tuple, the other colors represent neighbor tuples, and the red node is the one node of that neighbor tuple which is not in the center tuple. (B1) Example for $3$-E-DisGNN. With the green edge, two 3-tuples can form a 4-tuple. (B2) Example for 3-F-DisGNN. Four 3-tuples form a 4-tuple.
• Figure 3: An example of the augmentation. In this example, $(G^L_{ori}, G^R_{ori})$ are from Figure \ref{['fig:ce20']}, and $\mathcal{QR}_3 = \{(ori, r_r), (all, r_p), (com, r_g)\}$ (with $r, p, g$ representing red, purple and green respectively). Nodes with different colors indicate that they are derived from different $G_{Q,r}$ generations. The augmented graphs $(G^L, G^R) = {\rm AUG}_{\mathcal{QR}3}(G^L_{ori}, G^R_{ori})$ consist of all the colored nodes.
• Figure :
• Figure :

Theorems & Definitions (16)

• Theorem 3.1
• Definition 6.1
• Proposition 6.2
• Proposition 6.3
• Theorem 6.4
• Theorem 6.5
• Theorem A.1
• Lemma A.2
• proof : Proof of Lemma \ref{['lem:stable state']}
• Lemma A.3