Graph neural operators use less memory with functional tensor decomposition
FTD-GNO: Memory-Efficient Graph Neural Operators through Functional Tensor Decomposition of the Kernel
Machine LearningArtificial Intelligence
Summary
Graph Neural Operators are tools that help computers learn to solve complex math problems called partial differential equations. But they can use a lot of memory, especially when working with big or detailed graphs. The authors created a new method called FTD-GNO that breaks down the complex math parts into smaller, simpler pieces using a technique called tensor decomposition. This makes the process use less memory and run faster without losing accuracy. Their tests showed this method works better and scales well for big problems.
What this means in practice
- •For computational engineers: Train surrogate models for solving partial differential equations with less memory usage and faster times on large graphs.
- •For deep learning developers: Build memory-efficient graph layers for graph-based operator learning to handle high-resolution or high-dimensional data.
Authors
Xiaomin Zhang, Boyue Wang, Junbin Gao, Yongli Hu anbd Baocai Yin
Abstract
Graph Neural Operators (GNOs) provide flexible surrogate models for learning solution operators of partial differential equations (PDEs). However, standard GNOs typically parameterize the integral kernel with a monolithic neural network and evaluate kernel interactions over graph edges, leading to substantial computational and memory overhead at high resolutions or with large neighborhoods. To address these limitations, we propose Functional Tensor Decomposition Graph Neural Operator (FTD-GNO), a memory-efficient GNO framework that decouples the high-dimensional continuous integral kernel into low-dimensional mode-wise functions. By instantiating the kernel with classical tensor decomposition formats, including CP, Tensor-Train, and Tucker decompositions, FTD-GNO enables algebraic reconstruction of the integral operator without explicitly materializing full edge-wise kernel tensors. This factorized formulation reduces the memory footprint of kernel evaluation and aggregation while retaining the continuous operator-learning structure of GNOs. Theoretical complexity analysis shows that FTD-GNO substantially lowers parameter and activation-memory costs associated with high-dimensional kernel construction. Experiments show lower peak memory than the corresponding unfactorized graph-integral baselines, with shorter recorded training times. Fourier-graph experiments further demonstrate that FTD can improve the efficiency of a graph-integral layer within a hybrid operator and has good scalability.