Training neural networks under physical constraints using a stochastic augmented Lagrangian approach
Abstract
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.
Cite
@article{arxiv.2009.07330,
title = {Training neural networks under physical constraints using a stochastic augmented Lagrangian approach},
author = {Alp Dener and Marco Andres Miller and Randy Michael Churchill and Todd Munson and Choong-Seock Chang},
journal= {arXiv preprint arXiv:2009.07330},
year = {2020}
}