⚡ electrical engineering

VDW-GNNs: Vector diffusion wavelets for geometric graph neural networks

This paper introduces Vector Diffusion Wavelets (VDWs), a novel wavelet family inspired by vector diffusion maps, and demonstrates their effectiveness and desirable theoretical properties—such as frame stability and rotational/translational symmetries—when incorporated into geometric Graph Neural Networks (VDW-GNNs) for analyzing diverse data ranging from synthetic point clouds to real-world wind fields and neural activity.

Original authors: David R. Johnson, Alexander Sietsema, Rishabh Anand, Deanna Needell, Smita Krishnaswamy, Michael Perlmutter

Published 2026-06-23

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: David R. Johnson, Alexander Sietsema, Rishabh Anand, Deanna Needell, Smita Krishnaswamy, Michael Perlmutter

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a computer to understand shapes and movements in the real world, like the swirling patterns of wind or the firing patterns of neurons in a brain. Usually, computers are great at looking at simple lists of numbers (like a list of prices or temperatures). But when data has direction and geometry—like a wind blowing north at 10 mph, or a neuron firing with a specific vector—the standard tools often get confused. They might treat "North" and "South" as just two different numbers, missing the fact that they are opposite directions of the same thing.

This paper introduces a new tool called Vector Diffusion Wavelets (VDWs) and a neural network built around them, called VDW-GNN. Here is how it works, using simple analogies:

1. The Problem: The "Oversmoothing" and "Underreaching" Trap

Most current AI models for graphs (networks of connected points) work like a game of "Telephone."

The Message Passing Trap: In a standard model, a node (a point) asks its neighbors, "What do you think?" and averages their answers. If you do this too many times, everyone starts saying the exact same thing. The unique details get washed out. This is called oversmoothing.
The Short-Sighted Trap: If you stop the game too early, a node only knows about its immediate neighbors and misses the big picture. This is called underreaching.

2. The Solution: The "Multi-Scale Flashlight"

The authors use something called Diffusion Wavelets. Think of this as a special flashlight that can zoom in and out instantly.

Traditional Wavelets: These look at data at different "scales." One setting looks at the immediate neighborhood (close-up), another looks at the whole cluster (zoomed out), and another looks at the entire shape (wide-angle).
The Innovation: The authors realized that while we had these flashlights for simple numbers (scalar data), we didn't have them for directions and vectors (like wind or force). Their new Vector Diffusion Wavelets are flashlights that can handle directions. They don't just ask, "How hot is it?" they ask, "Which way is the wind blowing, and how does that direction change as we look at neighbors further away?"

3. How It Handles Rotation: The "Spinning Top" Analogy

One of the biggest challenges in geometry is rotation. If you take a picture of a wind map and spin the paper 90 degrees, the wind is still blowing the same way relative to the ground; only your view changed.

Old Models: If you rotated the input, many old models would get confused and give a completely different (and wrong) answer because they didn't understand that the relationship between the points stayed the same, even if the coordinates changed.
VDW-GNN: This new model is rotationally equivariant. Imagine a spinning top. If you rotate the table the top is sitting on, the top rotates with it, but it keeps spinning the same way relative to the table. Similarly, if you rotate the wind data, the VDW-GNN rotates its internal understanding of the wind to match, ensuring the answer remains consistent and correct.

4. How It Works: The "Local Map" Trick

To understand the direction of a vector at a specific point, the model builds a tiny, local map.

It looks at a point and its neighbors.
It uses a mathematical trick (called Singular Value Decomposition) to figure out the "local coordinate system" for that specific spot. It's like a hiker setting up a temporary compass based on the trees around them.
It then compares the directions of neighbors using these local compasses to create a smooth, multi-scale picture of how the vectors flow across the whole shape.

5. What They Tested It On

The authors didn't just talk about theory; they tested this on three specific things:

Synthetic Ellipsoids: They created fake 3D shapes (like stretched balls) and asked the model to guess their size or predict a direction on their surface. The model was much better at this than others, especially when the shapes were rotated.
Wind Fields: They used real-world data of wind blowing over the Earth. The goal was to guess missing wind speeds in areas where sensors failed. The model successfully filled in the gaps, even when the data was rotated, and did it with far fewer computer "brain cells" (parameters) than competitors.
Neural Activity: They looked at data from a monkey's brain while it moved its hand. The goal was to understand the pattern of neural signals. The model performed just as well as the best existing methods but was much faster to train and required less memory.

The Bottom Line

The paper claims that by inventing Vector Diffusion Wavelets, they have created a new way for AI to understand geometric data that involves direction. This new method avoids the common pitfalls of getting "too smooth" or "too short-sighted," respects the rules of rotation (so it works no matter how you turn the data), and does all this efficiently. It proves that you can build powerful, rotation-aware AI systems that are also lightweight and fast.

Technical Summary: VDW-GNNs: Vector Diffusion Wavelets for Geometric Graph Neural Networks

Problem Statement

Geometric Deep Learning (GDL) aims to extend deep learning to data with geometric structures, such as graphs and manifolds. A central challenge in this field is designing representations that respect the intrinsic symmetries and structures of the data (e.g., rotational and translational invariance/equivariance).

Existing Graph Neural Networks (GNNs) primarily rely on the message-passing paradigm, where node representations are updated via local averaging. While effective for tasks like node classification on homophilous networks, this approach suffers from two critical limitations:

Oversmoothing: As the number of layers increases, node representations become indistinguishable, limiting the depth of the network.
Underreaching: Local message passing struggles to capture long-range interactions necessary for certain tasks.

While diffusion wavelets and geometric scattering transforms have been proposed to address the oversmoothing/underreaching tradeoff by capturing multiscale graph geometry, they have historically been restricted to scalar-valued signals. Many real-world geometric data sets, however, involve vector-valued features (e.g., wind velocity vectors, neural activity velocities, or tangent bundle values on a manifold). Standard approaches often treat these vector components as independent scalar channels, failing to leverage the geometric relationships between vector components and the underlying manifold structure. Furthermore, existing methods for vector-valued data often lack rigorous theoretical guarantees regarding frame stability or equivariance under rigid motions.

Methodology

The authors introduce Vector Diffusion Wavelets (VDWs) and a corresponding family of Vector Diffusion Wavelet Graph Neural Networks (VDW-GNNs). This framework extends the theory of diffusion wavelets and geometric scattering to vector-valued signals defined on geometric graphs.

1. Vector Diffusion Matrix ( $Q$ )

Inspired by the Vector Diffusion Maps (VDM) algorithm, the core of the method is the construction of a vector diffusion matrix $Q \in \mathbb{R}^{nD \times nD}$ (where $n$ is the number of nodes and $D$ is the ambient dimension).

Local Frames: For each node $v_i$ , a local orthonormal basis (frame) $\{u_{i,1}, \dots, u_{i,D}\}$ is constructed using Singular Value Decomposition (SVD) of the relative distance matrix to its neighbors. This defines a local coordinate system.
Parallel Transport: To relate features between nodes, the method computes rotation matrices $O_{i,j} = U_i U_j^\top$ that map the local frame of node $v_j$ to that of $v_i$ . A sign-flipping technique is employed to ensure these matrices are well-defined and unique up to sign.
Construction: The matrix $Q$ is defined in block form as $Q[i, j] = P[i, j] O_{i,j}$ , where $P$ is the standard scalar diffusion matrix. This operator acts as a computationally efficient proxy for the heat semigroup associated with the Connection Laplacian ( $\Delta$ ), which governs the diffusion of vector fields on a manifold.

2. Vector Diffusion Wavelets and Scattering

Similar to scalar diffusion wavelets, the authors define a filter bank using powers of $Q$ :

Wavelets: $\tilde{\Psi}_j = Q^{2^{j-1}} - Q^{2^j}$ (for $j \ge 1$ ) and $\tilde{\Psi}_0 = I - Q$ .
Low-pass Filter: $\tilde{\Phi}_J = Q^{2^J}$ .
Scattering Transform: Non-linear features are extracted by alternating these wavelet filters with point-wise activation functions $\sigma$ . The authors define first-order coefficients $\tilde{U}[j]w = \sigma(\tilde{\Psi}_j w)$ and higher-order coefficients recursively.

3. VDW-GNN Architecture

The authors incorporate these wavelets into a GNN architecture that processes scalar and vector tracks separately but allows for interaction:

Scalar Track: Processes scalar node features using standard diffusion wavelets.
Vector Track: Processes vector features using VDWs.
Mixing: Learnable mixing layers combine scattering coefficients across scales.
Prediction Heads: The network is designed to be rotationally invariant for scalar targets (e.g., graph diameter) and rotationally equivariant for vector targets (e.g., wind reconstruction), ensuring the output transforms correctly under global rotations of the input.

Key Contributions

Introduction of VDWs: The authors propose a novel family of wavelets for graphs with vector-valued node features, extending diffusion wavelets to the tangent bundle setting.
Theoretical Guarantees:
- Frame Bounds: They prove that VDWs satisfy frame bounds similar to traditional scalar diffusion wavelets, ensuring stable inversion and robustness to noise.
- Equivariance: They prove that the wavelets and scattering coefficients are equivariant with respect to the Special Orthogonal group $SO(D)$ (rotations) and invariant to translations.
VDW-GNNs: They demonstrate the effective integration of these wavelets into neural networks, showing that the method captures multiscale geometry without oversmoothing.
Empirical Validation: The method is validated on synthetic data, real-world wind field data, and neural activity manifolds, showing competitive or superior performance to state-of-the-art equivariant GNNs (like EGNN and TFN) with significantly fewer parameters.

Experimental Results

The paper evaluates VDW-GNNs across three distinct settings:

Synthetic Ellipsoids:
- Task: Predicting the Euclidean diameter of random ellipsoids (graph-level) and regressing a vector field defined on the ellipsoid surface (node-level).
- Findings: Non-equivariant models failed significantly when test data was rotated, whereas VDW-GNNs maintained performance. VDW-GNNs achieved comparable accuracy to heavy equivariant baselines (EGNN, TFN) but with orders of magnitude fewer parameters (e.g., ~60k vs ~29M for TFN).
Wind Field Reconstruction:
- Task: Reconstructing masked wind velocity measurements (tangent bundle values on a sphere) from a real-world NOAA dataset.
- Findings: VDW-GNNs achieved the best test MSE and second-best rotated-test MSE among comparable models. They substantially outperformed DD-TNN (a tangent-bundle specific network) and were the fastest to train per epoch.
Multi-channel Neural Recordings:
- Task: Learning low-dimensional embeddings of macaque neural activity time-series to predict behavioral conditions (inductive setting).
- Findings: VDW-GNNs achieved a mean classification accuracy of 66%, slightly outperforming MARBLE (65%) and significantly outperforming CEBRA and LFADS. The model trained much faster and used fewer parameters than MARBLE.

Significance and Claims

The paper claims that VDW-GNNs provide a flexible, parameter-efficient, and theoretically grounded framework for learning on geometric data with vector-valued features. By leveraging the intrinsic low-dimensional structure of data (manifold hypothesis) and the symmetries of the ambient space, the method avoids the oversmoothing issues of standard message passing while respecting the geometric nature of vector fields.

The authors emphasize that their approach:

Bridges Theory and Practice: It connects the spectral theory of the Connection Laplacian with practical, scalable GNN implementations using sparse matrix operations, avoiding expensive eigendecompositions.
Offers Strong Inductive Bias: The rotational equivariance and multiscale nature of the wavelets serve as effective inductive biases, allowing the model to generalize well to rotated data and capture global structure with fewer layers.
Is Computationally Efficient: Despite handling vector-valued data, the use of sparse diffusion operators and the avoidance of explicit spectral decomposition makes the method scalable to large datasets.

The authors modestly note limitations, including reliance on local PCA for tangent space estimation (which may degrade in sparse neighborhoods) and sensitivity to hyperparameters regarding diffusion scales. They also acknowledge that the method assumes the manifold hypothesis holds; if data does not lie on a smooth low-dimensional manifold, the theoretical derivation may not hold.