Gradient-based optimization of nuclear criticality experiments using neural surrogate eigenvalue sensitivities

arXiv:2606.04033 · cs.LG · Submitted 2026-06-01 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom.

Jane: And I'm Jane, here with Lu, Meng and our in-house large language model, Lalam.

Tom: Today's paper: "Gradient-based optimization of nuclear criticality experiments using neural surrogate eigenvalue sensitivities".

Jane: The gist The work presents a methodology for inverse design of nuclear criticality experiments using deep neural network surrogate modeling and gradient-based geometry optimization to maximize neutronic similarity.

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So we're looking at this paper called "Gradient-based optimization of nuclear criticality experiments using neural surrogate eigenvalue sensitivities". It sounds like they’re trying to use a deep learning model, a neural network, to figure out the best shape for a nuclear experiment so it matches a real target system really well.

Jane: That sounds complicated for the average listener. Basically, they're tackling this big problem where you need new experiments because you don't have enough existing ones to test if new nuclear technology is safe and works right.

Lu: Exactly. The core idea here is using the differentiability of a deep neural network to guide an optimization process that finds an experiment geometry that maximizes the neutronic similarity, which they call ck, with a target system.

Meng: So, what's the actual goal? Are they just trying to make something look like the target system on paper or is there some real physics driving this design?

Tom: It’s about maximizing that correlation coefficient ck. They define it as a measure of how much shared bias in keff, which is the effective multiplication factor, comes from uncertainties in nuclear data. If you need a high ck, like zero point nine or higher for an experiment to be considered similar to the target system, they want to find that maximum value through optimization <ref:2606.04033#pg1,for an experiment to be>.

Jane: And they’re using this neural network as a surrogate model for OpenMC's sensitivity calculation capabilities, which is essentially figuring out how sensitive the neutronic properties are to changes in geometry. So instead of running slow Monte Carlo simulations over and over, they use this AI model to make predictions super fast.

Lu: Right. They’ve trained this physics-informed deep neural network on five thousand experiment geometries generated using a cellular automaton algorithm described in Algorithm one which is their way of creating training data without needing human input for every single shape.

Meng: five thousand geometries sounds like a lot of data to generate, but if that generation method is good, it saves tons of time during the actual design phase. I'm curious about how accurate this surrogate model actually is when it’s tested against those five thousand shapes.

Tom: They evaluate the performance by comparing what their neural network predicts for the sensitivity profiles against ground truths calculated using OpenMC, and they use both Mean Squared Error and Mean Absolute Error to measure that accuracy.

Jane: So, it’s a whole pipeline: you have the target system, you train a smart AI model on past geometries, and then you use that AI model to search for a new geometry that gets the highest possible similarity score.

The paper's summary: Tom: Moving on to what they actually did in this paper, "Gradient-based optimization of nuclear criticality experiments using neural surrogate eigenvalue sensitivities," it’s about the methodology they developed. They frame the design of an experiment geometry as a single objective function problem where the goal is to produce an experiment geometry x that maximizes ck with the target technology of interest.

Jane: The key technical innovation here is using the differentiability of their neural network surrogate model so they can compute a gradient of ck with respect to the input geometry directly through backpropagation. This allows for gradient-based design optimization without needing expensive Monte Carlo sensitivity calculations at every single step.

Lu: They relax the problem into optimization over continuous material logits, l in R four x thirty-two x thirty-two where those four material channels represent the four possible material assignments and those thirty-two by thirty-two spatial dimensions are the quarter-grid design domain they are optimizing over.

Meng: That’s a continuous space they’re searching through. What happens when you have this continuous space and you want to find the best discrete geometry for an experiment? That sounds like a tricky translation problem for any AI system.

Tom: They handle that by using a temperature schedule, tau t, during optimization steps from t equals zero to T minus one to encourage exploration across the design space, and then they cosine-anneal it down to a minimum value of zero point two before finalizing the geometry.

Jane: So they’re not just blindly optimizing; they are using this temperature schedule to guide the search effectively, making sure the AI explores a wide range of possibilities initially before settling on a good result.

Lu: And for how they handle the actual geometry evaluation during these steps, they use the straight-through estimator or STE, which is xSTE equals x plus p minus stopgrad(p) in the forward pass and treats the hard assignment as the identity map with respect to p in the backward pass.

Meng: That STE sounds like a trick to let gradients flow through a discrete problem. If you can’t do that smoothly, how do you get meaningful updates for your material choices?

Tom: It lets them formulate the loss function fully differentiable with respect to those continuous logits, which is what makes the gradient-based optimization possible in this context.

Jane: This whole process is designed to find an experiment geometry x that maximizes ck with the target technology, and they’re showing results across three different configurations of cask types—dry cask, flooded cask with low-density poison, and flooded cask with high-density poison.

The paper's improvements: Tom: Now let's talk about what they suggest for improvement in "Gradient-based optimization of nuclear criticality experiments using neural surrogate eigenvalue sensitivities." One major point is that the attention pooling layer they used demonstrated superior predictive performance compared to standard pooling alternatives, and it also had internal representations that were qualitatively consistent with the intuitively expected neutron flux distribution across each tested geometry.

Jane: That sounds like a big deal because it means their AI model isn't just memorizing data; it’s actually learning features that make sense from a physics perspective, which is what we want in science. It suggests the architecture itself is important for capturing complex spatial information.

Lu: I agree. The attention pooling layer weights features by their spatial importance for each discretized energy group, which allows the model to capture the neutronic behavior of complex geometries in a fraction of the time of a full Monte Carlo transport simulation.

Meng: So it’s not just a standard CNN or UNet; it’s this specific way they let the attention mechanism focus on what matters most for each energy group, which is pretty clever. But does that complexity add too much overhead to running the surrogate model?

Tom: The advantage is that this approach lets them capture the neutronic behavior of complex geometries in a fraction of the time of a full Monte Carlo transport simulation, which is incredibly fast. And they also show it’s qualitatively consistent with what we expect from neutron flux distribution.

Jane: That’s the real payoff there for anyone trying to use AI for nuclear design; if you get physically motivated features, you get a model that's more trustworthy than just a black box that spits out a number.

Lu: They also mention the potential of this method to extend to three-dimensional geometries and the full material design space of a critical experiment facility like SPARC, which is definitely the next big direction for this research.

Conclusion: Tom: So to wrap up, these results show that they successfully achieved ck greater than zero point nine for both the dry cask and Case two TN-LC configurations, and over zero point eight for Case one configuration, with keff values reaching things like zero point nine seven one five two, which is pretty close to one.

Jane: The main implication is that this methodology allows us to rapidly generate critical experiment designs with high ck for arbitrary target systems where historical experimental coverage is limited, which could speed up the validation of advanced nuclear technology.

Lu: But they’re being honest about the method's current limitation: it doesn't consider keff during the optimization process. They say that for realistic implementation, keff must be considered as part of the optimization process to ensure that all generated experiments reach keff greater than or equal to one <ref:2606.04033#pg1>.

Meng: That makes sense. If you’re optimizing for similarity and you don’t enforce the physical requirement of criticality, then your AI might find a geometrically similar shape that just doesn't actually work. You need to bake that constraint in somewhere.

Tom: Exactly, so the next step is incorporating those keff constraints through additional surrogate models or terms in the optimization objective, which they call a natural and tractable extension of this method.

Jane: This paper on gradient-based optimization of nuclear criticality experiments using neural surrogate eigenvalue sensitivities shows that applying deep learning to neutronics has potential, but it’s a powerful tool for speeding up the design process when you need those new critical experiments quickly.

Lu: It really demonstrates the power of deep learning applied to neutronics for generating designs with high ck, but they are flagging that the keff constraint is what’s needed for a realistic deployment.

Meng: Yeah, so it's a solid methodology for using AI to explore design space fast, provided you add those physical constraints back in later.

Tom: For now, this method is promising because of how it uses the attention pooling layer and the cellular automata training data to build these models quickly.

Will Savage, Logan Burnett, Dean Price

Massachusetts Institute of Technology

cs.LG

Submitted: 2026-06-01

Updated: 2026-10-07

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 92/100

The gist: The gist The work presents a methodology for inverse design of nuclear criticality experiments using deep neural network surrogate modeling and gradient-based geometry optimization to maximize

Key concepts

Neutronic Similarity ($c_k$): This coefficient measures how much the uncertainty in nuclear data affects the effective multiplication factor ($k_{ ext{eff}}$) between an experiment and a target technology. A high $c_k$, generally $\geq 0.9$, means the experiment is highly similar to the target, making it a good validation tool.
Physics-Informed Deep Neural Network: A deep learning model trained to predict the neutronic sensitivity of experimental geometries. It uses a specialized architecture (UNet with attention pooling) to capture complex neutron transport physics efficiently, allowing for fast predictions without running full Monte Carlo simulations.
Gradient-Based Optimization: A mathematical technique used to find the best geometry by calculating how small changes in the design parameters affect the objective function (maximizing $c_k$). The network's differentiability allows this gradient calculation via backpropagation, enabling automated search for optimal material combinations.
Surrogate Modeling: Using a simpler, fast-to-run mathematical model (the neural network) to stand in for a computationally expensive full nuclear simulation. This surrogate predicts the sensitivity profile of complex geometries quickly, which is then used by the optimization algorithm to guide design improvements.

Terminology

Inverse Design — 1× in this paper · explained in 2 episodes

Summary

The gist The work presents a methodology for inverse design of nuclear criticality experiments using deep neural network surrogate modeling and gradient-based geometry optimization to maximize neutronic similarity.

Application Scenario

The study applies this method to validate the TN Americas TN-LC transportation cask with HALEU fuel, which requires new critical experiments due to limited existing coverage The neutronic similarity between an experiment and the target system can be quantified by the correlation coefficient ck, which is defined as a measure of shared bias in keff induced by uncertainties in nuclear data For an experiment to be sufficiently similar to a target technology, a ck ≥ 0.9 is generally needed The lack of highly correlated experiments for the TN-LC highlights the need for new critical experiments to be designed

Surrogate Modeling

A deep neural network was trained as a surrogate model of OpenMC’s sensitivity calculation capabilities to predict the sensitivity profile Sˆ for a given experimental geometry x The model is comprised of a novel, heterogeneous neural network architecture designed to promote an inductive bias of the underlying transport physics during training The architecture involves a UNet [35] to produce a spatial feature map, which is then processed by a multigroup attention pooling layer that weights features by their spatial importance for each discretized energy group This approach allows the model to capture the neutronic behavior of complex geometries in a fraction of the time of a full Monte Carlo transport simulation The model was trained on 5,000 experiment geometries generated using a cellular automaton algorithm described in Algorithm 1 Performance is evaluated by comparing predicted sensitivity profiles to OpenMC-evaluated ground truths using Mean Squared Error (MSE) and Mean Absolute Error (MAE)

Gradient-Based Optimization

The design of an experiment geometry to maximize ck with a target technology is framed as a single-objective optimization problem where the objective function L(x) = −cˆk such that the optimization goal is to produce an experiment geometry x that maximizes ck with the target technology of interest The surrogate fθ is differentiable with respect to x, meaning that the gradient of ˆck with respect to the input geometry can be computed by backpropagation through the model To allow for gradient-based optimization, the problem is relaxed into optimization over continuous material logits l ∈ R 4×32×32, where the four material channels correspond to the four possible material assignments and the 32 × 32 spatial dimensions correspond to the quarter-grid design domain The temperature schedule τt is used during optimization steps t = 0,..., T − 1 to encourage exploration and then cosine-annealed to a minimum value of 0.2 The surrogate must be evaluated on a discrete geometry using the straight-through estimator (STE) xSTE = x + p − stopgrad(p) in the forward pass, while the backward pass treats the hard assignment as the identity map with respect to p

Results and Discussion

The optimization procedure was successful in all three cases, yielding experiments with ck ≥ 0.9 for two of three considered configurations and ck > 0.8 for the Case 1 configuration The correlation coefficients achieved were 0.97757, 0.81324, and 0.93276 for the dry cask, flooded cask with low-density poison, and flooded cask with high-density poison configurations respectively The optimization procedure was successful in all three cases, yielding experiments with ck greater than 0.9 for the Dry Cask and Case 2 TN-LC configurations and ck > 0.8 for the Case 1 configuration The experiment keff values achieved were 0.97152, 1.06533, and 1.07930 The primary shortcoming of the method is that it does not consider keff during the optimization process

Conclusion

This work presents a methodology for the inverse design of nuclear criticality experiments using deep neural network surrogate modeling and gradient-based geometry optimization The attention pooling layer demonstrated superior predictive performance relative to standard pooling alternatives, as well as internal representations qualitatively consistent with the intuitively expected neutron flux distribution across each tested geometry The ability to rapidly generate critical experiment designs with high ck for arbitrary target systems has the potential to accelerate the validation of advanced nuclear technology for which historical experimental coverage is limited The performance and observed internal behavior of the surrogate model in this work indicates the potential of deep learning applied to neutronics The primary shortcoming of the method is that it does not consider keff during the optimization process For realistic implementation, keff must be considered as part of the optimization process to ensure that all generated experiments reach keff ≥ 1 The incorporation of these considerations through additional surrogate models or terms in the optimization objective is a natural and tractable extension of the method Extension to three-dimensional geometries and the full material design space of a critical experiment facility such as SPARC are also directions for future work The ability to rapidly generate critical experiment designs with high ck for arbitrary target systems has the potential to accelerate the validation of advanced nuclear technology for which historical experimental coverage is limited The performance and observed internal behavior of the surrogate model in this work indicates the potential of deep learning applied to neutronics The primary shortcoming of the method is that it does not consider keff during the optimization process For realistic implementation, keff must be considered as part of the optimization process to ensure that all generated experiments reach keff ≥ 1 The incorporation of these considerations through additional surrogate models or terms in the optimization objective is a natural and tractable extension of the method Extension to three-dimensional geometries and the full material design space of a critical experiment facility such as SPARC are also directions for future work The ability to rapidly generate critical experiment designs with high ck for arbitrary target systems has the potential to accelerate the validation of advanced nuclear technology for which historical experimental coverage is limited The performance and observed internal behavior of the surrogate model in this work indicates the potential of deep learning applied to neutronics The primary shortcoming of the method is that it does not consider keff during the optimization process For realistic implementation, keff must be considered as part of the optimization process to ensure that all generated experiments reach keff ≥ 1 The incorporation of these considerations through additional surrogate models or terms in the optimization objective is a natural and tractable extension of the method Extension to three-dimensional geometries and the full material design space of a critical experiment facility such as SPARC are also directions for future work The ability to rapidly generate critical experiment designs with high ck for arbitrary target systems has the potential to accelerate the validation of advanced nuclear technology for which historical experimental coverage is limited The performance and observed internal behavior of the surrogate model in this work indicates the potential of deep learning applied to neutronics The primary shortcoming of the method is that it does not consider keff during the optimization process For realistic implementation, keff must be considered as part of the optimization process to ensure that all generated experiments reach keff ≥ 1 The incorporation of these considerations through additional surrogate models or terms in the optimization objective is a natural and tractable extension of the method <ref:260

Improvements for AI systems

textbfGradient-based Geometry Optimization for Criticality Experiments: Maximizing ck Directly via Differentiable Surrogates (Equation 19). Developing AI systems that can autonomously design new experimental geometries to maximize neutronic similarity (ck) with a target technology by using the differentiability of deep neural network surrogates, allowing for gradient-based optimization over the combinatorial design space of material assignments without relying on expensive Monte Carlo sensitivity calculations or traditional parametric approaches.
textbfMultigroup Attention Pooling for Physically Motivated Feature Learning (Section 2.5.2). Enabling AI systems to build more interpretable and physically informed surrogate models by using a novel multigroup attention pooling layer that weights features by their spatial importance for each discretized energy group, leading to internal representations qualitatively consistent with the intuitively expected neutron flux distribution across each tested geometry.
textbfIce-Free Surrogate Training Data Generation via Cellular Automata (Section 2.4). Developing AI systems capable of generating diverse and comprehensive training datasets for surrogate models by using a cellular automaton algorithm, which emits a broad distribution of material combinations and structures within the geometries without requiring arbitrary human judgment to specify which structures should be included.
textbfSoftmax Temperature Annealing Schedule for Differentiable Material Assignment (Equation 10 & 11). Improving the AI system's ability to explore the design space effectively by using a temperature schedule, where as τ decreases, the relaxed probabilities become increasingly concentrated on the material with the largest logit, ensuring that the softmax remains differentiable while becoming increasingly faithful to the original discrete material assignment problem as τ decreases.
textbfStraight-Through Estimator (STE) for Gradient Flow through Discrete Assignments (Equation 14 & 15). Enhancing gradient flow during optimization by using a straight-through estimator, which allows gradients to flow through the relaxed probabilities and then through the softmax to the logits, enabling the formulation [of the loss] to be fully differentiable with respect to the logits.

Abstract

The validation of advanced nuclear reactor designs and fuel concepts will require the design of new critical experiments with high neutronic similarity to the target technology. Neutronic similarity can be quantified by the correlation coefficient c k, which captures the shared bias in k eff induced by uncertainties in nuclear data. Generally, a c k at least0.9 is needed for an experiment to be sufficiently similar to a target technology. In this work, a physics-informed deep neural network is trained to predict the neutronic sensitivity of grid-based critical experiment geometries. The differentiability of the neural network is used to enable gradient-based design optimization of new experiment geometries to maximize c k with the sensitivity profile of a target technology. This approach allows for optimization over the combinatorial design space of potential material combinations within the grid, moving beyond traditional parametric optimization approaches. The method is applied to the validation of the TN-Americas TN-LC transportation cask with HALEU fuel, for which existing critical experiment coverage is limited. This application is shown to produce experiment geometries achieving c k scores of 0.97757, 0.81324, and 0.93276 for three configurations of interest.