Harnessing Low Dimensionality in Diffusion Models: From Theory to Practice: Lecture III: Diffusion Inverse Solvers for Scientific Applications

Harnessing Low Dimensionality in Diffusion Models: From Theory to Practice: Lecture III: Diffusion Inverse Solvers for Scientific Applications

Speakers: Liyue Shen, Postdoctoral Researcher, Google; Qing Qu, Assistant Professor, University of Michigan; Yuxin Chen, Assistant Professor, Princeton University

Conference: ICML 2025

YouTube: https://slideslive.com/39043830

Overview

This article delves into the third lecture of a comprehensive tutorial on diffusion models, focusing specifically on their application as inverse solvers for scientific problems. Presented by Liyue Shen, the talk bridges the theoretical foundations of diffusion models, previously covered in the first two lectures by Dr. Qing Qu and Dr. Yuxin Chen, with their profound practical implications. The core theme revolves around how the powerful data distribution priors learned by diffusion models can be leveraged to tackle challenging ill-posed inverse problems prevalent across various scientific and engineering domains.

Dr. Shen highlights the transformative potential of diffusion models, particularly in medical imaging, where they offer solutions to critical issues like reducing radiation dose in CT scans and accelerating MRI acquisition. Beyond merely showcasing existing applications, the lecture critically examines the current limitations of diffusion inverse solvers, categorizing them into efficiency, generalization, and controllability challenges. It then presents cutting-edge research, including Dr. Shen's own contributions, that aim to overcome these hurdles, pushing the boundaries of what diffusion models can achieve in real-world scientific applications.

The talk underscores the shift from conventional generative sampling to constrained posterior sampling, where diffusion models integrate learned data priors with observed measurements. This approach offers significant advantages over traditional and even supervised deep learning methods, often requiring less training data and demonstrating superior flexibility. By addressing the computational burden, enhancing adaptability to out-of-distribution data, and providing mechanisms for uncertainty quantification, this research paves the way for diffusion models to become indispensable tools for a broader spectrum of scientific discovery and practical deployment.

Background

▶ Watch: Introduction to talk: Diffusion models for scientific applications (0:00)

The concept of the inverse problem is fundamental across numerous scientific and engineering disciplines. It refers to the challenge of inferring underlying causes or properties from observed effects or measurements. Unlike forward problems, which predict effects from known causes, inverse problems are often ill-posed, meaning that a unique and stable solution is difficult or impossible to obtain without additional constraints or prior knowledge. Examples abound, from medical imaging (e.g., reconstructing internal body structures from external measurements) and black hole imaging to biological imaging and high-dimensional dynamic system solving.

In the context of medical imaging, two prominent examples illustrate the inverse problem:

• Computed Tomography (CT) imaging: The goal is to reconstruct cross-sectional images of internal volumetric structures from projections acquired at different view angles. A significant challenge arises from the desire to reduce patient exposure to radiation. Decreasing the number of projection samples, however, leads to severe artifacts when using conventional reconstruction methods like filtered back projection (FBP). The inverse problem here is to obtain high-quality CT reconstructions from sparse view projections.
• Magnetic Resonance Imaging (MRI): This modality reconstructs internal structures from measurements acquired in k-space (the Fourier domain). MRI scans are often time-consuming. To accelerate the process, researchers explore undersampling in k-space. Similar to CT, undersampling causes artifacts in the reconstructed images, posing the inverse problem of efficient and accurate reconstruction from limited frequency-domain data.

Mathematically, the inverse problem can be broadly formulated as y = Ax + n, where:

• y represents the observed measurements (e.g., k-space data in MRI, projection data in CT).
• A is the forward model or imaging system operator, mapping the desired image x to the measurement space. This often involves physics-based transformations like the Radon transform for CT or the Fourier transform for MRI, combined with sparse sampling.
• x is the high-quality image or underlying structure we aim to reconstruct.
• n denotes noise.

Historically, solving these ill-posed problems has evolved through several stages:

• Conventional Reconstruction: Relies on sufficient sampling according to sampling theory, often leading to high data acquisition requirements.
• Regularized Reconstruction: Techniques like compressed sensing emerged, introducing sparsity priors in a transform domain (e.g., wavelet sparsity) to regularize the optimization problem and enable reconstruction from undersampled data.
• Deep Learning Reconstruction: In recent years, deep learning methods have gained prominence, leveraging large datasets to learn complex mappings from y to x or to provide powerful data-driven priors. These approaches often involve supervised training with paired (y, x) data.

This sets the stage for diffusion models. As a recent family of generative models, diffusion models have demonstrated an exceptional ability to learn and represent complex data distributions. The core idea is that if a diffusion model can accurately capture the underlying distribution p(x) of high-quality, artifact-free images, this learned distribution can serve as a potent prior for solving inverse problems. This is particularly impactful because, unlike many supervised deep learning methods, the diffusion model can be trained solely on clean target data x, without requiring explicit knowledge or paired examples of the forward operator A or the degraded measurements y during its training phase. The challenge then becomes how to effectively integrate this powerful data distribution prior with the observed measurements y during the inference (sampling) process to achieve accurate and robust reconstructions.

Key Findings

▶ Watch: Inverse problems in CT imaging: reducing radiation dose (3:00)

The lecture highlights several pivotal findings and contributions in the application of diffusion models to inverse problems:

1. Powerful Data Priors: Diffusion models, trained to learn the score function of a data distribution, provide an exceptionally powerful and flexible prior for high-quality data. This prior can generalize beyond natural images to various domains, including medical images, time series, and graph structures.

1. Superior Performance for Ill-Posed Problems: Early research in 2022 demonstrated that diffusion models, when integrated as priors, can achieve state-of-the-art performance in medical imaging reconstruction (e.g., 2D CT and MRI) from sparse measurements. Crucially, they often outperform conventional optimization-based approaches and even previous supervised deep learning methods, despite the diffusion model itself being trained only on clean data x and not on (y, x) pairs or explicit knowledge of the forward operator A.

1. Flexible Inference-Time Integration: The ability to incorporate measurement consistency (y = Ax + n) during the diffusion sampling process allows a single pre-trained diffusion model to solve a wide variety of inverse problems (e.g., different sampling processes, anatomical sites, or modalities) at inference time without retraining.

1. Generalization to Diverse Inverse Problems: Methods like Diffusion Posterior Sampling (DPS) extended the application of diffusion priors to general noisy inverse problems, including linear and non-linear cases like super-resolution, deblurring with various kernels, and image inpainting, establishing standard benchmarks.

1. Identification of Core Challenges: Despite promising results, the talk identifies three significant, unsolved challenges hindering broader real-world adoption:

• Efficiency: High computational cost and memory burden, especially for high-dimensional (e.g., 3D CT with 512^3 voxels) and high-resolution data.
• Generalization: Adapting models trained on specific distributions to out-of-distribution (OOD) samples or different domains without extensive retraining.
• Controllability: The ability to reliably perform posterior sampling for tasks like uncertainty quantification, beyond just generating a single "best" reconstruction.

1. Novel Solutions for Efficiency:

• Diffusion Blend: A novel approach for 3D reconstruction using 2D diffusion priors that blends score functions from neighboring slices during training and incorporates positional encoding. This achieves better cross-sectional consistency and avoids external regularizations like Total Variation (TV), demonstrating effective 3D CT reconstruction (256x256x500 slices).
• Resample: An innovative method for integrating latent diffusion models (LDMs) into inverse problems. It ensures robust data consistency by blending prior consistency and measurement consistency in the latent space through a novel "stochastic resampling" step. This approach is more stable than latent DPS and achieved state-of-the-art results for CT reconstruction with significantly reduced training data (e.g., 2000 CT images).

1. Advancements in Generalization:

• Test-Time Domain Adaptation: Techniques like self-supervised training (similar to Deep Image Prior, DIP) and LoRA adaptation allow diffusion models to adapt to OOD data during inference, effectively overcoming artifacts without full retraining.
• Patch-based Diffusion Models: Training diffusion models on image patches with positional encoding improves efficiency and offers better generalization abilities for OOD cases compared to whole-image models.

1. Emerging Controllability: The Controllable and Constrained Sampling (CCS) method investigates the initial noise space to enable targeted control over sampled image properties (e.g., mean and variance) while adhering to measurements. This opens avenues for uncertainty estimation and generating diverse, yet consistent, solutions to ill-posed problems.

These findings collectively demonstrate the rapid evolution of diffusion models as inverse solvers, moving from theoretical promise to practical, high-impact applications while actively addressing their inherent limitations.

Technical Deep Dive

▶ Watch: General mathematical formulation of inverse problems (y=Ax+n) (5:39)

The technical foundation for diffusion inverse solvers relies on two core components from the broader diffusion model paradigm: the learning of score functions and the process of reverse denoising sampling. As elaborated in the preceding lectures, diffusion models are trained to learn the gradient of the log-probability density of data corrupted by Gaussian noise at various noise levels (the score function, ∇ log p(x_t)). This learned score function then guides a reverse stochastic differential equation (SDE) or ordinary differential equation (ODE) to denoise a pure noise sample back into a clean data sample, effectively generating new data from the learned distribution.

For inverse problems, the task shifts from unconditional generation to posterior sampling, where the generated sample x must be consistent with observed measurements y. This is formalized through Bayes' rule: p(x|y) ∝ p(y|x)p(x). Here, p(x) is the data distribution prior learned by the diffusion model, and p(y|x) is the likelihood term derived from the forward model y = Ax + n.

Early works, serving as a proof of concept (e.g., 2022 research for medical imaging), integrated this information by modifying each step of the reverse diffusion sampling process. Instead of purely following the learned score function, an additional optimization objective was introduced. This objective sought to constrain the current sample x_t to be consistent with the measurements y through the forward operator A. Specifically, it aimed to minimize the distance between Ax_t and y, effectively ensuring that the intermediate samples remained on or near the "hyperplane" defined by Ax = y. This robust integration allowed the model to leverage the powerful data prior while adhering to observational constraints.

A significant advancement came with Diffusion Posterior Sampling (DPS), which provided a more general and gradient-based guidance mechanism. DPS decomposes the posterior sampling into the prior part (p(x)) and the forward modeling part (p(y|x)). It estimates the posterior mean using Tweedie's formula, bridging the noisy latent variable x_t with the clean data estimation x_0. The key innovation is to use gradient-based guidance by backpropagating gradients from the likelihood term p(y|x) through the network. This gradient ∇ log p(y|x) is then added to the learned score function ∇ log p(x_t) to guide the sampling process at each time step t, steering the generation towards samples consistent with y. This approach is highly flexible and applicable to various linear and non-linear inverse problems.

Addressing the efficiency challenge for high-dimensional data, particularly 3D medical images (e.g., 512x512x512 CT volumes), has been a major focus. One approach involves using 2D diffusion priors for 3D reconstruction. Early attempts, like a 2023 work, reconstructed each 2D slice independently using a 2D prior and then enforced cross-sectional consistency with an external Total Variation (TV) regularization term on neighboring slices. While enabling 3D reconstruction, this introduced hyperparameter tuning and could compromise consistency.

To overcome these limitations, the speaker introduced Diffusion Blend. This method enhances the 2D diffusion prior by blending score functions from neighboring slices during the training process. It also incorporates positional encoding to inform the model about the relative Z-axis separation of slices. This allows the model to inherently learn and enforce cross-slice consistency without external regularizations, leading to high-quality 3D CT reconstructions (e.g., 256x256x500 slices).

For time efficiency, techniques like distillation and consistency models have been explored in general diffusion generation to reduce sampling steps. However, integrating these into inverse solvers is non-trivial because their training often "breaks" the theoretical consistency with the underlying SDE/ODE, making it hard to apply measurement constraints. A recent work addressed this with a two-stage constraint approach: a soft measurement constraint using ControlNet (a conditional diffusion architecture) and a hard measurement constraint through optimization, allowing efficient consistency models to be used for inverse problems.

A promising direction for both efficiency and performance is leveraging latent diffusion models (LDMs). LDMs operate in a lower-dimensional latent space, encoded by a separate encoder-decoder network, making training and sampling more efficient. The challenge, however, is guaranteeing data consistency with the measurements y due to the non-linearity introduced by the decoder. The speaker's Resample approach tackles this:

1. Unconditional Latent Estimation: At each diffusion step, an unconditional sample Z_0 (clean latent representation) is estimated from the current noisy latent Z_t.
2. Measurement-Consistent Optimization: Z_0 is then optimized to incorporate information from y. This involves passing Z_0 through the non-linear decoder to x_0, applying the forward operator A to get A(x_0), and minimizing the difference with y. The resulting optimized latent code captures measurement consistency.
3. Stochastic Resampling: This is the critical step. Instead of simply adding noise back to the optimized Z_0 (which can lead to unstable results and move off the noisy data manifold), Resample combines the information from the initial unconditional Z_0 and the measurement-optimized Z_0 to return to a noisy latent Z_t' that is both consistent with the prior and the measurements. The paper provides theoretical justification for balancing the ratio between these two sources of information. This method proved more stable than direct latent DPS.

Regarding generalization, research focuses on test-time domain adaptation. Techniques involve self-supervised training (similar to Deep Image Prior, DIP) to fine-tune the diffusion denoiser weights using only the test-time measurements. This can be combined with LoRA adaptation for parameter-efficient fine-tuning. Another approach is patch-based diffusion models, where the model is trained on image patches with positional encoding. This improves efficiency (smaller data for training) and enhances generalization to out-of-distribution (OOD) samples.

Finally, controllability explores the rich information within the diffusion model's initial noise space. The Controllable and Constrained Sampling (CCS) method investigates how to control properties of the generated samples (e.g., mean and variance) while adhering to measurements. By understanding how different initial noise vectors lead to diverse yet consistent reconstructions, CCS opens doors for uncertainty quantification and generating a set of plausible solutions, rather than just a single one, for inherently ill-posed problems.

Experimental Setup & Results

▶ Watch: Evolution of inverse problem solving: conventional to deep learning (6:30)

The experimental validation of diffusion inverse solvers has progressed significantly, demonstrating their capabilities across various scientific applications, primarily medical imaging.

Early Works (2022 - Proof of Concept):

• Application: 2D CT and MRI reconstruction from sparse/undersampled data.
• Training Data: Diffusion models were trained on datasets of clean, high-quality 2D CT lung images or 2D MRI images. Crucially, no paired (y, x) data (i.e., degraded measurements and corresponding clean images) was used during diffusion model training; only the clean target images x were required to learn the data prior.
• Baselines: Performance was compared against conventional reconstruction methods (e.g., FBP for CT, conventional k-space reconstruction for MRI) and prior deep learning-based supervised reconstruction approaches.
• Metrics: Standard image quality metrics like Peak Signal-to-Noise Ratio (PSNR) and Structural Similarity Index Measure (SSIM) were used.
• Headline Results: These initial works showed that diffusion inverse solvers achieved very promising results, often outperforming both conventional and supervised deep learning methods. A notable finding was the flexibility of a single diffusion model to solve different inverse problems (e.g., varying undersampling patterns or different anatomical sites) at inference time without retraining, purely by adjusting the measurement consistency step.

Diffusion Blend (3D CT Reconstruction):

• Application: High-resolution 3D CT reconstruction.
• Data: The method was tested on 3D CT volumes, targeting reconstructions of 256x256x500 slices.
• Baselines: Compared against methods using 2D diffusion priors with external Total Variation (TV) regularization for cross-slice consistency.
• Results: Diffusion Blend achieved superior cross-slice consistency in the Z-axis compared to TV-regularized approaches, as evidenced by visual inspection and potentially quantitative metrics (though specific numbers were not detailed in the talk, "better consistency" was highlighted). It demonstrated effective 3D reconstruction without the need for additional regularization terms and their associated hyperparameter tuning. While prior 2D-to-3D methods could take more than 20 hours for a single 3D volume, Diffusion Blend's efficiency improvements were implied by its novel training strategy, although exact time comparisons were not provided.

Resample (Latent Diffusion for Inverse Problems):

• Application: General inverse problems, including natural image tasks (super-resolution, deblurring) and medical imaging (CT reconstruction).
• Baselines: Compared against direct application of DPS to latent diffusion models (latent DPS), which often showed unstable reconstructions due to the non-linear decoder.
• Training Efficiency: A significant result was the ability to achieve state-of-the-art CT reconstruction performance by fine-tuning a latent diffusion model pre-trained on natural images using only 2000 CT images. This dramatically reduces the data requirements compared to training from scratch.
• Results: Resample demonstrated more stable and consistent reconstructions with measurements compared to latent DPS. It achieved state-of-the-art results for CT reconstruction, showcasing its effectiveness for medical images. The approach was also shown to be applicable to pixel-based diffusion models, yielding comparable results to existing state-of-the-art methods.

Generalization (Patch-based Diffusion and Test-Time Adaptation):

• Application: Overcoming out-of-distribution (OOD) challenges.
• Data: One example showed training on synthetic images of ellipses and testing on computed tomography (CT) images.
• Results: Without adaptation, significant artifacts were present. Through test-time adaptations (e.g., self-supervised fine-tuning of weights), the model was able to overcome these artifacts and produce high-quality CT reconstructions, demonstrating its ability to generalize across different domains.
• Patch-based diffusion models were shown to exhibit better generalization abilities for OOD cases compared to whole-image models, attributed to their efficiency in training and sampling.

Controllable and Constrained Sampling (CCS):

• Application: Controllable image generation and uncertainty estimation.
• Experimental Setup: Toy examples were presented where multiple samples were generated from the same diffusion model with the same condition and measurement constraint, but with different initial noise seeds.
• Results: The generated samples exhibited diversity while adhering to the constraints. By averaging these diverse samples, a mean image was obtained that could be precisely controlled to be similar to a target image, demonstrating the ability to control average properties and hint at uncertainty quantification by analyzing the variance across samples.

In summary, the experimental results consistently highlight the superior performance, flexibility, and increasing efficiency of diffusion models as inverse solvers. The ongoing research focuses on pushing these boundaries further, particularly for high-dimensional and challenging real-world scientific data.

Practical Implications

▶ Watch: Generalization of diffusion models to various data types (8:00)

The advancements in diffusion inverse solvers carry profound practical implications for various stakeholders in the AI/ML and scientific communities:

For Practitioners (e.g., Radiologists, Scientists, Engineers)

• Improved Image Quality with Reduced Data: Diffusion models enable the reconstruction of high-quality images from significantly sparse or undersampled measurements. This directly translates to reduced radiation dose in CT scans (fewer projections) and accelerated acquisition times for MRI (less k-space sampling), benefiting patient safety and clinical workflow efficiency.
• Versatile Reconstruction Tools: A single pre-trained diffusion model can be adapted at inference time to solve a diverse range of inverse problems, including varying imaging modalities, anatomical sites, or degradation types, without the need for specialized models for each scenario. This simplifies deployment and maintenance.
• Uncertainty Quantification: Emerging capabilities in controllable and constrained sampling offer the potential for uncertainty estimation. This is critical in medical diagnosis and scientific discovery, where understanding the confidence or range of plausible solutions for an ill-posed problem is as important as the solution itself.
• New Scientific Applications: Beyond medical imaging, the framework can be extended to other scientific domains like black hole imaging, dynamic systems, time series imputation (forecasting missing data), and molecular design (generating molecular structures with desired properties for drug discovery).

For Infrastructure Teams and Model Deployers

• Computational and Memory Demands: The primary challenge is the high computational cost and GPU memory requirements. Training powerful diffusion priors, especially for high-resolution 3D data (e.g., 512^3 CT volumes), demands substantial hardware resources. Efficient sampling algorithms are crucial for deployment in clinical or real-time scientific settings.
• Scalability: Deploying these models for large-scale, high-throughput applications (e.g., processing thousands of patient scans daily) necessitates highly optimized inference pipelines and potentially specialized hardware.
• Data Management: While diffusion models can be trained on clean data only, acquiring and curating large datasets of high-quality, artifact-free images for training remains a significant effort.

For Model Builders and Researchers

• Novel Architectural Designs: The need for efficiency drives innovation in architectures like latent diffusion models (LDMs) and patch-based diffusion models. Integrating these efficiently with inverse solvers, as demonstrated by Resample, is a key research direction.
• Data Consistency Mechanisms: Developing robust and theoretically sound methods to integrate measurement consistency into the diffusion sampling process (e.g., gradient-based guidance, optimization objectives, stochastic resampling) is central to the field.
• Generalization Strategies: Research into test-time domain adaptation and unified/versatile diffusion models is vital for making these tools robust to out-of-distribution (OOD) data, reducing the need for extensive retraining for every new application or data shift.
• Controllability and Interpretability: Moving beyond merely generating "good" images to understanding the posterior distribution and enabling controllable sampling enhances the scientific utility and trustworthiness of these models.

Tradeoffs and Limitations

• Computational Cost vs. Performance: More accurate and robust inverse solvers often come with higher computational burdens, necessitating a careful balance for practical deployment. Techniques like distillation or consistency models offer speedups but may require additional constraints to maintain data consistency.
• Hyperparameter Tuning: Methods like TV regularization or specific blending ratios (as in Resample) introduce hyperparameters that require careful tuning, potentially adding complexity to implementation.
• Generalization Gaps: While progress has been made, diffusion models still face challenges with highly OOD data. Test-time adaptation can mitigate this but might add to inference time.
• Unsupervised vs. Fine-tuning: While the core diffusion prior can be learned unsupervised, integrating it with specific inverse problems or efficient sampling techniques (like consistency models) might necessitate fine-tuning or additional constraints, blurring the line of "unsupervised" application.
• Complexity of Non-Linear Operators: Handling non-linear forward models, especially within latent spaces, adds significant complexity and requires sophisticated solutions like Resample to maintain stability.

In essence, diffusion inverse solvers represent a powerful paradigm shift, offering unprecedented capabilities for solving challenging scientific problems. However, their full potential hinges on continued innovation to address the practical hurdles of efficiency, generalization, and controllability, making them robust and accessible tools for a broader range of real-world applications.

Key Takeaways

• Diffusion models are powerful priors for inverse problems: They effectively learn complex data distributions, enabling high-quality reconstructions from sparse or undersampled measurements in scientific applications like medical imaging.
• Integration of prior and measurements is crucial: The key to diffusion inverse solvers is combining the learned data distribution (score function) with measurement consistency (forward model y=Ax+n) during the reverse sampling process, often via gradient-based guidance or optimization.
• Efficiency, generalization, and controllability are key challenges: Current diffusion inverse solvers face significant hurdles in computational cost for high-dimensional data, adaptability to out-of-distribution samples, and the ability to perform reliable posterior sampling for uncertainty quantification.
• Innovative solutions are emerging: Research efforts like Diffusion Blend (for 3D consistency), Resample (for stable latent diffusion integration), patch-based diffusion, and Controllable and Constrained Sampling (CCS) are actively addressing these challenges, pushing the boundaries of practical applicability.
• Latent Diffusion Models (LDMs) offer efficiency gains: Operating in a lower-dimensional latent space, LDMs can be more efficient, but their non-linear decoder requires sophisticated methods like Resample to ensure robust data consistency for inverse problems.
• Broad impact across scientific domains: Beyond medical imaging, these advancements hold promise for diverse applications, including dynamic system solving, black hole imaging, time series forecasting, and molecular design, opening new avenues for scientific discovery.

About the Speaker(s)

Liyue Shen is a key researcher in the field of diffusion models and their applications, particularly in scientific domains. In this specific lecture, as the primary speaker, she shared her recent research improvements and progress in leveraging diffusion models to solve inverse problems. Her work focuses on bridging the theoretical understanding of diffusion models with practical, real-world applications, especially in areas like medical imaging. She discussed her contributions to addressing critical challenges such as efficiency, generalization, and controllability in diffusion inverse solvers, including methods like Diffusion Blend and Resample.

The preceding two lectures in this tutorial series, which laid the theoretical groundwork for diffusion models, were given by Dr. Qing Qu and Dr. Yuxin Chen. Dr. Qu (Assistant Professor at the University of Michigan) and Dr. Chen (Assistant Professor at Princeton University) are recognized for their foundational contributions to the theoretical aspects of diffusion models, which Dr. Shen built upon in her discussion of practical applications.