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

This paper introduces JointDiff, a joint diffusion-based generative framework that unifies forward modeling, inverse inference, and output imputation to predict multi-modal simulation distributions from partial observations, demonstrating high accuracy and transferability to National Ignition Facility experiments for advancing inertial confinement fusion design.

The Gist

Imagine you are trying to solve a massive, 3D jigsaw puzzle, but you only have a few pieces and a blurry photo of the final picture. That is essentially the challenge scientists face in Inertial Confinement Fusion (ICF), a field trying to create clean energy by smashing tiny fuel pellets together.

Here is the problem:

• The Simulation (The "Perfect" World): Computer models can simulate the entire explosion in 3D. They know everything: the temperature, the pressure, the shape of the fuel, and they can "see" the explosion from every angle with perfect clarity.
• The Experiment (The "Real" World): When scientists actually run these experiments at the National Ignition Facility (NIF), they can only see a tiny fraction of that data. Some cameras get blocked, some sensors fail, and they can't measure things like the internal pressure directly. They have a "partial" picture.

The paper introduces a new AI tool called JointDiff to bridge this gap. Think of JointDiff as a super-smart, probabilistic detective that has studied millions of "perfect" computer simulations.

How JointDiff Works: The "All-in-One" Detective

Usually, AI models are like specialists: one is good at predicting the future (forward modeling), another is good at guessing the past (inverse modeling), and a third is good at filling in missing puzzle pieces (imputation).

JointDiff is different. It uses a technique called Joint Diffusion. Imagine a noisy, static-filled TV screen that slowly clears up to reveal a picture. JointDiff learns to "clean up" the noise for everything at the same time—numbers (scalars) and pictures (images). Because it learns the relationship between the numbers and the pictures together, it can do three things simultaneously:

1. The "Forward" Prediction: If you give it the starting conditions (like the pressure and shape of the fuel), it predicts what the explosion will look like and what numbers it will produce.
2. The "Inverse" Prediction: If you give it the results of an experiment (the blurry pictures and a few numbers), it works backward to guess what the starting conditions must have been.
3. The "Fill-in-the-Blank" (Imputation): If you have a picture but are missing a number (or vice versa), it can guess the missing piece based on the patterns it learned from the millions of simulations.

The "Magic" of Uncertainty

What makes JointDiff special is that it doesn't just give you one answer; it gives you a range of likely answers.

Think of it like a weather forecaster. A simple model might say, "It will rain at 2 PM." JointDiff says, "There is a 90% chance it rains between 1:45 PM and 2:15 PM, but if the wind shifts, it might be later."

In the paper, the authors tested this by hiding half the data (masking it) and asking JointDiff to guess the rest.

• The Result: Even when the AI was "blind" to 50% of the data, it could still guess the missing pieces with high accuracy.
• The Confidence: When the AI was unsure (because too much data was missing), it naturally gave a wider range of guesses. When it was confident, the guesses were tight. This helps scientists know when to trust the AI and when to be careful.

Testing on Real Life (The NIF Experiments)

The team didn't just test this on computer simulations; they tried it on real experiments from the National Ignition Facility.

• The Catch: They did not teach the AI any of the real experimental data. They only fed it the computer simulations.
• The Outcome: When they gave the AI real, messy experimental data (with missing pieces), it successfully guessed the starting conditions that would create those results.
• The Reality Check: The AI was very good at matching the general shape of the explosion and most numbers. However, it struggled with a few specific details (like a specific type of neutron scattering). This actually helped the scientists realize that their underlying computer physics model might need a small tweak to match reality better.

The Bottom Line

JointDiff is a flexible, all-in-one AI tool that acts as a bridge between perfect computer simulations and messy real-world experiments. It allows scientists to:

1. Predict what will happen before they build an experiment.
2. Figure out what went wrong after an experiment by working backward.
3. Fill in the gaps when their sensors fail.

It's like having a time machine that can show you the future, the past, and the missing pages of your diary, all based on the patterns of a million previous stories.

Technical Summary

Technical Summary: A Joint Diffusion Approach to Multi-Modal Inference in Inertial Confinement Fusion

Problem Statement
Inertial Confinement Fusion (ICF) research relies on a tight iterative loop between physics-based simulations and experimental campaigns, such as those conducted at the National Ignition Facility (NIF). A fundamental challenge exists in the data disparity between these two domains: simulations generate a rich, uniform set of multi-modal outputs (comprising both scalar quantities and 2D/3D images), whereas experiments yield only a limited, distinct subset of diagnostics. Furthermore, experimental diagnostics are often incomplete due to hardware failures, physical obstructions, or resolution limits. Traditional machine learning surrogates often struggle to model the complex, conditional joint distributions required to map between partial experimental observations and the full simulation state. Existing approaches, such as Markov Chain Monte Carlo (MCMC) for inverse inference, are computationally intensive, sensitive to prior choices, and may fail to converge. Conversely, standard generative models often fail to capture subtle inter-modal dependencies when not explicitly modeling the full joint distribution.

Methodology: The JointDiff Framework
To address these challenges, the authors introduce JointDiff, a generative framework based on joint diffusion probabilistic models. Unlike approaches that align pre-trained latent spaces or combine data prior to diffusion, JointDiff trains a single architecture to learn the full joint distribution of multi-modal data (scalars and images) simultaneously.

• Architecture: The model utilizes a multi-modal U-Net architecture. It employs convolutional networks for image data and fully connected networks for scalar data. Crucially, the architecture incorporates discrete masks that inform the model whether specific modalities (or channels within modalities) are provided as true data or are masked (missing/noisy). Time embeddings and mask embeddings are concatenated with input features at each convolution step.
• Training Objective: The model is trained on a large ensemble of over 443,000 three-dimensional Multi-Rocket Piston (RP) simulations. During training, noise is gradually added to both scalar and image modalities. The model learns to "denoise" the entire joint state. Random masking is applied to inputs and outputs during training, forcing the model to learn:
  1. Forward modeling: Predicting outputs given inputs.
  2. Inverse modeling: Predicting inputs given outputs.
  3. Imputation: Predicting missing outputs given partial observations.
• Data Modalities: The training data includes 28 scalar inputs (e.g., initial pressure, adiabat, drive symmetry modes) and a rich set of outputs: 12 scalar outputs (including Yield, Bang-time, inferred areal density $\rho R$, and residual kinetic energy RKE) and 3 image views, each with two energy channels (Primary and Down-scattered neutrons).

Key Contributions and Results
The paper evaluates JointDiff across three core prediction tasks, demonstrating high accuracy, statistical robustness, and transferability without fine-tuning on experimental data.

1. Forward Surrogate Modeling: When predicting simulation outputs from scalar inputs, JointDiff achieves excellent accuracy, with $R^2$ values ranging from 0.935 to 0.999 across scalar outputs. The model produces meaningful probability distributions rather than deterministic point estimates; approximately 92.8% of ground truth samples fall within two standard deviations of the predicted distribution. The model also successfully generates high-fidelity neutron images conditioned on inputs, capturing distinct intensity profiles across a three-order-of-magnitude yield range.
2. Inverse Modeling and Imputation: The model demonstrates robust capability in inferring inputs from outputs and imputing missing data.
   • Imputation: When specific scalar outputs (e.g., Yield, Ion temp) are masked, the model predicts their distributions with accuracy comparable to or slightly better than the forward model, leveraging correlations between outputs.
   • Inverse Inference: When predicting inputs from partial outputs, the model maintains high accuracy ($R^2$ 0.967–0.999). Sensitivity analysis reveals that specific inputs (e.g., symmetry modes) are highly dependent on specific image views (e.g., removing View 3 significantly degrades $l=2, m=1$ symmetry prediction), providing insight into diagnostic utility.
3. Round-Trip Consistency: The authors perform "round-trip" tests where predicted inputs are fed back into the forward model to reconstruct outputs. Even when 50% of the data (scalars and images) is masked, the distributions remain self-consistent. The reconstructed outputs center on the original conditioning values, verifying that the learned inverse and forward distributions are mutually consistent.
4. Transferability to NIF Experiments: The model was applied to eight recent NIF experimental shots without any fine-tuning on experimental data. Despite the RP physics model being a simplification, JointDiff successfully produced input distributions that, when passed through the forward model, reconstructed most experimental observables (scalars and images) with reasonable correlation. Discrepancies were noted in specific features (e.g., Down-scattered Ratio and subtle image details), which the authors attribute to limitations in the underlying RP physics model rather than the generative framework.

Significance and Claims
The paper claims that JointDiff establishes a flexible, robust generative surrogate for multi-modal scientific tasks characterized by partial observations. Its primary significance lies in unifying forward surrogate modeling, inverse inference, and output imputation into a single architecture capable of handling heterogeneous data types (scalars and images).

The authors assert that this approach enables:

• Tunable Uncertainty Quantification: Providing probability distributions rather than single values allows researchers to assess the confidence of predictions, particularly when diagnostics are missing.
• Diagnostic Constraint Analysis: The model can identify which specific diagnostics are critical for resolving specific physical parameters (e.g., the necessity of equatorial views for certain symmetry modes).
• Alignment of Simulation and Experiment: By generating input distributions that reproduce experimental outputs, the framework aids in understanding the constraints of current diagnostics and the limitations of underlying physics models.
• Acceleration of ICF Design: The framework offers a pathway to accelerate the design of higher-yield and more robust experiments by rapidly exploring input conditions and inferring missing data from partial experimental sets.

The work concludes that while the current implementation relies on the RP physics model, the JointDiff framework is generalizable and could be extended to higher-fidelity simulations and experimental data to further refine ICF design and understanding.