# Hyperbolic Machine Learning Moment Closures for the BGK Equations

arXiv: https://arxiv.org/abs/2401.04783 · HTML: https://mlsanity.com/paper/2401.04783/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 8 of its 47 references that are in the map.
- First posted on arXiv 9 Jan 2024.

## Builds on (most cited first, 8 of 8)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Cyclical Learning Rates for Training Neural Networks](https://mlsanity.com/paper/1506.01186/) (2015)
- [Machine learning moment closure models for the radiative transfer equation I: directly learning a gradient based closure](https://mlsanity.com/paper/2105.05690/) (2021)
- [Machine learning moment closure models for the radiative transfer equation II: enforcing global hyperbolicity in gradient based closures](https://mlsanity.com/paper/2105.14410/) (2021)
- [Machine learning moment closure models for the radiative transfer equation III: enforcing hyperbolicity and physical characteristic speeds](https://mlsanity.com/paper/2109.00700/) (2021)
- [Hybrid quadrature moment method for accurate and stable representation of non-Gaussian processes and their dynamics](https://mlsanity.com/paper/2110.01374/) (2021)

## Built on by (most cited first, 4 of 4)

- [Electron neural closure for turbulent magnetosheath simulations: energy channels](https://mlsanity.com/paper/2510.00282/) (2025)
- [Machine learning moment closure models for the radiative transfer equation IV: enforcing symmetrizable hyperbolicity in two dimensions](https://mlsanity.com/paper/2604.20143/) (2026)
- [Structure-preserving neural networks for the regularized entropy-based closure of the Boltzmann moment system](https://mlsanity.com/paper/2404.14312/) (2024)
- [Learning Closure of Dynamical Systems with Kernel Ridge Regression](https://mlsanity.com/paper/2610.02564/) (2026)

## Abstract (arXiv)

We introduce a hyperbolic closure for the Grad moment expansion of the Bhatnagar-Gross-Krook's (BGK) kinetic model using a neural network (NN) trained on BGK's moment data. This closure is motivated by the exact closure for the free streaming limit that we derived in our paper on closures in transport citeHuang2022-RTE1. The exact closure relates the gradient of the highest moment to the gradient of four lower moments. As with our past work, the model presented here learns the gradient of the highest moment in terms of the coefficients of gradients for all lower ones. By necessity, this means that the resulting hyperbolic system is not conservative in the highest moment. For stability, the output layers of the NN are designed to enforce hyperbolicity and Galilean invariance. This ensures the model can be run outside of the training window of the NN. Unlike our previous work on radiation transport that dealt with linear models, the BGK model's nonlinearity demanded advanced training tools. These comprised an optimal learning rate discovery, one cycle training, batch normalization in each neural layer, and the use of the textttAdamW optimizer. To address the non-conservative structure of the hyperbolic model, we adopt the FORCE numerical method to achieve robust solutions. This results in a comprehensive computing model combining learned closures with methods for solving hyperbolic models. The proposed model can capture accurate moment solutions across a broad spectrum of Knudsen numbers. Our paper details the multi-scale model construction and is run on a range of test problems.