Equivariant Geometric Scattering Networks via Vector Diffusion Wavelets

TL;DR

This paper introduces a novel vector-valued geometric scattering transform built on vector diffusion wavelets to deliver $SE(3)$-equivariant processing on geometric graphs with scalar and vector node features. It provides theoretical guarantees, including frame bounds and rotation-equivariance of both wavelets and scattering coefficients, under rotations of the input. The proposed Equivariant Scattering-based GNN (ESc-GNN) demonstrates competitive performance with a fraction of the parameters compared to existing equivariant GNNs on synthetic 3D ellipsoid data, addressing oversmoothing and underreaching by leveraging multiscale diffusion features. The work offers a principled, low-parameter approach for enforcing rotational symmetry in GNNs with vector-valued signals, with potential impact on molecular, physical-geometry, and point-cloud learning tasks.

Abstract

We introduce a novel version of the geometric scattering transform for geometric graphs containing scalar and vector node features. This new scattering transform has desirable symmetries with respect to rigid-body roto-translations (i.e., $SE(3)$-equivariance) and may be incorporated into a geometric GNN framework. We empirically show that our equivariant scattering-based GNN achieves comparable performance to other equivariant message-passing-based GNNs at a fraction of the parameter count.

Figures (3)

• Figure 1: Illustration of rotational equivariance of $\mathbf{Q}^m$, applied to a $2D$ vector field. (The field was generated with random uniform sampling of points $x_i \sim \mathcal{U}([-1,1]^2)$, angles $\theta_i \sim \mathcal{U}[0,2\pi)$, and magnitudes $m_i \sim \mathcal{U}[0.4, 1.0]$ such that $w_i = m_i [\cos \theta_i, \; \sin \theta_i]$). The central vector (the magnitude of which was made to be the largest, for illustrative purposes) is colored blue. The top row shows $\mathbf{Q}^m$ applied to vectors in an unrotated vector field (where we first used the sample points $x_i$ to construct a $k$-NN graph, $k = 3$). The bottom row shows the same system, rotated 90 degrees counter-clockwise, then diffused by $\overline{\mathbf{Q}}$. The black vectors highlight those involved in the next diffusion step. That is, the vectors from which the central vector will receive (indirect) diffusion messages: first, its (symmetric) $k$-nearest neighbor vectors, then neighbors of neighbors, and so on. After three diffusion steps, the top (unrotated) system is rotated 90 degrees counter-clockwise, like the bottom system was initially. (This panel has a light blue background). Notably, this panel is identical, (other than the background) to the panel immediately below it thereby demonstrating the equivariance of $\mathbf{Q}^m$, since diffusing and then rotating yields the same result as rotating and then diffusing.
• Figure 2: Illustration of rotational equivariance of vector diffusion wavelets applied to a vector field (which is defined the same way as in Figure \ref{['fig:2d_Q_equivar']}). Here, the top row shows vector diffusion wavelets constructed from $\mathbf{Q}$ applied to vectors in the unrotated vector field; the bottom row shows the same system, rotated 90 degrees counter-clockwise, then diffused using wavelets constructed from $\overline{\mathbf{Q}}$. The black vectors again highlight those from which the central vector will receive (indirect) messages in the next diffusion step. After three wavelet diffusion steps, the top (unrotated) system is rotated 90 degrees counter-clockwise, like the bottom system was initially (this panel has a light blue background). The fact that the diffused vector fields in the rightmost column are equivalent shows the rotational equivariance of the wavelets.
• Figure 3: Architecture the equivariant scattering-based GNN. Given a geometric graph, scalar features $\mathbf{X}$ are passed into the scalar track while vector features $\mathbf{W}$ are passed into the vector track. Upon computing the respective scattering coefficients, the updated scalar and vector features are concatenated and projected to form the final node-level embeddings, which can be aggregated to obtain graph-level embeddings.

Theorems & Definitions (17)

• Theorem 3.1
• Theorem 3.2: Wavelet Equivariance
• Theorem 3.3: Scattering Equivariance
• proof
• Lemma A.1
• Lemma A.2
• Lemma A.3
• proof : The proof of Theorem \ref{['thm:wavelet-equivariance']}
• proof
• Proposition B.1: Proposition 2.2 of perlmutter2019understanding