Graph Neural Dynamics via Learned Energy and Tangential Flows
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 non-increasing direction that guarantees 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 is designed to help alleviate 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.
Lay Summary
Many important data sets can be described as networks: molecules are networks of atoms, social platforms are networks of users, and images or text can be turned into networks of related parts. AI models that learn from networks often struggle when useful information must travel across many steps. Important signals can become diluted or squeezed before they reach the places where they are needed. This paper introduces TANGO, a new way to guide how information moves inside these network-based AI models. The method allows the model to follow a learned “energy landscape,” similar to how a physical system tends to move toward more stable states. At the same time, TANGO also allows the model to move sideways along this landscape, so it can keep exploring useful directions instead of getting stuck too early. This combination gives the model a controlled but flexible way to update information. In experiments across many network learning tasks, including molecular and graph prediction benchmarks, TANGO improves information flow and often performs better than strong existing models.