A joint diffusion approach to multi-modal inference in inertial confinement fusion

TL;DR

This work introduces a generative framework, called JointDiff, which enables predictions of conditional simulation input and output distributions from partial, multi-modal observations, and establishes JointDiff as a flexible generative surrogate for multi-modal scientific tasks.

Abstract

A combination of physics-based simulation and experiments has been critical to achieving ignition in inertial confinement fusion (ICF). Simulation and experiment both produce a mixture of scalar and images outputs, however only a subset of simulated data are available experimentally. We introduce a generative framework, called JointDiff, which enables predictions of conditional simulation input and output distributions from partial, multi-modal observations. The model leverages joint diffusion to unify forward surrogate modeling, inverse inference, and output imputation into one architecture. We train our model on a large ensemble of three-dimensional Multi-Rocket Piston simulations and demonstrate high accuracy, statistical robustness, and transferability to experiments performed at the National Ignition Facility (NIF). This work establishes JointDiff as a flexible generative surrogate for multi-modal scientific tasks, with implications for understanding diagnostic constraints, aligning simulation to experiment, and accelerating ICF design.

Figures (11)

• Figure 1: JointDiff is generative surrogate for forward, inverse, and imputation tasks. Simulations (upper row) produce a complete set of multi-modal outputs given a complete set of inputs. Surrogate models are trained to replace simulations and perform the forward or inverse predictions tasks. Experiments (bottom row) produce incomplete sets of outputs, and inputs to the capsule are unknown. In addition to forward and inverse surrogate capabilities, JointDiff predicts conditional distribution of inputs and outputs given partial multi-modal observations.
• Figure 2: Architecture of the JointDiff model. Outputs images contain primary (green) and down-scattered (pink) neutron intensities, which are encoded as separate image channels. Trapezoids represents data encoders and decoders, which are convolutional networks for images and fully connected networks for scalars data. Gray and blue boxes correspond to discrete masks which inform the model if true or noisy data is provided for each modality and are embedded by a fully connected network. Diffusion time is encoded by a sin/cos positional embedding. The time, mask, and input/output embeddings are concatenated to the image embeddings at each convolution step. Individual decoders predict noise to be removed at diffusion time $t$ for each image and set of input and output scalars.
• Figure 3: JointDiff predicts accurate distributions of RP simulation outputs. A) For each set of test inputs, 10 outputs predictions are made. Error bars show the mean and standard deviation across these 10 predictions. $R$2 metrics are computed between the predicted means and the simulated ground truth. The percent of ground truth samples within two standard deviations of the predicted distribution is also shown for each scalar. A subset of 8 outputs are shown with all inputs in Figure \ref{['fig:outputs_all']}. B) Primary neutron images (View 1) sampled from each quintile of yield in the test data. The first row shows the ground truth simulated images, the second row shows the mean model prediction across generated samples, and the third row shows MAE between each generated sample and the ground truth.
• Figure 4: Sensitivity of input predictions to removal of outputs image views and channels. Mean average error for six inputs when providing all outputs (blue) and when removing partial image information (each view and each channel). All inputs are shown in the Figure \ref{['fig:mae_all']}
• Figure 5: Input and round-trip distributions for single test sample. A) Distribution of inputs for 500 generations when providing model all outputs (blue) or half of outputs (orange). Ground truth inputs are shown by black dashed line. Distribution of all training data is shown in gray. B) Round-trip output distributions produced by running inputs from A through the forward model. Both round-trip distributions of given outputs (green shading) are similar and centered on the original outputs, despite originating from different inputs distribution. Round-trip distributions for masked outputs (orange shading) are less consistent but retain significant overlap. C) Round-trip images produced by running inputs from A through the forward model. Simulated images (first row) are compared to mean images of full outputs (second row) and partial outputs (third row).