TANGO: Graph Neural Dynamics via Learned Energy and Tangential Flows

TL;DR

TANGO introduces a novel graph neural network framework that combines learned energy-based dynamics with tangential flows, improving stability, flexibility, and performance in graph learning tasks.

Contribution

It proposes a new energy-based dynamical system with tangential flows for graph neural networks, enhancing stability and signal propagation in graph representation learning.

Findings

Achieves strong results on node and graph classification benchmarks.

Effectively mitigates oversquashing in graph neural networks.

Demonstrates compatibility with various GNN backbones.

Abstract

We introduce TANGO -- a dynamical systems inspired framework for graph representation learning that governs node feature evolution through a learned energy landscape and its associated descent dynamics. At the core of our approach is a learnable Lyapunov function over node embeddings, whose gradient defines an energy-reducing direction that guarantees convergence and stability. To enhance flexibility while preserving the benefits of energy-based dynamics, we incorporate a novel tangential component, learned via message passing, that evolves features while maintaining the energy value. This decomposition into orthogonal flows of energy gradient descent and tangential evolution yields a flexible form of graph dynamics, and enables effective signal propagation even in flat or ill-conditioned energy regions, that often appear in graph learning. Our method mitigates oversquashing and is compatible with different graph neural network backbones. Empirically, TANGO achieves strong performance across a diverse set of node and graph classification and regression benchmarks, demonstrating the effectiveness of jointly learned energy functions and tangential flows for graph neural networks.