Training neural networks under physical constraints using a stochastic augmented Lagrangian approach

• DOI: 10.48550/arxiv.2009.07330
• Publication Date: 2020-01-01
• Authors: Dener, Alp; Miller, Marco Andres; Churchill, Randy ...; Munson, Todd; Chang, Choong-Seock
• Keywords: FOS: Computer and information sciences; Computer Science - Machine Learning; FOS: Physical sciences; Machine Learning (stat.ML); Computational Physics (physics.comp-ph); 01 natural sciences; Physics - Plasma Physics; Machine Learning (cs.LG); Plasma Physics (physics.plasm-ph); Statistics - Machine Learning; Optimization and Control (math.OC); 0103 physical sciences; FOS: Mathematics; Physics - Computational Physics; Mathematics - Optimization and Control

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.