⚛️ lattice

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

This paper benchmarks two parameter-matched neural surrogate architectures, Transolver and BSMS-MGN, for approximating linear radiation transport on Lattice and Hohlraum benchmarks, revealing that optimal inductive bias choices are architecture-dependent and that downstream utility relies on per-quantity-of-interest sensitivity rather than a single field-level metric.

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

Published 2026-10-06

📖 5 min read🧠 Deep dive

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the vast, invisible world of nuclear engineering and astrophysics, scientists rely on a set of rules to track how energy moves through matter. Imagine a beam of light or a stream of particles traveling through a room filled with different materials—some that absorb the energy, some that scatter it, and some that let it pass through. Predicting exactly where that energy goes is essential for designing safer nuclear reactors, planning cancer treatments, and understanding how stars shine. However, calculating these paths is incredibly difficult. The particles move in three dimensions of space and three dimensions of momentum, all changing over time. To get a precise answer, computers must break the problem down into billions of tiny pieces and perform trillions of calculations. This process is so slow and expensive that engineers often cannot afford to run the simulations they need to test new designs or explore uncertainties. They are forced to use rougher, faster approximations that sacrifice accuracy for speed.

A team of researchers has now tested a new approach to this problem: using artificial intelligence to act as a fast, accurate stand-in for the slow, heavy calculations. They trained two different types of neural networks—computer systems designed to learn patterns from data—to predict the final state of particle movement in two specific, challenging scenarios. One scenario mimics a section of a nuclear reactor, a grid of alternating materials that absorb and scatter radiation. The other simulates a "hohlraum," a hollow chamber used in fusion research where radiation bounces off walls to heat a central target. The researchers generated thousands of high-quality examples using a powerful, traditional computer solver and fed this data to the neural networks. The goal was to see if the networks could learn to predict the outcome of a simulation in a fraction of the time it takes the traditional computer to solve it, without losing the precision needed for real-world engineering.

The results show that these neural networks can indeed become powerful tools, but they do not all learn in the same way. The researchers compared two distinct architectures. The first, called Transolver, uses a method that allows the computer to look at the entire problem at once, finding connections between distant parts of the simulation. The second, known as BSMS-MGN, works more like a local neighborhood watch, passing information from one tiny cell to its immediate neighbors and then up through a hierarchy of larger groups. When tested on the reactor-like grid, both networks performed well, reproducing the complex patterns of energy flow with high accuracy. However, when the researchers tested them on the fusion chamber, a clear difference emerged. The Transolver network, when given a specific type of coordinate encoding that helps it understand spatial positions, became significantly more accurate at predicting how much energy was absorbed by the central target. The graph-based network, by contrast, struggled with this specific detail when given the same coordinate help, but it performed best when the training process was adjusted to pay extra attention to the material boundaries.

This finding suggests that there is no single "best" design for an AI model in this field. The choices that make one network accurate can actually hurt another. The researchers found that the Transolver model thrived when it was given extra information about the geometry of the space, while the graph-based model performed better when the training process was tweaked to emphasize the importance of specific material regions. This means that engineers cannot simply copy a successful AI setup from one problem to another; they must tailor the design to the specific architecture they are using. Despite these differences, both networks achieved a remarkable speedup. Once trained, they could produce a solution in a fraction of a second, whereas the traditional computer solver took seconds or even minutes. The researchers calculated that for tasks requiring thousands of simulations, such as optimizing a reactor design or testing for safety margins, the time saved by using these AI models would far outweigh the initial time spent training them.

The study also revealed a limitation in how these models handle different levels of detail. The networks were trained on very fine, high-resolution grids that captured every small detail of the particle flow. When the researchers asked them to run on a coarser, lower-resolution grid without retraining, the accuracy dropped, particularly for the graph-based model, which produced overly smooth results that missed sharp changes in energy. This indicates that for the highest fidelity, the AI should be run on the same level of detail it was trained on. However, even with this constraint, the speed gains are substantial. The researchers estimate that for the fusion chamber problem, the AI models would need to be run roughly 800 to 1,000 times to pay back the time spent generating the training data and training the model. After that point, every additional simulation is a massive win in efficiency.

Ultimately, the work demonstrates that neural surrogates are ready to join the toolkit of nuclear engineers and physicists, provided they are used with care. They are not a magic replacement for all physics calculations, but they are a highly effective tool for the specific, repetitive tasks that currently bottleneck design and analysis. By showing that different network structures require different training strategies, the researchers have provided a roadmap for building better models in the future. The key takeaway is that the path to accurate, fast simulation is not about finding a single perfect algorithm, but about understanding the unique strengths and weaknesses of each approach and matching them to the specific problem at hand. As these models are refined and applied to more complex three-dimensional scenarios, they promise to unlock new possibilities for designing cleaner energy sources and understanding the universe.

Technical Summary: Efficient Neural Surrogates for Linear Radiation Transport on the Lattice and Hohlraum Benchmarks

Problem Statement
Linear radiation transport equations (RTEs) are fundamental to simulations in nuclear engineering, inertial confinement fusion, medical imaging, and astrophysics. However, resolving the high-dimensional phase space (up to six dimensions plus time) required for engineering fidelity is computationally prohibitive. Traditional high-resolution solvers, such as discrete-ordinates ( $S_N$ ) methods, demand significant memory and floating-point operations, while Monte Carlo approaches suffer from statistical noise. Consequently, outer-loop workflows like design optimization, uncertainty quantification (UQ), and parameter sweeps are often budget-constrained. While surrogate models offer a potential solution by amortizing simulation costs across thousands of queries, architectural choices and inductive biases that yield accuracy in one transport problem do not necessarily transfer across different model families.

Methodology
This work benchmarks two parameter-matched neural surrogate architectures as end-to-end approximations of the final-time particle concentration for the two-dimensional linear RTE. The models are trained on data generated by the KiT-RT solver, a high-resolution, GPU-accelerated discrete-ordinates/finite-volume solver.

Benchmarks: The study utilizes two canonical problems from [SH25]:
- Lattice Benchmark: A heterogeneous nuclear reactor assembly mock-up featuring a checkerboard pattern of absorbing, scattering, and void materials with a central source.
- Hohlraum Benchmark: A linearized, symmetric variation of a thermal radiative transfer problem involving a central capsule irradiated from the sides, surrounded by shielding materials. This benchmark includes design parameters that alter the geometry (shielding locations and capsule displacement), requiring dynamic meshing.
Surrogate Architectures:
- Transolver: A transformer-based architecture utilizing "Physics-Attention." Instead of standard all-pair attention, it learns soft assignments of mesh cells to physics-aware slices (tokens). This reduces computational complexity from quadratic to linear in the number of cells while maintaining global coupling.
- Bi-Stride Multi-Scale MeshGraphNet (BSMS-MGN): A hierarchical graph neural network (GNN) operating on the cell-dual graph of the unstructured mesh. It employs a U-Net-style multi-scale processor with bi-stride pooling to propagate information across long distances (streaming regimes) that standard local message passing cannot resolve efficiently.
Training and Evaluation:
- Models are trained to map static material configurations and source fields to the particle concentration field $\Phi$ at a fixed final time $T$ .
- The loss function combines a region-weighted Mean Squared Error (MSE) with a Quantity of Interest (QoI) alignment loss, specifically targeting radiation absorption in specific subdomains.
- An ablation study investigates the impact of Fourier features (positional encoding) and region-weighted loss (assigning higher weights to material regions vs. voids) on both architectures.
- Evaluation is performed on both the high-resolution mesh (training resolution) and a coarser low-resolution mesh to test cross-mesh transfer capabilities.

Key Results

Performance: Both surrogates achieve significant inference speedups over the high-resolution KiT-RT reference solver, ranging from approximately 90× to 4400×, depending on the benchmark, architecture, and evaluation mesh.
Accuracy:
- On the Lattice benchmark, both models outperform the low-resolution KiT-RT baseline in terms of absorption QoI error. BSMS-MGN shows a slight edge in QoI accuracy, while Transolver performs marginally better on field RMSE.
- On the Hohlraum benchmark, the architectures exhibit divergent behaviors based on design choices. Transolver achieves its best performance with Fourier features enabled, significantly reducing field RMSE and QoI errors (often below 1%). BSMS-MGN, conversely, performs best with region-weighted loss and no Fourier features; adding Fourier features degrades its performance on the central capsule QoI.
Inductive Bias: The ablation study reveals that inductive biases are architecture-dependent. Fourier features, which help Transolver disambiguate cells with similar material properties in different spatial regimes, act as a distractor for BSMS-MGN, whose cell-dual graph topology already encodes spatial structure.
Amortization: The analysis quantifies the "break-even" point where the cost of training and data generation is recovered by the speedup in inference. For the Hohlraum benchmark, Transolver amortizes after ~781 high-resolution evaluations, while BSMS-MGN requires ~930. Both are viable for outer-loop workflows requiring thousands of queries.
Cross-Mesh Transfer: Evaluating models on a coarser mesh without retraining degrades accuracy, particularly for BSMS-MGN, which produces over-diffused fields. The authors recommend matching inference mesh resolution to training resolution.

Significance and Claims
The paper claims that neural surrogates can effectively replace high-fidelity transport solvers in outer-loop workflows, provided that:

Architecture-Specific Design: Design choices (such as feature engineering and loss weighting) cannot be imported across model families; they must be revisited for each specific architecture.
Utility: The models are useful for downstream tasks like multi-fidelity optimization, UQ, and design-space exploration where repeated high-resolution simulation is prohibitive.
Reproducibility: The work provides a complete pipeline, including model training recipes, training data, and evaluation code, to support reproduction and transfer to related transport problems.

The authors maintain a modest stance, noting that while the surrogates outperform coarse conventional solvers, they do not yet match the absolute fidelity of high-resolution solvers in all metrics without careful architectural tuning. Future work is identified as extending inputs to include inflow boundary data, broadening QoIs to time-resolved quantities, and scaling to three-dimensional benchmarks.

More like this