A Sub-linear Low-Rank Solver for Poisson's Equation using Machine Learning Frameworks for GPU Acceleration

2026-07-07

AI Summary

This work addresses the computational bottleneck in solving high-dimensional, large-scale low-rank Poisson equations by proposing an efficient low-rank solver implemented in PyTorch. The method integrates the Cross-DEIM framework with statistical leverage score–driven adaptive index selection and discrete sine transforms (DST), marking the first incorporation of a machine learning framework into low-rank PDE solvers. GPU-accelerated parallelism is achieved through batched FFTs without global transposition, enabling significant performance gains. Evaluated on A100 GPUs and AMD EPYC CPUs, the solver demonstrates substantial speedups and successfully tackles problem scales previously deemed infeasible, offering both algorithmic novelty and practical engineering value.

Abstract

In this paper we explore a fast Poisson solver for problems with a solution that is known to be low-rank. We use an adaptive and warm started cross approximation called Cross-DEIM that iterates between index selection and and cross approximation to generate a low-rank solution. This paper focuses on leveraging a modern machine learning framework, PyTorch, as a general purpose array language to implement low-rank solvers based on Cross-DEIM. PyTorch enables native access to GPUs and accelerators but with a user-friendly high-level interface. We investigate statistical leverage scores for the index selection for the cross approximation due to the cost associated with the pivoted algorithms used with the discrete empirical interpolation methods (DEIM and QDEIM) which are historically preferred. The cross approximation is naturally paired with a Discrete Sine Transform (DST) Poisson solver. This allows the Fast Fourier Transform (FFT) to be evaluated in batches along dimensions independently without any global transpose even in higher dimensions. We present performance results running on a A100 GPU and AMD EPYC CPU demonstrating the usefulness of the approach that enables problems sizes that previously were not feasible.

Problem

• Poisson's Equation
• Low-Rank Solver
• GPU Acceleration
• Cross Approximation
• Machine Learning Frameworks

Innovation

• Cross-DEIM
• low-rank solver
• Poisson's equation
• GPU acceleration
• statistical leverage scores