Towards a scalable ML-accelerated space-time multigrid solver for nonlinear equations

High-fidelity computational fluid dynamics (CFD) simulations play a critical role in various scientific and engineering domains. Reducing wall-clock time without loss of accuracy is particularly important for time-sensitive applications affecting human life, such as digital twins for biomedicine. Even with the advent of supercomputers, effective speedups remain challenging due to the saturation of spatial parallelization. Reduced order models and machine learning (ML) surrogates reduce wall-clock time at the cost of introducing error. To overcome this issue, we propose a novel framework that leverages neural operators (NO) to accelerate the convergence of space-time multigrid methods. We extend these techniques to non-linear systems and explore different strategies for employing NOs on coarse grid levels in a way that does not alter the solution, i.e., preserves numerical accuracy. To evaluate the effectiveness of this approach, we study canonical nonlinear 1D and 2D PDEs, including the viscous Burgers and Kuramoto–Sivashinsky equations, as representative test cases. We investigate the extent to which the proposed method can accelerate convergence while maintaining physical fidelity, which highlights a path for the reliable integration of ML into conventional CFD pipelines in the future. This framework is non-intrusive and generalizable to more complex 3D systems.