💻 computer science

A Plug-in Interpretation of Conditioning in Score-Based Diffusion Models

This paper proposes a plug-in conditioning mechanism for score-based diffusion models that separates conditioning from unconditional dynamics via a multi-speed joint diffusion framework, deriving explicit conditional samplers and introducing log-Fokker-Planck residual regularization to enhance deterministic ODE sampling quality.

Original authors: Libo Chen, Souvik Ghosh, Teo Deveney, Chris Budd, Vinay P. Namboodiri

Published 2026-08-21

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Libo Chen, Souvik Ghosh, Teo Deveney, Chris Budd, Vinay P. Namboodiri

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a world where computers can paint, compose music, or reconstruct blurry photographs by learning the hidden patterns of data. For years, the most powerful tools for this task have been generative models that work by starting with pure chaos—random static noise—and slowly, step by step, refining it into a clear, meaningful image. This process is like watching a blurry photograph slowly come into focus, but in reverse: the computer starts with the blur and learns the rules to sharpen it. However, real-world problems rarely ask for just any image. They ask for a specific kind of image: a face that matches a specific description, a photo where a missing part is filled in correctly, or a low-resolution picture turned into a high-definition masterpiece. This is called conditional generation. The challenge has always been how to guide the computer to create exactly what is needed without losing the natural quality of the image or having to retrain the entire system for every new request.

A team of researchers at the University of Bath and the International Institute of Information Technology Hyderabad has proposed a new way to solve this puzzle. Instead of trying to teach the computer a completely new set of rules for every specific condition, they developed a method that acts like a universal adapter. Their approach allows the computer to learn the general rules of how images form once, and then, at the very moment of creation, it simply plugs in a correction based on the specific request. This "plug-in" mechanism separates the general learning from the specific instruction, offering a clear window into how the computer decides to shape the image. The result is a system that can fill in missing parts of a photo or sharpen a blurry one with high precision, all while using a single, pre-trained model that doesn't need to be re-learned for every new task.

To understand why this is significant, consider how these systems usually work. Traditional methods often try to learn a specific path for every possible condition. If you want to fill in a missing eye in a face, the model must learn the specific relationship between the rest of the face and that missing eye. If you want to sharpen a different kind of blur, it must learn a different relationship. This makes the system rigid and difficult to adapt. Other methods try to guide a pre-trained model by constantly checking its progress against the desired outcome, but this requires heavy, repeated calculations that slow everything down. The researchers' new method takes a different path. They teach the computer to understand the relationship between two things at once: the final image it wants to create and the condition it is given, such as a partial view of a face or a low-resolution sketch. They let both the image and the condition evolve through the noise together, but at different speeds. The condition, which is usually more stable or known, changes very slowly, while the target image changes rapidly.

Once the computer has learned this joint dance of noise and structure, the real innovation happens during the creation phase. Instead of forcing the computer to relearn the rules, the researchers apply a simple, mathematical correction at the end of the process. This correction acts like a steering wheel that gently nudges the generation toward the specific condition provided. Because this correction is calculated directly from the rules of how the noise was added in the first place, it is transparent and easy to understand. The computer knows exactly how the condition influences the final result. This approach allows the system to use a single, powerful model trained on general data and then apply it to specific tasks like image inpainting (filling in missing parts) or super-resolution (making small images big) without needing to retrain the model from scratch.

The researchers tested this idea on several challenging tasks. In one experiment, they asked the system to fill in missing squares of a face, where up to 25% of the image was hidden. In another, they asked it to turn a tiny, 16-by-16 pixel image into a crisp 128-by-128 pixel portrait. The results were impressive. The new method produced images that were not only clearer and more accurate than previous approaches but also stayed truer to the original conditions. For instance, when filling in a missing part of a face, the system didn't just guess a random face; it created a face that perfectly matched the visible parts and the specific shape of the missing area. In tests comparing the quality of the generated images, the new method consistently scored higher on measures of realism and consistency than other leading techniques. It managed to balance the need for a sharp, detailed image with the need to stay faithful to the input data better than systems that tried to learn everything from scratch or those that relied on heavy, repeated calculations.

Perhaps even more importantly, the researchers showed that this method works even when using a model that was already trained by someone else. They took a powerful, pre-existing model designed to generate faces and applied their plug-in correction to it. Without changing the model's internal brain, they were able to make it perform complex tasks like repairing noisy, damaged photos. In these tests, their method outperformed other popular techniques that required the computer to constantly recalculate gradients—a complex mathematical process that slows down the generation. Their approach achieved better results in terms of image clarity and how well the image matched the original damaged version, all while running faster because it avoided those heavy calculations.

The team also looked at the mathematical stability of their method. They found that by adding a specific type of regularizer—a term that encourages the system to stay consistent with the laws of probability—they could make the deterministic version of their process (which produces the same result every time) match the performance of the more chaotic, random version. This is a subtle but important finding. It means the system can be made more reliable and predictable without sacrificing the quality of the images it produces. The researchers demonstrated this by showing that the gap between the two versions of the process narrowed significantly when they applied this extra layer of training, leading to more consistent and higher-quality outputs.

In the end, this work offers a clearer, more flexible way to guide artificial intelligence in creating specific content. By treating the condition as a simple, analytic correction rather than a fundamental part of the learning process, the researchers have created a framework that is both powerful and transparent. It suggests that we do not need to build a new engine for every new task; instead, we can build a robust engine and simply plug in the right guidance when we need it. This approach not only improves the quality of the images but also provides a deeper understanding of how the computer makes its decisions, turning a black box of complex mathematics into a process that can be seen, understood, and trusted.

Technical Summary: A Plug-in Interpretation of Conditioning in Score-Based Diffusion Models

Problem Statement

Score-based diffusion models have achieved state-of-the-art performance in unconditional generation tasks. However, many practical applications require conditional generation, where samples must be consistent with specific observations or constraints (e.g., image inpainting, super-resolution, data assimilation).

Existing approaches to conditional generation generally fall into two categories, both of which have limitations:

Fixed-condition estimators: These learn a conditional score network directly from paired data by keeping the condition fixed during the forward diffusion process. While effective, they tightly couple the model to a specific conditioning interface, making adaptation to new inverse problems difficult. Furthermore, they offer limited insight into how observations explicitly influence the reverse-time dynamics.
Joint-diffusion estimators: These jointly diffuse the target and the condition, learning a joint score field. While more flexible, existing methods often rely on implicit schedule matching and lack an explicit formulation of the conditional reverse-time Stochastic Differential Equation (SDE) or Ordinary Differential Equation (ODE). Consequently, the mechanism by which observations guide generation remains opaque.
Posterior-guided sampling (e.g., DPS, DiffPIR): These use pretrained unconditional priors and impose observations during sampling. While avoiding retraining, methods like Diffusion Posterior Sampling (DPS) require iterative likelihood-gradient evaluations at every step, increasing computational cost.

The core challenge is to develop a unified framework that provides a transparent interpretation of conditioning, separates unconditional dynamics from observation-consistency corrections, and enables principled conditional samplers without requiring per-step backpropagation through the score network.

Methodology

The authors propose a unified conditional diffusion framework based on multi-speed joint diffusion with an analytic plug-in correction.

1. Multi-Speed Joint Diffusion and Plug-in Conditioning

Instead of learning a separate conditional score for every observation, the method jointly diffuses the target $X_t$ and the condition $Y_t$ under potentially different noise schedules.

Forward Process: The pair $(X_t, Y_t)$ is diffused with independent Brownian motions but different diffusion intensities ( $g_X(t)$ and $g_Y(t)$ ). The condition channel is typically diffused more slowly (multi-speed design) to encourage a "screening" effect where the high-SNR corrupted observation $Y_t$ renders the original observation $Y_0$ nearly independent of $X_t$ given $Y_t$ .
Score Approximation: A single unconditional joint score network $s_\theta(X_t, Y_t, t)$ is trained to approximate $\nabla_{(x,y)} \log p_t(X_t, Y_t)$ .
Plug-in Correction: Conditioning is not learned implicitly but imposed explicitly at inference time via an analytic correction term derived from the forward corruption process. Using Bayes' rule and the screening approximation, the conditional joint score is decomposed as:
$\nabla_{(x,y)} \log p_t(X_t, Y_t | Y_0) \approx \nabla_{(x,y)} \log p_t(X_t, Y_t) + (0, c_y(t; Y_0, Y_t))$
where $c_y(t; Y_0, Y_t) \approx \nabla_y \log p_t(Y_0 | Y_t)$ . Under Gaussian forward corruption, this term has a closed-form expression: $c_y = (Y_0 - Y_t) / \sigma_Y^2(t)$ .
Pretrained-Prior Instantiation: This principle extends to pretrained unconditional models (e.g., for inpainting). Here, the condition is a partial observation of the target ( $Y_0 = M \odot X_0$ ). The plug-in correction is applied directly to the observed pixels of the noisy state $X_t$ , effectively treating the observed region as the condition channel.

2. Explicit Conditional Reverse Dynamics

Based on the plug-in correction, the authors derive explicit formulations for both stochastic and deterministic sampling:

Conditional Reverse SDE: The reverse-time SDE includes the plug-in term in the drift of the condition channel, guiding the generation of $X_t$ through the learned joint score's dependence on $Y_t$ .
Conditional Probability-Flow ODE: A matched approximate probability-flow ODE is derived by substituting the plug-in corrected score into the standard ODE formulation. This provides a deterministic sampler with interpretable conditioning dynamics.

3. Log-Fokker–Planck Residual Regularization

To address the discrepancy between the deterministic ODE sampler and the stochastic SDE sampler (which can degrade ODE sample quality), the authors introduce a log-Fokker–Planck (log-FP) residual regularization.

They define a log-FP operator $F_{Y_0}[u]$ associated with the corrected score dynamics.
The training objective is augmented with a term minimizing the $L_2$ residual of this operator: $R_{LFP}(\theta, t)$ .
Theoretical Guarantee: The authors prove that a small log-FP residual bounds the Wasserstein-2 distance between the ODE and SDE distributions ( $W_2(p^{ODE}, p^{SDE}) \leq C \delta$ ). Thus, minimizing this residual ensures the deterministic ODE sampler closely tracks the reverse-time SDE distribution.

Key Contributions

Plug-in Conditioning Interpretation: The paper provides a new interpretation of conditional diffusion where conditioning is an analytic plug-in term added to a learned unconditional joint score. This separates unconditional generation dynamics from observation-consistency corrections, offering transparency. It demonstrates that this formulation works with both jointly trained models and pretrained priors (e.g., for inpainting).
Explicit Conditional Dynamics: The authors derive explicit conditional reverse-time SDEs and matched approximate probability-flow ODEs. The conditioning contribution appears in closed form, clarifying how observations influence both stochastic and deterministic trajectories.
ODE–SDE Discrepancy Control: A theoretical connection is established between the conditional ODE–SDE discrepancy and a controllable log-Fokker–Planck residual. The proposed regularization mechanism improves the quality of deterministic ODE sampling by reducing this gap.

Experimental Results

The framework was evaluated on conditional image generation tasks (inpainting and super-resolution) across CelebA, FFHQ, and ImageNet.

Conditional Generation Performance:
- On CelebA inpainting, the method achieved the best PSNR (25.38), SSIM (0.876), and JFID (15.99), outperforming fixed-condition (CDE) and joint-diffusion baselines (CDiffE, CMDE). Notably, it improved JFID significantly over CMDE, indicating better joint realism and consistency with the observation.
- On CelebA super-resolution (8x), the method achieved the best PSNR (25.12), LPIPS (0.098), and JFID (14.56), demonstrating superior conditional consistency compared to HCFLOW and other diffusion baselines.
Pretrained-Prior Application:
- In noisy random inpainting on FFHQ and ImageNet256, the method was compared against DPS and DiffPIR using the same pretrained prior.
- The proposed method outperformed DPS in both reconstruction fidelity (PSNR) and conditional consistency (JFID) without requiring the per-step U-Net backpropagation that DPS demands.
- It remained competitive with DiffPIR, achieving the best PSNR and JFID on ImageNet256 and the best LPIPS and PSNR on FFHQ.
ODE–SDE Gap Reduction:
- On MNIST conditional inpainting, the log-FP residual regularization (CLOSE) reduced the sliced Wasserstein-2 distance between the ODE and SDE samplers ( $W_2(ODE, SDE)$ ) from 0.0329 to 0.0323.
- Crucially, this reduction in discrepancy did not degrade sample quality; the regularized model also improved alignment with ground truth ( $W_2(ODE, GT)$ and $W_2(SDE, GT)$ ).

Significance and Claims

The paper claims that its primary significance lies in providing a transparent and modular view of conditioning. By treating conditioning as an analytic correction to a learned score, the method:

Decouples the learning of generation dynamics from the specific observation operator.
Avoids the computational overhead of iterative likelihood-gradient evaluations (unlike DPS).
Offers a principled way to derive both stochastic and deterministic conditional samplers.
Provides a theoretical mechanism (log-FP residual) to control the discrepancy between ODE and SDE samplers, thereby improving deterministic sampling quality.

The authors acknowledge limitations, noting that the closed-form correction relies on linear-Gaussian forward kernels and that the screening approximation may be less accurate for highly multimodal posteriors or high noise levels. They also note that while the method is conceptually complementary to Classifier-Free Guidance (CFG), combining them is left for future work.