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

Resumen Técnico: Diseño de Experimentos Críticos Inversos mediante Optimización de Gradientes y una Red Neuronal Basada en Atención de Multigrupos

Planteamiento del Problema
La validación de diseños de reactores nucleares avanzados y conceptos de combustible, particularmente aquellos que utilizan Uranio de Bajo Enriquecimiento de Alta Contenido (HALEU), requiere experimentos críticos con una similitud neutrónica elevada respecto a la tecnología objetivo. Esta similitud se cuantifica mediante el coeficiente de correlación, $c_k$, que mide el sesgo compartido en el factor de multiplicación efectiva ($k_{eff}$) inducido por las incertidumbres en los datos nucleares. Generalmente se requiere un $c_k \geq 0.9$ para que un experimento sea considerado suficientemente similar a un sistema objetivo. Sin embargo, para aplicaciones específicas como el contenedor de transporte TN-Americas TN-LC, los experimentos de referencia de experimentos críticos existentes a menudo no logran alcanzar este umbral (por ejemplo, ningún experimento existente logró un $c_k \geq 0.9$ para ciertas configuraciones de TN-LC). Los métodos de diseño tradicionales están limitados por el alto costo computacional de las simulaciones de Monte Carlo necesarias para evaluar geometrías candidatas, lo que hace que la exploración del espacio completo de diseño combinatorio sea intratable.

Metodología
Los autores proponen una metodología para el diseño inverso de experimentos críticos que combina el modelado de sustitución (surrogate modeling) mediante redes neuronales profundas (DNN) con la optimización de gradientes no paramétricos.

1. Modelado de Sustitución:

   • Generación de Datos: Se creó un conjunto de datos de 5,000 geometrías de experimentos 2D generadas procedimentalmente utilizando un algoritmo de autómatas celulares. Estas geometrías fueron modeladas en OpenMC para calcular los vectores de sensibilidad para cuatro reacciones ($^{235}$U fisión, $^{12}$C dispersión elástica, $^{1}$H dispersión elástica y $^{58}$Ni absorción) a través de 252 grupos de energía.
   • Arquitectura: El modelo de sustitución emplea una arquitectura heterogénea que combina un codificador-decodificador convolucional U-Net con una novedosa capa de Agrupación por Atención de Multigrupos (Multigroup Attention Pooling).
     • La U-Net extrae características espaciales de la geometría de entrada (codificada como un cuarto de cuadrícula de 32x32 de asignación de materiales).
     • La capa de Agrupación por Atención de Multigrupos aprende mapas de atención distintos para cada grupo de energía, ponderando las características espaciales según su importancia para ese grupo específico. Esto introduce un sesgo inductivo que refleja las diferentes dependencias espaciales del transporte de neutrones en diversas escalas de energía.
     • Un regresor convolucional 1D procesa los descriptores de grupo para predecir el perfil de sensibilidad final.
   • Entrenamiento: El modelo se entrena para predecir perfiles de sensibilidad (en lugar de $c_k$ directamente) utilizando una pérdida de error cuadrático medio (MSE) normalizada, lo que permite una evaluación detallada a través de reacciones y grupos de energía.

2. Optimización Basada en Gradientes:

   • La diferenciabilidad del modelo de sustitución permite la optimización directa de la asignación de materiales en la cuadrícula de la geometría para maximizar $c_k$.
   • Para manejar la naturaleza discreta de las asignaciones de materiales, el problema se relaja optimizando los logits de materiales continuos ($\ell$). Una función softmax escalada por temperatura convierte estos logits en probabilidades, que luego se discretizan utilizando un Estimador de Paso Directo (Straight-Through Estimator, STE) para la pasada hacia adelante (forward pass), manteniendo la diferenciabilidad para la pasada hacia atrás (backward pass).
   • La optimización minimiza el coeficiente de correlación predicho negativo ($-\hat{c}_k$) utilizando el optimizador AdamW. Se aplica un esquema de recocido de coseno (cosine annealing) a la temperatura del softmax para transicionar de la exploración a la explotación.

Contribuciones Clave

• Marco de Diseño Inverso: Este trabajo establece un marco para generar geometrías de experimentos críticos que maximizan la similitud neutrónica ($c_k$) con un sistema objetivo mediante la optimización directa de las asignaciones de materiales vía descenso de gradiente.
• Mecanismo de Atención de Multigrupos: La introducción de una capa de agrupación por atención de multigrupos permite que la red capture las distintas dependencias espaciales de las sensibilidades de los neutrones en diferentes grupos de energía. Los autores demuestran que esta arquitectura supera a los métodos tradicionales de agrupación máxima y promedio en términos de precisión de predicción (menor Error Absoluto Medio) y proporciona un comportamiento interno interpretable consistente con la intuición física (por ejemplo, los neutrones térmicos permanecen localizados en el combustible, mientras que los neutrones rápidos interactúan con las regiones circundantes).
• Modelos de Sustitución Diferenciables para Neutrónica: El estudio demuestra la viabilidad del uso de modelos de sustitución de aprendizaje profundo diferenciables para navegar el espacio de diseño combinatorio completo de los experimentos críticos, una tarea previamente dificultada por el costo computacional de los cálculos de sensibilidad de Monte Carlo.

Resultados
La metodología se aplicó para diseñar experimentos para tres configuraciones de un contenedor de transporte TN-LC: un contenedor seco, un contenedor inundado con veneno de baja densidad (Caso 1) y un contenedor inundado con veneno de alta densidad (Caso 2).

• Desempeño: La optimización generó con éxito geometrías de experimentos con altos puntajes de $c_k$:
  • Contenedor Seco: $c_k = 0.97757$
  • Caso 1: $c_k = 0.81324$ (superando el umbral de 0.8 requerido para evitar penalizaciones de $k_{eff}$)
  • Caso 2: $c_k = 0.93276$ (un logro significativo, ya que ningún experimento previo con $c_k \geq 0.8$ existía para esta configuración).
• Características de la Geometría: Las geometrías optimizadas fueron altamente heterogéneas y de diseño "ajeno", priorizando UO$_2$ y aluminio borado sobre acero inoxidable y polietileno. El optimizador se centró principalmente en igualar la magnitud y la forma del perfil de sensibilidad de la fisión de $^{235}$U, que es el contribuyente dominante a la incertidumbre de $k_{eff}$.
• Precisión del Modelo: El modelo de sustitución alcanzó un Error Absoluto Medio (MAE) de prueba de 41.19 pcm, mostrando la configuración de agrupación por atención de multigrupos el MAE más bajo en comparación con los métodos de agrupación tradicionales.

Significancia y Reivindicaciones
El artículo afirma que este enfoque demuestra el potencial del aprendizaje profundo y la optimización de gradientes para acelerar el desarrollo y la validación de tecnologías nucleares avanzadas. Específicamente, aborda la brecha crítica en la cobertura experimental para escenarios de transporte de HALEU donde el benchmarking tradicional falla.

Los autores señalan explícitamente las limitaciones y las reivindicaciones modestas:

• Restricción de $k_{eff}$: El método actual no restringe el factor de multiplicación efectiva ($k_{eff}$) durante la optimización. Si bien los experimentos resultantes tuvieron valores de $k_{eff}$ cercanos o superiores a 1, asegurar la subcriticidad o márgenes de criticidad específicos requiere restricciones o modelos de sustitución adicionales, los cuales se identifican como trabajo futuro.
• Alcance: El estudio se centra en geometrías 2D y una paleta de materiales específica. La extensión a geometrías 3D y al espacio completo de diseño de materiales de instalaciones como SPARC se propone como una dirección para futuras investigaciones.
• Interpretabilidad: Aunque los mapas de atención proporcionan información cualitativa consistente con la física de neutrones, los autores reconocen que se necesita un análisis de interpretabilidad más riguroso para cuantificar completamente el comportamiento interno del modelo.

En resumen, este trabajo constituye una prueba de concepto de que los modelos de sustitución diferenciables pueden utilizarse para resolver problemas de diseño inverso en ingeniería nuclear, generando nuevas configuraciones de experimentos que alcanzan una alta similitud neutrónica donde antes no existían.