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

Abstract: The validation of advanced nuclear reactor designs and fuel concepts requires critical experiments with high neutronic similarity to the target technology. Neutronic similarity is quantified by the correlation coefficient $c_k$, which captures the shared bias in $k_\text{eff}$ induced by uncertainties in nuclear data. Generally, a $c_k\geq0.9$ is needed for an experiment to be sufficiently similar to a target technology. This work presents a methodology for the inverse design of critical experiments. Deep neural network surrogate modeling and nonparametric gradient optimization are used to generate experiment geometries that maximize $c_k$. A deep neural network is trained on OpenMC-calculated sensitivity vectors for grid-based critical experiment geometries. The model architecture combines a U-Net convolutional encoder-decoder with a novel multigroup attention pooling layer, introduced to capture the differing spatial dependencies of sensitivities. Multigroup attention pooling is shown to achieve better performance than traditional pooling, as well as interpretable internal behavior. The differentiability of the surrogate enables gradient-based optimization of the full combinatorial design space, allowing $c_k$ to be maximized by directly changing the material assignment of each position in the geometry grid. 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. The optimization procedure is shown to produce experiment geometries achieving $c_k$ scores of 0.97757, 0.81324, and 0.93276 for three configurations of interest. This approach demonstrates the potential of deep learning and gradient optimization to accelerate the development of advanced nuclear technology.

• The paper introduces a gradient-based optimization method integrating a multigroup attention neural network to surrogate neutronic sensitivity profiles.
• The paper utilizes a U-Net encoder-decoder with energy-specific attention pooling, reducing sensitivity prediction error to under 50 pcm.
• The paper demonstrates improved critical experiment designs, achieving c_k values above 0.9 for reactor validation across multiple configurations.

Inverse Design of Critical Experiments Using Gradient Optimization and Multigroup Attention Neural Networks

Motivation and Background

Validation of novel nuclear reactor designs mandates the creation of critical experiments with maximal neutronic similarity to the deployment target. The statistical measure of this similarity, $c_k$, quantifies shared sensitivities in the effective neutron multiplication factor ($k_\text{eff}$) to nuclear data uncertainties. For rigorous validation transfer, $c_k\geq0.9$ is typically required. However, existing experimental coverage often fails to provide sufficient $c_k$ for advanced technologies, particularly where HALEU fuels and complex cask environments are involved (see the TN Americas TN-LC case).

Conventional strategies for experiment design, such as Constrained Bayesian Optimization or particle swarm methods, are hindered by the computational expense of Monte Carlo neutronics and the high-dimensional, combinatorial design space. Recent developments in deep learning have enabled surrogate modeling of complex, nonlinear physical phenomena, providing differentiability that opens the design space to efficient gradient-based search.

Problem Formulation and $c_k$ Optimization

The core technical objective is to maximize $c_k$ between a candidate experiment and a target reactor/cask configuration. The $c_k$ value is computed from sensitivity profiles $\mathbf{S}_\text{TARGET}$ and $\mathbf{S}$ (an experimental candidate) as:

$$c_k = \frac{\mathbf{S}_\text{TARGET}^\top C_{XX} \mathbf{S}}{\sqrt{(\mathbf{S}_\text{TARGET}^\top C_{XX} \mathbf{S}_\text{TARGET})(\mathbf{S}^\top C_{XX} \mathbf{S})}}$$

where $k_\text{eff}$0 is the nuclear data covariance submatrix over $k_\text{eff}$1-dimensional reaction-energy sensitivity coefficients.

Full-order sensitivity calculations via OpenMC or Serpent are employed to produce ground-truth target profiles and training data for surrogate modeling. Experiment geometries are parameterized as $k_\text{eff}$2 voxel grids with a fixed palette (HALEU $k_\text{eff}$3, SS304, AlB, polyethylene), capturing the arrangement as a discrete input tensor.

Figure 1: Renderings of the simplified TN-LC model in OpenMC showing the arrangement of HALEU canisters (red) within the cask.

Figure 2: Correlation matrix for uncertainties in the $k_\text{eff}$4U fission cross-section using the SCALE-252 group structure.

Surrogate Model Design: U-Net with Multigroup Attention

A novel surrogate model architecture is introduced to predict $k_\text{eff}$5 for arbitrary material assignments. The motivation is both predictive performance and the imposition of physically-informed inductive bias.

• Encoder/Decoder: Uses a U-Net backbone. Skip connections preserve high-resolution spatial information lost in hierarchically downsampled representations.
• Multigroup Attention Pooling: For each energy group, an independent attention mechanism learns the spatial importance of local features, modeling the physics that neutrons of different energies interact with material distributions on different length scales. This enables the model to parse heterogeneous and long-range interactions more effectively than fixed pooling.
• Group Regressor: Post-attention, a residual 1D CNN operates over the group dimension, projecting descriptors to scalar reaction-energy sensitivities.

  Figure 4: Surrogate model architecture: geometry is processed by a shared U-Net, followed by energy/reaction-specific attention pools and group regressors for sensitivity prediction.

Empirical error analysis demonstrates that multigroup attention pooling yields a lower MAE compared with naively averaged or max-pooled alternatives, with the model achieving $k_\text{eff}$6 pcm test MAE in sensitivity prediction.

Dataset Generation and Geometry Representation

Training datasets are generated algorithmically using stochastic cellular automata, ensuring diversity and minimal human bias in material patterns while maintaining physical plausibility. Key hyperparameters ($k_\text{eff}$7, number of steps; $k_\text{eff}$8, neighbor weighting) control heterogeneity and segregation.

Figure 5: Procedural geometry generation over 40 cellular automaton steps.

Figure 6: Effect of different $k_\text{eff}$9 and $c_k\geq0.9$0 settings on generated geometry diversity from the same random seed.

Gradient-Based Optimization in Combinatorial Design Space

The crux of the inverse experiment design strategy is leveraging the differentiability of the surrogate to perform direct gradient-based optimization over the discrete geometry space. Material assignment logits are continuously relaxed via softmax-temperature annealing; a straight-through estimator (STE) enables gradient flow through hard assignments. The optimizer (AdamW) updates the logits to maximize $c_k\geq0.9$1, ultimately returning a discrete, physically interpretable geometry.

Batch initialization permits wide exploration; only the most promising candidates (by predicted $c_k\geq0.9$2) are evaluated in OpenMC for final $c_k\geq0.9$3 confirmation, thus reducing reliance on surrogate fidelity.

Figure 3: Optimization trajectory—mean and maximum $c_k\geq0.9$4 across 500 batch-initialized geometries for the TN-LC Case 2 target.

Figure 9: Snapshots of the best geometry during optimization for the TN-LC Case 2 target, demonstrating progressive structural adaptation.

Results: TN-LC Validation and Quantitative Benchmarks

The approach is evaluated on three configurations of the TN-LC cask: dry, flooded/low-$c_k\geq0.9$5B (Case 1), and flooded/high-$c_k\geq0.9$6B (Case 2). In all cases, the optimization yields discretized geometries with high $c_k\geq0.9$7—specifically:

Configuration	Initial Mean $c_k\geq0.9$8	Optimized $c_k\geq0.9$9 (OpenMC)	$c_k$0 (exp)
Dry	0.86083	0.97757	0.97152
Case 1	0.67963	0.81324	1.06533
Case 2	0.70278	0.93276	1.07930

• Case 2 result is salient: Prior work reported no experiments with $c_k$1; this method delivers $c_k$2 for Case 2.

  Figure 11: Optimized experiment geometry for the dry TN-LC configuration.

  

  Figure 13: Optimized experiment geometry for the TN-LC Case 1 configuration.

  

  Figure 15: Optimized experiment geometry for the TN-LC Case 2 configuration.

  

  Figure 17: Full sensitivity profiles for the three optimized experiment geometries.

Model Interpretability

Analysis of the learned attention maps reveals that the model develops spatial group selectivity aligning with physical transport behavior:

• Low-energy attention concentrates in fuel, consistent with localized fission and low-energy neutron mean free paths.
• Fast-group attention emphasizes geometry boundaries and reflector regions, reflecting the transport scale of fast neutrons.

  Figure 18: Normalized attention pool activations for $c_k$3U fission in thermal and fast groups, confirming spatial alignment with transport physics.

Implications and Future Perspectives

This framework demonstrates a scalable, automated approach to inverse nuclear experiment design, leveraging deep learning to bridge the gap between computationally expensive full-order simulation and high-dimensional, nonconvex optimization. By achieving $c_k$4 in previously uncovered configurations, the method offers practical pathways for targeted benchmark experimentation necessary for advanced reactor deployment and regulatory approval.

The principal limitation is the absence of $c_k$5 or engineering constraints in the optimization loop. Extension to multi-objective optimization, incorporating separate surrogates for operationality ($c_k$6), cost, and constructability, is a direct pathway for enhancing real-world applicability.

Architecturally, the observed advantage of multigroup attention suggests further opportunities for physics-informed network designs in computational neutronics. Extension to 3D representations, variable palette sizes, and application to facility-scale critical experiments (e.g., full SPARC material design space) will pose interesting research challenges.

Conclusion

The integration of a deep surrogate with group-wise attention pooling and batch gradient-based optimization establishes an effective pipeline for the inverse design of critical nuclear experiments. Achieving validated, high-similarity benchmarks for reactor/cask validation becomes a tractable endeavor, offering substantial impact for nuclear technology adoption. The results underscore the promise of tailored neural architectures and differentiable combinatorial optimization in physical sciences, with immediate application to advanced nuclear validation contexts.