Differentiability in Unrolled Training of Neural Physics Simulators on Transient Dynamics

TL;DR

It is shown that these behaviors are invariant to the physical system, the network architecture and size, and the numerical scheme, which motivates integrating non-differentiable numerical simulators into training setups even if full differentiability is unavailable.

Abstract

Unrolling training trajectories over time strongly influences the inference accuracy of neural network-augmented physics simulators. We analyze this in three variants of training neural time-steppers. In addition to one-step setups and fully differentiable unrolling, we include a third, less widely used variant: unrolling without temporal gradients. Comparing networks trained with these three modalities disentangles the two dominant effects of unrolling, training distribution shift and long-term gradients. We present detailed study across physical systems, network sizes and architectures, training setups, and test scenarios. It also encompasses two simulation modes: In prediction setups, we rely solely on neural networks to compute a trajectory. In contrast, correction setups include a numerical solver that is supported by a neural network. Spanning these variations, our study provides the empirical basis for our main findings: Non-differentiable but unrolled training with a numerical solver in a correction setup can yield substantial improvements over a fully differentiable prediction setup not utilizing this solver. The accuracy of models trained in a fully differentiable setup differs compared to their non-differentiable counterparts. Differentiable ones perform best in a comparison among correction networks as well as among prediction setups. For both, the accuracy of non-differentiable unrolling comes close. Furthermore, we show that these behaviors are invariant to the physical system, the network architecture and size, and the numerical scheme. These results motivate integrating non-differentiable numerical simulators into training setups even if full differentiability is unavailable. We show the convergence rate of common architectures to be low compared to numerical algorithms. This motivates correction setups combining neural and numerical parts which utilize benefits of both.

Figures (35)

• Figure 1: Illustration of data shift and gradient divergence over unrolling; gradients of non-differentiable unrolling diverge from the true gradient landscape (blue), gradients of differentiable simulators are prone to explosions over long horizons (orange); shaded areas mark potential benefits over one-step learning; full gradient calculation performs best.
• Figure 2: Left: Illustration of unrolling for a horizon $m=2$, with numerical solver $\mathcal{S}$ and neural network $f_\theta$, visualized for correction and prediction chains; importantly, it highlights differences in the gradient flow shown with dotted lines in the backward pass for non-differentiable (blue, no-gradient (NOG)) and differentiable (orange, with-gradient (WIG)) setups: note how the gradients do not flow through a time step of $\mathcal{S}$ for , instead they contain two shorter, disconnected backpropagation chains; Right: Contributions of networks and numerical solvers in possible prediction and correction paradigms
• Figure 3: Visualizations of our physical systems, from left to right: KS equation state over time, KOLM vorticity field, WAKE vorticity field, AERO local Mach number field visualization for different upstream Mach numbers M$_\infty$
• Figure 4: Inference accuracy measured in $\mathcal{L}_\mathrm{rel}$ for correction setups on , , and systems; displayed models were trained with (brown), (blue), (orange); across network architectures (graph networks for KS, conv-nets otherwise), and network sizes has lowest errors; one 32k model diverged in the case that and kept stable
• Figure 5: Exemplary inference trajectories for the system on extrapolation data for $\mathcal{X}=70.4$; trajectories of CNNs on the left, errors of CNNs and GCN in the middle, and on the right respectively for all training variants