scieee AI-readable full text Open interactive document viewer

Accelerating Fokker-Planck Simulations by Substituting the Moment Closure with a GPU-Native Deep Neural Network

Roohi, Ehsan; Mahdavi, Amir Mehran

Abstract

Particle-based Fokker-Planck (FP) models represent a high-fidelity method for simulatingrarefied gas dynamics, but they suffer from a severe computational bottleneck: the“closure problem.” This step requires the expensive, cell-wise calculation of high-ordermoments and the solution of a 9 × 9 linear system at every simulation time step. Thispaper introduces a new computational methodology designed to eliminate this bottleneckby substituting the physics-based solver with a Deep Neural Network (DNN) surrogatedeployed via a novel, high-performance strategy. Our workflow makes a critical distinctionbetween a complex offline training phase (where a 16-256-256-256-256-9 DNN is trained)and a lightweight online inference phase. Crucially, for online deployment, we avoid allframework overhead and I/O bottlenecks by extracting the raw parameters (weights andbiases) and executing the model’s forward pass as a simple, batched matrix-multiplicationfunction written natively in CuPy, ensuring all operations remain on the GPU.We validatethis approach through a rigorous, multi-stage test campaign. First, for 1D Couette flow, amodel trained on a Knudsen number sweep (Kn ≈ 0.0015−0.3) demonstrates outstandingaccuracy in both interpolations (Kn = 0.05 and Kn = 0.09) and significant extrapolation(Kn = 0.7). To test fundamental generalization, we deployed this 1D-trained modelto the 2D cavity geometry. This test yielded excellent agreement for velocity and densitystructures but produced minor, localized errors in the temperature field, confirmingthat representative multi-dimensional data is required for full thermal accuracy. Consequently, a robust 2D cavity model, trained on a lid velocity sweep (50 m/s to 600 m/s),proves capable of extreme extrapolation, accurately predicting the complex, high-energyphysics of a hyper-velocity 800 m/s case. The primary finding of this work is a fundamentalshift in the computational paradigm for this method. Performance benchmarksshow a 1.63x–1.73x speedup, but critically, a strong-scaling analysis proves this accelerationreaches the theoretical maximum predicted by Amdahl’s Law. This result provides adefinitive insight: the GPU-native surrogate is, for all practical purposes, “infinitely fast”(zero-cost) relative to the remaining tasks, and the true computational bottleneck hasbeen decisively shifted from the physics solver to the particle moment-gathering processitself.

Full text

Accelerating Fokker-Planck Simulations by Substituting the Moment Closure with a GPU-Native Deep Neural Network Ehsan Roohia,∗, Amirmehran Mahdavib aMechanical and Industrial Engineering, University of Massachusetts Amherst, 160 Governors Dr., Amherst, MA 01003, USA bDepartment of Mechanical Engineering, Hakim Sabzevari University, Sabzevar, Iran Abstract Particle-based Fokker-Planck (FP) models represent a high-fidelity method for simulating rarefied gas dynamics, but they suffer from a severe computational bottleneck: the “closure problem.” This step requires the expensive, cell-wise calculation of high-order moments and the solution of a 9×9linear system at every simulation time step. This paper introduces a new computational methodology designed to eliminate this bottleneck by substituting the physics-based solver with a Deep Neural Network (DNN) surrogate deployed via a novel, high-performance strategy. Our workflow makes a critical distinction between a complex offline training phase (where a 16-256-256-256-256-9 DNN is trained) and a lightweight online inference phase. Crucially, for online deployment, we avoid all framework overhead and I/O bottlenecks by extracting the raw parameters (weights and biases) and executing the model’s forward pass as a simple, batched matrix-multiplication function written natively in CuPy, ensuring all operations remain on the GPU. We validate this approach through a rigorous, multi-stage test campaign. First, for 1D Couette flow, a model trained on a Knudsen number sweep (Kn ≈0.0015−0.3) demonstrates outstanding accuracy in both interpolations (Kn = 0.05 and Kn = 0.09) and significant extrapolation (Kn = 0.7). To test fundamental generalization, we deployed this 1D-trained model to the 2D cavity geometry. This test yielded excellent agreement for velocity and density structures but produced minor, localized errors in the temperature field, confirming that representative multi-dimensional data is required for full thermal accuracy. Conse- ∗Corresponding author. Email address: [email protected] (Ehsan Roohi) 2 quently, a robust 2D cavity model, trained on a lid velocity sweep (50 m/s to 600 m/s), proves capable of extreme extrapolation, accurately predicting the complex, high-energy physics of a hyper-velocity 800 m/s case. The primary finding of this work is a fundamental shift in the computational paradigm for this method. Performance benchmarks show a 1.63x–1.73x speedup, but critically, a strong-scaling analysis proves this acceleration reaches the theoretical maximum predicted by Amdahl’s Law. This result provides a definitive insight: the GPU-native surrogate is, for all practical purposes, “infinitely fast” (zero-cost) relative to the remaining tasks, and the true computational bottleneck has been decisively shifted from the physics solver to the particle moment-gathering process itself. Keywords: Fokker-Planck simulation, Surrogate Model, Closure Problem, Deep Neural Network (DNN), GPU-Native, CuPy, Rarefied Gas Dynamics, Amdahl’s Law 1. Introduction High-fidelity simulation of fluid dynamics, particularly in rarefied or non-equilibrium regimes, is a cornerstone of modern engineering and physics [1,2]. While traditional Computational Fluid Dynamics (CFD) solvers are efficient, their underlying assumptions (e.g., the Navier-Stokes equations) break down when the gas mean free path becomes comparable to the characteristic length scale, as described by the Knudsen number (Kn). In these regimes, particle-based methods such as the Direct Simulation Monte Carlo (DSMC) method are the gold standard [1,3]. However, DSMC becomes computationally intractable at low to moderate Knudsen numbers due to the high number of required particle collisions. Hybrid methods, which couple a particle-based kinetic solver with a fluid description, offer a promising compromise. One such approach is the particle-based Fokker-Planck (FP) model [4], which replaces the stochastic DSMC collision operator with a deterministic drift and diffusion process in velocity space. This approach is particularly effective in reducing the stochastic noise associated with DSMC. There is a continued interest in improving the Fokker-Planck method in recent years [5,6,7,8,9,10,11,12, 13,14,15,16,17]. 3 Despite their advantages, particularly the physical fidelity of advanced formulations like the cubic-Fokker-Planck model [5], particle-FP methods introduce a severe computational bottleneck: the moment closure problem. The coefficients governing the particle’s velocity evolution (drift and diffusion) are not known a priori. They must be re-computed at every time step and for every cell in the domain by calculating high-order velocity moments (up to the 4th or 5th order) from the particle ensemble, and then solving a dense, complex linear system (e.g., a 9×9system for the cubic-FP model [9]). This repeated, cell-wise, and expensive calculation represents a primary obstacle to the scalability of high-fidelity FP simulations. The immense computational burden of high-fidelity particle methods has catalyzed a recent and rapidly growing interest in machine learning-based acceleration. Deep Neural Networks (DNNs), in particular, have emerged as a powerful paradigm for creating surrogate models capable of learning the complex, non-linear mappings inherent in fluid physics, often at a fraction of the computational cost. This trend is especially prominent in rarefied gas dynamics, where data-driven surrogates have been developed to replicate DSMC solutions [18]. More advanced architectures, such as DeepONets [19] and its modern variants [20], are being applied to learn complex operator mappings for challenging micro-nozzle flows [21]. Furthermore, the integration of physical constraints, such as shock-aware or zonal loss functions, is enhancing the accuracy of these surrogates for flows with discontinuities, like those over micro-steps [22]. This overall approach of using physics-enforced neural networks to learn the fundamental properties of rarefied gas dynamics represents a vibrant and active frontier of research [23]. Our work contributes to this field by targeting the specific, deterministic bottleneck of the moment closure solver within the Fokker-Planck framework. On modern hardware, this "closure solver" represents a primary computational bottleneck, consuming a significant fraction of the total runtime (as high as 40% in the 2D cavity test cases considered in this paper). This paper hypothesizes that this expensive, repeated, and deterministic calculation is an ideal target for replacement by a machine learning surrogate. The solver is, in effect, a complex, high-dimensional, non-linear func- 4 tion. We propose that a DNN can be trained to learn this mapping, taking readilyavailable low-order moments (e.g., density, velocity, temperature, and stress tensor) as input and directly predicting the required closure coefficients as output. This approach would entirely bypass the costly construction and solution of the linear system, as well as the need to gather the expensive high-order moments required only by the physics solver. Furthermore, this paper addresses a critical, and often overlooked, performance pitfall of hybrid ML-physics simulations: the I/O bottleneck. While the particle data (potentially billions of points) resides in GPU memory, A ’naive surrogate implementation’ that calls a standard framework’s inference function from the CPU-based control script would thus require a costly data transfer from GPU-to-CPU (for the input moments) and CPUto-GPU (for the output coefficients). A ’naive’ surrogate implementation would thus require a costly data transfer from GPU-to-CPU (for the input moments) and CPU-toGPU (for the output coefficients) for every cell at every time step. This data transfer overhead would negate, or even reverse, any computational gains. We demonstrate that this pitfall is overcome by developing a "GPU-Native Surrogate": we extract the trained model’s parameters (weights and biases) and re-implement the entire neural network forward pass using pure GPU-native matrix operations via the CuPy library. This approach ensures the entire solver replacement executes on the GPU, completely avoiding CPUGPU communication within the simulation loop. This paper details a rigorous, multi-stage validation strategy across canonical benchmarks to rigorously test the accuracy, robustness, and generalization capabilities of our GPU-native surrogate. The methodology hinges on a two-part workflow: a comprehensive offline training phase, where data from full-physics simulations is used to train a deep network (16-256-256-256-256-9) via Keras, and a highly optimized online inference phase, where the trained weights are used in a simple, lightweight CuPy function to replace the solver during a live simulation. We validate this strategy by systematically increasing complexity: (1) First, we establish the methodology in 1D Couette flow, training a model on a rich dataset spanning a wide Knudsen number sweep (Kn ≈0.0015 to 0.3) and prove its high fidelity in both interpolations (Kn = 0.05 and Kn = 0.09) and extrap- 5 olation (Kn = 0.7) tests. (2) Second, we test the limits of generalization by deploying this 1D-trained model directly into a 2D cavity flow, revealing that while it captures primary velocity structures, it fails on complex thermal fields, proving that the training data must be physically representative of the problem for accurate prediction of the higherorder moments. (3) Finally, we create a robust 2D model, training it on data from a lid velocity sweep (50 m/s to 600 m/s), and demonstrate its power through challenging extrapolation tests (e.g., 200 m/s and 800 m/s), where it accurately predicts hyper-velocity physics far outside its training distribution. This rigorous approach not only confirms a significant speedup (1.63x–1.73x) but also provides a powerful analysis that, via Amdahl’s Law, proves our optimization has hit its theoretical limit and has successfully shifted the computational bottleneck away from the solver to the particle-based moment-gathering operations. 2. Background of the Fokker Planck Model The governing equation for a rarefied gas is the Boltzmann equation, which provides a high-fidelity description of the velocity distribution function f(x,v, t)[24]. However, the binary collision integral Q(f, f)is a complex, non-linear, high-dimensional operator, making the direct solution of the Boltzmann equation computationally prohibitive for most practical engineering problems. The Direct Simulation Monte Carlo (DSMC) method, introduced by Bird, has become the benchmark computational tool for rarefied gas dynamics [1]. DSMC is a stochastic particle method that simulates the physics of the Boltzmann equation, offering excellent accuracy. Despite its success, DSMC suffers from a significant computational constraint: its spatial and temporal discretizations (∆x,∆t) must resolve the local mean free path and mean collision time. In the near-continuum flow regimes (i.e., at low local Knudsen numbers), this requirement leads to computational stiffness, where the simulation cost becomes excessively high. This "near-continuum bottleneck" of DSMC has motivated the development of alternative kinetic models that can efficiently bridge the gap between the rarefied and continuum regimes. Among the most promising alternatives are kinetic models based 6 on the Fokker-Planck (FP) approximation [4]. The FP approach replaces the discrete, jump-based Boltzmann collision operator with a continuous Markovian stochastic process, described by a drift (friction) vector and a diffusion (stochastic) tensor. The resulting particle method, governed by an equivalent Stochastic Differential Equation (SDE) or Langevin equation, is "collisionless" in a numerical sense. The particle evolution is decoupled from the microscopic collisional scales, allowing for time steps ∆tand cell sizes ∆xsignificantly larger than those required by DSMC. While computationally efficient, the simplest linear Fokker-Planck model (often called the Langevin model) possesses a critical physical deficiency. This model, which was foundational to the particle-FP algorithm, uses a single relaxation time scale derived from viscosity, which results in a fixed Prandtl number of Pr = 3/2[4]. This value is incorrect for a monatomic gas, where the physical value is Pr = 2/3. This discrepancy makes the linear model unsuitable for flows where heat transfer is significant, such as in high-Mach or non-isothermal problems. To address this "wrong Prandtl number" problem, higher-order, non-linear FokkerPlanck models have been developed. This line of research has primarily diverged into two main approaches. The first is the Cubic Fokker-Planck (cubic-FP) model, developed by Gorji et al. [9], which introduces a non-linear cubic drift term to provide the necessary degrees of freedom to match the heat flux relaxation rate, successfully recovering Pr = 2/3 [5]. The second is the Ellipsoidal Shakhov Fokker-Planck (ES-FP) model [25], which retains a linear drift but employs an anisotropic diffusion tensor based on the pressure tensor. Recent research has focused on the systematic comparison and validation of these advanced models [11,13,26]. These recent works have highlighted the fundamental tradeoff between these two approaches. The cubic-FP model demonstrates superior accuracy for transport properties in the near-continuum regime due to its correct Prandtl number. However, this comes at the cost of sacrificing a guaranteed H-theorem, which can lead to instabilities in far-from-equilibrium regions. Conversely, the ES-FP model preserves the entropy law (H-theorem), offering better stability in strong shock waves, but may be 7 less accurate in near-continuum heat transfer predictions [27]. Understanding this tradeoff is critical for the development of robust and accurate hybrid simulation methods for multi-scale rarefied gas flows. This line of research has continued to evolve, addressing the limitations of the original monatomic gas models. More recent work has focused on two key extensions: first, the development of kinetic-model-based FP algorithms for polyatomic gases (including diatomic gases), which account for rotational and vibrational energy modes [17,14]. Second, the formulation of high-order Fokker-Planck models that move beyond the cubic approximation to capture non-equilibrium phenomena with even greater fidelity [12]. These advanced models represent the current frontier in developing particle-based FP methods that are both computationally efficient and physically comprehensive. 3. The Fokker-Planck Approximation 3.1. Context: The Boltzmann Equation and its Computational Challenge For a dilute monatomic gas, the statistical state is fully described by the velocity distribution function F(V, X, t). Its evolution is governed by the Boltzmann equation, which in the absence of external forces is given by [1]: ∂F ∂t +Vi ∂F ∂xi =SBoltz(F)(1) where SBoltz(F)is the non-linear binary collision operator. While physically accurate across the entire range of Knudsen numbers (Kn), direct numerical solution of the Boltzmann equation is computationally prohibitive. The Direct Simulation Monte Carlo (DSMC) method has become the benchmark stochastic approach, simulating the evolution of computational particles that mimic the Boltzmann equation [1,3,2]. The "near-continuum issue" of DSMC, discussed in the previous section, is the primary motivation for developing alternative kinetic models that can efficiently bridge the gap between the hydrodynamic and rarefied regimes. The FokkerPlanck approximation of the collision operator provides such an alternative. 8 3.2. The General Fokker-Planck Operator The Fokker-Planck (FP) equation replaces the discrete, jump-based Boltzmann collision operator SBoltz with a continuous Markovian stochastic process, specifically a diffusive one. This transformation yields a significant computational advantage: the resulting stochastic paths of particles are continuous and independent, allowing for time steps and cell sizes much larger than the collisional scales [5]. The generic form of the Fokker-Planck collision operator is: SFP(F) = ∂F ∂t FP =−∂ ∂Vi (AiF) + 1 2 ∂2 ∂Vi∂Vj (DijF)(2) Here, Airepresents the drift vector, which describes the systematic friction-like force on a particle, and Dij is the positive-definite diffusion tensor, which describes the stochastic (random) forces. The central task of any FP kinetic model is to define Aiand Dij. These coefficients are chosen to ensure consistency between the FP operator and the Boltzmann operator for a set of relevant polynomial velocity moments Ψα: ZR3 SFP(F)Ψαd3V=ZR3 SBoltz(F)Ψαd3V(3) This moment-matching ensures that the model correctly reproduces the relaxation rates of the corresponding macroscopic quantities (e.g., mass, momentum, energy, stress, heat flux). 3.3. The Baseline: The Linear Drift (Langevin) Model and its Deficiency The simplest non-trivial FP model is the linear drift model, which corresponds to the classic Langevin equation. This model was employed by Jenny et al. and is defined by the following coefficients [4]: Ai=−1 τ(Vi−Ui)(4) Dij =θ τδij (5) 9 where Uiis the bulk velocity and θ=kT/m is the temperature (with kbeing the Boltzmann constant and mthe molecular mass). The single time scale τis set to recover the correct viscosity µin the hydrodynamic limit, given for a Maxwell-type interaction by: τ=2µ p(6) where pis the ideal gas pressure. This linear model is computationally efficient and rigorously admits an H-theorem for the entropy functional H(f) = Rflog fd3V, ensuring relaxation to the correct equilibrium. However, it possesses a critical physical deficiency: the use of a single time scale τfor both momentum relaxation (viscosity) and energy relaxation (thermal conduction) results in an incorrect, fixed Prandtl number of Pr = 3/2. For a dilute monatomic gas, the physical value is Pr = 2/3. This "wrong Prandtl number" makes the linear model unsuitable for accurate simulations of flows where heat transfer is significant, such as high-Mach flows or non-isothermal boundary layers. This deficiency is the direct motivation for the development of higherorder FP models, such as the cubic-FP model [5]. 3.4. Formulation of the Cubic Fokker-Planck (cubic-FP) Model The cubic-FP model, as developed by Gorji et al. , is a direct generalization of the linear model. It is designed specifically to correct the Prandtl number by introducing additional degrees of freedom into the drift coefficient, allowing for the independent matching of heat flux relaxation rates. The cubic-FP model begins by positing a more complex, non-linear form for the drift coefficient Ai, while retaining the simple, isotropic diffusion tensor Dij of the linear model. The drift force is expanded as a polynomial series in the peculiar velocity, v′ i=Vi−Ui, using a Hermite basis: Ai=c(0) i+c(1) ij v′ j+c(2) ijkv′ jv′ k+. . . (7) The coefficients c(n)are macroscopic quantities that depend on the moments of the dis- 16 Phase 0: Baseline Physics Sim Run highfidelity solver. Collect 16-feature inputs (X) and 9-feature outputs (Y). Training Data Phase 1: Offline Training) Load data. Normalize with StandardScaler. Train DNN (Keras/TF). Trained Model (fp_model.keras) & Scalers (scalers.pkl) Phase 2: Parameter Extraction Load model & scalers. Extract all weights (W), biases (b), means (µ), scales (σ). Raw Parameters Phase 3: GPU-Native Inference Load parameters to GPU. Run simulation using a pureCuPy forward pass. Final Results (1.7x Speedup) Figure 1: The four-phase pipeline for data generation, offline training, parameter extraction, and GPUnative deployment of the surrogate model. 4.2. Phase 1-2: Data Generation and DNN Training We first generated a training dataset from the 1D Couette flow problem, running 20 simulations with randomized boundary conditions. At each sampling step, we saved the inputs and outputs of the physics solver. •Input Features (X∈R16): We selected 16 low-order moments and properties: ρ, T, Ux, Uy, Uz, Pxx, Pxy, Pxz, Pyy, Pyz, Pzz, Qx, Qy, Qz, DM2, ν. 17 Algorithm 1 Baseline Physics Simulation Loop Initialize Npparticles on GPU (p_data) Initialize Ncgrid cells on GPU (grid_gpu,coeffs_gpu,linsys_gpu) for nt = 1 to Nsteps do Move_Particles_2D(p_data,DT) Apply_Boundary_Cavity(p_data) sort_and_calc_moments_FULL(p_data,grid_gpu)▷Calculates M1-M5 if nt > Nss then average_results(avg_grid_gpu,grid_gpu,nt,Nss) end if build_linear_systems(grid_gpu,linsys_gpu) solve_linear_systems(linsys_gpu,coeffs_gpu)▷Target Bottleneck evolve_velocities(p_data,grid_gpu,coeffs_gpu,DT) end for Transfer avg_grid_gpu to CPU for plotting. •Target Labels (Y∈R9): The 9 closure coefficients (Axx, . . . , Bz). This dataset was used to train a 4-layer MLP using Keras. The model fp_model.keras and scalers scaler_X.pkl,scaler_y.pkl were saved. 4.3. Phase 3: GPU-Native Deployment This is the most critical phase for performance. A naive implementation that calls model.predict() in the simulation loop is unacceptably slow due to CPU-GPU data transfer. We therefore developed a GPU-native deployment strategy. 1. Parameter Extraction: A one-time script (extract_params.py) was run to load the Keras/SKLearn files and dump all parameters—weights (W), biases (b), means (µ), and scales (σ)—into a single NumPy .npz file (model_params_for_cupy.npz). 2. Optimization (LITE Moments): We created a new moment function, ...LITE, which skips the expensive bincount operations for M3, M4, M5, λ, etc., and only computes the 16 features required by the ML model. 3. GPU-Native Solver: At the start of the new simulation, the .npz file is loaded and all parameters are transferred to GPU memory as CuPy arrays (e.g., W1_gpu, b1_gpu,X_mean_gpu). We then replace the physics solver (Steps 10-11) with a single function call, predict_coeffs_cupy_native, which executes the entire DNN forward pass on the GPU. 18 The resulting accelerated simulation loop is shown in Algorithm 2. Algorithm 2 Accelerated ML Simulation Loop Initialize Npparticles on GPU (p_data) Initialize Ncgrid cells on GPU (grid_gpu,coeffs_gpu) Load_ML_Params_to_GPU(model_params.npz)→ml_params_gpu for nt = 1 to Nsteps do Move_Particles_2D(p_data,DT) Apply_Boundary_Cavity(p_data) sort_and_calc_moments_LITE(p_data,grid_gpu)▷Calculates M1-M3 (density, velocity, stress, heat flux) required by DNN if nt > Nss then average_results(avg_grid_gpu,grid_gpu,nt,Nss) end if predict_coeffs_cupy_native(grid_gpu,coeffs_gpu,ml_params_gpu)▷New Solver evolve_velocities(p_data,grid_gpu,coeffs_gpu,DT) end for Transfer avg_grid_gpu to CPU for plotting. The performance of Algorithm 1(Physics) versus Algorithm 2(Fast ML) is the central comparison of this study. 5. Results and Discussion We validate our approach on two canonical flow problems. All tests were run on an NVIDIA A100-SXM4-80GB GPU. 5.1. Case Study 1: 1D Couette Flow The 1D Couette flow serves as the first validation case. The working gas for all simulations was Argon (mAr = 66.3×10−27 kg), modeled as a Maxwell molecule (ω= 1.0) with a reference viscosity µ0= 2.117 ×10−5Pa s and a specific heat ratio γ= 5/3. The baseline parameters, including an initial gas temperature (Tin) and wall temperature (Tw) of 273.15 K, were set to establish a baseline Knudsen number of Kn ≈0.15. Table 2details the simulation parameters for the 1D Couette flow benchmark. This case is used to compare the performance of the baseline physics solver against the GPUnative ML surrogate. 19 Table 2: Simulation specifications for the 1D Couette flow benchmark (Case 1). Parameter Value Description Physical Constants & Gas Properties Gas Argon Argon Mass (mAr)66.3×10−27 kg Boltzmann (kB)1.380 ×10−23 J K−1 Viscosity Model Maxwell (ω= 1.0) Ref. Viscosity (µ0)2.117 ×10−5Pa s Specific Heat Ratio (γ)5/3 Flow & Boundary Conditions Base Density (ρin)≈1.747 ×10−5kg m−3(Yields Kn ≈0.15) Wall Temp (Tw1, Tw2)273.15 K Inlet Temp (Tin)273.15 K Wall Velocity (Uw1)−50.0 m s−1(in Y-direction) Wall Velocity (Uw2)50.0 m s−1(in Y-direction) Domain & Discretization Domain Length (Lx)0.001 m Grid Cells (Nx) 100 Total Cells (NC) 100 Particle & Time Parameters Target Particles/Cell 30,000 Total Particles (NP) 3,000,000 Time Step (∆t)≈9.36 ×10−8s Simulation Run Total Steps 4,000 Steady State Steps (NTSS) 1,000 (Steps before averaging begins) 5.1.1. Accuracy Validation Figure 2compares the final time-averaged profiles for velocity, temperature, and density from both the physics baseline and the ML surrogate. The results demonstrate excellent agreement. The ML surrogate (red, dashed line) perfectly captures the linear velocity profile, the parabolic temperature profile (caused by viscous dissipation), and the corresponding U-shaped density profile. This confirms that the ML model, trained on randomized 1D data, has successfully learned the underlying closure physics. 20 Figure 2: Comparison of 1D Couette flow profiles. The high-fidelity Physics solver (blue, solid) and the GPU-Native ML surrogate (red, dashed) show excellent agreement. 21 5.1.2. Performance Benchmark Table 3details the wall-clock time for both simulations. The ML-driven simulation, by replacing the full moment calculation and linear solver with the "LITE" moment calculation and native ML inference, achieves a 1.56x speedup. This confirms that for this problem, the solver component (Tsolver = 262.12 −167.73 = 94.39 s) represents a significant 36% of the total computational cost. Table 3: Performance: 1D Couette Flow (4000 steps, Nc= 100,Np= 3M) Solver Solver Method Time (s) Speedup Physics Baseline ..._FULL +cp.linalg.solve 262.12 1.0x Fast ML ..._LITE +predict_native 167.73 1.56x 5.1.3. Extension to a Wide-Range of Knudsen Numbers In the next stage, the aim is to train our DNN over a wide range of Knudsen number for the one-dimensional planar Couette flow problem. Simulation Parameters. To ensure the surrogate model is robust across different flow regimes, we performed simulations at five distinct Knudsen (Kn) numbers. The baseline simulation (Kn ≈0.15) was modified by a set of factors, kn_factors = [0.01, 0.1, 0.5, 1.0, 2.0], resulting in a training dataset spanning the following Knudsen numbers: •Kn ≈0.0015 (Slip-flow regime) •Kn ≈0.015 (Slip-flow regime) •Kn ≈0.075 (Slip-flow regime) •Kn ≈0.15 (Transition regime) •Kn ≈0.3 (Transition regime) All simulations were conducted using a 1D grid with NC= 300 cells. To ensure low statistical noise, a high particle count of NP= 9,000,000 (corresponding to an average of 30,000 particles per cell) was used. 22 Convergence and Data Extraction. A critical challenge in rarefied gas dynamics, particularly at high Kn, is the long simulation time required to reach a steady state. Preliminary tests revealed that simulations at Kn > 0.5 required significantly more steps for the macroscopic profiles (temperature, density) to converge. To generate a clean, high-fidelity dataset, we employed a very long simulation run for all cases. Each simulation was run for a total of Nsteps = 40,000 iterations. We defined a steady-state convergence threshold of NTSS = 30,000 steps, ensuring that all data was collected only after the system had fully stabilized, even for the highest Kn cases. Post-Processing and Feature Engineering. Data was sampled for Nsteps −NT SS = 10,000 iterations for each of the 5 Kn simulations. At each of these steady-state steps, we performed a post-processing operation to extract the input features and their corresponding output labels for the neural network. Input Features (X): For each of the 300 cells, we computed a 16-dimensional feature vector of macroscopic properties. This vector includes density (ρ), temperature (T), the 3 velocity components (U), the 6 unique components of the pressure tensor (Π), the 3 components of the heat flux vector (q), the second-order moment (DM2), and the collision frequency (ν). Output Labels (y): Simultaneously, we saved the 9-dimensional solution vector (coefficients Aand B) computed by the full-physics FP linear system solver. This process resulted in a comprehensive training dataset of 15,000,000 samples (10,000 steps ×300 cells ×5 Kn). This dataset forms a direct mapping from the macroscopic flow state to the required microscopic collision coefficients. 5.1.4. Neural Network Architecture The surrogate model is a fully-connected, feed-forward Deep Neural Network (DNN), also known as a Multi-Layer Perceptron (MLP). The architecture was designed to be sufficiently deep and wide to capture the complex, non-linear relationships between the macroscopic moments and the FP coefficients. The network architecture is as follows: Input Layer: 16 neurons, corresponding to the 23 16 input features. Hidden Layers: A deep architecture consisting of four hidden layers, with each layer containing 256 neurons. Activation Function: The Rectified Linear Unit (ReLU) activation function is applied to the output of each hidden layer. Output Layer: 9 neurons, corresponding to the 9 output coefficients. Activation Function: A linear activation function is used for the output layer, as this is a regression problem. This architecture, shown in Figure 3, summarized as 16-256-256-256-256-9, contains a total of 204,041 trainable parameters, providing significant expressive power. This architecture was selected after a preliminary hyperparameter study to find an optimal balance between expressive power and computational cost. The mapping from the 16-dimensional moment space to the 9-dimensional coefficient space is highly nonlinear. A shallower network (e.g., 1-2 hidden layers) was found to be insufficient, failing to capture the complex inter-dependencies (underfitting). Conversely, a significantly wider (e.g., 1024 neurons) or deeper network introduced a risk of overfitting, increased offline training time, and—most importantly—would incur a higher computational cost during the online inference phase, diminishing the potential speedup. The 16-256-256-256-256-9 architecture, with four hidden layers, provides the necessary depth for hierarchical feature extraction to model the underlying physics, while the 256-neuron width offers sufficient capacity without becoming computationally burdensome for the GPU-native forward pass. As this configuration achieved the excellent validation loss (<2×10−4) and generalization performance shown in this paper, a more exhaustive hyperparameter search was deemed unnecessary. The selection of these 16 features is a cornerstone of our optimization strategy. This vector represents the complete set of low-order macroscopic moments (up to the 3rd order) and fundamental properties that define the local thermodynamic state of the gas in a cell: density (ρ), temperature (T), the velocity vector (Ui), the six unique components of the pressure/stress tensor (Pij), the three components of the heat flux vector (Qi), the second-order moment (DM2, which is inherently linked to temperature), and the collision frequency (ν). The central hypothesis of this work is that the 9 target closure coefficients (which, 24 in the full-physics solver, are non-linear functions of expensive-to-compute high-order moments like M3, M4, and M5) can be directly inferred, or "closed," from this fundamental 16-dimensional state vector. This choice is deliberate: by using only these low-order, readily-available moments as input, our surrogate model (the DNN) entirely bypasses the need to compute the high-order moments, which constituted the primary computational burden of the original physics-based solver. Input Layer 16 Features H.L. 1 256 Neurons H.L. 2 256 Neurons H.L. 3 256 Neurons H.L. 4 256 Neurons Output Layer 9 Coefficients ReLU AC ReLU AC ReLU AC Linear AC Figure 3: Schematic of the 16-256-256-256-256-9 feed-forward Deep Neural Network (DNN) architecture used as the surrogate model. "H.L." stands for Hidden Layer. "AC" stands for Activation. It is critical to distinguish this complex training architecture from the final "inference" model deployed within the simulation. The trained parameters (weights, biases, and normalization scalers) are extracted as simple numerical arrays (e.g., numpy or cupy arrays). The "network" implemented in the live Fokker-Planck code is therefore not a complex, backpropagation-enabled object from a heavy library like TensorFlow. Instead, it is a simple, forward-pass-only function composed of a series of highly-optimized, batched matrix multiplications (X @ W + b) and basic ReLU activation functions (cp.maximum(0, X)), all executed natively on the GPU via CuPy. This lightweight, native implementation is what makes the surrogate "infinitely fast" (as shown in the Amdahl’s Law analysis) and avoids all CPU-GPU communication overhead. 5.1.5. Training and Validation The 15-million-sample dataset was shuffled and split into a 90% training set (13.5 million samples) and a 10% validation set (1.5 million samples). Data Standardization. Prior to training, both the input features (X) and output labels (y) were standardized using the StandardScaler from Scikit-learn. The scaler was fit only on the training data, and the resulting mean and standard deviation were then applied to normalize both the training and validation sets. This ensures all inputs and outputs 25 are centered around zero with a standard deviation of one, which is essential for stable and efficient training. Training. The model was trained using the Adam optimizer with a learning rate of 1×10−4. The loss function selected was the Mean Squared Error (MSE), which is standard for regression tasks. The training was monitored using an EarlyStopping callback (monitoring val_loss with a patience of 10 epochs) and a ModelCheckpoint callback to save only the best-performing model. The model trained for 50 epochs before ‘EarlyStopping‘ terminated the run, having achieved a final validation loss of 1.8785 ×10−4at epoch 40. This extremely low validation loss indicates a highly accurate and well-generalized model. Offline Computational Cost. A valid consideration for this methodology is the one-time, upfront computational cost of the offline phase. This phase includes both data generation and model training. The data generation, which involved running five separate, highfidelity physics simulations for the 2D cavity case (at lid velocities of 50, 100, 200, 400, and 600 m/s) for 35,000 time steps each, required approximately 2.5 GPU-hours in total on the NVIDIA A100-SXM4-80GB GPU. The subsequent DNN training, which processed the resulting 13.5 million data samples for 50 epochs, was also completed in approximately one hour on the same hardware. This entire offline investment is amortized; it is performed only once. This modest, one-time cost of roughly 4 GPU-hours is negligible when contrasted with the cumulative and recurring online performance gains (e.g., the 1.63x-1.73x speedup) that are realized every time the lightweight, GPU-native surrogate solver is executed. 5.1.6. Results and Discussion To evaluate the true generalization capability of the trained DNN surrogate, we performed validation tests on two Knudsen numbers that were not included in the training data. The model was tested on: Kn = 0.05 (Interpolation test, as 0.015 <0.05 <0.075), Kn = 0.09 (Interpolation test, as 0.075 <0.09 <0.15), Kn = 0.7 (Extrapolation test, as 0.7>0.3). The second 32 5.1.7. Robustness to Flow Regime (Knudsen Sweep) To test the performance gains across different flow regimes, we performed a parameter sweep on the 1D Couette flow problem. The baseline Knudsen number (Kn) of ≈0.15 (transitional flow) was varied over three orders of magnitude, from the near-continuum regime (Kn ≈0.0015) to the mid-transitional regime (Kn ≈1.5). This was achieved by decreasing (or increasing) the base density (ρ), while all other parameters (geometry, Np=9M, Nc=300) were held constant. The objective was to determine if the computational cost of the physics solver is sensitive to the collision frequency ν(which is proportional to ρ), and if the ML model’s speedup is maintained in these different physical contexts. 5.1.8. Performance Analysis Table 4summarizes the performance of both solvers across the four tested Knudsen numbers. The results are highly insightful and confirm a key aspect of this computational model: Constant Physics Time: The runtime of the Physics Baseline solver remained nearly constant at ≈448 ±5seconds, regardless of the 1,000-fold change in Knudsen number. This conclusively demonstrates that the computational cost of the physics solver (calculating M3-M5 and solving 300×(9×9) systems) is not sensitive to the values of the fluid properties (like νor ρ). On the NVIDIA A100-SXM4-80GB GPU, this computebound task is dominated by the number of operations (a function of Nc), not their specific values. Constant ML Time: The runtime of the Fast ML solver also remained nearly constant (≈270 ±15 seconds). This is expected, as its runtime is dominated by the Tparticles_LITE component, which depends on the total number of particles (Np= 9M), not the physical properties. Consistent Speedup: As a result, the speedup achieved by the ML surrogate remained stable and significant across the entire 1,000-fold range of Knudsen numbers, consistently performing ≈1.63x to ≈1.73x faster than the physics solver. 33 Table 4: Performance Comparison for 1D Couette Flow Across Varying Knudsen Numbers (4000 steps, Nc=300, Np=9M) Knudsen No. (Approx.) TimePhysics (s) TimeML (s) Speedup Kn ≈0.0015 448.89 s 266.96 s 1.68x Kn ≈0.015 448.27 s 259.26 s 1.73x Kn ≈0.15 449.73 s 269.41 s 1.67x Kn ≈1.50 449.24 s 275.09 s 1.63x This test confirms that our ML surrogate is robust, generalizable to different flow regimes, and provides a consistent performance benefit regardless of the gas density. This is a critical finding, as it validates the model as a general-purpose replacement for the physics solver, not one tuned to a specific flow condition. 5.2. Case Study 2: 2D Lid-Driven Cavity To test the generality and scalability of our DNN model, we apply it to the 2D liddriven cavity problem at the same Knudsen number of Kn=0.15. This is a more complex flow featuring a large central vortex and significant gradients. 5.2.1. Model Generality Crucially, we did not retrain the model. We used the exact same model_params_for _cupy.npz file that was trained on 1D Couette flow data. This tests the hypothesis that the DNN learned the local, cell-level physics, which is independent of the global 1D or 2D geometry. Figure 13 plots 1D slices from the 2D domain, comparing the physics solver and the ML solver. The ML model (red, dashed line) again shows excellent agreement, correctly capturing the velocity and temperature profiles, including the non-trivial S-curve of the U-velocity profile. This strongly indicates that the ML surrogate is generalizable and not overfit to the 1D training geometry. 5.2.2. Performance and Scalability (Amdahl’s Law) We conducted two strong-scaling experiments by keeping the particle-per-cell count low and increasing the number of cells in the cavity geometry. The total runtime is 34 Ttotal =Tparticles +Tsolver, where Tparticles is the non-optimizable part (particle move, ...LITE moments) and Tsolver is the optimizable part (extra moments + linear solve). The "...LITE moments" routine refers to an optimized version of the moment-gathering function, which computes only the low-order moments (e.g., density, velocity, stress tensor) required as input for the DNN. This optimized routine entirely skips the computationally expensive calculation of the high-order moments (3rd, 4th, 5th) that were only necessary for the full physics-based solver. Table 5summarizes the results. 25x25 Test: With 625 cells, the T_solver component took 35.14 seconds, representing 41.4% of the total runtime. Our ML model eliminated this, achieving a 1.71x speedup. 100x100 Test: We increased the cell count 16-fold to 10,000, while keeping the particle count roughly constant. The ‘T_solver‘ component (34.37s) remained surprisingly constant, indicating the ‘cp.linalg.solve‘ operation is heavily parallelized and limited by kernel-launch overhead, not compute, on the NVIDIA A100-SXM4-80GB GPU. This second test is the most revealing. The non-optimizable portion (T_particles) was 47.35s, or 57.9% of the total physics time. According to Amdahl’s Law, the maximum possible speedup, even with an infinitely fast solver, is: Smax =1 (1 −P)=1 0.579 ≈1.727x where Pis the fraction of the code that is parallelizable (or in this case, optimizable). Our measured speedup was 1.73x. This remarkable result indicates that our GPUnative ML solver is, for all practical purposes, infinitely fast—its computational cost is so low that it is completely masked by the remaining T_particles time. We have successfully hit the theoretical performance limit defined by Amdahl’s Law for this optimization strategy. 5.2.3. Model Generality and the Impact of Training Data Distribution A key objective of this research was to determine if a surrogate model, trained on data from a simple 1D flow, could generalize to predict the physics of a more complex 2D flow. 35 Table 5: Strong Scaling Performance: 2D Cavity (5000 steps, Np≈630k) Grid (NC) T_PHYSICS T_ML T_SOLVER Solver % Speedup (Total) (Particles) (Physics – ML) (Optimizable) 25×25 (625) 84.96 s 49.82 s 35.14 s 41.4% 1.71× 100×100 (10,000) 81.72 s 47.35 s 34.37 s 42.1% 1.73× To test this, we first deployed our original model (trained on 1D Couette flow data) into the 2D cavity simulation and this time consider the entire flow field of the cavity. The results, shown in Figure 14, are insightful. The model demonstrated a surprising ability to capture the primary flow structures. The Speed Comparison and Density Comparison plots show that the ML surrogate’s predictions (red dashed lines) align remarkably well with the physics baseline (black solid lines), capturing the main vortex and density gradients. However, a failure was observed in the Temperature Comparison. While the general shape is similar, the ML-predicted contours (red dashed) slightly deviate from the physics baseline (black) in the lower half of the cavity. This demonstrates an extrapolation error. The 1D Couette training data, while physically accurate, did not contain the complex, multi-dimensional physics of a 2D recirculating flow, such as stagnation points and the associated complex thermal transport. The model was forced to predict physics it had never seen, and it slightly failed on the most complex field, i.e., temperature. 5.2.4. Validation with a 2D-Trained Model To correct this, we followed our methodology a second time: 1. We ran the 2D physics baseline solver (Algorithm 1) for the cavity problem and generated a new, dedicated training dataset (ml_training_data_cavity.npz). 2. We trained a new surrogate model (fp_model_cavity.keras) using this 2D-specific data. 3. We extracted its parameters into a new file (model_params_cavity_for_cupy.npz) and re-ran the 2D ML simulation. 36 The results from this correctly-trained model are shown in Figure 15. The improvement is dramatic and conclusive. The ML surrogate’s predictions (red dashed lines) now show perfect agreement with the physics baseline (black solid lines) across all three fields: Speed, Temperature, and Density. The previous discrepancy in the temperature field is completely resolved. This two-part experiment confirms a critical finding: while the local physics is general, a surrogate model’s predictive accuracy is fundamentally bound by the complexity and richness of its training data. For high-fidelity reproduction of complex, multi-dimensional flows, the model must be trained on data that adequately samples the physical phenomena of that regime. This result validates that our methodology is sound, and when provided with appropriate data, the GPU-native surrogate can successfully and accurately accelerate complex 2D flows. 5.2.5. Generalization Test: Extrapolation to Unseen Flow Conditions A critical measure of success for any machine learning surrogate model is not its ability to interpolate within its training data, but its capacity to generalize—or extrapolate—to conditions it has never encountered. To evaluate this capability, we performed a rigorous generalization test using the 2D driven cavity simulation. 5.2.6. Test Methodology The surrogate "Fast ML" model used in this test was the network trained exclusively on data from the 100 m/s lid velocity simulation. We then configured a new simulation case with the lid velocity doubled to 200 m/s, a regime outside the training distribution. This test is designed to answer a key question: has the neural network merely "memorized" the specific flow patterns of the 100 m/s case, or has it successfully learned the underlying, non-linear relationships between the 16 macroscopic input features (density, temperature, velocity, stresses, etc.) and the 9 required cubic-FP coefficients (Aij, Bi)? The "Physics" solver (full cubic-FP) was run for this new 200 m/s case to establish the high-fidelity "ground truth" solution. The "Fast ML" solver, using the 100 m/s-trained network, was run under the identical 200 m/s flow conditions. The resulting steady-state 37 contour fields are compared in Figure 16. The results of the extrapolation test, as described in Figure 16, are outstanding. The analysis confirmed a near-perfect match between the ground-truth "Physics" solver and the "Fast ML" prediction. The surrogate model, despite having no prior exposure to this flow regime, accurately captures all major flow features: The location and shape of the primary recirculation vortex, the complex thermal gradients, especially the hot region near the top-right corner and the cool region near the top-left, and the density field, including the low-density zone in the vortex core and the high-density accumulation in the corners. This result is highly significant. It provides strong evidence that the neural network has not simply performed a trivial curve-fitting exercise. Instead, it has successfully learned a high-dimensional, non-linear function that robustly approximates the core physics of the cubic-FP closure model. The model’s ability to extrapolate its learned relationships to a 100% increase in the primary driving velocity confirms its robustness and its value as a true surrogate for the physics solver. 5.2.7. Robust Training Dataset Generation for 2D Cavity over a Wide Range of Lid Speed To ensure the DNN surrogate for the 2D lid-driven cavity is robust and capable of handling a wide range of flow conditions, a composite training dataset was generated. Instead of training on a single simulation, we combined data from five independent, highfidelity physics-based simulations. The five simulation runs were differentiated by the lid velocity (Ulid), spanning a significant range to capture diverse flow physics: Ulid = 50 m s−1,Ulid = 100 m s−1,Ulid = 200 m s−1,Ulid = 400 m s−1,Ulid = 600 m s−1. Instantaneous flow-field data (inputs) and the corresponding solved closure coefficients (targets) were sampled from all five runs after they reached a steady state. This composite dataset, containing ml_training_data_ROBUST_cavity.npz, ensures the final surrogate model is exposed to a rich variety of flow regimes, from low-speed incompressible recirculation to high-speed flows with significant compressibility and viscous heating effects, 38 thereby enhancing its generalization and extrapolation capabilities. 5.2.8. Extrapolation Test: 800 m/s Lid Velocity To further probe the limits of the model’s robustness, we conducted a final, extrapolation test. The same DNN model trained on the 2D cavity dataset (which included velocities up to 200 m/s) was deployed without any modification to simulate a hypervelocity case at 800 m s−1. This represents a 4x extrapolation beyond the maximum velocity seen during training and introduces significantly more intense non-equilibrium, compressibility, and viscous heating effects. The results are shown in Figures 17a through 17c. Speed Field. As shown in Figure 17a, the ML surrogate’s prediction (dashed red line) for the speed field is remarkably accurate. It perfectly captures the location, size, and shape of the primary recirculation vortex, the secondary vortices in the bottom corners, and the intense shear layer near the top lid. The agreement with the full-physics solver (solid black line) is visually near-perfect, even in this extreme regime. Temperature Field. The temperature field (Figure 17b) provides the most stringent test, as the peak temperature now exceeds 670 K due to intense viscous dissipation and compression, far beyond the thermal range in the training data. Despite this, the ML model’s prediction is outstanding. It correctly identifies the location and magnitude of the hot spot in the top-right corner and accurately reproduces the complex, non-linear contour shapes throughout the domain. Density Field. Similarly, the density field comparison in Figure 17c shows excellent agreement. The surrogate model correctly predicts the low-density region associated with the primary vortex core and the significant high-density accumulation in the top-right corner. While minor deviations are visible in the sharpest gradient region, the overall structural integrity and quantitative accuracy are exceptional. This successful 4x extrapolation provides powerful evidence that the DNN has learned the underlying, generalized physics of the closure model, rather than simply curve-fitting its training data. 39 5.2.9. L2 Error Metric The prediction error magnitude is characterized by the Euclidean norm of the coefficient vectors: Estress =∥A∥2=v u u t 6 X j=1 A2 j(19) Eheat flux =∥B∥2=v u u t 3 X j=1 B2 j(20) Etotal =qE2 stress +E2 heat flux (21) The error contours identify regions where the model exhibits high confidence (low σ) and regions with large magnitude predictions (high E), guiding the application of physics-informed corrections or hybrid approaches. Figure 18 depicts the absolute L2 error for the extrapolation test case considered in the previous section. A critical observation is that the error is not distributed randomly but is almost entirely localized in the top-right corner of the domain. This is the most physically extreme region, characterized by the intense shear layer, high compression, and severe viscous heating generated by the 800 m s−1lid meeting the stationary wall. The model’s ability to remain stable and highly accurate across the other 95% of the domain, despite this being an extreme 4x extrapolation test (trained up to 200 m s−1), demonstrates significant robustness. The large absolute error values (e.g., 1.5×106) are expected, as they correspond to a small relative error of very large physical quantities (stresses) in this singular region. 5.2.10. Uncertainty Quantification and Error Metrics To quantify model uncertainty, we employ Monte Carlo Dropout by performing N stochastic forward passes through the neural network with dropout applied at inference time. For each spatial cell, the predictions are: 40 Y={y1, y2, . . . , yN}where yi=NN(x;maski)(22) where NN denotes the neural network and maskirepresents independent random dropout masks with rate pdropout = 0.2. The mean prediction is computed as: µ=1 N N X i=1 yi(23) The uncertainty (standard deviation) is defined as: σ=v u u t 1 N N X i=1 (yi−µ)2(24) The model outputs 9 components: stress tensor A= [A11, A12, A13, A22, A23, A33]and heat flux B= [Bx, By, Bz]. Plots shown in Fig. 19 show the model’s own predictive uncertainty, quantified by the standard deviation (Std Dev) of predictions from 50 forward passes with Monte Carlo Dropout enabled. High values indicate high model uncertainty. The most significant finding is the direct correlation between these maps and the L2 Error maps in Figure 18. The model reports its highest uncertainty (brightest colors) in the exact same top-right corner where the L2 error is largest. Conversely, it reports very high confidence (darkest colors) over the vast majority of the domain, where the L2 error is near-zero. This alignment is a critical validation, demonstrating that the model has not only learned the physics but has also learned the limits of its own knowledge, successfully identifying and flagging the regions of extreme physical extrapolation. 6. Conclusion In this paper, we successfully designed, trained, and deployed a GPU-native Deep Neural Network (DNN) surrogate to replace the computationally expensive moment closure solver in a Fokker-Planck particle simulation. Our methodology makes a clear distinction 41 between the offline training pipeline, which uses Keras/TensorFlow to learn the complex 16-input to 9-output mapping from a high-fidelity dataset, and the online inference solver. Crucially, the online solver does not load a heavy framework like TensorFlow; it is a simple, lightweight CuPy function that executes the network’s forward pass as a series of batched matrix multiplications and ReLU activations. This native-GPU implementation eliminates all CPU-GPU I/O overhead and is so computationally efficient that its cost is negligible, as proven by our performance analysis. Our key findings are three-fold: 1. Performance and Bottleneck Identification: The GPU-native surrogate achieved a consistent and significant speedup of 1.63x to 1.73x. A strong-scaling analysis on the 2D cavity case confirmed this speedup reaches the theoretical maximum predicted by Amdahl’s Law. This result is paramount, as it proves the DNN solver is, in effect, "infinitely fast" (zero-cost) relative to the remaining simulation tasks and that the true computational bottleneck has been decisively shifted from the solver to the particle-based sort_and_calc_moments (i.e., moment-gathering) routine. 2. Robustness in Interpolation and Extrapolation: The model demonstrated exceptional robustness. For the 1D Couette flow, a model trained on a Knudsen number sweep (Kn ≈0.0015 −0.3) showed outstanding accuracy in both interpolation (e.g., Kn = 0.09) and significant extrapolation (Kn = 0.7). For the 2D cavity flow, a model trained on a lid velocity sweep (up to 600 m/s) successfully predicted the complex, hyper-velocity physics of an 800 m/s case. This confirms the network learned the underlying physical relationships rather than memorizing specific flow patterns. 3. Methodology Generalizability and Data-Dependency: Our experiments confirmed that while the methodology (learning the local moment-to-coefficient mapping) is general, the resulting model is data-dependent. The test of the 1D-trained model on the 2D cavity (Fig. 14) failed to capture complex 2D thermal fields, whereas the model trained on 2D data (Fig. 15) achieved perfect agreement. This confirms that the training data must be representative of the multi-dimensional 48 (a) Stress Tensor L2 Error (b) Heat Flux L2 Error (c) Total L2 Error Figure 18: L2 (Absolute) Error Analysis for the 800 m s−1Extrapolation Test. These plots show the absolute L2 error between the Physics baseline and the ML surrogate for the (a) stress tensor, (b) heat flux, and (c) total combined coefficients. 49 (a) Stress Tensor Uncertainty (Std Dev) (b) Heat Flux Uncertainty (Std Dev) (c) Total Combined Uncertainty (Std Dev) Figure 19: Uncertainty Quantification (MC Dropout) for the 800 m s−1Extrapolation Test. 50 References [1] G. A. Bird. Molecular Gas Dynamics and the Direct Simulation of Gas Flows. Clarendon Press, Oxford, 1994. [2] Ehsan Roohi, Houman Akhlaghi, and Stefan Stefanov. Advances in Direct Simulation Monte Carlo: From Micro-scale to Rarefied Flow Phenomena. Springer, 2025. [3] G.A. Bird. The DSMC method. CreateSpace Independent Publishing Platform, USA, 2013. [4] Patrick Jenny, Manuel Torrilhon, and Stefan Heinz. A solution algorithm for the fluid dynamic equations based on a stochastic model for molecular motion. Journal of computational physics, 229(4):1077–1098, 2010. [5] Mohammad H Gorji, Maniel Torrilhon, and Patrick Jenny. Fokker–planck model for computational studies of monatomic rarefied gas flows. Journal of fluid mechanics, 680:574–601, 2011. [6] Vahid RezapourJaghargh, Amirmehran Mahdavi, and Ehsan Roohi. Shear-driven micro/nano flows simulation using fokker planck approach: Investigating accuracy and efficiency. Vacuum, 172:109065, 2020. [7] Amirmehran Mahdavi and Ehsan Roohi. A novel hybrid dsmc-fokker planck algorithm implemented to rarefied gas flows. Vacuum, 181:109736, 2020. [8] Amirmehran Mahdavi and Ehsan Roohi. A study on micro-step flow using a hybrid direct simulation monte carlo–fokker–planck approach. Physics of Fluids, 34(6), 2022. [9] M Hossein Gorji and Patrick Jenny. An efficient particle fokker–planck algorithm for rarefied gas flows. Journal of Computational Physics, 262:325–343, 2014. [10] M Hossein Gorji and Patrick Jenny. Fokker–planck–dsmc algorithm for simulations of rarefied gas flows. Journal of Computational Physics, 287:110–129, 2015. 51 [11] Sanghun Kim, Hossein Gorji, and Eunji Jun. Critical assessment of various particle fokker–planck models for monatomic rarefied gas flows. Physics of Fluids, 35(4), 2023. [12] Sanghun Kim, Woonghwi Park, and Eunji Jun. A second-order particle fokker-planck model for rarefied gas flows. Computer Physics Communications, 304:109323, 2024. [13] Sanghun Kim and Eunji Jun. A particle fokker–planck method for rarefied gas flows of monatomic mixtures. Physics of Fluids, 37(1), 2025. [14] Sanghun Kim and Eunji Jun. A stochastic particle method based on the fokker– planck master equation for rarefied gas flows of diatomic mixtures. Physics of Fluids, 37(3), 2025. [15] Baoyan Sun. Exponential stability for the kinetic ellipsoidal fokker–planck equation in weighted sobolev spaces. Mathematical Methods in the Applied Sciences, 2025. [16] Fei Fei, Zhaohui Liu, Jun Zhang, and Chuguang Zheng. A particle fokker-planck algorithm with multiscale temporal discretization for rarefied and continuum gas flows. Communications in Computational Physics, 22(2):338–374, 2017. [17] Ziqi Cui, Kaikai Feng, Qihan Ma, and Jun Zhang. A multiscale stochastic particle method based on the fokker-planck model for nonequilibrium gas flows. Journal of Computational Physics, 520:113458, 2025. [18] Ehsan Roohi and Amir Shoja-sani. Data-driven surrogate modeling of dsmc solutions using deep neural networks. Aerospace Science and Technology, 168(A):110785, 2024. [19] Ali Peyvan, Vismay Oommen, Ameya D Jagtap, and George E Karniadakis. Riemannonets: Interpretable neural operators for Riemann problems. Computer Methods in Applied Mechanics and Engineering, 426:116996, 2024. [20] Ali Peyvan, Vivek Kumar, and George E Karniadakis. Fusion-deeponet: A dataefficient neural operator for geometry-dependent hypersonic and supersonic flows. Journal of Computational Physics, 544:114432, 2026. 52 [21] Ehsan Roohi and Amirmehran Mahdavi. Shock-aware physics-guided fusiondeeponet operator for rarefied micro-nozzle flows. arXiv preprint arXiv:2510.17887, 2025. [22] Ehsan Roohi and Amirmehran Mahdavi. Analysis of the rarefied flow at micro-step using a deeponet surrogate model with a physics-guided zonal loss function. arXiv preprint arXiv:2509.17254, 2025. [23] Ehsan Roohi, Amir Shoja-Sani, B. Goshayeshi, and A. Peyvan. Learning rarefied gas dynamics with physics-enforced neural networks. arXiv preprint arXiv:2509.06231, 2025. [24] Carlo Cercignani. Rarefied gas dynamics: from basic concepts to actual calculations, volume 21. Cambridge university press, 2000. [25] Julien Mathiaud and Luc Mieussens. A fokker–planck model of the boltzmann equation with correct prandtl number. Journal of Statistical Physics, 162(2):397–414, 2016. [26] Fei Fei, Haihong Liu, Zhaohui Liu, and Jun Zhang. A benchmark study of kinetic models for shock waves. AIAA Journal, 58(6):2596–2608, 2020. [27] Eunji Jun, Marcel Pfeiffer, Luc Mieussens, and M Hossein Gorji. Comparative study between cubic and ellipsoidal fokker–planck kinetic models. AIAA Journal, 57(6):2524–2533, 2019.