Parameter Inference and Uncertainty Quantification with Diffusion Models: Extending CDI to 2D Spatial Conditioning

TL;DR

This work extends the Conditional Diffusion Model-based Inverse Problem Solver (CDI) from 1D temporal signals to 2D spatial conditioning, enabling probabilistic parameter inference directly from CBED images. By introducing a cross-modal transformer that jointly tokenizes parameters and 2D observations, the method performs diffusion-based posterior sampling $p(\boldsymbol{\theta}|\mathbf{y})$, yielding well-calibrated uncertainties that reflect parameter identifiability. Results on a physics-based CBED dataset with 13 parameters show CDI provides sharp posteriors for well-constrained quantities (e.g., geometry and structure factors) and broad, honest uncertainties for weakly informative parameters (e.g., Debye-Waller factors), outperforming regression in revealing when parameters are truly ambiguous. The findings demonstrate CDI’s capacity to generalize uncertainty-aware inference from temporal signals to spatial inverse problems, with implications for robust parameter estimation in materials science and other physics domains.

Abstract

Uncertainty quantification is critical in scientific inverse problems to distinguish identifiable parameters from those that remain ambiguous given available measurements. The Conditional Diffusion Model-based Inverse Problem Solver (CDI) has previously demonstrated effective probabilistic inference for one-dimensional temporal signals, but its applicability to higher-dimensional spatial data remains unexplored. We extend CDI to two-dimensional spatial conditioning, enabling probabilistic parameter inference directly from spatial observations. We validate this extension on convergent beam electron diffraction (CBED) parameter inference - a challenging multi-parameter inverse problem in materials characterization where sample geometry, electronic structure, and thermal properties must be extracted from 2D diffraction patterns. Using simulated CBED data with ground-truth parameters, we demonstrate that CDI produces well-calibrated posterior distributions that accurately reflect measurement constraints: tight distributions for well-determined quantities and appropriately broad distributions for ambiguous parameters. In contrast, standard regression methods - while appearing accurate on aggregate metrics - mask this underlying uncertainty by predicting training set means for poorly constrained parameters. Our results confirm that CDI successfully extends from temporal to spatial domains, providing the genuine uncertainty information required for robust scientific inference.

Figures (9)

• Figure 1: Representative CBED patterns from the simulated dataset. The patterns exhibit intensity variations that encode crystal thickness, tilt angles, Debye-Waller factors, and structure factors.
• Figure 2: CDI architecture for CBED parameter inference. A ResNet-based encoder processes the input CBED pattern and extracts spatial features, which are tokenized by the feature map tokenizer (FMT). During reverse diffusion, the transformer encoder iteratively denoises parameter tokens from random noise ($t=T$) to clean estimates ($t=0$). At each timestep, the transformer processes noisy parameter tokens conditioned on the image tokens and timestep embedding. Parameter encoder and decoder ($\mathcal{P}^E/\mathcal{P}^D$) map between parameter space and transformer embedding space.
• Figure 3: Parameter distribution comparison for unconstrained and well-constrained parameters. Left: Debye-Waller factor $u_{12}$ (Mg) distributions from ground truth (uniform), regression predictions, and CDI samples. Right: Structure factor $F_{001}$ distributions from the same three sources. Distributions are constructed from 100,000 test samples.
• Figure 4: Calibration plots for CBED parameters showing empirical coverage frequencies vs predicted confidence levels. Diagonal lines indicate perfect calibration. (Top left) Debye-Waller Mg, (top right) Debye-Waller B, (bottom left) structure factors, (bottom right) sample geometry parameters.
• Figure 5: Coverage-sharpness trade-off analysis for representative CBED parameters. Each plot shows coverage probability (y-axis) versus prediction interval sharpness (x-axis) normalized over the parameter range at various confidence levels. (Top left) Debye-Waller Mg, (top right) Debye-Waller B, (bottom left) structure factors, (bottom right) sample geometry parameters.