Neurons for Neutrons: A Transformer Model for Computation Load Estimation on Domain-Decomposed Neutron Transport Problems

TL;DR

This paper introduces a Transformer-based model that accurately predicts computation loads in domain-decomposed neutron transport problems, reducing the need for time-consuming simulations and enabling faster, scalable analysis.

Contribution

The paper presents a novel Transformer model with 3D embeddings tailored for neutron transport problems, achieving high accuracy and eliminating the need for repeated small-scale simulations.

Findings

Achieves 98.2% accuracy in load prediction

Can bypass small-scale simulations entirely

Demonstrates robustness across different geometries and parameters

Abstract

Domain decomposition is a technique used to reduce memory overhead on large neutron transport problems. Currently, the optimal load-balanced processor allocation for these domains is typically determined through small-scale simulations of the problem, which can be time-consuming for researchers and must be repeated anytime a problem input is changed. We propose a Transformer model with a unique 3D input embedding, and input representations designed for domain-decomposed neutron transport problems, which can predict the subdomain computation loads generated by small-scale simulations. We demonstrate that such a model trained on domain-decomposed Small Modular Reactor (SMR) simulations achieves 98.2% accuracy while being able to skip the small-scale simulation step entirely. Tests of the model's robustness on variant fuel assemblies, other problem geometries, and changes in simulation parameters are also discussed.