A Plug-in Interpretation of Conditioning in Score-Based Diffusion Models
TL;DR
A conditioning mechanism for diffusion models based on multi-speed joint diffusion of the target and the condition learns an unconditional joint score network and enforces conditioning at inference via a plug-in correction term, and derives explicit conditional reverse-time SDEs and approximate probability-flow ODEs.
Abstract
We propose a conditioning mechanism for diffusion models based on multi-speed joint diffusion of the target and the condition. The mechanism learns an unconditional joint score network and enforces conditioning at inference via a plug-in correction term. The plug-in term separates the conditioning contribution from the learned unconditional dynamics, offering a transparent view of how the condition steers generation of the target distribution. Building on this, we derive explicit conditional reverse-time SDEs and approximate probability-flow ODEs, enabling principled and directly comparable conditional samplers. To reduce the induced ODE--SDE discrepancy, we introduce a log-Fokker--Planck residual regularization that improves ODE sampling quality. Experiments on conditional image generation tasks demonstrate competitive performance and support the effectiveness of the plug-in conditioning view. Additional ODE--SDE comparison experiments show that the log-Fokker--Planck residual regularization improves deterministic ODE sampling.