Blind Drifting: Diffusion models with a linear SDE drift term for blind image restoration tasks
Abstract
In this work, we utilize the high-fidelity generation abilities of diffusion models to solve blind image restoration tasks, using JPEG artifact removal at high compression levels as an example.
We propose a simple modification of the forward stochastic differential equation (SDE) of diffusion models to adapt them to such tasks.
Comparing our approach against a regression baseline with the same network architecture, we show that our approach can escape the baseline’s tendency to generate blurry images and recovers the distribution of clean images significantly more faithfully, while also only requiring a dataset of clean/corrupted image pairs and no knowledge about the corruption operation.
By utilizing the idea that the distributions of clean and corrupted images are much closer to each other than to a Gaussian prior, our approach requires only low levels of added noise, and thus needs comparatively few sampling steps even without further optimizations.
Introduction
Diffusion models have taken the world of machine learning by storm due to their unprecedented ability to generate high-fidelity images and great flexibility to condition on a variety of user inputs.
Previous works on diffusion models have largely concentrated on unconditional and conditional image generation.
Recently, Bansal et al.
[1] and Daras et al.
[2] have proposed to extend and improve diffusion models by adding known deterministic corruptions, but have only evaluated their ideas for unconditional generation tasks, rather than for faithfully inverting corruptions to restore plausible original images.
Here, we try to go a different route, and explicitly modify and train diffusion models to restore plausible images from corrupted ones.
In contrast to other recent works [2, 3], our approach does not require the underlying corruption operator to be known, linear nor differentiable and instead requires only a dataset of (clean image, corrupted image) pairs.
We also propose an alternative way of combining a deterministic corruption with noise.
In this work, we consider JPEG compression with low quality levels (10–20 % ), as an example corruption with a nonlinear and nondifferentiable corruption operator that is unknown to the restoration procedure at inference time.
We propose a simple modification of the forward stochastic differential equation (SDE) of diffusion models to adapt them to such tasks, building upon previous work from the speech processing literature [4, 5] but extending their modification to a more general form and newly applying these ideas to a nonlinear inverse problem in the image domain.
Methods 2.1 A
Following [6], the forward process of a diffusion model can be interpreted as a dynamical system following a stochastic differential equation
[Equation]
where in this work, x ∈ R C × H × W is the current image and w is a standard Wiener process of the same dimensionality as x , and the process runs forward from t =
t ε until t =
T := 1 , with t ε ⪆ 0 for numerical reasons [6].
Each image in the training dataset then represents the initial value x 0 of a particular realization of this SDE.
Song et al.
[6] showed that previous diffusion models in the discrete-time domain can be interpreted to follow either the so-called Variance Exploding (VE) SDE or the so-called Variance Preserving (VP) SDE .
Both SDEs have the aim of progressively turning images into Gaussian white noise, thereby turning the intractable image distribution into a tractable prior.
To generate images, one then samples from this prior and numerically solves the corresponding reverse SDE [7],
[Equation]
where the only unknown term is the score ∇ x t log p t ( x t , t ) .
A deep neural network called a score network S θ ( x t , t ) is then trained to estimate this score, given the current process state x and time t .
Rather than turning the clean image distribution into pure noise, here we aim to turn the clean image distribution into a noisy version of the corrupted image distribution.
This has two purposes:
Instead of pure noise, this uses the corrupted image (plus tractable noise) as the initial value of the reverse SDE, thus achieving the task adaptation through the formulation of the process itself, as opposed to only providing the corrupted image as conditioning information.
Since the added Gaussian noise is white, it functions as a continual source of all possible spatial frequencies throughout the reverse process.
The trained score model then filters these frequencies appropriately, generating plausible clean image estimates without a loss of high-frequency detail
Note that the distribution of noisy corrupted images is still “tractable” for the purposes of the restoration task, as a sample from it can be drawn by taking a corrupted image and adding Gaussian noise.
We now propose the following family of linear forward SDEs to realize our idea:
[Equation]
where y is the corrupted image corresponding to x 0 , γ is a stiffness hyperparameter controlling how strongly x t is pulled towards y , and α ∈ R ≥ 0 controls the shape of the curve pulling x t towards y .
ν is a normalization factor determined to ensure that σ T ≈ σ max , where σ t is the closed-form variance of the Gaussian process described by (3).
Intuitively, our family of SDEs combines the diffusion g of the VE SDE with an added drift term f that pulls x t towards the corrupted image y .
2.2 The two considered SDEs
In this work, we will only consider α ∈ { 0 , 1 } for simplicity.
The case of α = 0 was previously proposed for similar tasks in the speech processing literature [4], and we refer to it as the OrnsteinUhlenbeck Variance Exploding (OUVE) SDE.
We also newly propose using α = 1 , which we call the t-squared Decay Variance Exploding (TSDVE) SDE.
To allow for efficient forward sampling in order to perform denoising score matching [6], we determine closed-form expressions of the mean µ t and variance σ t of the Gaussian processes described by each SDE, see for instance [8].
Since the
SDEs are linear in the state x t , the means of the two SDEs follow
[Equation]
Both expressions describe a linear interpolation between x 0 and y , with the interpolation parameter controlled by an exponential (OUVE) or half-Gaussian-shaped (TSDVE) decay over time.
We relegate the somewhat involved closed-form expressions for the variance to Appendix A.
For both SDEs, µ t ̸ = y for all finite t , which may seem like an issue since the aim was to have the process move towards y .
However, letting z ∼ N (0 , I ) , it is only required that the distributions of ( µ T + σ T z ) and ( y + σ T z ) are similar, so that the latter can function as a plausible initial value for the reverse sampling process.
We can control how well the distributions of these two expressions match, either by increasing the stiffness γ at the cost of potentially destabilizing the reverse process, or by increasing σ max to further smooth the density functions of both distributions at the cost of more reverse iterations.
Here, we choose a set of parameters that empirically work well, and leave further optimization of them to future work.
2.3 Dataset
For x 0 , we use the CelebA-HQ dataset [9], resized to 256x256 and split into 24000 images for training, 1500 for validation and 4500 for testing.
To generate y , we first randomly sample a JPEG quality value from 0 to 30 for each image x 0 and training iteration, and then apply JPEG compression to x 0 .
During evaluation, we set the JPEG quality values to a constant value across all compared images and models (but do not provide this quality value to any model).
2.4 Network training and process parameterization
We utilize the NCSN++ architecture [6], both for the regression baseline and as a score network.
We train the regression baseline R θ to recover x 0 given y via a simple L 2 loss:
[Equation]
(5) where we pass a constant “dummy” value of t = 1 to the time embedding layers of NCSN++ to avoid making any changes to the DNN that may affect the qualitative behavior of each layer.
To train the score models S θ , we use the idea that diffusion models based on such linear SDEs can still be trained by the same denoising score matching target as in [6]:
[Equation]
with z ∼ N (0 , I ) , and image pairs ( x 0 , y ) sampled from the dataset.
The only changes in this objective are that the expectation is also calculated over y , and that y is provided to the score network.
We provide y as an input to S θ by concatenating x t and y along the channel dimension, for which we change the number of input channels of NCSN++ from 3 to 6.
For training, we use the
AdamW optimizer [10], and set the hyperparameters of the process and training as listed in Table 1.
For evaluation, we use each method’s checkpoint with minimum loss on the validation set.
In line with the diffusion model literature [6], we update an exponential moving average (decay of 0.999) of all network parameters after each training step, and use these parameters for evaluation.
[Table]
Results and Discussion
We generate reconstructions from our diffusion models with the Euler-Maruyama sampler discretized to N = 100 steps, and retrieve reconstructions from the regression baseline via a single pass.
In Figure 1, we compare the regression baseline against our SDE-based diffusion models for JPEG quality level 10.
Qualitatively, the baseline reconstructs major features well, but fails to produce plausible high-frequency details.
This is particularly visible for hair, facial hair and skin textures, and results in a painting-like look.
In contrast, our approach subjectively results much more natural-looking images and, for both proposed SDEs, shows a remarkable ability to reconstruct natural skin and hair textures that are plausible given the corrupted image.
Further example images can be found in Appendix B. Both SDEs generally result in perceptually very similar images.
In Table 2, we compare the distributions of the corrupted and reconstructed images of each method against the distribution of the ground-truth images, using the metrics FID [11] and KID [12].
We also list the average SSIM [13] and LPIPS [14] values, and evaluate all metrics on our test set of 4500 images.
Judging from the distribution-based metrics, both of our diffusion-based approaches model the clean image distribution of CelebA-HQ significantly more faithfully than the baseline, which even performs worse than the corrupted (compressed) images in this regard.
On the other hand, the baseline achieves a significant SSIM improvement whereas our proposed models do not.
This is to be somewhat expected due to the generative nature of our approach.
We argue that SSIM does not match human perception well here: the blurry look of the baseline images is clearly visible, but does not seem to be strongly penalized.
Indeed, for the perceptual LPIPS metric our methods consistently outperform the baseline.
Both proposed SDEs perform very similarly in all regards.
[Table]
Conclusion
Based on work from the speech processing literature, we propose a simple change to SDE-based diffusion models to adapt them to image restoration tasks.
Our approach does not require the corruption operator to be available in closed form and does not impose strong restrictions on its nature, but only requires a dataset of paired images.
Compared to a regression baseline using the same architecture, our approach restores images of perceptually higher quality and models the ground-truth image distribution significantly more faithfully.
While currently requiring 100x more DNN passes than the baseline, we expect that recent progress on efficient sampling for diffusion models will allow this number to drastically decrease and make the approach competitive in terms of runtime.
Acknowledgments and Disclosure of Funding
We acknowledge the support by DASHH (Data Science in Hamburg - HELMHOLTZ Graduate School for the Structure of Matter) with the Grant-No. HIDSS-0002.
A Closed-form variance solutions of the OUVE and TSDVE SDEs
In the following, we provide the expressions for the diffusion normalization factors ν and the solved variance of the Gaussian process, for the OUVE SDE and the TSDVE SDE.
We utilized the software
Mathematica 12.1 to solve the mean and variance ODEs [8] for the initial value σ 0 = 0 .
This choice is in contrast to [6], where the initial value σ 0 = σ min was used.
Our reasoning for this change is that due to the influence of our added drift terms, the choice by Song et al.
may result in nonmonotonous functions σ t , particularly depending on the choice of γ and the relative scale of σ max in comparison to σ min .
We admit that the convenient name of σ min may be misleading under this assumption, since it is not a “minimum sigma” under our assumptions.
Nonetheless, we keep this parameter as a way to control the shape of the variance curve, and also keep its name in line with the previous literature.
For the OUVE SDE, we follow [6] to determine an approximate normalization factor:
[Equation]
resulting in the variance
[Equation]
One may observe that the condition σ 1 = σ max is only approximately fulfilled here, but the error is small when σ min is small and γ is reasonably large: for our parameter choice, it is equal to
[Equation]
and we therefore use this factor ν OUVE due to its relative simplicity.
For the TSDVE SDE, we first solved for the unnormalized variance expression using ν = 1 : σ 2 t, TSDVE,
[Equation]
where D is the Dawson function.
Solving for ν such that ν 2 σ 2 T = σ max exactly, we retrieve
[Equation]
We then multiply the diffusion term g ( t ) of the forward SDE by this ν TSDVE , finally resulting in
[Equation]
Note that we are now already dealing with unwieldy functions such as the Dawson function just by setting α = 1 (TSDVE SDE).
This was the principal reason for us to not investigate α =and higher, though exploring a solution or approximation of σ t for these cases may be helpful:
Increasing α keeps the process mean close to x 0 for longer early in the forward process, which may help the reverse process to result in higher-quality reconstructions; however, in practice we did not find significant differences between α = 0 and α = 1 in this work.
B More example images
[Figure] [Figure]