Diff-NekRS: A Scalable Differentiable Framework for Multi-Timestep Solver-in-the-Loop Training
Abstract
Hybrid physics-machine-learning solvers improve under-resolved simulations by embedding trainable corrections into the time integration. During autoregressive inference, repeated solver-model interactions can amplify small errors, motivating multi-timestep solver-in-the-loop training. However, production solvers rarely expose the derivatives needed to backpropagate through such rollouts. We introduce Diff-NekRS, a scalable differentiable framework that embeds neural corrections directly in the GPU-accelerated NekRS incompressible-flow solver. NekRS computes the authoritative forward trajectory, a manually implemented exact discrete adjoint differentiates the supported fully discrete timestep, and LibTorch supplies neural vector-Jacobian products and parameter gradients. End-to-end Taylor and centered finite-difference tests verify the assembled gradient for two-dimensional cylinder flow (2Dcyl) and the three-dimensional Taylor-Green vortex (3DTGV) across five horizons and 12-1,020 MPI ranks. At 1,020 ranks, optimizer-enabled post-setup training updates retain 54.5%-78.0% and 80.7%-81.9% weak-scaling efficiency for 2Dcyl and 3DTGV, respectively. In 200-step autoregressive inference, the M = 50 model reduces the three-seed median terminal relative L2 velocity error by 59.2% for 2Dcyl and 12.1% for 3DTGV relative to the uncorrected coarse-grid P = 2 baseline, and retains wall-clock speedups of 5.38x and 2.49x, respectively, relative to the corresponding P = 7 configurations for equal simulated-time intervals. These results establish a verified and scalable path for multi-timestep solver-in-the-loop training that improves coarse-grid trajectory accuracy while retaining a speed advantage over the high-order reference
Related
- DiffSWE2d: a differentiable Shallow Water Equations solver for end-to-end flood and tsunami modelling
- A Differentiable Brain Tumor Mass Effect Solver
- AGNI: A differentiable MHD stability solver & optimizer for magnetic confinement fusion devices
- A scalable sample-based quantum diagonalization solver for general Pauli-sum Hamiltonians
- Beyond the Thin-Layer Limit: Differentiable Volumetric Training for Visible-Range Diffractive Neural Networks
BibTeX
@misc{jung2026diffnekrs,
title = {{Diff-NekRS: A Scalable Differentiable Framework for Multi-Timestep Solver-in-the-Loop Training}},
author = {Junoh Jung and Riccardo Balin and Bethany Lusch and Emil Constantinescu},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.23208},
eprint = {2609.23208},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.23208},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Diff-NekRS: A Scalable Differentiable Framework for Multi-Timestep Solver-in-the-Loop Training
AU - Junoh Jung
AU - Riccardo Balin
AU - Bethany Lusch
AU - Emil Constantinescu
PY - 2026
DA - 2026/09/19
PB - arXiv
DO - 10.48550/arxiv.2609.23208
UR - https://arxiv.org/abs/2609.23208
AB - Hybrid physics-machine-learning solvers improve under-resolved simulations by embedding trainable corrections into the time integration. During autoregressive inference, repeated solver-model interactions can amplify small errors, motivating multi-timestep solver-in-the-loop training. However, production solvers rarely expose the derivatives needed to backpropagate through such rollouts. We introduce Diff-NekRS, a scalable differentiable framework that embeds neural corrections directly in the GPU-accelerated NekRS incompressible-flow solver. NekRS computes the authoritative forward trajectory, a manually implemented exact discrete adjoint differentiates the supported fully discrete timestep, and LibTorch supplies neural vector-Jacobian products and parameter gradients. End-to-end Taylor and centered finite-difference tests verify the assembled gradient for two-dimensional cylinder flow (2Dcyl) and the three-dimensional Taylor-Green vortex (3DTGV) across five horizons and 12-1,020 MPI ranks. At 1,020 ranks, optimizer-enabled post-setup training updates retain 54.5%-78.0% and 80.7%-81.9% weak-scaling efficiency for 2Dcyl and 3DTGV, respectively. In 200-step autoregressive inference, the M = 50 model reduces the three-seed median terminal relative L2 velocity error by 59.2% for 2Dcyl and 12.1% for 3DTGV relative to the uncorrected coarse-grid P = 2 baseline, and retains wall-clock speedups of 5.38x and 2.49x, respectively, relative to the corresponding P = 7 configurations for equal simulated-time intervals. These results establish a verified and scalable path for multi-timestep solver-in-the-loop training that improves coarse-grid trajectory accuracy while retaining a speed advantage over the high-order reference
ER -