SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jung Yeon Park, Lawson L.S. Wong†, Robin Walters†
{park.jungy@northeastern.edu, lsw@ccs.neu.edu,
r.walters@northeastern.edu}
Data over non-Euclidean manifolds, often discretized as surface meshes, naturally arise in computer graphics and biological and physical systems. In particular, solutions to partial differential equations (PDEs) over manifolds depend critically on the underlying geometry. While graph neural networks have been successfully applied to PDEs, they do not incorporate surface geometry and do not consider local gauge symmetries of the manifold. Alternatively, recent works on gauge equivariant convolutional and attentional architectures on meshes leverage the underlying geometry but underperform in modeling surface PDEs with complex nonlinear dynamics. To address these issues, we introduce a new gauge equivariant architecture using nonlinear message passing. Our novel architecture achieves higher performance than either convolutional or attentional networks on domains with highly complex and nonlinear dynamics. However, similar to the non-mesh case, design trade-offs favor convolutional, attentional, or message passing networks for different tasks; we investigate in which circumstances our message passing method provides the most benefit<sup>1</sup>.
Surfaces embedded in 3D space appear in many domains such as computer graphics [1], structural biology [2, 3], neuroscience [4], climate pattern segmentation [5, 6], and fluid dynamics [7]. Such objects are naturally Riemannian manifolds but are often discretized into meshes for computational tractability. Unlike other 3D representations such as point clouds or voxels, meshes encode both the geometry and topology of the surface.
One important task that naturally arises over meshes is solving partial differential equations (PDEs) over surfaces and non-Euclidean manifolds [11-14]. A classic approach to solving surface PDEs is to model the domain as a mesh and use the finite element method and the variational formulation [15, 16]. Recently, there has been increasing success in applying deep learning methods to accelerate solving PDEs and building differentiable surrogates for dynamics models [17-20]. While there are many effective approaches for gridded data or point clouds, fewer methods exist for leveraging the geometric structure of meshes. Several approaches simply treat the mesh as a graph by approximating patches of the manifold as Euclidean [21-23]. However, these approaches do not exploit the geometric properties of meshes, and Verma et al. [24], de Haan et al. [9] have shown that not retaining this information leads to poor performance. One way to express such information is to specify the local
$\dagger$  Equal advising.
<sup>1</sup>Project website with code: https://jypark0.github.io/hermes
(a)Mesh
GemCNN
Figure 1: Mesh example. The neighbors of vertex  $b$  are projected onto the tangent plane to compute their local orientation  $(\theta_{a},\theta_{c})$  after choosing a reference neighbor (gauge). There are 3 flavors of gauge equivariant message passing, cf. Fig. 17 in Bronstein et al. [8]. GemCNN [9] computes linear convolutions using the neighbor node features, EMAN [10] uses linear messages with attention, while Hermes performs nonlinear message passing. Self-interactions are omitted for clarity.
EMAN
gauge symmetries and explicitly model operations that are equivariant to these transformations [6, 9], by using parallel transport and applying a gauge equivariant filter.
Although previous work on gauge equivariance for meshes including convolutional (GemCNN [9]) and attentional (EMAN [10]) methods have shown good performance on shape correspondence and mesh classification tasks, we find they are inadequate for modeling complex dynamics on meshes, such as solving surface PDEs, which can be highly nonlinear. In this work, we introduce a new architecture, Hermes, that combines gauge equivariance with nonlinear message passing. As Hermes can express linear convolutions or attention, Hermes is strictly more expressive than GemCNN and EMAN. Hermes completes the 3 flavors of gauge equivariant message passing, analogous to graph neural networks [8]. See Figure 1 for a comparison.
We evaluate Hermes on several linear and nonlinear partial differential equations, and also on shape correspondence, object interaction system, and cloth dynamics. Our experiments show that Hermes outperforms convolutional or attentional counterparts in most domains, particularly on nonlinear surface PDEs. We find that Hermes is more robust to both the fineness of the mesh and surface roughness compared to both GemCNN and EMAN. We further investigate when the added model complexity of nonlinear message passing outweighs the cost. With this work, we hope to help guide practitioners on when to use nonlinear message passing schemes over linear schemes.
Our contributions are summarized as follows: 1) we propose a novel gauge equivariant, nonlinear message passing architecture, Hermes, for learning on meshes, 2) when evaluated on complex and nonlinear dynamics such as surface PDEs, Hermes outperforms both convolutional and attentional architectures and can generate realistic prediction rollouts, and 3) we investigate in which situations nonlinear message passing should be preferred over convolutional or attentional counterparts.
In this work, we focus on learning signals over meshes. Similar to de Haan et al. [9], our goal is to define a model in coordinates that are intrinsic to the 2D mesh instead of using the extrinsic 3D coordinates of the embedding space. This avoids dependence on the 3D embedding, which is irrelevant. However, in order to encode data over the mesh, it is still necessary to make a choice of local coordinate frame at each vertex. Since this coordinate frame is arbitrary, it follows that the model should be invariant to change of coordinates frame at each vertex, i.e. gauge equivariant.
2.1 Data over Meshes
The Mesh Datum A mesh  $M$  consists of  $(\mathcal{V},\mathcal{E},\mathcal{F})$ , where  $\mathcal{V}$  is the set of vertices,  $\mathcal{E} = \{(i,j)\}$  is the set of ordered vertex indices  $i,j$  connected by an edge, and  $\mathcal{F} = \{(i,j,k)\}$  is the set of ordered vertex indices  $i,j,k$  connected by a triangular face. Let  $\mathcal{N}_v$  denote the set of vertices connected to a vertex  $v$ . We assume that the mesh represents a discrete 2-dimensional manifold embedded in  $\mathbb{R}^3$  (i.e. a manifold mesh). Let  $x_b\in \mathbb{R}^3$  be the coordinates of vertex  $b$ . We assign a normal vector  $n_b$  to each vertex  $b$  equal to the normalized average of the surface normals of the faces that contain the vertex. The associated tangent plane  $T_{b}M$  at vertex  $b$  is the 2-dimensional affine space through  $x_{b}$  orthogonal to the normal  $n_b$ .
Local Coordinate Frame We adopt the strategy of [9] by defining the local coordinate frame at each vertex in terms of a reference neighboring vertex  $d \in \mathcal{N}_b$ . See Figure 1a for a visualization. The reference neighbor defines a positively oriented orthonormal basis  $\{e_1^b, e_2^b\}$  for the tangent space  $T_bM$ . The orientation of every neighbor  $v \in \mathcal{N}_b$  can be represented with polar angles, where  $\theta_v$  is the angle between  $v$  and the reference neighbor after projection to the tangent plane. In Figure 1a,  $\theta_d = 0$  as  $d$  is itself the reference neighbor and we show the  $\theta_v$  for the other neighbors with respect to the reference neighbor  $d$ . See Appendix A.1 for more information.
Change of Gauge A different choice of reference neighbor will result in a different positively oriented orthonormal basis. Any two such bases will be related by an element  $g\in \mathrm{SO}(2)$ . Here  $\mathrm{SO}(2)$  is called the gauge group and  $g$  is the gauge transformation. Note that such a change is local, and a different change  $g^{p}$  can be performed at each vertex  $p\in \mathcal{V}$ . For example, if we switch from reference neighbor  $d$  to  $c$ , then we induce a gauge transformation  $\mathrm{Rot}(\phi)\in \mathrm{SO}(2)$  where  $\phi = \theta_{c} - \theta_{d}$  is the angle between  $d$  and  $c$ . The orientations are then updated  $\theta_q\mapsto \theta_q - \phi ,\forall q\in \mathcal{N}_p$
Feature Fields on the Mesh Input, output, and hidden features are encoded as data over the vertices  $\{f_p\}_{p\in \mathcal{V}}$  or edges of the mesh  $\{e_{pq}\}_{(p,q)\in \mathcal{E}}$ . We consider directed graphs and so for the edge feature  $e_{pq}$ ,  $p$  is the source and  $q$  is the target node. We situate  $e_{pq}$  to be at  $T_{p}M$  for simplicity as they are often chosen to be vectors relative to  $p$  or the edge distance [25, 18]. A scalar field, such as temperature, is invariant to change of gauge, and thus  $f_{p}\in \mathbb{R}$  transforms according to the trivial representation  $\rho_0$  of SO(2). A vector field over the mesh, representing e.g. a flow across the manifold would have  $f_{p}\in \mathbb{R}^{2}$  and transform according to the standard rotation representation  $\rho_{1}$  of SO(2) by  $2\times 2$  matrices. In general, feature fields may be any representation  $\rho_{n}$  of the gauge group SO(2).
Parallel Transport Let  $f_{p}$  and  $f_{q}$  be feature vectors at nodes  $p, q$ , respectively, where  $q$  is adjacent to  $p$ . As features at different nodes live in different tangent spaces and have different bases, in order for the anisotropic kernel to be applied to  $f_{q}$  for each  $q \in \mathcal{N}_{p}$ , each  $f_{q}$  must be parallel transported to  $T_{p}M$  and be in the same gauge. The parallel transporter  $g_{q \to p}$  first aligns the tangent plane at  $q$  to the tangent plane at  $p$  by a 3D rotation and then transforms the gauge at  $q$  to the gauge at  $p$  by a planar rotation. Acting on  $f_{q}$  with  $g_{q \to p}$  transports  $f_{q}$  to the gauge at  $p$ . For the edge features  $e_{pq}$ , no parallel transport is required as the features live in  $T_{p}M$ . For more details, see Appendix A.2 and Appendix A.2 of [9].
2.2 Flavors of Message Passing
Meshes are graphs with additional geometric information. Neural networks designed for meshes may be understood by analogy to graph neural networks (GNNs). As outlined in Bronstein et al. [8], there are largely three flavors of GNNs: convolutional, attentional, and nonlinear message-passing.
The three images on the right of Figure 1 show the three different GNN flavors in the context of gauge equivariant networks over meshes. To date, both convolutional and attentional GNNs have gauge equivariant analogues for learning over meshes. Here, we introduce a gauge equivariant mesh network in the nonlinear message-passing flavor. Let  $p$  be the source node ( $b$  in Figure 1a),  $q \in \mathcal{N}_p$  be the target nodes (one-hop neighborhood), and  $f$  the signals over nodes.
Convolutional GemCNN [9] uses convolutions where the kernel  $K_{\mathrm{neigh}}$  is applied (anisotropically) to features at the target nodes. It is comparable to graph convolutional networks [26] but with anisotropy and gauge equivariance. As convolutions are linear, the messages passed to the source node are linear, and these messages are aggregated in a permutation invariant way. GemCNN also has a kernel  $K_{\mathrm{self}}$  to model self-interactions before updating the feature at the source node. The gauge equivariant convolution (GemCNN layer) is defined as
$$
f _ {p} ^ {\prime} = K _ {\text {s e l f}} f _ {p} + \sum_ {q \in \mathcal {N} _ {p}} K _ {\text {n e i g h}} \left(\theta_ {q}\right) \rho_ {\text {i n}} \left(g _ {q \rightarrow p}\right) f _ {q}, \tag {1}
$$
where  $K_{\mathrm{neigh}}(\theta_q)$  is the kernel for  $q$  and  $\rho_{\mathrm{in}}(g_{q\to p})f_q$  is the feature vector at  $q$  parallel transported to the gauge at  $p$ . Input features  $f_{q}$  transform according to  $\rho_{\mathrm{in}}$  and output features  $f_{q}^{\prime}$  according to  $\rho_{\mathrm{out}}$ .
The filter  $K_{\mathrm{neigh}}$  is anisotropic, i.e., it depends on the orientation  $\theta_q$ . This is strictly more expressive than an isotropic filter (as in a GNN) which would simply use the same weights for each neighboring vertex. However, the dependence on orientation means that without constraints, the convolution depends on the choice of local coordinate frame. To fix this and impose gauge equivariance,  $K_{\mathrm{neigh}}$  must satisfy  $K = \rho_{\mathrm{out}}(g)^{-1}K\rho_{\mathrm{in}}(g)$ . We also define the kernel to only depend on the orientation and
not on the radial distance of neighboring nodes, as including the radius in the parameterization was not beneficial [9] and verified in our experiments.
Attentional EMAN [10] extends GemCNN by adding an attention mechanism to the messages before the aggregation step, analogous to graph attention networks [27]. Gauge equivariant attention with the self-interaction term is defined as
$$
f _ {p} ^ {\prime} = \alpha_ {p p} K _ {\text {v a l u e}} ^ {\text {s e l f}} f _ {p} + \sum_ {q \in \mathcal {N} _ {p}} \alpha_ {p q} K _ {\text {v a l u e}} ^ {\text {n e i g h}} \left(\theta_ {q}\right) \rho \left(g _ {q \rightarrow p}\right) f _ {q}, \tag {2}
$$
where  $\alpha_{pp},\alpha_{pq}$  are the attention weights for the self-interaction and neighbor interaction terms (see Equation (12) in Appendix B.2). Multi-headed attention can be used by concatenating the results of the individual heads. The attention weights are combined linearly with the neighboring node features and can be seen as weighted local averaging. Note that due to the separate kernels for the keys, queries, and values, an EMAN layer uses many more parameters than an equivalent GemCNN layer with the same dimensions.
3 Gauge Equivariant Nonlinear Message Passing
We propose gauge equivariant nonlinear message passing for meshes, complementing the convolutional and attentional gauge equivariant methods. As is the case for nonlinear message passing GNNs, our network is designed to be better suited for tasks with complex local interactions. Our network, named Hermes, consists of a sequence of message passing blocks, each of which contains an edge network  $\phi_{e}$ , an aggregation (we use sum throughout), and a node network  $\phi_{n}$ .
The edge network models neighboring interactions and takes as inputs the source node features  $f_{p}$ , target node features  $f_{q}$  where  $q \in \mathcal{N}_{p}$ , and edge features  $e_{pq}$  where  $(p,q) \in \mathcal{E}$ . The target node features are parallel transported to the gauge at  $p$  as  $\rho(g_{q \to p})$ . The network  $\phi_{e}$  consists of  $N_{e}$  gauge equivariant convolutions followed by regular nonlinearities [9] which also preserve gauge equivariance. The nonlinear messages  $m_{pq}$  are then aggregated to  $m_{p}$ . The node network models self-interactions and takes as input  $m_{p} \oplus f_{p}$ , where  $\oplus$  represents the direct sum, and then applies  $N_{n}$  gauge equivariant convolutions with  $K_{\mathrm{self}}$  kernels and regular nonlinearities. Equations (3)-(7) define gauge equivariant nonlinear message passing,
$$
h _ {p q} = f _ {p} \oplus \rho \left(g _ {q \rightarrow p}\right) f _ {q} \oplus e _ {p q}, \quad \forall (p, q) \in \mathcal {E} \tag {3}
$$
$$
m _ {p q} = \phi_ {e} \left(h _ {p q}\right) = \sigma \circ K _ {\text {n e i g h}} ^ {N _ {e}} \left(\theta_ {q}\right) \circ \dots \circ \sigma \circ K _ {\text {n e i g h}} ^ {1} \left(\theta_ {q}\right) \left(h _ {p q}\right), \tag {4}
$$
$$
m _ {p} = \sum_ {q \in \mathcal {N} _ {p}} m _ {p q}, \quad \forall p \in \mathcal {V} \tag {5}
$$
$$
h _ {p} = m _ {p} \oplus f _ {p} \tag {6}
$$
$$
f _ {p} ^ {\prime} = \phi_ {n} \left(h _ {p}\right) = \sigma \circ K _ {\text {s e l f}} ^ {N _ {n}} \circ \dots \circ \sigma \circ K _ {\text {s e l f}} ^ {1} \left(h _ {p}\right), \tag {7}
$$
where  $h_{pq}$  and  $h_p$  are inputs to  $\phi_e$  and  $\phi_n$ ,  $\sigma$  is the regular nonlinearity, and  $\circ$  denotes composition.
Figure 2 shows the Hermes architecture for one of the datasets used in experiments (wave PDE, see Section 4.1). For comparison, we also include the GemCNN and EMAN architectures in Appendix B. A residual connection is included at the end of each HermesBlock for more expressivity, see Section 5.5 for ablation results.
Proof of Gauge Equivariance We explicitly show that Hermes is equivariant to local gauge transformations in Proposition 1.
Proposition 1. Hermes is equivariant to local gauge transformations  $g^p$  for any  $p \in \mathcal{V}$  of a mesh.
The proof, provided in Appendix B.3, uses the fact that all of the kernels used in Hermes satisfy the constraints for gauge equivariance. The nonlinearities are exactly gauge equivariant as the number of intermediate samples  $N$  used in the discrete Fourier transform approaches infinity and approximately gauge equivariant for finite  $N$  (see Section 4 in [9] for more details). Thus Hermes is also gauge equivariant (in the limit).
Hermes combines both nonlinear message passing and gauge equivariance, generalizing GemCNN and EMAN, and completes the full picture with the 3 flavors of message passing (cf. Figure 1). An
Figure 2: Hermes network architecture for the Wave PDE dataset. There are two message blocks, each with 2 layers for the edge network  $\phi_e$  and 1 layer for the node network  $\phi_n$ . To illustrate the computations within the edge and node networks, we use the example mesh from Figure 1a as input. In the edge network, we show only the computations for node  $b$  for clarity.
important point to note is that nonlinear message passing decouples the effect of the number of layers in the network (depth) with the receptive field of the graph. For graph convolutions using 1-hop neighbors, the receptive field of messages is exactly equal to the network depth. In contrast, nonlinear message passing computes messages using features from neighbors at an arbitrary hop distance away from the source node before the message is passed to the source node. We hypothesize that this decoupling is beneficial for tasks with complicated interaction dynamics between node neighbors, and previous works with graph message passing networks [28, 29] have shown good performance for abstract graphs (non-spatial graphs) involving objects. Although meshes are inherently spatial, we show that nonlinear message passing is an important tool for solving PDEs on surfaces and can outperform either convolutional or attentional mechanisms.
4 Experiment Design
To evaluate gauge equivariant nonlinear message passing, we consider several different tasks. Our primary domains are partial differential equations on meshes, but we also consider shape correspondence, object interactions, and cloth simulation. Throughout, we consider triangular meshes as they are the most common. Further details and visualizations of the domains are provided in Appendix C.
Partial Differential Equations on meshes We consider three linear and nonlinear surface partial differential equations (PDEs), where the dynamics occur on the surface of an object, represented as a two-dimensional mesh embedded in 3D space. Since solving PDEs on meshes, e.g. heat diffusing over a surface, naturally depends on the intrinsic mesh geometry, gauge-equivariant nonlinear message passing is a promising solution. For all three equations, we use example meshes in the PyVista library [30] and generate 5 trajectories with different initial conditions, see Appendix C for more details.
Heat/Wave equation: The heat and wave equations are second-order linear partial differential equations (see Appendix C.1). Solving the heat equation on a surface mesh can correspond to modeling the dynamics when an external hot object makes contact at certain points on a thin hollow object. The wave equation describes how acoustic waves propagate on a thin hollow surface. The dynamics of both equations are highly dependent on the local mesh geometry.
Cahn-Hilliard equation: The Cahn-Hilliard equation [31] describes phase separation in a binary fluid mixture and is often used to model spinodal decomposition. It is a fourth-order, nonlinear,
time-dependent PDE and is often factored into two coupled second-order equations. The surface Cahn-Hilliard equation can model real-world applications such as cell proliferation [32] and two-component vesicles [33]. As a nonlinear PDE, we expect our nonlinear message passing method to exhibit greater performance than other linear message passing flavors.
Shape correspondence As a standard mesh benchmark, we use the same FAUST dataset used in previous work [9, 10]. The dataset consists of 80 train and 20 test high resolution scans of 10 humans in 10 different poses. The task is to determine shape correspondence between different meshes. As the vertices are all registered and represent the same position on the human body, this task is equivalent to classifying the correct label for each vertex.
Object interactions Inspired by interaction systems [28, 29, 34], we consider complex dynamics of interacting objects on a mesh. On a coarse triangular mesh with random hills, each object occupies a vertex and is oriented towards a neighboring vertex. An action can either turn the object left (changing its orientation), move the object forward, or turn right. Objects cannot move forward if there is another object at the destination vertex, giving rise to complex interacting dynamics. Furthermore, we consider the geometry of the mesh such that an object cannot move forward if the height difference between the current vertex and destination vertex is too high, or if the angle between the vertex normals is too large. If an object is able to move forward, we parallel transport its orientation and then choose the nearest neighboring node as its new orientation.
FlagSimple We also include the FlagSimple dataset from [20], which simulates cloth dynamics of a flag with self-collisions. Unlike the other datasets, the mesh is dynamic where the node positions change over time. The dataset was created using ArcSim [35] with regular meshing over 400 timesteps. See Appendix A.1 of [20] for more information.
4.2 Training Details
For the PDE datasets, we report test root mean squared error (RMSE) of the prediction at the next timestep given the previous 5 timesteps. We use three separate test datasets and evaluate generalization to future timesteps (test time), unseen initial conditions (test init), and unseen meshes (test mesh). For test time, we train on timesteps  $T = 0, \dots, 149$  and test on  $T = 150, \dots, 200$ . For test init, we test on trajectories with new initial conditions. For test mesh, we evaluate on completely unseen meshes to see whether a method overfits to specific mesh geometries.
For baselines, our main comparison is with GemCNN and EMAN to gauge how beneficial nonlinear message passing is over convolutional or attentional flavors. We also compare against 1) a SOTA message passing, mesh-aware method MeshGraphNet [20], 2) an E(3)-equivariant, non mesh-aware message passing network EGNN [25], 3) a non-equivariant, mesh-aware baseline SpiralNet++ [36], and 4) standard non-equivariant GNNs: graph convolutional networks (GCN) [26] and message passing networks (MPNN) [37, 38]. We tune each architecture to use a similar number of trainable parameters for a fair comparison, see Table 9 in Appendix D. A more detailed feature comparison of each method is provided in Table 8 and all other training details are relegated to Appendix D.
Table 1 shows the results on the PDE datasets and other results are given in Table 10 in Appendix E. On Heat, we see that EMAN outperforms GemCNN and Hermes outperforms EMAN significantly. This coincides with our expectation as attention mechanisms can express convolutions using constant weights, and Hermes generalizes both linear convolution and attention mechanisms. However, this does not hold for the Wave dataset, where EMAN is noticeably worse than GemCNN while Hermes still performs well. On the nonlinear Cahn-Hilliard dataset, we see that EMAN cannot generalize well to unseen meshes. Overall Hermes achieves an RMSE approximately 3 times lower than that of GemCNN, and between 2 to 8 times lower than that of EMAN.