Training neural networks under physical constraints using a stochastic augmented Lagrangian approach
This addresses the problem of improving neural network training under physical constraints for kinetic simulations in fusion research, representing an incremental advancement.
The paper tackled physics-constrained training of a neural network for approximating the Fokker-Planck-Landau collision operator in kinetic fusion simulation, proposing a stochastic augmented Lagrangian approach that achieved higher accuracy than a fixed penalty method, with accuracy sufficient for practical use.
We investigate the physics-constrained training of an encoder-decoder neural network for approximating the Fokker-Planck-Landau collision operator in the 5-dimensional kinetic fusion simulation in XGC. To train this network, we propose a stochastic augmented Lagrangian approach that utilizes pyTorch's native stochastic gradient descent method to solve the inner unconstrained minimization subproblem, paired with a heuristic update for the penalty factor and Lagrange multipliers in the outer augmented Lagrangian loop. Our training results for a single ion species case, with self-collisions and collision against electrons, show that the proposed stochastic augmented Lagrangian approach can achieve higher model prediction accuracy than training with a fixed penalty method for our application problem, with the accuracy high enough for practical applications in kinetic simulations.