[论文解读] Training neural networks under physical constraints using a stochastic augmented Lagrangian approach

Alp Dener, M.A. Miller+3 more|arXiv (Cornell University)|Sep 15, 2020

Nuclear reactor physics and engineering参考文献 17被引用 12

一句话总结

该论文提出了一种随机增广拉格朗日方法，用于在物理约束（具体为质量、动量和能量守恒）下训练神经网络，以近似融合模拟中的福克-普朗克-兰道碰撞算子。通过将随机梯度下降与惩罚参数和拉格朗日乘子的启发式更新策略相结合，该方法实现了低于10⁻⁶的约束违反，且预测精度优于固定惩罚方法，从而可实现与XGC模拟工作流的集成。

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.

研究动机与目标

• 解决XGC模拟中涉及多种等离子体组分时福克-普朗克-兰道碰撞算子的O(n²)计算瓶颈问题。
• 开发一种机器学习模型，以高精度近似碰撞算子，同时严格遵守质量、动量和能量守恒等物理守恒定律。
• 克服固定惩罚方法的局限性，后者在约束执行方面表现不佳，且需要繁琐的超参数调优。
• 通过将约束违反控制在每次碰撞步骤低于10⁻⁵，实现ML模型在XGC模拟工作流中的实际集成。

提出的方法

• 使用增广拉格朗日方法将物理解析训练表述为带有非线性守恒约束的约束优化问题。
• 利用PyTorch原生的随机梯度下降高效求解内部无约束最小化子问题。
• 在外部增广拉格朗日循环中实施惩罚参数和拉格朗日乘子的启发式更新策略，以提升收敛性。
• 通过惩罚项将约束强制直接集成到损失函数中，以惩罚质量、动量和能量守恒的违反。
• 在从XGC FPL算子生成的合成数据上训练编码器-解码器神经网络，确保学习过程中的物理解析一致性。
• 通过小批量数据打乱和每轮打乱执行10次外部增广拉格朗日迭代，共30个训练周期，以提升泛化能力和稳定性。

实验结果

研究问题

• RQ1 随机增广拉格朗日方法能否在融合模拟算子的神经网络训练中有效强制执行物理解析约束？
• RQ2 与固定惩罚方法相比，该方法在约束违反和模型精度方面的表现如何？
• RQ3 更严格的守恒定律强制是否能提升ML模型的泛化能力和预测精度？
• RQ4 在无需完整Hessian信息的情况下，惩罚参数和乘子的启发式更新策略是否能实现充分收敛？
• RQ5 考虑到每次碰撞步骤的误差容限为10⁻⁵，该模型的精度是否足以实现与XGC模拟工作流的集成？

主要发现

• 随机增广拉格朗日方法在质量、动量和能量守恒定律上均实现了至少10⁻⁶的中位数约束违反，满足XGC集成所需的精度阈值。
• 该方法显著降低了模型预测误差，优于固定惩罚方法，表明更严格的约束强制引导网络向更好的泛化方向发展。
• 所有三项守恒定律的约束违反均得到均匀强制，任一物理量的性能均未出现退化。
• 训练收敛速度和稳定性均优于固定惩罚方法，损失函数和约束范数图谱显示了这一点。
• 在分布外测试中，采用增广拉格朗日方法训练的模型优于无约束和固定惩罚基线模型，表现出更强的鲁棒性。
• 结果表明，强制执行物理约束不仅能确保模拟保真度，还能提升所学算子的内在精度。