Learning Differentiable Weak-Form Corrections to Accelerate Finite Element Simulations

Abstract

We present a differentiable weak-form learning approach for accelerating finite element simulations. Rather than introducing black-box source terms in the strong form of the governing equations, we augment the momentum equation directly in the variational (weak) form with parameterized bilinear operators. The coefficients of these operators are learned from high-resolution simulations so that unresolved small-scale dynamics can be represented on coarse grids. Applying the correction at the weak-form level aligns the learned model with the finite element discretization, preserving key numerical structure and better respecting the fundamental properties of incompressible flow. In the same setting, the approach yields solutions that are more accurate and more stable over long time horizons than comparable strong-form corrections. We implement the proposed method in the Firedrake finite element solver and evaluate it on benchmark problems, including the one-dimensional convection-diffusion equation and the two-dimensional incompressible Navier-Stokes equations. End-to-end differentiable training is enabled by coupling PyTorch with the Firedrake adjoint framework. Across these tests, the learned variational operators improve long-term accuracy while reducing computational cost. Overall, our results suggest that weak-form learning provides a principled, structure-preserving route to accurate and stable coarse-grid simulations of incompressible flows.