Efficient Neural Surrogates for Linear Radiation Transport on the Lattice and Hohlraum benchmarks

Carmelo Gonzales; Steffen Schotthöfer; Cory D. Hauck

cs.AIOct 3, 2026

Efficient Neural Surrogates for Linear Radiation Transport on the Lattice and Hohlraum benchmarks

Authors: Carmelo Gonzales, Steffen Schotthöfer, Cory D. Hauck

Organizations: NVIDIA, Santa Clara, US · Seamless Labs, Mountain View, US · Oak Ridge National Laboratory

Abstract

Linear radiation transport equations (RTEs) form the simulation foundations underpinning design and analysis tasks in nuclear engineering, inertial confinement fusion, medical imaging, and astrophysics, but resolving the high-dimensional phase space at engineering fidelity remains expensive enough that outer-loop workflows, such as design optimization, uncertainty quantification, and parameter sweeps, are routinely budget-bound on traditional solvers. Neural surrogates promise to relax this bottleneck by amortizing simulation cost across thousands of downstream queries, but the architectural choices and engineered inductive biases that make a surrogate accurate on one transport problem do not transfer straightforwardly across model families. We benchmark two parameter-matched neural surrogate architectures, the physics-attention Transolver and the multi-scale graph network Bi-Stride Multi-Scale MeshGraphNet (BSMS-MGN), as end-to-end approximations of the final-time particle concentration for the two-dimensional linear RTE on the canonical Lattice and Hohlraum benchmarks. An ablation across Fourier features and region-weighted training loss exposes strongly architecture-dependent inductive-bias preferences, indicating that design choices common to physics-informed surrogate workflows must be revisited per architecture rather than imported across model families, and that downstream utility depends on per-QoI sensitivity rather than a single field-level score. The model training recipe, training data, and evaluation pipeline are released alongside this paper to support reproduction, transfer to related transport problems, and evaluation as amortized forward-model components in larger outer-loop simulation workflows.

Figures & tables

Panel of Figure 1 — Figure 1: The Lattice benchmark: (a) geometric layout; (b) material parameters; (c) particle concentration $\Phi$ from a high-resolution simulation, using the method described in Section 4 . The default values are $\sigma_{\mathrm{s}}^{W}=1$ for the scattering cross section in the white material and $\sigma_{\mathrm{a}}^{B}=10$ for the absorption cross section in the blue material.

Panel of Figure 2 — Figure 2: The Hohlraum benchmark: (a) geometric layout The labels in (a) are the eight design parameters of Section 3.2.2 : $\mathrm{ulr},\mathrm{llr}$ and $\mathrm{urr},\mathrm{lrr}$ are the $y$ coordinates of the upper and lower edges of the left and right red strips ( DP-1 ); $\mathrm{hlr}$ and $\mathrm{hrr}$ are the distances of the interior edges of the left and right red strips from the vertical symmetry axis, i.e., the interior edges lie at $x=-\mathrm{hlr}$ and $x=\mathrm{hrr}$ ( DP-2 ); and $(c_{x},c_{y})$ is the center of the capsule ( DP-3 ). The default configuration is $\mathrm{ulr}=\mathrm{urr}=0.4$ , $\mathrm{llr}=\mathrm{lrr}=-0.4$ , $\mathrm{hlr}=\mathrm{hrr}=0.6$ , and $(c_{x},c_{y})=(0,0)$ .; (b) material parameters; (c) particle concentration $\Phi$ from a high-resolution simulation, using the method described in Section 4 .

Resolution	Geometry	char. length	Min Cells	Max Cells	quad order	quadrature points	time steps	max. degrees of freedom
High-res	Lattice	0.015	79,862	79,862	6	144	978	11,500,128
High-res	Hohlraum	0.0075	69,897	81,775	6	144	1728	11,775,600
Low-res	Lattice	0.05	7,364	7,364	4	64	176	471,296
Low-res	Hohlraum	0.02	10,196	11,958	6	144	662	1,721,952

Table 1: Mesh resolution across datasets. The choices for feasible low-resolution discretizations are inferred from [ 51 ] . The characteristic length is the target cell edge length passed to the mesh generator; the realized cell size follows from the cell count and the domain area.

	Lattice	Hohlraum
Train	494	592
Validation	106	126
Test	107	128

Table 2: Dataset distribution for training and evaluation.

Panel of Figure 3 — Figure 3: Transolver training and validation loss curves (see Equation 11 ) on the training and validation split of Table 2 . Solid lines show the median across three random seeds and shaded bands show the seed-wise min-max envelope; we select the lowest-validation-loss checkpoint per seed for evaluation. The Lattice panel shows the baseline configuration, the only configuration trained on Lattice. The Hohlraum panel overlays all four configurations of the Hohlraum ablation validation loss curves ( Section 6.3 ): Base, +RW, +FF, and +FF+RW.

Panel of Figure 4 — Figure 4: BSMS-MGN training and validation loss curves (see Equation 11 ) on the training and validation split of Table 2 . Solid lines show the median across three random seeds and shaded bands show the seed-wise min-max envelope; we select the lowest-validation-loss checkpoint per seed for evaluation. The Lattice panel shows the baseline configuration, the only one trained on Lattice. The Hohlraum panel overlays all four configurations of the Hohlraum ablation validation loss: Base, +RW, +FF, and +FF+RW.

Panel of Figure 5 — Figure 5: Lattice benchmark, high-resolution evaluation, baseline configuration for both architectures (no Fourier features, no region-weighted loss), representative test sample with design parameters $\sigma_{\mathrm{a}}^{B}=40.0$ , $\sigma_{\mathrm{s}}^{W}=1.1$ . Per row: KiT-RT reference solution (left), surrogate prediction (middle), and absolute error (right). The particle-concentration field is shown in log scale. We observe that both surrogates reproduce the qualitative features of the benchmark, with absolute errors that are visually concentrated along the strongest gradients between absorbing and streaming subdomains. The reference is computed with the high-resolution discretization of Table 1 : characteristic mesh length $0.015$ ( $79{,}862$ cells), $144$ ordinates, and $978$ time steps. The streak-like pattern that radiates from the source through the scattering medium is a ray effect of the discrete-ordinates angular discretization [ 51 ] .

Method	Speedup vs high-res KiT-RT	Time (s)	RMSE $\Phi$	$E_{\textup{QoI},\textup{rel}}$ (%)
KiT-RT (high-res)	$1.00\times$	$13.62\pm 0.455$	$-$	$0.00$
KiT-RT (low-res)	$59.2\times$	$0.23\pm 0.007$	$-$	$6.60$
Transolver	$197\times$	$0.069\pm 0.0007$	$\bm{0.0828{\pm}5.20{\times}10^{-3}}$	$1.06{\pm}0.0620$
BSMS-MGN	$90.4{\times}$	$0.1506\pm 0.0045$	$0.0835{\pm}2.90{\times}10^{-5}$	$\bm{0.276{\pm}0.210}$

Table 3: Lattice benchmark, high-resolution surrogate evaluation (baseline configuration - no Fourier features, no region-weighted loss). Columns report the speedup over the high-resolution KiT-RT solver, per-sample inference time, field RMSE ( Equation 26 ), and the relative error in the absorption QoI (

E_{\textup{QoI},\textup{rel}}

, Equation 27 ) at

t=T

. The high-resolution KiT-RT reference and the cheap low-resolution KiT-RT baseline are shown for comparison. Surrogate rows report mean

\pm

1 standard deviation across

n=3

random seeds. The best surrogate mean per column is shown in bold. Full high- and low-resolution results are given in .

Panel of Figure 6 — Figure 6: Hohlraum benchmark with high-resolution evaluation, full (+FF+RW) configuration for both architectures (Fourier features and region-weighted MSE loss enabled), sample with design parameters $\mathrm{ulr}=0.3$ , $\mathrm{llr}=-0.4$ , $\mathrm{urr}=0.5$ , $\mathrm{lrr}=-0.4$ , $\mathrm{hlr}=0.5$ , $\mathrm{hrr}=0.6$ , $c_{x}=0.1$ , $c_{y}=0.0$ . Per row: KiT-RT reference solution (left), surrogate prediction (middle), and absolute error (right). Both neural surrogates reproduce the main qualitative features of the high-resolution KiT-RT solution, including the inflow-driven radiation field, attenuation through the red side shields, and absorption around the central capsule. The Transolver full-configuration prediction is visibly closer to the high-resolution reference, while the BSMS-MGN full-configuration prediction captures the streaming pattern correctly but exhibits noticeably larger field errors, especially above and below the central absorbing capsule.

Method	Speedup vs high-res KiT-RT	Time (s)	RMSE $\Phi$	$E_{\textup{QoI},\textup{rel}}$ Center (%)	$E_{\textup{QoI},\textup{rel}}$ Vert. (%)	$E_{\textup{QoI},\textup{rel}}$ Horiz. (%)
KiT-RT (high-res)	$1{\times}$	$35.95\pm 0.495$	$-$	$0.00$	$0.00$	$0.00$
KiT-RT (low-res)	$18.5{\times}$	$1.94\pm 0.043$	$-$	$10.6$	$8.05$	$4.00$
Transolver Base	$637{\times}$	$0.05645\pm 0.0067$	$0.0559{\pm}1.20{\times}10^{-3}$	$0.860{\pm}0.0660$	$0.551{\pm}0.0420$	$1.17{\pm}0.190$
Transolver +RW	$618{\times}$	$0.05815\pm 0.0055$	$0.0517{\pm}2.20{\times}10^{-4}$	$0.959{\pm}0.0570$	$0.684{\pm}0.0720$	$1.37{\pm}0.130$
Transolver +FF	$639{\times}$	$0.05625\pm 0.0085$	$\bm{0.0244{\pm}2.80{\times}10^{-4}}$	$\bm{0.624{\pm}0.0380}$	$\bm{0.280{\pm}0.0130}$	$\bm{0.422{\pm}5.60{\times}10^{-3}}$
Transolver +FF+RW	$637{\times}$	$0.05645\pm 0.0058$	$0.0252{\pm}3.70{\times}10^{-4}$	$0.649{\pm}0.0660$	$0.335{\pm}9.10{\times}10^{-3}$	$0.672{\pm}0.0450$

Table 4: Hohlraum benchmark, high-resolution surrogate evaluation across the four ablation configurations (Base, +RW, +FF, +FF+RW; defined in Table 5 ). Columns report the speedup over the high-resolution KiT-RT solver, per-sample inference time, field RMSE ( Equation 26 ), and the relative errors in the three per-region absorption QoIs ( Equation 27 ): center

S=G\cup B

, vertical-strip shield

S=R

, and horizontal-strip shield

S=K

. The high-resolution KiT-RT reference and the cheap low-resolution KiT-RT baseline are shown for comparison; surrogate rows report mean and standard deviation across

n=3

random seeds. The best surrogate mean per column within each architecture is shown in bold. Full high- and low-resolution results are given in .

Configuration	Fourier features (FF)	Region-weighted loss (RW)
Base	$-$	$-$
+FF	✓	$-$
+RW	$-$	✓
+FF+RW	✓	✓

Table 5: The four Hohlraum ablation configurations, obtained by toggling two engineered factors: Fourier features (FF) and region-weighted MSE loss (RW). A checkmark denotes that the factor is enabled.

Model	Benchmark	Parameters	Wall time (h)	GPU-hours	s/epoch
BSMS-MGN	Lattice	5,842,669	$2.06$	$16.5$	$14.8$
BSMS-MGN	Hohlraum	5,842,513	$3.30$	$26.4$	$23.7$
Transolver	Lattice	5,951,105	$1.57$	$12.5$	$11.3$
Transolver	Hohlraum	5,950,593	$1.86$	$14.9$	$13.4$

Table 6: Training cost for the Base configuration of each benchmark and surrogate model on a single DGX H100 node with 8 H100 80 GB GPUs. Transolver performs global attention over a much smaller latent representation (physics tokens), while BSMS-MGN performs explicit message passing on the cell-dual graph at the finest level and an additional bi-stride U-Net pass over the multi-level hierarchy, leading to longer per-epoch wall times. GPU-hours equal wall time multiplied by the 8 GPUs used per task. Parameter counts and wall times are reported for the Base configuration. The other Hohlraum ablation configurations differ in trainable parameter count by less than

0.1\%

and in wall time by less than

1\%

(see and ).

Figure 7: Amortization time and amortized speedup of the surrogate models against the high-resolution KiT-RT solver. The Transolver model requires 917 evaluations on Lattice and 781 evaluations on Hohlraum to amortize training and data-generation cost. BSMS-MGN requires approximately 1,050 and 930 evaluations, respectively; the difference to Transolver reflects its longer training and inference times.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Meaning
$X$	Physical material domain in two spatial dimensions.
$\partial X$	Boundary of the physical domain.
$\mathbf{x}=(x,y)$	Spatial coordinate.
$\mathbf{n}(\mathbf{x})$	Outward unit normal on $\partial X$ .
$\mathbb{S}^{2}$	Unit sphere in angular space.
$\mathbb{S}_{+}^{2}$	Upper half of the unit sphere in angular space.

Appendix

Table 7: Notation for the continuous radiation transport equation.

Symbol	Meaning
$h$	Maximum mesh-cell diameter.
$N_{h}$	Number of finite-volume mesh cells.
$i$	Mesh-cell index.
$\mathbf{x}_{i}$	Centroid of mesh cell $i$ .
$\Phi_{i}$	Reference particle concentration at cell $i$ and final time.
$N_{q}$	Tessellation quadrature order.

Appendix

Table 8: Notation for the finite-volume and discrete-ordinates reference solver.

Symbol	Meaning
$\mathcal{F}_{\theta}$	Learned surrogate operator mapping cell-wise inputs to final-time particle concentration.
$\theta$	Trainable model parameters.
$\bar{\mathbf{x}}_{i}$	Normalized cell-centroid coordinate used as a model input.
$\mathbf{p}_{i}$	Physical input vector for spatial cell $i$ , $\mathbf{p}_{i}=(\sigma_{\textup{a},i},\sigma_{\textup{s},i},\sigma_{\textup{t},i},Q_{i})$ . $Q_{i}$ is excluded when the source is everywhere zero.
$\mathbf{f}_{i}$	Fourier features of cell $i$ , $\mathbf{f}_{i}=\mathbf{f}(\bar{\mathbf{x}}_{i})$ , ( 15 ); optional input, ablation factor FF.
$\mathbf{f}(\bar{\mathbf{x}})$	Fourier feature map ( 15 ), applied to a normalized coordinate $\bar{\mathbf{x}}\in[-1,1]^{2}$ .

Appendix

Table 9: Notation for the generic surrogate-model formulation and training loss.

Symbol	Meaning
$E_{\mathrm{in}}$	Pointwise input embedding from raw cell features to hidden states.
$E_{\mathrm{out}}$	Pointwise output embedding from hidden state to particle-concentration prediction.
$\mathbf{h}_{i}^{\ell}$	Hidden feature of cell $i$ after Transolver layer $\ell$ .
$\mathbf{H}^{\ell}$	Matrix of the $N_{h}$ hidden features after layer $\ell$ , $\mathbf{H}^{\ell}=[\mathbf{h}_{1}^{\ell},\ldots,\mathbf{h}_{N_{h}}^{\ell}]^{T}$ .
$L$	Number of Transolver layers.
$C$	Hidden channel dimension.

Appendix

Table 10: Notation for the Transolver architecture. Multi-indices are comma-separated (

i

cell,

a

head,

j

token,

\ell

layer) and every description names the indices it depends on.

Symbol	Meaning
$l$	Mesh-level index in the bi-stride hierarchy, $l\in\{0,\ldots,L_{g}\}$ , with $l=0$ the finest and $l=L_{g}$ the coarsest.
$L_{g}$	Number of bi-stride coarsening levels, set to $L_{g}=5$ in the final models (distinct from the Transolver layer count $L$ ).
$\mathcal{G}^{(l)}=(\mathcal{V}^{(l)},\mathcal{E}^{(l)})$	Directed graph at level $l$ of the bi-stride hierarchy; $\mathcal{G}^{(0)}$ is the cell-dual graph of the finite-volume mesh.
$\mathcal{V}^{(l)}$	Set of graph nodes at level $l$ .
$\mathcal{E}^{(l)}$	Set of directed graph edges at level $l$ .
$(i,j)$	Directed edge from source node $i$ to destination node $j$ .

Appendix

Table 11: Notation for the Bi-Stride Multi-Scale MeshGraphNet (BSMS-MGN) architecture.

Explore similar work

When a neural surrogate cannot accelerate a solver: runtime share, closed-loop drift, and the economics of uncertainty gating in a stiff coupled simulation
Aug 24, 2026L. Thümmler, T. KurodaNeural Surrogate Modeling PDE Surrogate Modeling
Department of Mathematics, ETH Zürich, Rämistrasse 101, 8092 Zürich, Switzerland · Max Planck Institute for Gravitational Physics (Albert Einstein Institute), Am Mühlenberg 1, D-14476 Potsdam, Germany
Read, Write, Relax: Why Neural PDE Surrogates Need Both Global and Local Processing
Aug 21, 2026Anuj Kumar, Heiko Zimmermann, Josiah Bjorgaard +4PDE Surrogate Modeling Neural PDE Solvers
Pasteur Labs, Brooklyn, NY, USA · Cornell University, Ithaca, NY, USA · Institute for Simulation Intelligence, New York, NY, USA
PTNO: Training Neural Operators with Noisy Monte Carlo Estimates for Particle Transport Problems
Sep 30, 2026Yubo Cao, Xi Deng, Mengqi Xia +3Neural Surrogate Modeling Noisy-Label Learning
California Institute of Technology · Yale University · UK Atomic Energy Authority +1

Aug 24, 2026J/K move · Enter open · S save

When a neural surrogate cannot accelerate a solver: runtime share, closed-loop drift, and the economics of uncertainty gating in a stiff coupled simulation

L. Thümmler, T. Kuroda

Department of Mathematics, ETH Zürich, Rämistrasse 101, 8092 Zürich, Switzerland · Max Planck Institute for Gravitational Physics (Albert Einstein Institute), Am Mühlenberg 1, D-14476 Potsdam, Germany

Neural Surrogate Modeling PDE Surrogate Modeling Physics-Based Simulation

Open paper AskPDF arXiv

Learned surrogates for expensive inner solver blocks are a widely pursued route to faster multiphysics simulation. We report a controlled, end-to-end negative result and identify three structural barriers, none of them a deficiency of the network we trained. The testbed is the implicit Newton solve coupling energy-dependent neutrino radiation to matter in a general-relativistic radiation-hydrodynamics code, its most expensive physics routine per call. First, per-call cost and share of runtime are different quantities, and only the second bounds acceleration. An exclusive self-time profile puts the target block at 16.9% of critical-rank wall clock, capping any surrogate at ~1.2x by Amdahl's law. A surrogate 5.8x cheaper per call merely ties the solver, and the configuration stable enough to run without fallback reaches only parity. Second, offline accuracy cannot rank surrogates for deployment: across fourteen networks the pooled Spearman error-versus-survival correlation (rho = +0.73) is a between-family confound that vanishes under control (rho = -0.04). Third, a correct out-of-distribution gate cannot accelerate a loop that leaves its training distribution. We give the break-even deferral fraction in closed form: because the visited states sit 73x off the data manifold, the gate defers 96.8 to 99.7% of cells, almost invariant to surrogate quality. Including its own cost, the gated loop is a 0.94 to 0.96x slowdown. We further separate stability from fidelity: a never-crashing gated run accumulates a linear -19.9% density bias over 6000 steps. The error is a directed, ballistically accumulating bias, not the variance-driven divergence the autoregressive literature targets.

Abstract

Explore similar work

When a neural surrogate cannot accelerate a solver: runtime share, closed-loop drift, and the economics of uncertainty gating in a stiff coupled simulation

Read, Write, Relax: Why Neural PDE Surrogates Need Both Global and Local Processing

PTNO: Training Neural Operators with Noisy Monte Carlo Estimates for Particle Transport Problems