PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2610.04665v1cs.AIhep-lat

TL;DR

High-fidelity radiation-transport simulations are costly because they resolve high-dimensional phase space and multiple scales. This paper benchmarks two neural surrogates on Lattice and Hohlraum and finds architecture-dependent preferences across its ablations. The results support their use as amortized forward models when mesh-transfer limits and feature choices are considered.

  • Problem

    High-fidelity RTE simulation is computationally costly because accurate solutions require resolving high-dimensional phase space and transport across broad spatial and temporal scales.

    Source: §1 · pp. 1–2

  • Method

    The paper benchmarks Transolver and BSMS-MGN as approximations of final-time particle concentration on Lattice and Hohlraum, using high-resolution simulation data.

    Source: §1 · p. 2, §5.1 · p. 7, §5.4.1 · p. 11

  • Results

    At matched parameter counts, Transolver outperforms BSMS-MGN on Hohlraum center-capsule QoI with Fourier features, while BSMS-MGN benefits from region weighting and is competitive on Lattice field RMSE.

    Source: §7 · p. 20

  • Takeaways & Limitations

    The results support trained neural surrogates as amortized forward models in outer-loop transport workflows when architecture-specific feature choices and mesh-transfer limitations are accounted for.

    Source: §7 · p. 20

    Low-resolution surrogate evaluation is competitive in inference time but not accuracy, and the study focuses on final-time concentration rather than all-time or kinetic-level quantities.

    Source: §3.1.1 · p. 4, §6.2 · p. 16

Abstract

from arXiv · show

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.

1 Introduction

High-fidelity RTE simulations are costly because they resolve high-dimensional phase space and transport across multiple scales. This work benchmarks neural surrogates on two transport problems and examines their accuracy and cost for downstream queries.

1 Introduction

  • The RTE's high-dimensional phase space and broad spatial and temporal scales make accurate grid-based simulations computationally expensive, while Monte Carlo methods can be noisy.The equation underpins high-fidelity particle-transport modeling across several application areas.Source: §1 · pp. 1–2
  • The study builds on neural operator learning, comparing a transformer-based Transolver with the hierarchical multi-scale graph network BSMS-MGN.Both models are trained and tested using PhysicsNeMo.Source: §1 · p. 2
  • The paper benchmarks Transolver and BSMS-MGN against a discrete-ordinates/finite-volume method on the Lattice and Hohlraum benchmarks, which include design parameters and quantities of interest.The Lattice is a toy reactor-assembly mock-up, and Hohlraum is a linearized, symmetric variant of the original problem.Source: §1 · p. 2
  • Using high-resolution KiT-RT data, the surrogates map material properties to final-time particle concentration and outperform the coarse simulation in almost all cases.The paper also evaluates how many input samples are needed to recoup training and data-generation costs in outer-loop workflows.Source: §1 · p. 2
  • Training, inference, reference-simulation code, and pre-generated data are publicly available.The passages identify the code and dataset repositories.Source: §1 · p. 2

2 The radiation transport equation

The paper formulates particle transport in a two-dimensional spatial domain with angular and temporal dependence. Its model assumes unit-speed particles, isotropic scattering, and a known isotropic source in a static medium.

2 The radiation transport equation

  • The benchmarks treat their geometries as extrusions into a third direction, reducing the usual three-dimensional spatial domain to two dimensions.The text refers readers to and for details of this reduction.Source: §2 · p. 3
  • The equation defines inflow and outflow on the boundary by the sign of the outward normal's dot product with the travel direction.The initial particle density and inflow boundary data are given.Source: §2 · p. 3
  • Particles travel at unit speed through a static material, with scattering and absorption rates given by spatially varying cross-sections.Scattering is assumed isotropic, and the particle source is known and isotropic.Source: §2 · p. 3
  • Particle concentration is the particle density at position x with respect to spatial measure at time t.The paper denotes the phase-space particle density by Ψ and the concentration by Φ.Source: §2 · p. 3

3 Benchmark problems

The section specifies the Lattice and Hohlraum transport benchmarks, their quantities of interest, and the design parameters used to vary their configurations.

3.1 Lattice benchmark

  • The Lattice benchmark models a source emitting particles into checkerboard-arranged absorbing blue regions, a red source region, and white void quadrants.The source is initially the only particle source; the default blue-material absorption cross-section is σW_a = 10.Source: §3.1 · p. 3, Figure 1 · p. 4

3.1.1 Quantities of interest

  • The Lattice quantity of interest is absorption in a subdomain at final time, while other quantities requiring all time steps or kinetic-level Ψ are out of scope.The subdomain absorption functional is also used for Hohlraum quantities of interest and the loss and error metrics.Source: §3.1.1 · p. 4

3.1.2 Design Parameters

  • Lattice design parameters vary the benchmark configuration, with the full Cartesian product of uniformly spaced values yielding 800 configurations.Each configuration is simulated once and assigned to training, test, or validation splits.Source: §3.1.2 · p. 4

3.2 Hohlraum benchmark

  • Hohlraum is a linearized, symmetric variant of a nonlinear thermal radiative-transfer problem, with shielding around a capsule irradiated from both sides.The linearization assumes no coupling between radiation and material; the domain includes red and black shielding, a green capsule hull, a blue checkerboard interior, and white void.Source: §3.2 · pp. 4–5

3.2.1 Quantities of interest

  • The Hohlraum quantities of interest measure final-time absorption in the capsule, red shielding, and black shielding regions.The capsule region is G ∪ B; the other regions are R and K.Source: §3.2.1 · p. 5

3.2.2 Design parameters

  • Hohlraum design parameters vary red-strip edges and thicknesses and displace the capsule center, spanning up to eight parameters.The parameters are treated as random variables when an outer-loop workflow places a probability distribution on them.Source: §3.2.2 · pp. 5–6

4 Numerical discretization and training data generation

The authors generate high- and low-resolution KiT-RT data for training and evaluation, with benchmark-specific discretizations and parameterized sample sets.

4 Numerical discretization and training data generation

  • KiT-RT uses discrete ordinates in angle, a second-order finite-volume scheme on an unstructured triangular mesh in space, and explicit second-order Runge-Kutta time stepping.The CFL number is 0.7 for Lattice and 0.8 for Hohlraum.Source: §4 · p. 6

4.1 Data generation for surrogate model training

  • For both benchmarks, high- and low-resolution discretizations are defined, and low-resolution test samples use the high-resolution test parameter values except for mesh-generation parameters.The low-resolution runs provide a cheaper conventional-solver baseline for comparing accuracy and computational cost.Source: §4.1 · pp. 6–7, Table 2 · p. 7
  • The Lattice design set contains 800 configurations, of which 707 yielded valid simulations; 1,000 of 6,561 Hohlraum configurations were sampled, and 846 yielded valid simulations.Incomplete runs resulted from scheduling, timeouts, or other infrastructure failures, not known solver failures.Source: §4.1 · p. 7
  • The stored target is cellwise particle concentration Φ at final time, together with cell-centroid coordinates and cell areas.Hohlraum geometries require remeshing for each design-parameter set, so their meshes have varying cell counts.Source: §4.1 · pp. 6–7

5 Neural network-based surrogates

The section formulates surrogates that map static material configurations and source fields to final-time particle concentration, using a composite field-and-QoI objective. It instantiates Transolver and BSMS-MGN on mesh data, with both architectures designed for linear-in-mesh-size inference at fixed architecture settings.

5.1 Surrogate Model Formulation

  • Surrogates predict final-time particle concentration from mesh coordinates and material/source inputs, trained with region-weighted field error and an absorption-QoI alignment loss.The field-loss weights depend on material type; the QoI loss compares surrogate and high-resolution finite-volume absorption values.Source: §5.1 · pp. 7–8

5.2 Input feature design

  • Fourier features append multi-resolution sinusoidal coordinate encodings to raw coordinates, with their inclusion ablated alongside region-weighted loss.The encoding uses three frequencies, 2, 4, and 8.Source: §5.2 · p. 8

5.3.1 Overview

  • Transolver uses learned soft assignments to aggregate cells into physics-aware tokens, applies attention across tokens, then dispatches updated information back to cells.This replaces all-pairs cell attention with communication through M learned tokens, where M is much smaller than the number of cells.Source: §5.3.1 · pp. 8–9, §5.3.2 · p. 9

5.3.3 Computational cost

  • At fixed architecture parameters, Transolver inference scales linearly with mesh-cell count, while BSMS-MGN inference is measured as two to three times longer.Transolver attends over learned physics tokens rather than all cell pairs; BSMS-MGN's larger constant and irregular gather/scatter access partly explain its longer measured times.Source: §5.3.3 · p. 9, §5.4.3 · p. 12

5.3.4 Architecture choices

  • Transolver uses eight layers, width 256, and 128 physics tokens; training selects the lowest-validation-loss checkpoint and uses benchmark-specific QoI loss weights.The weights are 5 × 10^-3 for Lattice and 10^-2 for Hohlraum; Fourier features and region-weighted loss are ablated.Source: §5.3.4 · p. 10, §5.3.5 · p. 10

5.4.1 Overview

  • BSMS-MGN addresses local message-passing limits for streaming transport with a U-Net-style multiscale processor over coarsened graphs.Coarse-level hops span multiple mesh cells, extending communication across the domain within a forward pass.Source: §5.4.1 · p. 11

5.4.2 BSMS-MGN surrogate instantiation

  • BSMS-MGN represents mesh cells as graph nodes and shared faces as directed edges, then builds parameter-free coarsened levels through bi-stride pooling.Each coarsening level has roughly half as many nodes, and a hop at level l corresponds to 2^l hops on the mesh graph.Source: §5.4.2 · p. 11

5.4.4 Architecture choices

  • BSMS-MGN uses 16 finest-level message-passing blocks, five coarsening levels, and width 156 to match Transolver's parameter count; it shares the composite loss and training setup.Its fixed-setting forward pass is linear in cell count, and Fourier features and region-weighted loss are also ablated.Source: §5.4.3 · p. 12, §5.4.4 · p. 12, §5.4.5 · p. 12

6 Numerical results

The benchmarks compare high-resolution neural surrogates with KiT-RT and test cross-mesh transfer, architecture-specific ablations, and amortization costs.

6.1.1 Lattice

  • On Lattice, both surrogates beat the low-resolution KiT-RT baseline on absorption-QoI error; BSMS-MGN is about four times more accurate on that QoI, while Transolver has slightly lower field RMSE.Both reproduce the qualitative streaming and absorption patterns, with visible errors concentrated along strong gradients between absorbing and streaming regions.Source: §6.1.1 · pp. 13–14, Table 3 · p. 15

6.1.2 Hohlraum

  • On Hohlraum, both surrogates reproduce the main radiation and shielding features, but Transolver’s full configuration is visibly closer to the reference than BSMS-MGN’s.BSMS-MGN’s larger field errors are concentrated around the central absorbing capsule.Source: §6.1.2 · p. 15, Figure 6 · p. 16

6.2 Cross-mesh transfer

  • Cross-mesh evaluation without retraining degrades both models; BSMS-MGN degrades more severely, so the paper recommends matching inference and training mesh resolutions, especially for BSMS-MGN.Transolver largely retains field RMSE but loses absorption-QoI accuracy, whereas BSMS-MGN produces over-diffused fields.Source: §6.2 · p. 16

6.3 Ablations

  • Across Hohlraum, Transolver performs best with +FF, while BSMS-MGN favors +RW; Fourier features worsen BSMS-MGN’s field and center-capsule errors.For Transolver, +FF is best on every metric; for BSMS-MGN, +RW is best except on the horizontal-shield QoI, where Base is statistically tied within one standard deviation.Source: §6.1.2 · p. 15, Table 4 · p. 17, §6.3 · p. 18
  • The ablation associates Transolver’s gains with Fourier positional features, while BSMS-MGN’s cell-dual graph already encodes spatial structure and benefits more consistently from region weighting.For BSMS-MGN, Fourier features increase field RMSE and center-capsule error; the authors do not retune baselines in this study.Source: §6.3 · pp. 17–18

6.4 Inference efficiency and amortization

  • Surrogate training and data-generation costs are amortized after about 917 Lattice solves for Transolver and 1,050 for BSMS-MGN, and 781 and 930 Hohlraum solves, respectively.These are break-even counts for high-resolution KiT-RT evaluations when the surrogates are evaluated on high-resolution meshes.Source: Figure 7 · p. 19, §6.4 · p. 19

7 Conclusion and future work

The benchmark comparison finds architecture-dependent strengths: Transolver leads on the Hohlraum center-capsule QoI with Fourier features, while BSMS-MGN benefits from region-weighted training and is competitive on Lattice field RMSE. The conclusion supports architecture-specific design choices and use as amortized forward models when mesh-transfer limits are considered.

7 Conclusion and future work

  • The paper tests the hypothesis that surrogates trained on a few hundred high-resolution SN solutions can replace the transport solver when QoI accuracy and per-query cost meet the stated target.The target is accuracy matching the high-resolution solver on engineering QoIs at a per-query cost below a coarse conventional solve.Source: §7 · p. 20
  • At matched parameter counts, Transolver outperforms BSMS-MGN on the Hohlraum center-capsule QoI with Fourier features, while BSMS-MGN benefits from region-weighted training and is competitive on Lattice field RMSE.The ablation suggests feature and loss choices should be revisited for each architecture; future work extends inputs to include inflow data, broadens QoIs, and targets three-dimensional benchmarks.Source: §7 · p. 20
  • Surrogate performance relative to wall-time cost supports their use as amortized forward models in outer-loop transport workflows, provided architecture-specific feature choices and mesh-transfer limitations are accounted for.The cited workflows include multi-fidelity optimization, uncertainty quantification, design-space exploration, and inverse analysis.Source: §7 · p. 20

A Notation

  • The notation appendix provides symbol references for the continuous transport equation, reference solver, and generic surrogate formulation and training loss..Source: Table 7 · p. 25, Table 8 · p. 26, Table 9 · p. 26
  • Architecture-specific notation tables cover Transolver and BSMS-MGN, including the latter’s bi-stride U-Net passes and decoder.The BSMS-MGN notation describes per-level message-passing blocks and maps the final latent node state to particle concentration.Source: A Notation · pp. 27–28, Table 10 · p. 27, Table 11 · p. 28

B.1 Definition Transolver

Transolver represents the surrogate with learned physics-aware tokens, applying attention to those tokens rather than directly across mesh cells. The updated information returns to cells for final-time concentration prediction.

B.1 Definition Transolver

  • Transolver implements Fθ as a transformer that communicates through learned physics-aware tokens rather than direct all-to-all attention over Nh mesh cells.The same operator can be evaluated on meshes different from the training mesh.Source: §B.1 · pp. 29–30
  • Each cell’s inputs include normalized centroid coordinates, auxiliary features, and material/source quantities; Fourier features are an engineered input factor, while Hohlraum omits Q.A pointwise embedding maps the cell-wise inputs into the hidden representation.Source: §B.1 · p. 29
  • Feature-dependent slice assignments allow spatially separated cells with similar learned states to contribute to the same token.Source: §B.1 · p. 29
  • Self-attention operates over the M tokens, and updated token information is dispatched back to mesh cells using the same slice weights.In the reference implementation, token-attention query, key, and value projections are shared across heads.Source: §B.1 · p. 29
  • The model concatenates head outputs, projects them back to cell representations, and uses a final pointwise embedding to predict particle concentration.The Transolver surrogate factors the operator into its input embedding, per-layer transforms, and output embedding.Source: §B.1 · p. 30

B.2 Implementation details of BSMS-MGN

BSMS-MGN processes mesh-cell inputs on a hierarchy of graphs, combining finest-level message passing with a bi-stride U-Net pass. Its decoder predicts final-time particle concentration, and the graph and positions are supplied as inputs.

B.2 Implementation details of BSMS-MGN

  • BSMS-MGN represents the operator with a hierarchy of graphs, using the cell-dual graph at the finest level and a latent-space, multi-grid-like U-Net across levels.Source: §B.2 · pp. 30–31
  • Finest-level node inputs combine normalized coordinates, Fourier features, and material/source properties; Hohlraum omits Q, while edges encode relative offsets and their norms.Source: §B.2 · p. 30
  • An encoder maps node and edge inputs to a shared latent space, then Kmp finest-level message-passing blocks update edges before nodes with residual connections.Source: §B.2 · p. 31
  • A bi-stride U-Net pass transfers information across mesh levels using per-level graph message passing and non-parametric contribution tables for downsampling and upsampling.Per-level edge messages are recomputed from positional offsets and sender/receiver latents.Source: §B.2 · pp. 31–32
  • The decoder maps the finest-level U-Net output to final-time particle concentration, while graph structure, hierarchy, and node positions remain inputs that permit evaluation on different meshes.Source: §B.2 · p. 32
  • With Kmp=16 and Lg=5, one forward pass reaches about 3.3 length units on Lattice and about 0.7 on Hohlraum; both are upper bounds on Euclidean reach.The Lattice reach is comparable to the central-source-to-boundary distance but short of the domain diagonal.Source: §B.2 · p. 32
  • Cg=156 matches the trainable parameter count to the Transolver baseline, and the Lattice model uses Base configuration with void weight 1.Source: Table 15 · p. 35

C Hohlraum ablation visualizations

The Hohlraum visualizations show contrasting ablation sensitivities: Fourier features reduce visible Transolver errors, while BSMS-MGN’s capsule absorption region is closer to the reference without Fourier features. Panels compare reference, prediction, and absolute error across four configurations.

C Hohlraum ablation visualizations

C.1 Transolver ablation grid

In the Transolver Hohlraum ablation grid, Fourier features reduce visible error near material interfaces and the central capsule, while region weighting alone leaves it essentially unchanged.

C.1 Transolver ablation grid

  • Fourier-equipped Transolver configurations show noticeably lower visible error near material interfaces and in the central capsule region.Source: §C.1 · p. 36
  • Region weighting alone leaves the visible Transolver error magnitude essentially unchanged in the Hohlraum grid.Source: §C.1 · p. 36
  • The figure stacks Base, +RW, +FF, and +FF+RW configurations from top to bottom, making the Fourier-feature comparison visible across rows.Source: §C.1 · p. 36

C.2 BSMS-MGN ablation grid

In the BSMS-MGN ablation grid, adding region-weighted loss leaves absolute-error panels visibly similar to the base configuration, while Fourier features degrade central-capsule absorption.

C.2 BSMS-MGN ablation grid

  • The Base and +RW configurations have visibly similar absolute-error panels.Figure 9 stacks the four BSMS-MGN configurations in the same order.Source: §C.2 · p. 36
  • Central-capsule absorption degrades when Fourier features are added.Source: §C.2 · p. 36

D Cross-mesh (low-resolution) surrogate evaluation

Evaluating high-resolution-trained surrogates on coarser meshes reduces accuracy, with substantially greater degradation for BSMS-MGN. The authors recommend keeping inference resolution close to training resolution, especially for BSMS-MGN.

D Cross-mesh (low-resolution) surrogate evaluation

  • The experiment applies the same high-resolution-trained checkpoints to a coarser mesh without retraining; point-set and graph inputs allow evaluation at mesh centroids.Source: D Cross-mesh (low-resolution)… · p. 36
  • Transolver largely retains field RMSE, but its absorption-QoI error rises above the low-resolution KiT-RT baseline; on Hohlraum, +FF remains best for field RMSE and horizontal-shield QoI.Source: D Cross-mesh (low-resolution)… · p. 36
  • Cross-mesh evaluation is competitive on inference time but not accuracy; the authors recommend keeping inference resolution close to training resolution, especially for BSMS-MGN.Source: D Cross-mesh (low-resolution)… · p. 36
  • On low-resolution evaluation, BSMS-MGN field RMSE roughly triples to quadruples and most QoIs worsen by an order of magnitude or more, while both architectures lose accuracy.Source: D Cross-mesh (low-resolution)… · p. 36, Table 16 · p. 37, Table 17 · p. 37
  • On the low-resolution Lattice evaluation, both models have larger errors; Transolver captures high-absorbing regions, whereas BSMS-MGN over-diffuses.Source: Figure 10 · p. 40
  • On low-resolution Hohlraum, errors increase especially near inflow; Transolver retains qualitative features, whereas BSMS-MGN exhibits nonphysical behavior.Source: Figure 11 · p. 41

Related papers

Loading 2610.04665v1…