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

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.