🤖 machine learning

Inverse Critical Experiment Design via Gradient Optimization and a Multigroup Attention-Based Neural Network Architecture

This paper introduces a deep learning framework combining a multigroup attention-based neural network with gradient-based optimization to inversely design critical experiments that maximize neutronic similarity ( $c_k$ ) for validating advanced nuclear technologies, such as HALEU fuel casks, achieving high similarity scores where traditional methods fall short.

Original authors: Will Savage, Logan Burnett, Dean Price

Published 2026-06-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Will Savage, Logan Burnett, Dean Price

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

The Big Picture: Designing a "Twin" for a Nuclear System

Imagine you have a brand-new, high-tech nuclear reactor (or a fuel container) that you want to build. Before you build the real thing, you need to prove it's safe. To do this, scientists usually run "criticality experiments"—small, low-power test versions of the system.

However, there's a catch: The test version must be a perfect "twin" of the real thing. If the test version behaves differently than the real one, the results are useless. Scientists measure this "twin-ness" using a score called $c_k$ .

$c_k = 1.0$ : Perfect twin.
$c_k \ge 0.9$ : Good enough to trust.
$c_k < 0.9$ : Not similar enough; the test results might be misleading.

The problem? For some new technologies (like a specific container for transporting advanced nuclear fuel called TN-LC), scientists couldn't find any existing experiments that were similar enough. They were stuck.

This paper presents a solution: Instead of looking for an existing experiment, they used a computer to invent a new one from scratch that is a perfect twin.

The Tools: A "Crystal Ball" and a "Tuning Knob"

To invent this new experiment, the authors used two main tools:

1. The Crystal Ball (The Neural Network)

Running a real nuclear simulation is like trying to predict the weather by building a miniature atmosphere in a lab and waiting a month to see what happens. It takes a long time and costs a lot of computing power.

The authors built a Deep Neural Network (a type of AI) to act as a "Crystal Ball."

Training: They showed the AI thousands of different experiment designs and the results of real, slow simulations.
The Trick: They didn't just teach the AI to guess the final result. They taught it to understand the physics of how neutrons move through different materials (like Uranium, steel, or water).
The Innovation: They added a special layer called "Multigroup Attention." Think of this like a spotlight. Neutrons of different energies (speeds) care about different parts of the experiment. Fast neutrons might care about the edges, while slow neutrons care about the center. The "Attention" layer lets the AI focus its spotlight on the right areas for the right type of neutron, just like a human expert would.

2. The Tuning Knob (Gradient Optimization)

Usually, if you want to design the perfect experiment, you might try random shapes until something works. That's like trying to find a specific song by randomly pressing piano keys.

Because their "Crystal Ball" (the AI) is mathematically smooth, the authors could use Gradient Optimization.

Imagine you are in a dark room trying to find the highest peak of a hill.
Instead of stumbling around randomly, you can feel the slope under your feet.
The AI calculates exactly which direction to move the materials (e.g., "move a bit of steel here, add more uranium there") to instantly climb higher on the "Similarity Hill."
This allows them to tweak the design thousands of times in seconds, finding the absolute best shape that maximizes the $c_k$ score.

The Experiment: The "TN-LC" Suitcase

The team applied this method to a specific problem: validating a transport cask (a heavy-duty suitcase) for HALEU (a type of advanced nuclear fuel).

They looked at three scary scenarios:

Dry Cask: Normal conditions.
Case 1: Water leaks in, and there's a light poison coating.
Case 2: Water leaks in, and there's a heavy poison coating.

The Problem: Previous searches found zero existing experiments that were similar enough ( $c_k \ge 0.9$ ) for Case 2.

The Solution:
The AI started with a blank grid and used its "Tuning Knob" to rearrange materials (Uranium, steel, borated aluminum, plastic) until it found a shape that matched the TN-LC cask perfectly.

The Results:

Dry Cask: Achieved a similarity score of 0.97 (Excellent).
Case 1: Achieved 0.81 (Good enough to be useful, though not perfect).
Case 2: Achieved 0.93 (Excellent). This is huge because no previous experiment existed for this scenario.

What Did the AI Invent?

The resulting designs looked a bit "alien." They didn't look like the neat, symmetrical experiments humans usually build.

The AI realized that to match the physics, it needed to use mostly Uranium and Borated Aluminum, with very little steel or plastic.
The shapes were complex and patchy, which is exactly what the math demanded to make the neutrons behave the same way as they would in the real transport cask.

The Catch (Limitations)

The paper is honest about one major limitation: The AI only cared about making the twin ( $c_k$ ). It didn't care if the resulting experiment was actually safe to build or if it would explode (reach a critical state where $k_{eff} > 1$ ).

In the real world, you need the experiment to be stable.
The authors note that in the future, they could add a "safety rule" to the AI so it designs twins that are also safe to run.

Summary

This paper shows that by combining a smart AI (that understands how neutrons move) with a mathematical "tuning knob," scientists can reverse-engineer the perfect experiment. Instead of hoping to find a matching test in a library, they can now generate a custom-made test for any new nuclear technology, even for scenarios where no test currently exists.

Technical Summary: Inverse Critical Experiment Design via Gradient Optimization and a Multigroup Attention-Based Neural Network

Problem Statement
The validation of advanced nuclear reactor designs and fuel concepts, particularly those utilizing High-Assay Low-Enriched Uranium (HALEU), requires critical experiments with high neutronic similarity to the target technology. This similarity is quantified by the correlation coefficient, $c_k$ , which measures the shared bias in the effective multiplication factor ( $k_{eff}$ ) induced by uncertainties in nuclear data. A $c_k \geq 0.9$ is generally required for an experiment to be considered sufficiently similar to a target system. However, for specific applications such as the TN-Americas TN-LC transportation cask, existing critical experiment benchmarks often fail to meet this threshold (e.g., no existing experiments achieved $c_k \geq 0.9$ for certain TN-LC configurations). Traditional design methods are limited by the high computational cost of Monte Carlo simulations required to evaluate candidate geometries, making the exploration of the full combinatorial design space intractable.

Methodology
The authors propose a methodology for the inverse design of critical experiments that combines deep neural network (DNN) surrogate modeling with nonparametric gradient optimization.

Surrogate Modeling:
- Data Generation: A dataset of 5,000 procedurally generated 2D experiment geometries was created using a cellular automaton algorithm. These geometries were modeled in OpenMC to calculate sensitivity vectors for four reactions ( $^{235}$ U fission, $^{12}$ C elastic scattering, $^{1}$ H elastic scattering, and $^{58}$ Ni absorption) across 252 energy groups.
- Architecture: The surrogate model employs a heterogeneous architecture combining a U-Net convolutional encoder-decoder with a novel Multigroup Attention Pooling layer.
  - The U-Net extracts spatial features from the input geometry (encoded as a 32x32 quarter-grid of material assignments).
  - The Multigroup Attention Pooling layer learns distinct attention maps for each energy group, weighting spatial features based on their importance to that specific group. This introduces an inductive bias reflecting the differing spatial dependencies of neutron transport at various energy scales.
  - A 1D convolutional regressor processes the group descriptors to predict the final sensitivity profile.
- Training: The model is trained to predict sensitivity profiles (rather than $c_k$ directly) using normalized mean squared error (MSE) loss, allowing for fine-grained evaluation across reactions and energy groups.
Gradient-Based Optimization:
- The differentiability of the surrogate model enables the direct optimization of the material assignment in the geometry grid to maximize $c_k$ .
- To handle the discrete nature of material assignments, the problem is relaxed into optimizing continuous material logits ( $\ell$ ). A temperature-scaled softmax converts these logits into probabilities, which are then discretized using a Straight-Through Estimator (STE) for the forward pass while maintaining differentiability for the backward pass.
- The optimization minimizes the negative predicted correlation coefficient ( $-\hat{c}_k$ ) using the AdamW optimizer. A cosine annealing schedule is applied to the softmax temperature to transition from exploration to exploitation.

Key Contributions

Inverse Design Framework: The work establishes a framework for generating critical experiment geometries that maximize neutronic similarity ( $c_k$ ) to a target system by directly optimizing material assignments via gradient descent.
Multigroup Attention Mechanism: The introduction of a multigroup attention pooling layer allows the network to capture the distinct spatial dependencies of neutron sensitivities across different energy groups. The authors demonstrate that this architecture outperforms traditional max and average pooling in terms of prediction accuracy (lower Mean Absolute Error) and provides interpretable internal behavior consistent with physical intuition (e.g., thermal neutrons remaining localized in fuel, fast neutrons interacting with surrounding regions).
Differentiable Surrogates for Neutronics: The study demonstrates the feasibility of using differentiable deep learning surrogates to navigate the full combinatorial design space of critical experiments, a task previously hindered by the computational cost of Monte Carlo sensitivity calculations.

Results
The methodology was applied to design experiments for three configurations of the TN-LC transportation cask: a dry cask, a flooded cask with low-density poison (Case 1), and a flooded cask with high-density poison (Case 2).

Performance: The optimization successfully generated experiment geometries with high $c_k$ $c_{k}$ scores:
- Dry Cask: $c_k = 0.97757$
- Case 1: $c_k = 0.81324$ (exceeding the 0.8 threshold required to avoid $k_{eff}$ penalties)
- Case 2: $c_k = 0.93276$ (a significant achievement, as no prior experiment with $c_k \geq 0.8$ existed for this configuration).
Geometry Characteristics: The optimized geometries were highly heterogeneous and "alien" in design, prioritizing UO $_2$ and borated aluminum over stainless steel and polyethylene. The optimizer focused primarily on matching the magnitude and shape of the $^{235}$ U fission sensitivity profile, which is the dominant contributor to $k_{eff}$ uncertainty.
Model Accuracy: The surrogate model achieved a test set Mean Absolute Error (MAE) of 41.19 pcm, with the multigroup attention pooling configuration showing the lowest error compared to traditional pooling methods.

Significance and Claims
The paper claims that this approach demonstrates the potential of deep learning and gradient optimization to accelerate the development and validation of advanced nuclear technologies. Specifically, it addresses the critical gap in experimental coverage for HALEU transport scenarios where traditional benchmarking fails.

The authors explicitly note limitations and modest claims:

$k_{eff}$ Constraint: The current method does not constrain the effective multiplication factor ( $k_{eff}$ ) during optimization. While the resulting experiments had $k_{eff}$ values near or above 1, ensuring subcriticality or specific criticality margins requires additional constraints or surrogate models, which are identified as future work.
Scope: The study focuses on 2D geometries and a specific material palette. Extension to 3D geometries and the full material design space of facilities like SPARC is proposed as a direction for future research.
Interpretability: While the attention maps provide qualitative insights consistent with neutron physics, the authors acknowledge that more rigorous interpretability analysis is needed to fully quantify the model's internal behavior.

In summary, the work provides a proof-of-concept that differentiable surrogate models can be used to solve inverse design problems in nuclear engineering, generating novel experiment configurations that achieve high neutronic similarity where none previously existed.