Conditional Variational Diffusion Models
Authors: Gabriel Della Maggiora
Luis Alberto Croquevielle
Nikita Deshpande
Harry Horsley
Thomas Heinis
Artur Yakimovich
Venue: ICLR 2024 poster (2024)
Keywords: Denoising Diffusion Probabilistic Models
Inverse Problems
Generative Models
Super Resolution
Phase Quantification
Variational Methods
Primary area: generative models
TL;DR: We introduce a method to automatically learn the schedule for diffusion models during training, improving results in various applications without the need for manual schedule fine-tuning.

Abstract: Inverse problems aim to determine parameters from observations, a crucial task in engineering and science. Lately, generative models, especially diffusion models, have gained popularity in this area for their ability to produce realistic solutions and their good mathematical properties. Despite their success, an important drawback of diffusion models is their sensitivity to the choice of variance schedule, which controls the dynamics of the diffusion process. Fine-tuning this schedule for specific applications is crucial but time-consuming and does not guarantee an optimal result. We propose a novel approach for learning the schedule as part of the training process. Our method supports probabilistic conditioning on data, provides high-quality solutions, and is flexible, proving able to adapt to different applications with minimum overhead. This approach is tested in two unrelated inverse problems: super-resolution microscopy and quantitative phase imaging, yielding comparable or superior results to previous methods and fine-tuned diffusion models. We conclude that fine-tuning the schedule by experimentation should be avoided because it can be learned during training in a stable way that yields better results. The code is available on https://github.com/casus/cvdm

Published as a conference paper at ICLR 2024
CONDITIONAL VARIATIONAL DIFFUSION MODELS
Gabriel della Maggiora Luis Alberto Croquevielle
Nikita Deshpande Harry Horsley", Thomas Heinis", Artur Yakimovich 125.
I Center for Advanced Systems Understanding (CASUS), Görlitz, Germany
2 Helmholtz-Zentrum fresden-Rossendorffs e. V. (HZDR), Dresden, Germany
3 Department of Computing, Imperial College London, London, United Kingdom
Bladder Infection and Immunity Group (BIIG), UCL Centre for Kidney and Bladder Health,
Division of Medicine, University College London,
Royal Free Hospital Campus, London, United Kingdom
5 Institute of Computer Science, University of Wroclaw, Wroclaw, Poland
correspendence: g. waa.elllmmagiiioraaaaaeeesseeeee
aac6228ic. ac. uk, a.yakimovichhhrdr.cco
ABSTRACT
Inverse problems aim determine parameters from observations, crucial task in engineering and science. Lately, generative models, especially diffusion models, have gained popularity in this area for their ability produce realistic solutions and their good mathematical properties. Despite their success, an important drawback of diffusion models is their sensitivity to the choice of variance schedule, which controls the dynamics of the diffusion process. Fine-tuning this schedule for specific applications is crucial but conseonnnuingggg and does not guarantee an optimal result. We propose novel approach for learning the schedule as part of the training process. Our method supports probabilistic conditioning on data, provides high-quality solutions, and is flexible, proving able to adapt to different applications with minimum overhead. This approach is tested in two unrelated inverse problems: super-resolution microscopy and quantitative phase imaging, yielding comparable or superior results to previous methods and fine-tuned diffusion models. We conclude that fine-tuning the schedule by experimentation should be avoided because it can be learned during training in a stable way that yields better results
INTRODUCTION
Inverse problems deal with the task of finding parameters of interest from observations. Formally, for mapping A y X' and x E X (the data), the inverse problem is find input y E y such that A(y)= = x. Examples abound in computer science. For instance, single-image super-resolution can be formulated this setting. Inverse problems are usually ill-posed, which means that observations underdetermine the system or small eтrors in the data propagate greatly to the solution.
Deep networks trained with supervised learning have recently gained popularity for tackling inverse problems in contexts where input-data samples are available. Depending on the application, the optimization might target the L1 or L2 norm, or pixel-wise distance measures image-based tasks. These methods are effective in many scenarios, but the supervised learning approach can introduce undesired artifacts (Lchtinen et al. (0188) and it does not yield uncertainty estimates the solutions.
Since most inverse problems are ill- ill-posed, it makes sense model the uncertainty explicitly. This can done by considering y and as realizations of random variables Y and X, respectively, and learning conditional probability distribution Pylx. Several methods available to learn et al.
a distribution from pairs of samples (Isola et al. 2016 Peng Li 2020
"Equal contribution
'The code is available on https://githhb. http://gtthuuu.cccc//ss//s//////

---

Published as a conference paper at ICLR 2024
2015 Kohl al., 2018). In this work, diffusion models, likelihood- method with ofate-of-the-arr generation capabilities (Dhariwal & Nichol 2021 Saharia et al. 2021).
Diffusion models produce high quality samples and offer stable training (Ho et al. 2020; Kingma
ct al. 2023). Despite these advantages, they have few drawbacks, like the number of steps required
to generate samples (Song al. 2022) and their sensitivity to the choice of variance schedule (Saharia
et a 2021 Nichol & Dhariwal 2021). ). The variance schedule controls the dynamics of the diffusion
process, and it is usually necessary to fine-tune with hyperparameter search for each application.
This is time-consuming and leads to suboptimal performance (Chen et al. 2020).
In this work, introduce the Conditional Variational Diffusion Model (CVDM), flexible method to learn the schedule that involves minimum fine-tuning, Our detailed contributions are:
Following Kingma et al. (023)) we learn the schedule as part of the training, extending their approach to the conditioned case. Furthermore, we allow for learning different schedule for each element the output (e.g., a pixel-wise for images). These extensions require several technical novelties, including a separation-off-varibbe
strategy for the schedule.
We prove that the rate of convergence of the discrete-time diffusion loss to the continuous- time case depends strongly on the derivatives of the schedule. This shows that the finding in KKingma al. (2023) that continuous-ting diffusion loss is invariant under choice of schedule may not hold in practice. Based on this result, we introduce novel regularization term that proves be critical for the performance of the method in our general formulation.
We implement the schedule by replacing the architecture in Kingma al. (2023) with two networks, one required to be positive for the conditioning variable and one monotonic- convolutional network. This allows us to test our model with inputs of different resolutions without retraining. Moreover, thanks to our clean formulation, our method does not need the post-processing of the schedule that Kingma et al. al.] (2023) introduces, nor the preprocessing of the input. This further contributes to streamlining our implementation.
We test CVDM in three distinct applications. For super-resolution microscopy, our method shows comparable reconstruction quality and enhanced image resolution compared to previous methods. For quantitative phase imaging, it significantly outperforms previous methods. For image super-resolution, reconstruction quality is also comparable to previous methods. CVDM proves to be versatile and accurate, and it can be applied to previously unseen contexts with minimum overhead.
2 RELATED WORK
Diffusion probabilistic models (DPMs) DSchl-Dickstein et al. 2015)
subset of variational
autoencoders methods (Kingma & Welling 2022] Kingma et al. 2023). Recently, these models
have demonstrated impressive generation capabilities (Ho et al. 2020; Dhariwal & Nichol [02I],
performing better than generative adversarial networks (Isola et al. (016) while maintaining greater training stability. Fundamentally, diffusion models operate by reversing process wherein the structure of an imput is progressively corrupted by noise, according to variance schedule.
Building on the DPM setup, conditioned denoising diffusion probabilistic models (CDDPMs) have been introduced to tackle challenges like image super-resolution (Saharia et al. 2021). inpainting, decompression (Saharia al., 2022), and image deblurring (Chen al. (023). All the cited works use different variance schedules to achieve their results, which is to be expected, because fine-tuning the schedule is usually a prerequisite to optimize performance (Saharia et al. 2021).
Efforts have been made to understand variance schedule optimization more broadly. Cosine-shaped schedule functions (Nichol & Dhariwal (02T) are now the standard in many applications (Dhariwal & Nichol 2021 Croitoru et al. 2202)), showing better performance than linear schedules in various settings. Beyond this empirical observation, Kingma et al. (023) developed a framework (Variational Diffusion Models, or VDMs) learn the variance schedule and proved several theoretical properties. Their work supports the idea that learning the schedule during training leads to better performance.
Specifically, Kingma et al. (203) formulate Gaussian diffusion process for unconditioned distribu- tion sampling. The latent variables are indexed in continuous time, and the forward process diffuses the input driven by learnable schedule. The schedule must satisfy a few minimal conditions, and its
2

---

Published as a conference paper at ICLR 2024
parameters are defined monotonic network. By forcing the schedule to be contained within range, the model can be trained by minimizing a weighted version of the noise prediction loss. Their work also introduces Fourier features to improve the prediction of nigh-frequency details.
In this work, we extend VDMs to the conditioned case. To achieve this, we define the schedule function of two variables, time t and condition x. This requires careful consideration because the schedule should be monotonic in & and not necessarily with respect to x. To solve this issue, we propose novel factorization of the schedule and derive several theoretical requirements for the schedule functions. By incorporating these requirements into the loss function, we can train the model straightforward way and dispense with the schedule post-processing of Kingma et al. (2023). We adapt the framework proposed by (Saharia et al. (202I). which learns a noise prediction model that takes the variance schedule directly as input. To streamline the framework, we eliminate the noise embedding and directly concatenate the output of our learned schedule to the noise prediction model.
3 METHODS
3.1
CONDITIONED DIFFUSION MODELS AND INVERSE PROBLEMS
For brief introduction inverse problems, see Appendix A For a formulation of non-conditioned DPMs, see see Ho et al. (2020) Kingma et al. (2023). In conditioned DPMs we are interested in sampling from a distribution Pylx=x yy) which we denote p(y[x). The samples from this distribution can be
considered as solutions to the inverse problem.
Following Kingma et al. (2023), we use continuous-time parametrization of the schedule and latent variables. Specifically, the schedule is represented by function >(,, x) and for each E [o. 1] there is a latent variable z₁. We formulate the diffusion process using a time discrctization, with continuous-time time introduced as a limiting case. This formulation is better for the introduction of key concepts like the regularization strategy in Section 3.4. To start with, for a finite number of steps T, let i/T for i (0, 1,.., T), and define the forward process by
q( ly, = -N (Vi(t,, (), x)y, (ti, x)1)
(I)
We use the prreeerrvgg condition (t.,, x) simplify notation. The whole process is conditioned on y: it starts at Zı, and it gradually injects noise, so that each subsequent Zt, a noisier version of y, and Zty is almost a pure Gaussian variable. The forward process should start y in L = O, so for all we enforce the condition 70, 1.
(t,,, x) but sometimes write (1,, x) to
We follow Kingma et al. (2023) learning the schedule instead of fine-tuning it as a hyperparameter, with some key differences. If has vector representation, we learn a different schedule for each component. For example, in the case of images, each pixel has a schedule. This means that functions of y apply element- wise and expressions like (1t,, x)y represent a Hadamard product.
Similarly to the non-conditioned case, the diffusion process is Markovian and each step 0 < i T is characterized by a Gaussian distribution qZz, 2t.... x) (see Appendix B.I]. The posterior q(z,_,, Lz, y, x) also distributes normally, and the parameters ultimately depend on the function (see Appendix B.2.. 2. The reverse process chosen to take the natural shape
(zt, Zt. x) g(zt. Ze., y Sr(Ze, 44, x),x),
such that only differs from the posterior in that y is replaced by a deep learning prediction 8r. The forward process should end in an almost pure Gaussian variable, so g(ze, x).x) ~ (O( I) should hold. Hence, for the reverse process we model Zt, with
Po (2t xx) = (O,
(3)
and without dependence on y, can use them to sample y. Specifically, we can sample from pr (zz, x)) and then use relation (2) repeatedly until reach 2ι, y. y. All the relevant distributions Gaussian, this procedure computationally feasible.
for all x. The reverse process is defined by equations
In other words, equations (2) and completely define the reverse process, such that sample any 2t, conditioned on x, including Z,, = y. we are able to learn the p. (zt,., Zt, x), we should have a good proxy for p(y[x) from which we can sample.
3

---

Published as a conference paper at ICLR 2024
3.2 DEFINING THE SCHEDULE
We now describe in more detail our approach to learning the schedule. Recall that the latent variables 2, are indexed by a continuous time parametert E [0, 1]. and we introduced the diffusion process using time discretization {4, We now consider the non-discretized case, starting with the forward process. In this case, the mechanics described by equation can be extended straightforwardly to the continuous case (z).
The continuous-time version the forward Markovian transitions and the posterior distribution are more complicated. Consider the forward transitions, which in the discretized version we denote by q(z, Zt; x). To extend this idea to the continuous case, Kingma et al. (2023) consider g x) for < t. We use different approach and focus on the infinitesimal transitions ((z x).. This idea can be formalized using stochastic calculus, but we do not need that framework here. From Appendix B.1 the forward transitions are given by
g(Zz, (x, = -
Bry,,r(I,, 1r1
(T x)y,
where Br (t; =1 (t,, x)/r(tt-1,x).. As defined, the values (t,, x) control the change in the latent variables over short period of time. Now, consider the continuous-time limit T - The intuition is that there is function / such that the change of the latent variables over an infinitesimal period of time dt is given by (t, x).t. In the discretization, this becomes ( This idea leads to the following relation between y and 3 (details in Appendix P
(t,, x) xt),,)))
(x, (t,,xx.
In view of this relation, our approach is as follows. First, we make the assumption that / can be decomposed into two independent functions, respectively depending the time and the data x. We write this as ((4, x) (x))))x(x), where both т and A are learnable positive functions. This assumption takes inspiration from many separable phenomena in physics, and it proves to be general enough in our experiments. Moreover, (tt, x) should be decreasing in t, since the forward process should start at in 0 and gradually inject noise from there. This monotony condition is much simpler to achieve in training if the and x variables are separated in the schedule functions. Replacing this form of /3 in equation and integrating with initial condition (0, x) 1, we get
x(,x)) = Ax(x) Ig ro(s)ds
(5)
Equation (5) could be used to compute r during training and inference, but we follow a different approach to avoid integration. Since To is defined to be positive, its integral is increasing function. This motivates the parametrization of 7 7(1, x) (0))))((() where Px is learnable function which is increasing by design. Summarizing, we parametrize 3 and 7 such that they share the same function (x(x) and separate denendependent functions. A priori, Px and the integral of тy could not coincide. So, to ensure that equation (4) holds, we use the Deep Galerkin Method (Sirignano & Spiliipouloss 2018) by including the following term part of the loss function:
[r(t,x) x) Ot
Cx(x) ((_
+ 3t, ((t,x))(ttx)))
+ ll((0, x) - 11l2 (1, -
where U(A represents the continuous uniform distribution over set A. The last two terms codify the soft constraints 700, x) and >1, = which help to ensure that the forward process starts at y (or close) and ends in a standard Gaussian variable (or close). Noise injection is gradual because defined by deep learning model as a smooth function (see details about architecture in Appendix G. In the rest of Section we provide the remaining details about the loss function.
3.3 NON-REGULARIZED LEARNING
The schedule functions 7, define the forward process and determine the form of the posterior distribution and the reverse process (Appendix B). On the other hand, y, helps to define the reverse process Pv- To learn these functions we need access dataset of input-data pairs (y, x). The
4

---

Published as a conference paper at ICLR 2024
standard approach for diffusion models is to these samples to minimize the Evidence Lower Bound (ELBO), given in our case by (details in Appendix
(yLBO x) = (y, + Laitlusion (y.
The term pprir = DKL (q(z), Et,, x) |l Pv (Ze, ()) helps to ensure that the forward process ends in almost pure Gaussian variable. As described in Appendix D Ldiffusion (y, x) is analytic and differentiable, but is computationally inconvenient. To avoid this problem, this term can be rewritten as an expected value LT (y, [y,x) x) which has a simple and efficient Monte Carlo estimator (Appendix For reasons outlined in the following section, we work in the continuous-
E.
case, i.c. the oc limit. Under certain assumptions, we get the following form of the ELBO when -
LELBO x) = DxL (g(21ly,3) (z x)) + (y,
In the above expression, represents the latent variable at time t = 1, such that (Z1 (x) is standard Gaussian and gizily, x) (1, x ((1,x))) On the other hand, Coo is continuous- time estimator for Ldiffusion and takes an integral form. See Appendix E.2 for the full details.
In summary. İELOO provides a differentiable and efficient-to -o-mpppttt- version of the ELBO, and we use it as the core loss function to learn the forward and reverse processes correctly. In the next section, we describe two important modifications we make to the loss function and the learning process.
3.4 REGULARIZED DIFFUSION MODEL
As mentioned above, work with diffusion process that is defined for continuous time. The schedule and the model prediction y, depend on a continuous variable / E [0, 1], and the latent variables Z, are also parametrized by E [0, 1]. We also derived continuous-time version of the diffusion loss, Loo. In our final implementation, we get better results by replacing Coo with
1 (y, = -
2 Ee.. [|I- (z(2)(e), x)]1]]
where é, is a noise prediction model. See details in Appendix E.3] This provides natural Monte Carlo estimator for the diffusion loss, by taking samples (0, and U[[0, 1)).
As shown in Kingma et al. (2023), increasing the number of timesteps T should reduce the diffusion setup is easier to
loss, so makes sense to work the T - oc limit. Also, a continuous- implement. Importantly, Kingma et al. (2023) prove that the continuous time diffusion loss is invariant to the choice of variance schedule. However, we argue this is not necessarily the case in practice. Since all computational implementations are ultimately discrete, we look for conditions on y that make the discrete case as close as possible to the continuous one.
As explained in Appendix F.2, one way of achieving this is to keep the Euclidean norm of SNR" (t. x) low, where SNR(t, x) = y(t, x) (t,, We use f'(t, x) to represent the partial derivative function (t, x) with respect to the time t. A natural way of incorporating this condition would be to include regularization term the loss function, with a form like
LsNR (x) = |SNR (.NR"(x))l (xx))lIz((()))) where by definition SNR(L, x) 7l, x1,
7l,x))
(t, 1 1- (t,
From this, we can see that Lsnr can be complicated to implement. involves a fraction and a second derivative, operations that can be both numerically unstable. Moreover, for t = 0 it should hold that 7(t, x) se 1, which makes SNR more unstable around t = 0. Since SNR(t, x) xs (t, x) for values of closer 1, we replace LsNR with the more stable
have mentioned before,
(x) E.-0((01 [II. (t,))
This regularization term is actually key for the performance of our method. To see this, recall that variable 2, ~ 2 ly, x) can be reparametrized as z, = (t, + (a(1, x)e with e ~ N((0, I). make z, e, so that the noise prediction model é (ze(e),
This means that 7 = o can perfectly predict and make É. = Now, 7 = is not compatible with the L, loss term, but any function that starts at 1 for / = 0 and then abruptly drops to 0 is permitted. L.y prevent this type of undesirable solution. Once we include this term, the full loss function takes the form
z,
LCVDM = xxyx))x((xxx
[xx()) + Dxr. (g(z)ly x) || (2) + Loo (y, x) ax,(x)) (6)
5

---

Published as a conference paper at ICLR 2024
where controls the weight of the regularization term and p(y. x) is the joint distribution of y and x. We optimize a Monte Carlo estimator of LCVDM by using the available (y, x) samples. For the KL divergence term, optimize the analytical form of the KL divergence between two Gaussian distributions with the log-variance (Kingma & Welling, 2022 Nichol & Dhariwal, 2021).
4 EXPERIMENTS AND RESULTS
We assess the performance of our model on three distinct benchmarks. First, we evaluate the model's ability recover high-spatial frequencies using the BioSR super-resolution microscopy benchmark (Qiao et al.) 2021). Second, we examine the model's effectiveness in retrieving the phase of a diffractive system with synthetic data and real brightfield image stacks from a clinical sample assessed by two experienced microscopists. The last benchmark is image super-resolutie ImageNet 1K (Deng et al. 2009). For the first two, performance is measured using two key metrics: multi-scale structural similariyy index measure (MS (Wang et al. 2003) and Mean Absolute Error (MAE), both detailed in Appendix G In the case of ImageNet, performance is measured using SSIM and peak hignal-tonoiss ratio (PSNR). For BioSR, the resolution of the reconstruction (a metric related to the highest frequency in the Fourier space) is additionally evaluated as per Qiao et al. | (2021), using a parameter-free estimation (Descloux al. 2019). We evaluate against methods developed for each benchmark, as well as CDDPM. In the case of CDDPM, we follow the implementation shown in Saharia al. (202I). The specific implementation of our model described in Appendix G
on
4.1
MICROSCOPY
Super-resolution microscopy aims overcome the diffraction limit, which restricts the observation of fine details in images. It involves reconstructing a high-resolution image y from its diffraction- version x, expressed mathematically as K y + 17. where K is the point spread function (PSF) and 1) represents inherent noise. Convolution of the PSF with the high-resolution image leads to the diffraction- limited image x, which complicates y recovery due to information loss and noise. In this context, we utilize the BioSR dataset (Qiao al. 2021).
BioSR consists of pairs widefield and structured illumination microscopy (SIM) images which encapsulate varied biological structures and signal-to-noise ratio (SNR) levels. The structures present the dataset have varying complexities: clathrin-coated pits (CCPs), endoplasmic reticulum (ER), microtubules (MTs), and F-actin filaments, ordered by increasing structural complexity. Each image pair is captured over ten different SNR levels. Our results are compared with DFCAN, a regression- based method implemented as in Qiao et al.) al.] (2021), and CDDPM trained in Saharia al. (2021). For diffusion methods during inference, best results are found at T = 500, resulting in similar inference time for CVDM and CDDPM. For CDDPM, fine-tuning results optimal performance for a linear schedule ranging from 0.0001 to 0.03. All models are trained for 400,000 iterations.
4.1.1 RESULTS
Table I shows the enhanced resolution achieved in the BioSR benchmark, with our model surpassing other methods in the ER and F-actin structures. Our approach consistently delivers comparable or superior performance than CDDPM across all structures, and improves markedly over DFCAN in the resolution metric. Figure (6b] facilitates a visual inspection of these achievements, comparing our reconstructions to those of DFCAN, and Figure (a] which underscores the contrast in quality between our model and the CDDPM benchmark. Additional comparative insights are detailed Appendix
4.2 QUANTITATIVE PHASE IMAGING
Quantitative phase imaging (QPI) has gained prominence in diverse applications, including bio- imaging, drug screening, object localization, and security scanning (Zuo al.) 2020). The Transport of Intensity Equation (TIE) method (Teague 1983 Streibi 1984) is a notable approach to phase retrieval, linking the diffraction intensity differential along the propagation direction to the lateral phase profile. This relationship, under the paraxial approximation, is formulated as
OT O1 y: = V (2.a)) (I(2, yi V (x 2)
0z

---

Published as a conference paper at ICLR 2024
Table 1: Performance on BioSR Structures. Bold represents statistically significant best result.
Underline represents statistically significant second place. Statistical significance is determined by comparing the sample mean errors of two different methods over the dataset, using hypothesis test.
(a) CCP structures.
(b) ER structures.
Metric/ Model
DFCAN
CDDPM
CVDM (ours)
Metri Model
DFCAN
CDDPM
CVDM (ours)
MS-SSIM (t)
0.957
0.952
0,955
MS-SSIM
0.928
0,920
0,934
MAE (1)
0.006
0.907
0.007
MAB (L)
0.053
0.033
0.032
(um) (
100
Resolution (mm) (1)
165
(c) MT structures.
(d)
F-actin structures.
Model
DFCAN
CDDPM
CVDM (ours)
Model
DECAN
CVDM (ours)
MS-SSIM (?)
0.901
0.857
0.887
MS-SSIM (?)
0,853
0,831
0.863
MAE
0.033
0,042
0.04
MAE (J)
0.049
0.049
0.043
Resolution (nm) (L)
127
101
97
Resolution (nm) (U)
151
104
Table 2: Performance on QPI Data. Bold represents a statistically significant best result.
(a) Performance on QPI on synthetic НCCCC.
(b) MS-SSIM on QPI Brightfield images.
Metric / Model
US-TIE
CDDPM
CVDM (ours)
Sample
US-TIE
CDDPM
CVDM (ours)
MS SSIM
0.907
0,881
0,943
0.740
0.863
0.892
MAB
0,110
0.134
0,073
2
0.625
0.653
0.742
where is is the wavenumber, I is the intensity, is the phase of the image, and x, y, are the coordinates. In practice, the intensity derivative AI(x, y; 2)/Oz at = is approximated via finite difference calculations, using 2-shot measurements z d and z -d:
(xr, 22
I_ Ia 2d
82
We train the model with a synthetic dataset created by using ImageNet to simulate Ta and I_a from grayscale images. The process, informed by the Fresnel diffraction approximation, is expressed
eike e (2P+0))))))))
Ia VToe
iXz
with 2 = d. the defocus distance, fixed at 2um during the training phase. Noise was added using (), distribution, where fluctuates randomly between 0 and 0.2 for every image pair to approximate Poisson noise.
To evaluate the effectiveness of method, compare against two baselines: CDDPM and US-TIE (Zhng et al. 2020), an iterative parameter-free estimation method for the QPI problem. For diffusion methods during inference, best results found at T 400, resulting in similar inference time for CVDM and CDDPM. For CDDPM, fine-tuning results in optimal performance for linear schedule ranging from 0.0001 to 0.02. Both our model and CDDPM undergo 200,000 iterations of training.
4.2.1 RESULTS FOR THE SYNTHETIC QPI DATASET
We estimate the phase of images using the HCOCO dataset (Cong al. 2019). These images simulated at a distance of +2pmm Table 2a presents the model's performance on the provided test split of HCOCO, which is conducted without noise. Figure Ta Ta] visually compares the methods. For more detailed comparisons, please refer to Appendix
4.2.2 RESULTS FOR QPI GENERATED FROM CLINICAL BRIGHTFIELD IMAGES
We evaluate our method on microscope brightfield images, consisting of three stacks with varying defocus distances, obtained from clinical urine microscopy samples. The phase ground truth is established by computing (I y: 2)/0,, using a 20th-order polynomial fitting for each pixel within the stack, following Waller al. (2010). This fitting is performed distances d = +2kkm, with k
7

---

Published as a conference paper at ICLR 2024
OODEN
(a)
(b)
Figure 1: QPI methods evaluated in the synthetic dataset (a) and brightfield clinical microscopy showing epithelial cells (b). From left to right: first column displays the defocused image at distance d (Ia). with the respective ground truth (GT) situated directly below. Second, third, and fourth columns represent each different method, with the reconstruction on top and the error image the bottom.
ranging from I to 20, and the gradient of the polynomial is employed to apply the US-TIE method to generate ground truth data. All methods are evaluated using 2-shot approach at d =2pm. The diffusion models are evaluated in a zero-shot scenario from the synthetic experiment. Figure [b illustrates the reconstructions of all stacks, while Table 2b presents the MS-SSIM for each sample and method, with sample numbers corresponding to the figure rows.
4.3 ABLATIONS AND ADDITIONAL EXPERIMENTS
In addition to super-resolution microscopy and QPI, we also test our method for the problem of image super-resolution over ImageNet. For this task, use the architecture of Saharia al. (202I), equipped with our schedule-learning framework. Without any fine-tuning of the schedule, our results are comparable to Saharia al. (2021) and slightly better than Kawar al. (2022). another iiffusion-basee method. See Appendix for more details and reconstruction examples.
We also evaluate the impact of our design decisions. First, the regularization term (Section 3.4 is critical in preventing the schedule from converging to meaningless solution. See Appendix K for detailed analysis. Second, the separation of variables for 3 (Section 3.2) ensures that the monotonic behavior of 7t,, x) with respect to / is not extended to This is key for achieving competitive performance, so we do not include a detailed ablation study. Finally, we compare learning a pixel-wise schedule to a single, global schedule using the synthetic QPI dataset. Our experiment shows that performance drops with a single learned schedule (see details Appendix K
5 ANALYSIS AND DISCUSSION
Accuracy. Our model competes favorably with DFCAN on the BioSR dataset, particularly excelling in image resolution (Table []. It shows the versatility of diffusion models in generating realistic solu- tions across diverse phenomena. Additionally, it outperforms CDDPM in more complex biological structures, and it advances significantly in the QPI problem by overcoming noise challenges near the singularity (Wu et al.) 2022). While not designed specifically for these tasks, our approach shows promise as flexible tool for solving inverse problems with available input-data pair samples.
Schedule, As explained in Section I the schedule is a key parameter for diffusion models, so we aim to understand whether our method yields reasonable values. Based on the relation between 7 and 3
8

---

Published as a conference paper at ICLR 2024
Sample mean
Semple standard deviation
Tremings
(a)
(b)
Figure 2: Schedules and sample mean and deviations for the represented images from the BioSR dataset. (a) Schedule (3) values for a microtubule image. The graph shows the average of the pixels in the respective region. (b) Mean and standard deviations for microtubule (top row) and endoplasmic reticulum (bottom row). The images were reconstructed using 20 samples obtained with CVDM.
the forward model can be parametrized as (details in Appendix D): e Jo ((xx)d))
q2,,y,, = (e-
2 p(exx)dy
We can see that for large values of the latent variable z, gets rapidly close (N I) distribution. For small values of 3, on the other hand, remains closer to for longer. In general, steeper graph of I results in a diffusion process that adds noise more abruptly around the middle timesteps, resulting in more difficult inversion of the process. In our image-based data applications, the pixel- wise dedicated schedule supports this analytical insight. In BioSR (Figure 2a), structure pixels (i.e.. pixels with high-freet frequeryy information) are more difficult denoise, which is reflected in the steeper graph. In contrast, background pixels (i.e., pixels with low-frequency information) are easier to resolve. This is consistent with the diffraction limit in optical microscopy, which mostly consists of low frequencies. Conversely, in QPI background pixels have a steeper / graph, a phenomenon linked the amplification of low-frequency noise around the singularity point in k-space (Wu et al. 2022)
Uncertainty. our experiments, we measure the uncertainty of the reconstruction by the pixel-wise variance the model predictions. This uncertainty is theoretically tied to 3. As described in Song et al. (021), the reverse diffusion process can be characterized by a stochastic differential equation with a diffusion coefficient of (t, æ) (in the prriance--reserrving case). Hence, higher values of introduce more diffusion into the reverse process over time, leading to more varied reconstructions. For BioSR, structure pixels exhibit higher values of 3 and consequently higher reconstruction variance, as seen in Figure 2b For QPI, the converse phenomenon is true (Figure 12). For further analysis and illustration of uncertainty, see Appendix M
6 CONCLUSION
We introduce a method that extends Variational Diffusion Models (VDMs) to the conditioned case, providing new theoretical and experimental insights for VDMs. While theoretically, the choice of schedule should be irrelevant for a continuous-time VDM, we show that it is important when a discretization is used for inference. We test our method in three applications and get comparable or better performance than previous methods. For image super-resolution and super-resolution microscopy we obtain comparable results in reconstruction metrics. Additionally, for super-resolution microscopy we improve resolution by 4.42% when compared against CDDPM and 26.27% against DFCAN. For quantitative phase imaging, we outperform 2-shot US-TIE by 50.6% in MAE and 3.90% in MS-SSIM and improve the CDDPM benchmark by 83.6% in MAE and 7.03% in MS-SSIM. In the wild brightfield dataset, we consistently outperform both methods. In general, our approach improves over fine-tuned diffusion models, showing that learning the schedule yields better results than setting it as a hyperparameter. Remarkably, our approach produces convincing results for a wild clinical microscopy sample, suggesting its immediate applicability to medical microscopy.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
This work was partially funded by the Center for Advanced Systems Understanding (CASUS) which is financed by Germany's Federal Ministry of Education and Research (BMBF) and by the Saxon Ministry for Science, Culture, and Tourism (SMWK) with tax funds on the basis of the budget approved by the Saxon State Parliament.
The authors acknowledge the financial support by the Federal Ministry of Education and Research of Germany and by Sächsische Staatsministerium für Wissenschaft, Kultur und Tourismus in the programme Center of Excellence for Al-research "Center for Scalable Data Analytics and Artificial Intelligence Dresden/Leipzig" project identification number: Sc.DS
The authors also acknowledge the support of the National Agency for Research and Development (ANID) through the Scholarship Program (DOCTORADO BECAS CHILE/2023 72230222)
REFERENCES
Martin Benning and Martin Burger. Modern regularization methods for inverse problems. Acta numerica, 27:1-111, 2018.
Sayantan Bhadra, Varun Kelkar, Frank Brooks, and Mark A Anastasio. On hallucinations in tomographic image reconstruction. IEEE transactions medical imaging, 40011):3249-32260 2021.
Daniela Calvetti and Erkki Somersalo. Inverse problems: From regularization to bayesian inference. Wiley Interdisciplinary Reviews: Computational Statistics, 10(3):e1427, 2018.
Nanxin Chen, Yu Zhang, Heiga Zen, Ron J Weiss, Mohammad Norouzi, and William Chan. Wavegrad: Estimating gradients for waveform generation. arXiv preprint arXiv:2009.000 00711, 2020.
Zheng Chen, Yulun Zhang, Ding Liu, Bin Xia, Jinjin Gu, Linghe Kong, and Xin Yuan. Hierarchical integration diffusion model for realistic image deblurring, 2023.
Wenyan Cong, Jianfu Zhang, Li Niu, Liu Liu, Zhixin Ling, Weiyuan Li, and Liqing Zhang. Image harmonization datasets: Hcoco, hadobe5k, hffickr, and hday2night. CoRR, abs/ 10s/1998.10526, 2019. URL httptttpp//aii...aaaap www.orr//////oooo//or/..../r..ee 908. 10526
Florinel-Alin Croitoru, Vlad Hondru, Radu Tudor Ionescu, and Mubarak Shah. Diffusion Models in Vision: A Survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(9); 18850-10869, September 2023. ISSN 0828, 21 0292,,9999222202 1939-3539. doi: 10. 109/TPAMI. 2023.3261988. BRRLhhtp:///axiiii wttp://gg/abb//aaaaaee 00722229900077777 04747 MXXi:2209.00747 [es].
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hier- archical image database. In 2009 IEEE Conference on Computer Vision and Puttern Recognition, pp. 248-255, 2009. doi: 10.1 109/C/0R. 20..2000955.444
A. Descloux, K. S. Grußmayer, and A. Radenovic. Parameter-free image resolution estimation based on decorrelation analysis. Nature Methods, 16(9):918 924,, September 2019. ISSN 1548- 7105. doi: 10. (1003///1152201-00515555 s41592- 919-019---00 0515-1
URL http:tt::: / www www.uratue.... com/ wwt.:/cco/aariiiie articles/
Number: 9 Publisher: Nature Publishing Group.
Prafulla Dhariwal and Alex Nichol. Diffusson models beat gans on image synthesis. CoRR,
Abs/2105.052333 2021. URL httpps://aaxiv..rrrrr.e..ss wawp://////aooao/o.o......ossssa 0atp://aaaaa11555533333
Juan Luis Fernández-Martinnz, Z Fernández-Mufizz JLG Pallero, and Luis Mariano Pedruelo- González. From bayes tarantola: new insights understand uncertainty in inverse problems. Journal of Applied Geophysics, 98:62-72, 2013.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models, December 2020. http: //arxiv...iiviv/a www.si..org//aae. / /wtp:///w..////////0..1.1.... 11239 arXiv:2006.1 11239 [cs, stat].
Phillip Isola, Jun Yan Zhu, Tinghui Zhou, and Alexei A. Efros. fmage-to-imagee translation with
conditional adversarial networks. CoRR, htbb//611..7000 2016. URL http: http::/tt//txti///t /ttp:///.xvv./////// abs/ 07004
10

---

Published as a conference paper at ICLR 2024
Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models, 2022.
Diederik P. Kingma and Max Welling. Auto-Encoding Variational Bayes, December 2022. URL http://arxi:: org/ abs 131111111.0 6114 1oXXv:1312.6114 [cs, stat].
Diederik P. Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational Diffusion Models, April wttp//aa///g/o////// 00630 arXiv:2107.000 000600 [cs, stat].
2023. URL httt::tp/aaaatti
Simon Kohl, Bernardino Romera-Paredes Clemens Meyer, Jeffrey De Fauw Joseph R. Led- Klaus Maier-Hein, S. M. Ali Eslami, Danilo Jimenez Rezende, and Olaf Ron- neberger. A Probabilistic U-Net for Segmentation of Ambiguous Images. In S. Ben- gio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL https: https::/pproceeddl
neur cc/paper loirr_fi ftless/ppeer/2018/ Paper.pppppp.
Yury Korolev and Jonas Latz. Lecture notes, michaelmas term 2020 university of cambridge. 2021.
Jaakko Lehtinen, Jacob Munkberg, Jon Hasselgren, Samuli Laine, Tero Karras, Miika Aittala, and Timo Aila. Noise2noise: Learning image restoration without clean data, 2018.
Michael Lustig, David Donoho, and John M. Pauly. Sparse MRI: The application of compressed sensing for rapid MR imaging. Magnetic Resonance in Medicine, 58(6): 182-1195, December 2007. ISSN 0740-3194. doi: 10. 10.10002/mm.21391
David JC MacKay. Information theory, inference and learning algorithms. Cambridge university
press, 2003.
Ahmad Mousavi, Mehdi Rezace, and Ramin Ayanzadeh. A survey on compressive sensing: Classical results and recent advancements, 2020.
Alex Nichol and Prafulla Dhariwal. Improved Denoising Diffusion Probabilistic Models, February 2021. URL HuRLpttp://arxiil wwtp:////raaaaaaa/aa/aa///////// ab220002. C9672 arXiv:2102.09672 [cs, stat].
Richard Nickl. Bayesian non non-linear statistical inverse problems. 2023.
Lukas Owens. Exploring the rate convergence approximations to the riemann integral. Preprint, 2014.
Shichong Peng and Ke Li. Generating unobserved alternatives, 2020.
Anna Pyzara, Beata Bylina, and Jaroslaw Bylina. The influence of matrix condition number on iterative methods' convergence. 201 Federuted Conference on Computer Science and Information Systems (FedCSIS), pp. 459-464. IEEE, 2011.
Chang Qiao, Di Li, Yuting Guo, Chong Liu, Tao Jiang, Qionghai Dai, and Dong Li. Evaluation and development of deep neural networks for image super-resolution optical microscopy. Nature Methods, 18(2):1944 202, February 2021. ISSN 1548-7105. doi: 10. 00013/444522222-01100 URL https: / www. nwwuuraaeee... com/
Number:
Publisher: Nature Publishing Group
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation, 2015.
Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David J. Fleet, and Mohammad Norouzi. Image Resolutnno via Iterative Refinement, June 2021. URt.tttp::/aciii.o0/// 2104.07633 arXiv:2104.07636 [cs, eess].
Chitwan Saharia, William Chan, Huiwen Chang, Chris A. Lee, Jonathan Ho, Tim Salimans, David J. Fleet, and Mohammad Norouzi. Palette: Image-to-image diffusion models, 2022.
Justin Sirignano and Konstantinos Spiliopoulos. DGM: A deep learning algorithm for solving partial differential equations. Journal of Computational Physit 375:1339 1364, dee 2018. doi: 10. 1016/ j-jep.2018 (atpp2201.....000 0.08.029. URL htttt::s///..oo//tt/tt 10000.100666000 jep 2018 08. 029

---

Published as a conference paper at ICLR 2024
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. International conference on machine learning. pp. PMLR, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models, 2022.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-Based Generative Modeling through Stochastic Differential Equations, February 2021.
wttp///gg//a/oooooaeoeeee 13456
arXiv:201 1. 13456 [cs, stat].
N. Streibl. Phase imaging by the transport equation intensity. Optics Communications, 49(1):
6-10, 1984. ISSN 0030-4018. doi: http://bbbooo////000000000004460 | https?
/ www.scieenedilrcttccooo cttp:/t.occccccci/eeccce http:///w........o.....oooooo//a
Michael Reed Teague.
Deterministic phase retrieval: a green's 's function solution. J. Opt. Soc. Am..
43311):1434-1441 Nov 1983. doi: 10.1164400SA..3.001144. URL htRtps:::///p optica
bog/abstract. www.mmmm........cccc ?URI josa 73 11314444
I Laura Waller, Lei Tian, and George Barbastathis. Transport of intensity phase-amplitude imaging with higher order intensity derivatives. Opt. Express, 18(12) Jun 2010. doi:
10.1364/OE. 18.012558 URL
cfm?URI
0e 18- 12 12552
Zhou Wang, Eero Simoncelli, and Alan C Bovik. Multiscale structural similarity for image quality assessment. In The Seventh Asilomar Conference Signals, Systems & Computers, 2003, volume 2, pp. 1398-1402. Icee, 2003.
Xiaofeng Wu, Ziling Wu, Sibi Chakravarthy Shanmugavel, Hang Z. Yu, and Yunhui Zhu. Physics- informed neural network for phase imaging based on transport of intensity equation. Opt. Express, 30(24):43398-33 43416, Nov 2022. doi: 10. 00.136400F.462444
URL /tttp::///ppp opt ca
org/oe/abstrace w oe 30- 24- 43398
Jialin Zhang, Qian Chen, Jiasong Sun, Long Tian, and Chao Zuo. On a universal solution to the rransport-of-intensity equation. Optics Letters, 45(13):364,, July 2020. ISSN 0146- 9592, 1539-4794, doi: 10. (0.133//01..99122 URL /all/tttp.///rri.orrr/bbbbs arXiv:1912.07371 [physics].
Chao Zuo, Jiaji Li, Jiasong Sun, Yao Fan, Jialin Zhang, Linpeng Lu, Runnan Zhang, Bowen Wang, Lei Huang, and Qian Chen. Transport of intensity equation: a tutorial. Optics and Lasers in Engineering. 135:106187, 2020. ISSN 0143-8166. doi: 2020. 1061 URL 0714381661932885
www. http://www..aaaaoocooooooooooooa
12

---

Published as a conference paper at ICLR 2024
A INVERSE PROBLEMS
We briefly describe a formalization of inverse problems, and mention some of the relevant ideas in this area. Consider two spaces X, y and a mapping A : y X. For any particular called the data), we are interested in finding an input y E 3 such that A(y) = x. In most interesting cases A is non-invertible, and even when it is, it may be ill-conditioned, in the sense that the inverse operation will amplify the noise or measurements errors in the data (Fernández: (Ferández-Marrinez et al.) 2013).
X' (usually
for invertible matrix A. If the condition number of the matrix is large, even this relatively simple problem may be dominated by numerical errors in practice (Pyzara et al. (01T). The situation is much more numerically complex when the problem involves mapping between rather than inverting the operator A, inverse problems are more commonly solved with a least-squares approach (Lustig et al. 2007 Mousavi et al.
Consider the case of solving Ay
spaces and there is need for discretization. As consequence,
2020
min ||(() - x||.
yey
The solution of the above optimization problem defines generalized inverse for A (Korolev & Latz 202T). However, if the problem ill-posed, this approach will still be highly sensitive to the noise in the data. This has motivated the use of regularization (Calvetti & Somersalo 2018 Benning & Burger 2018), which changes the problem into
minl || A(y) x112 + +x((), yεy
where the regularization term usually takes the form J(y) = |lyl| and a determines how much weight is assigned to the second term (prior knowledge y) as compared to the first (data fitting). Depending on several conditions, regularization can guarantee a stable optimization problem (Koroley) & Latz [021] Calvetti & Somersalo 2018). Still, further methods for solving inverse problems have been explored. One reason for this is that the theoretical understanding of regularization is not complete for nonlinear mappings (Benning & Burger 2018).
Moreover, as we mentioned already, one of the challenges for inverse problems is the presence of noise in the data. This is usually modeled as x = A(y) + 17, where r) the realization of a random variable. If the noise plays an important role, it would be desirable to provide uncertainty estimation for the solution. This has motivated the use of stochastic methods for inverse problems, one of the most prominent being the use of the Bayesian approach (Benning & Burger 2018 Nickl 2023).
B FORWARD PROCESS AND POSTERIOR DISTRIBUTION
B. FORWARD PROCESS MARKOVIAN TRANSITIONS
Using the result in Appendix A of Kingma et al. (2023), for each step 0 <is < T we have
( [x x)= =N (v x(1,x)
u (o(ts:x) x(,,,)
x() ((tx-1,))
We are working with the variance-preserving case, that is, xtt,x) 1 xt.x). Looking at the variance in the expression above, notice that
7(4, x) o (1o-1,,)) = (1 - x)),
7l1s x) (1 - --((t-1,,x))
ol,,,)
x)
7(1-11,
xl,,
1 - ((1,,)
7(4-11, + +7(1ix
= 7(1,x)
7(1-1,,))
To simplify these expressions, we introduce the following notation:
Br(t,.) (r(t,x):=1 = 1- (t,,x)
13

---

Published as a conference paper at ICLR 2024
Notice that By depends on both f4-1 and f,, but we have only parametrized it in terms of of We can do this because, for any fixed number of steps T, we can easily recover f,-1 from t,. Our notation will be convenient when consider the limit case T + oc, as explained in Section 3.2 With this notation, we can restate the result more concise way:
q(z,, Ze, [ = N (V1 (x
B.2 POSTERIOR DISTRIBUTION
By Bayes' theorem, we know that
g(Zt,_, Zz.. a(zz, Ze,-11 y, x)q((t,_. [x, ((z,, IZc, x)q(( Zt. ly,x)
(z x)
(zz, ly,x)
Notice that both the prior g(ze; ly, x) and the likelihood (Z,, Ze x) are Gaussian, which means
that the posterior will be normally distributed too (MacKay 2003 The computation of the parameters is somewhat involved, but it has already been done the diffusion models literature. Following
Kingma et al. (2023), get
g Ze. y, x)
where
(((--1,x)
We; + (r((i-1,,) 3r(t,,x)
1al2,,44y = V Br(li, x
O x)
oli,,
(1,-1,,
Br 3,,1,, x)
BT(,i, x)
be, + v /((li-1,x)
1 (1,, x)
711c, x)
and
OB(4,,x) = Br x) o(1--1,8)
o x)
(ti-1,,x)
= Br((,, x) (1,,
Notice that, similar to Br, both μæ and og require t,-1 as input besides t,. However, if the number of steps T is fixed, then i can be used to find f,-1- So, same as Pr, we write these functions as depending only on t, like (t,, x), rather than including t;-1 as an extra parameter.
C EVIDENCE LOWER BOUND
In this appendix, we derive the most important aspects of the loss function. Most of these details correspond to the derivations in Sohl- (nSSol--Dickseei et al. (2015); et al. (2020) and other references, but we include them for completeness and point out some key differences in our setup. For simplicity of notation, we denote Zt, as z; in most of this section. Let x be the data. We want the learned (ylx) distribution be close possible the true distribution p(y Hence, it makes sense to maximize the following log-likelihood term: E(y.)) [log (yl
Actual implementations of diffusion models take the equivalent approach of minimizing the negative log-likelihood, so we focus on the following term:
L
log Pa (y(x)
log Pv (Zul xx)
T
- log / Pe (2r xx
II
Pe Pe 12 x)dz1 ... dzr
1=1
- log / (x)())
T Po 1x
glz, 1, x)dzz dzr
li alz
2 xx
1=1
14

---

Published as a conference paper at ICLR 2024
Since the process is Markovian, IT ( 1, x) q 9(211...,, Z,r 20: x) and we get
I 1=1 Pv (Z 1 x)
- log E uy (r) (mg(()) ar ko,x)
pu(zr|x)
i=1 q 1i- x)
By Jensen's inequality:
I Pv Pu Ls, x) q ((z,|zz 1z x)
E Me aT ( my xo,x))
log
Pu(zr (x)
For simplicity of notation, in the appendices, we compute the loss terms before application of B(yxx)~p(yy x) [-]. This means, for instance, that L depends on (y, x) and should be written as L(y, x). write the loss terms without
Once again, in the interest of simplifying notation in the appendices, explicit dependence on y and x. So, for instance, we write Eprior in the appendices, but in the main body of the paper, we use the more precise piior (y, x).
In the following, we denote E (Z (........)))) (r) wr (x),x) [1] by just Eg [1] We continue by taking the logarithm of the whole product:
Ih Pa (u. (2 25, x)
C S E _x/xo.
log (u(l)))) x)
1= q 12 x)
T
g(z xx
-
log
Pv (2-_ (2,,))
T
g(z; xx)
+ g(21| /20, x)
=
log
log pe (zr (x)
Pe (2, x)
Pv (z) 21, x)
T
Iu,, x) )² ((2 (
(1(2-1
|i, 20, x) (( (221/20: x)
((2 (221|20.3 x) Pe (zu (2)
=
log
loog Pe (zr (x)
T
((2;-)
/i, 20, x)
T
log ((2) 20. x)
((2) 1(21|20, x)
= -
log
logg Per (2r(xx)
Pe (2 x2
q
Zu. 120, x)
Pu (zo (2,,
((Z
(i,, 20, x)
((2T| 20, x)
a(2) 20, 20,x)
E.gloox
log
log
log pe (2rlx)
Pe (z--1 |2,,
( I20 x)
(zo x),
T Σ log ((z;- 1 I2i, Z0, x)
log q 2o, (2,,)) Po(2r x)
= -Eglzo.x
logpo x))
Pe [2; [2;,
= Cprior + Eresomstrutioon + Luitfussinn
We take a look at each one of these terms in turn. To be consistent with the main body of the paper, we now return to the 2t; notation instead of using 2;- The prior loss term corresponds to
Lprice = DKL. (q [to, (Zto x) || Po (2 (x)),
and it helps ensure that the forward process ends with similar distribution to that with which the reverse process begins. In our setup, (ztz (x) is fixed as a normal distribution for all 2y and has no trainable parameters. In many DPM setups (Ho et al. 2020 Nichol & Dhariwal, (02T) the schedule is fixed and hence ( |2, x) has no trainable parameters either, so Eprior is discarded altogether for training. However, in our framework we need the key flexibility of learning the schedule function y, so we do not discard this term.
There are known ways to estimate and minimize Epeiee since Pa (Zt, x) is standard normal variable and (Zt, |,,,, x) has a Gaussian distribution too. Hence, the KL divergence between the two has an analytical expression that is differentiable and can be minimized with standard techniques (see Section 3.4). On the other hand, consider the reconstruction loss, given by
Creconstrnction Ee,, (z), Re. x) |- log Po (2t0 (2t, x)]:
15

---

Published as a conference paper at ICLR 2024
which helps ensure that the last step of the reverse process Pa (Zta Z6y gives the correct converse
operation with respect to the forward process (zz Zio x). In most DPM setups, this term is discarded or included in an altered form (Ho et al. 2020 Nichol & Dhariwal 202I
We choose to discard it for training, because its effect is very small when compared to Luifuusions and because the last step of the reverse process is not especially important. Consider the following: diffusion process has to produce change in a gradual way, and in the end pa (z(2,, (x) should be almost identical Pu(Zto x), especially when the number of steps T is high even infinite (we discuss the continuous-tiin case in Appendix D). The quality of pu is ultimately determined by the diffusion
loss whole rather than by the last step.
Finally, consider the diffusion loss, given by
T
log g(Z,, |2,,, 2o x)
Ldiffusion
E Miy Mep )mllay - mig meo x))
Pv Pu(Ze,_) Ze, x)
2=
E Bnima(an, xx)) [DL (g(z),_, Σ
Ze, Zeg, x))
Pu (Zt,-) Ze, x1]]
L4,
Each Lt, term simplifies greatly, as we show in Appendix reconstruction loss, the Evidence Lower Bound loss becomes
In the end, after discarding the
T
LELBO =
Lpeior + Lditfusion Cptio
+
Σ
Cu
D DIFFUSION LOss AND TRAINING
From Appendix
C
we know that one of the terms in the Evidence Lower Bound is
T
Ldiffusion
Σ
Cr
E Zz, (2,, (aeg x) [DLL. (g(Zt;-1 (2: Btg: x) Po (Zt, |24, x))]
1=2
We now simplify this term, essentially following Kingma et al. (023) We include this material for completeness and to highlight any differences that may come about from conditioning on Each Lt, includes the KL divergence between two Gaussian distributions, which has an analytic expression. Moreover, both distributions have the same variance. From Section 3.1 recall that by definition
Pv (z Zt,
q(zz,_. Zi,, Zip (2,, x),,,,
Now, from Appendix B.2 recall that the variance aB of q(2t,-) [2t, Zto x) only depends on 4 and and it does not depend on y. Hence, as mentioned before, both g(2t, / Zio x) and Pa (z Zt,, x) are Gaussian distributions with exactly the same variance. This means that the KL divergence between them simplifies
DKL. (q(zt. 2t,, Zio x) || Pe (2t., 24,, x)) = 20 (ts, x)) |1 1
where μ and uu are the means of q(Zz, 12,, Zeo x) and (Zr, 12,, x), respectively, Now, from Appendix B.2, we know that
# = uB v, V Br(4,, x) x( ne, + v7(4i-1, x) - x(t, y, (1,, x)
Br(t,,x)
Me = (Bl /1 1 (3. 1 - -((i-1,,
Br((;, 1- x-(,xxx)) 8v
21. + V)(ti-1.x)
1 (t,,
3r(,ix
1 Pv V7(li_1,x) x)
(y }w).
//,,, x)
16

---

Published as a conference paper at ICLR 2024
Since (t,,,, x) (t,, x1) (1,-1, x))\ ((t,:))) conclude that
DKL (((z))_ Rec, Zeoo x) Pv (Zt Zec x))
1
|1
x)
--(4;; x)
(t(ti_1, (T(
1 (t,, (y -
(tr x1 (t(--1xx))
(t;-1,, 3) (ti, x)
11y 1
2(1 - 1; - (t,, x))
x)) 1
7(4 2(14-1,x)) 1xx)
11y - -8u112
2(1 - -1. x:) (t,, x))
114-11 x) -, x)
11y S2112
2(1 - ,(1-1, x)) 7t,, x))
1
> (t)-1.
(t,,
) 11x - -
2
- (1i-1,,
1 - 1-7(1ix
1
(l4-1,x)
[((,xx) ||y 82112
2
044-1.
x14,
where 8. = Sz x) The fraction (t, x)/o x) is what Kingma et al. 20033) call the signal- to-noise ratio and should be decreasing function of time, so the above expression makes sense as nonnegative loss. Since we are in the srrtance-preeerviin case where (t, x) = I - (t, x), it sufficient that (t, x) is a decreasing function of time to guarantee that the above expression is nonnegative. Recapping, we have that
L1 Fzz, (a(z, la xx) ( (2144111))
7li,,) x)
1yy (2 x)12
(t,, x)
By the definition the forward process g(2,, 1et x) in Section 3.1 a variable Ze, g(z,, |2o x) can be reparametrized as Zı, (c) = xt), + (ot,,, x)e with e NO,,I). This means that we write the i-th term of the diffusion loss
( (2(4-11x) x)
[(1x [1y - (c),, ()
Lr 2 Fen
(t (f;,
Putting everything together, the diffusion loss is given by
Ldiffusion
Lu
( (2(t(-1,)) [((xx)) 1ly - to, x)11
Be-A((.))
(0(4111))))
xl1;,))
T
[1 (((-11,)) (2
1- 762, x) x3) ) |1x (z(zs, ()) (t)
7(1;,x] x))
E.-N((.))
2
(1-7(4-11,)
E ESTIMATORS FOR DIFFUSION LOSS
E. MONTE CARLO ESTIMATOR FOR DIFFUSION LOss
From Appendix D we know that
( [(1, ) 11y (x, x44,x))
Lditlusion 2 E
17

---

Published as a conference paper at ICLR 2024
For large T, computing the sum becomes computationally expensive. By linearity of expectation:
T
x))
7(1,, x)
BO.N(O.))
Σ
- -Su(aaz, (c), x)]1
2
(t.-11,))
0 (t,,,x)
T E No.N)
7(11,x) x))
1y Su(2,, (4), x)1
T- t= (t;-1,,))
(ts,
We can recognize the expression in brackets as an expected value, where is chosen uniformly at random from ( with probability 1/(T - 1). In other words, rewrite Ldiffusion
T-1 2 Vo..
[ (2(1-11)) (11_1,,)) xl1,
I(lix)
Ly(y,x)= =
|y
where UA represents the discrete uniform distribution over finite set A. As mentioned in Section 3.3 the advantage of writing the loss term like this is that it provides a straightforward Monte Carlo estimator, by taking samples (0, I) and Ty.
E.2 CONTINUOUS-TIME DIFFUSION LOSS
We mention again that throughout the appendices, we sometimes omit writing function parameters to simplify notation, mainly in the case the loss terms. For instance, we write Lr instead of Cr x).
Recall from Section 3.1 that Zr; i/T, intuitively the limit case T + takes us into continuous-time diffusion process, which can described by stochastic differential equation. This fact has been noticed before in the literature and it allows to present both diffusion models and score-based generative models as part of the same framework (Song et al. 2021).
In case, we are mostly interested in how the loss term Lr changes when T goes infinity. Kingma et al. 2023) give the following result when taking T - 0o:
[ SNR'
L = lim Loo =
E.
N(..)
(t,x))ly StR'
(z) t1),
(7)
T T-000
2
2 F.n.VooD,, (sNR'(1, x)]ly (,(2((e),,,
1
where SNR(t, x) represents the signal-to-noise ratio at time t and is defined as 7t, (//o.)///ttt x). Throughout this section, for simplicity of notation we use SNR' (1, x) to denote the partial derivative of SNR(,, x) with respect to the time I, and analogously for SNR" (1, x).
We are interested in better understanding this result, in particular regarding sufficient conditions under which the above equality holds, and its rate of convergence. This provides an interesting regularization idea for the training process, which we discuss in Section (..4] Starting from the expressions for LT and Lditfusion given in Appendices D and E.1
we know that
LT Lditfusion
x1 2(,,,))
2 Bo-N(O.1)
11 sz (e), (e), x)1
o(1111,))
o(4i,
1
Bo.N((..)) 2
(SNR(t-11:)) -SN((t,)))))) |yy (z), (e), x3 x311
Eo.N(o..)
(SNR(t,,x) x) SNR((t_1,,)) |1 (z (c) t4. x))
BENAiCO.) [ST],
18

---

Published as a conference paper at ICLR 2024
where we have denoted the sum inside the brackets Sy. Now, we denote hy 1/T and rewrite ST the following way, such that a derivative-like expression appears:
ST [SNR(,, x) (NN sz, (,,)22
SNR(4 - -SN 118 8z (c), tox x)
he
where the function f is defined
f (t) SNR((, SNR(L - No, ly - (su(z((e),,
ht
Now, need understand the behaviour of ST as T goes to infinity. Assuming that both SNR and Yv are continuously differentiable, it easy to see that f (t) converges pointwise:
(=) T T-+00 g(l) = SNR' (1, x)lly (z(z)(e), x)11
We have established pointwise convergence of f to g. but what we are really interested is to understand the convergence of Lr to Loo. defined as in equation (7). Notice that
Cr --- [ST], Loe - EE Lo s(()dt)
1
Denote the integral f/ g()dd by T. The above expression strongly suggests that we should understand the relation between ST and 7 when 7 goes to infinity, in order establish the convergence of Lr to Loo. With that in mind, for any given positive integer T', the difference between the sum Sy and the integral I can be bounded as follows:
L. g(t)dt
|St 11
(tt)TT
T
T
T glt - L. s(t)dt + +Ig(t))hg
$
(Ttt)/
Σ h(t
+
Σ
1=2
Es
E1
E2
We have assumed that both SNR and ÿr are continuously differentiable, which makes continuously differentiable too. Given the smoothness of g, for large T we have
Eg T
On the other hand, Eg is just the difference between Riemann Sum for g and its integral. For uniform partition such as ours (i.e. 1, = i/T) the Riemann Sum converges as O(1/T). particular, since is differentiable, for large 7 we have (Owens) 2014)
lg(1) .(0)
Eg
T
On the other hand, notice that
E
(fT) -g(t:) hr
SNR(I;,))
SNNR(1 l1r,)) SNR' (t,, _
x) 118 (c) x)1
hr
D:
19

---

Published as a conference paper at ICLR 2024
Assuming further that SNR is twice differentiable, the forward difference approximates the derivative:
D; SNR(L, x) SNR (t:
hr, x) SNR' S 2T |SNR
(,,x)]l
luT
Replacing the bound for E1, we get
T
E =
Σ
Dilly Fzlze, (c), x)1
Dillly - (u( 14, x)
1=2
|S (x) 82(2e, (c), x) 11
2T
SNR" ( x) L ( (c), t1
2T
i=2
Approximating the Riemann Sum by its integral
SNR" (t, x) - (z),((E), (,)
~
2T J
Putting everything together, we get
ST I
$ E3 + E2 + E1
lo(0))
19(1) (0)
1
/ (INR"(,))|IIL (,)
2T
Replacing 9 and using Cauchy-Schwand inequality:
(0, (z
T
SNR' (1, x)Il - (, (0, (z(zo(e), x)
+
T
+ (NI) (x) (
1/2
1/2
If the model is good enough, would expect lyy (r(zu( 0, x))1
ISNR' (1, x)llyy - (,((z((e),,, (,,
( (x))
1/2
T
2T
Summarizing, the main assumptions so far are that SNR is twice differentiable and that 8, is continuously differentiable. With that in mind, we can use the last expression to prove that Sr converges to I when T goes to infinity. Adding an extra condition of boundedness for yv. we can use the Dominated Convergence Theorem to conclude that
T - Bo.. [Sr] -E lim Sr - E..N... [ooooaat -im
lim CT lim T-000
We thus establish sufficient conditions for the convergence of Ly to Loo. Our analysis shows that, in some way, this convergence of order (1/T)) and its speed depends on the magnitude of SNR' and SNR". Also, it depends the approximation quality of the neural network model ÿ.. which is
20

---

Published as a conference paper at ICLR 2024
consistent with the finding in Kingma al. (2023) that increasing the number of steps T decreases the error on the condition that the model is good enough.
This provides an interesting insight. Kingma et al. (2023) show that in the continuous-time limit, the diffusion loss is invariant to the choice of SNR function, as long as fulfills a few basic conditions. However, in any computational implementation, continuous time cannot truly exist, the rate of convergence to the continuous-time case does matter. This means that for choices of SNR with ill-behaved derivatives, for a model yu that is not good enough, the "invariance under the choice of SNR" does not necessarily hold. This gives intuition for the regularization introduced in Section 3,4
E.3
NOISE PREDICTION MODEL
From Appendix
we know that the continuous-ting diffusion loss given by
Be..vio. [SNR' (t, x) S=(zo(e),,,x)11
By the definition of the forward process (Z,,,y, x) in Section 3.1 variable Z, ~ ((x,yy, x) can be reparametrized as Z, (t(t,x), + (t, x)e with E E~ (0, I). Rearranging the terms:
y Vr(t,x) (ze -o((t,))
(8)
Our model for predicting is x). From equation ( notice that ý. receives, explicit input, all the values necessary to compute y excepting c. As consequence, we can follow Ho et al. (2020) and repurpose y, into noise prediction model & by parametrizing
8,
(ze Vo),,
Vt(,,)
This means that the diffusion loss takes the form
L.0
E. ((10
SNR' (1, Su (,) x1)1
SNR'(1, x) 011, x)
E Be..N.o.),.-00001111)
(t,xx) ||e (l)
SNR' x) SNR(L,
E Be.-N(0.1.
( 4, x)11
This is the form of the diffusion loss that In experiments, we get better results by dropping the rational term SNR' (1, x) (L, x), which is consistent with the approach in Ho al. (2020). Otherwise, the training process can converge to trivial solutions that minimize Loo by making SNR SNR((t, so 0 for all t, x. In the end, the actual form of the diffusion loss for our method is given by
as part of the loss function in the end (see Section 3.4).
Lx - 2 - (
F CONTINUOUS-TIE TIME SCHEDULE
In Appendix B.I for all i e T,....,TT we define
7(41, x)
(r( =
Now, we want to study the relationship between functions and B7 when T' goes to infinity and we move into a continuous-time framework. Since we want diffusion process be smooth and free of sudden jumps, we require that y be least continuous on [0. 1] and continuously differentiable on (0, 1). By definition, notice that for any E {1,..., T}
7(4, x) 7(4-1, x) Br x)).
21

---

Published as a conference paper at ICLR 2024
Recall that the discretization uniform partition of [0. 1] defined by {t where t; = =i/T. Now, take any t E [0, 1] and define t' [t /T. Notice that t' is element of the partition {t,} defined above, and that t' + f when
using assumes number of steps T E N, and use the
oc. Then, equation
(9) implies
x(, > hr, x) (1-Br((',x)
Now, assume there exists continuous function (t,, x) such that
Br(t',
(10)
for all x. A function like this would allow to calculate Br for any discretization, interested in learning more about B. Denoting hr = 1/T, we have
7('', = 7('' ho, x) (1 B(t', x)hT)
7(' hT, x) 77(t -TT, - (t,B((())))) x)
(t',,) (('' hr, -t'' hr, 3)1 xi)rr
(1', 7(1' hr, x)
-t --r,,r x),(()(('x))
hT
We can now take the limit T + oo (which means hy + and e' + t). With our assumption that ? continuously differentiable, the left-hand side of the equation above converges to @rytt, x)/Ot. On the right-hand side, 7(t' - hr, x) converges to (t, x) and, with our assumption that is continuous, B(t', x) converges to (t,, x). In conclusion, we get the equation
@-(1, x) -1, (1, Øt
During training, we learn both the and 3 functions, and we enforce the above differential equation by adding corresponding term to the loss function. This ensures that the correct relation between ? and / is preserved during training. Afterwards, having learned these functions successfully, we can use equation 110 to discretize 3 into 3y. and then sample from the forward and the reverse process using the equations derived in Appendix
G IMPLEMENTATION DETAILS
G.1 METRICS DEFINITION
We use two main metrics to measure our results (Section 4.. The first is MAE, defined as:
MAE
Σ yr Ypl. PEP
where p indexes the set of pixel positions P in the images, y stands for the ground truth image, and y is the predicted image. The second metric is MS-SSIM, which measures the structural similarity between images and defined as:
M MS-SSIM(y,) = aM
where lj. Cj. and 8j are the measures of luminance, contrast, and structure corresponding to scale j.
We use five scales for this evaluation and set a, 3g = % such that Σ) 73 = 1.
G.2 ARCHITECTURAL DETAILS
Following Kingma et al. 2023), to learn the functions 7a (t) and P. (t) parametrize them using monotonic neural network. This network is composed of a residual block with three convolutional layers. The first and third layers employ linear activation and linked by skip connection, while
22

---

Published as a conference paper at ICLR 2024
the second layer uses sigmoid activation and is equipped with 1024 filters. All layers adopt x 1 convolutions. The decision to use convolutional layers over a dense network stems from the desire to facilitate the model's operation at various resolutions without retraining. Additionally, constrain the weights of the network to be positive. To satisfy the condition [80, x) 0, we multiply the network's s output 1.
For A (x), we parametrize it using a U-Net architecture (Ronneberger et al. 2015). represented in Figure The model incorporates 5 scales, doubling the number filters each scale while concurrently halving the resolution. Each block consists of two sequences: convolution, activation, and instance normalization. Upscaling is achieved using transposed convolutions ensure differen- tiability. Softplus activations are used throughout the architecture. Mirrored filters are concatenated, and the output's positivity is guaranteed by final softplus activation.
6 1024
87. 10,1024
61
128
129
128
As(a)
Figure 3: Architecture (x).
Our denoising model (Figure closely mirrors this implementation with two notable distinctions: first, input comprises the concatenation of x, yt, x), and z,, and the predicted z,-1 output has linear rather than softplus activation. In line with common practices, the network predicts the noise the corresponding timestep. We determine z by using the algorithm detailed in Appendix
81. 1024
87. 1024
647
144 1.28
1287 (64
Figure 4: Architecture of the score predictor used in BioSR and QPI.
H TRAINING AND INFERENCE ALGORITHMS
Figure presents high-level overview of the training algorithm and the inference process for CVDM. Algorithm [ describes the training of the denoiser using a learnable schedule, while Algorithm 2
demonstrates the inference process and how the learned schedule during this procedure.
23

---

Published as a conference paper at ICLR 2024
Algorithm 1: Training of Denoising Model Eu (z))
repeat
(x, y) ~ y)
<~ (0,1)
Take gradient descent step
(c x,, e) where w are the parameters of the joint model X. Φ, 0, vs equation 6
until converged;
Algorithm& 2: Inference in T timestops
for to do 9: xx)7xx) a; B
7: + II, act
for = T to do
if then
L <- 0
else
L c~ I(0,1)
7
1 var
2s - d (x + Voze
Figure 5: Training and Sampling algorithms for CVDM.
ADDITIONAL EXPERIMENTAL RESULTS FOR BIOSR AND QPI
The following figures further highlight results from our method in comparison to its counterparts in both the BioSR and QPI benchmarks. The reconstructions in Figures 7 and depict the differences between the CDDPM benchmark and our approach. Similarly, Figure contrasts an F-actin sample as reconstructed by our method and DFCAN.
Additionally, in Figure 10 we provide more examples of QPI evaluated in the synthetic benchmark.
CVDM
Ground Truth
Widerield
(a)
Figure 6: Results from super-resolution methods evaluated in BioSR. (a) Fine-grained comparison for CDDPM and CVDM on a microtubule (MT) sample. (b) Comparison of DFCAN and CVDM over different BioSR structures.
24

---

Published as a conference paper at ICLR 2024
CDDPM
CVDM
Ground Truth
Widefield
Figure 7: Comparison between CDDPM reconstruction and method F-actin sample.
CDDPM
CVDM
Ground Truth
Widefield
Figure Comparison between a CDDPM reconstruction and our method on a ER sample.
DFCAN
CVDM
Ground Truth
Widefield
Figure 9: Comparison between DFCAN reconstruction and method F-actin sample.
25

---

Published as a conference paper at ICLR 2024
lg / GT
US TIE
CDDPM
CVDM
AMG SCIO
STOP
STOP
STOP
STOP
CLES
COEF FUMMI WORNI TEET
RARKING
PORK
STOP
STOP
STOP
STO
Figure 10: Comparison of Phase Retrieval Methods on simulated HCOCO. From left to right: the first column displays the defocused image at distance (Ig), with the respective ground truth (GT) situated directly below. The second, third and fourth columns represent each different method, with the reconstruction on top and the error image the bottom.
J THE SCHEDULE AND THE FORWARD PROCESS IN EXPERIMENTS
As described in Section 3.1 the forward process can be characterized by the function ?
qzz,yx,) / [/(
There is an alternative, equivalent parametrization of the forward process in terms of the function B.
From Section 3.2 equation 4 we know that the following relation holds:
(x x) ((1, (),
This differential equation can be integrated:
0y(s, Olog Mogg((s,x) = -ff B(s, x)ds = Yt, x) log (0, x).
xs,
7(sxx)
Øs
Øs
As we describe in Section the initial condition 70, x) = should hold for all x in a diffusion process. Replacing in the above equation and applying exponentiation, get
xt,x) = - f3 8(e.x)ds
(11)
Figures and 17 show the averages of 3 and y for different groups of pixels (structure and background) in different images. Notice that the shapes of and consistent with equation (11)
26

---

Published as a conference paper at ICLR 2024
Stlucture Beskground
0.95
0.6
Cumma
80
0.92
0.00
0.0
25
50
75
100
125
150
175
200
0
25
50
75
1.00
125
150
175
700
Timaskase
Timestops
(a)
(b)
Figure 11: Schedule for the represented image from the clinical brightfield dataset. The graph shows the average of the pixels in the respective region (structure and background). (a) Schedule function B this image. (b) Schedule function y for this image.
Sample standard deviation
Sample mean
624
232
Figure 12: Mean and standard deviations for an image in one of the QPI datasets were reconstructed using 20 samples obtained with CVDM.
Moreover, notice that higher values lead values of y that decrease more rapidly to zero, we would expect. Now, using the variance- prieacc-preserving condition of,,x) = - r(t, x), the distribution of the forward process takes the form
xz) = N (e 1
(lxx)day (1 (6.x) Moxx)d) 1)
This is the relation we use Section to analyze the results of the BioSR experiments. note that it is equivalent to do the analysis in terms of y. If decreases fast, the latent variable z, gets rapidly close to a a.V(O. (0, I) distribution. This corresponds to pixels (or parts the image) that are harder to invert, and is reflected in the variance of those pixels in the reconstructed images. The converse is true for pixels where ry decreases more slowly. This relation between the learned schedule and the uncertainty (represented by the variance) is exemplified by Figures IT and 12 which respectively show the schedule and the variance of the reconstruction for an image.
K ABLATION STUDY
First, we study the behavior of our method without the regularization strategy. Figure 13 shows the learned schedule x), (averaged over all pixels) for BioSR, learned with the same specifications as CVDM but removing the regularization term. As can be observed, under these conditions, the schedule is mostly meaningless. Under this schedule, the input is steeply transformed the beginning
27

---

Published as a conference paper at ICLR 2024
of the diffusion process, then remains mostly equal for most of the time, and experiences another abrupt change at the end. There is no gradual injection of noise.
As explained in Section 3.4 regularization is important to prevent this type of result. Recall that variable 2, sampled from the distribution q( 2 y. x) be written 2, (t, xjy (f x)c, where e follows a Gaussian distribution, / (0, I). This reparametrization implies that setting y = and 1 gives 2, e, allowing the noise prediction model E (z.(), I, x) to perfectly predict and yield 2 0. It is true that y 0 conflicts with the Lg loss term, but any function y that starts at for = 0 and then abruptly decreases to O is admissible. In Figure 13 we can see similar behavior, with y starting at 1 and then abruptly dropping to low value. The regularization term by prevents undesired solutions of this type and ensures the gradual nature of the diffusion process.
MT
0.175
0.150
Average Schedule
0.100
0.075
0.050
0.000
250
1500
Timesteps
Figure 13: Value of y(t,x x) for each timestep, averaged over all pixels.
Additionally, we perform an ablation study on the impact of learning a pixel-wise schedule instead of a single, global one. Learning single schedule can be implemented method, which we denote as CVDM-simple. For the experiment we use the synthetic QPI dataset, introduced in Section Table 3 summarizes the results. As can be seen, CVDM-simple shows improved performance when compared to CDDPM, supporting the idea that learning the schedule can yield better results than fine-tuning it. At the same time, the results are not on par with CVDM, showing there is value in learning a pixel-wise schedule. This supports the idea that some regions of the image are intrinsically more difficult to reconstruct than others, and the schedule captures part of that difficulty. Figure 14 allows for visual comparison of the reconstructions given by these methods.
special case of our
Metric Model CDDPM CVDM-simple CVDM
MS-SSIM MAE
0.881
0.915
0.943
0.134
0.094
0.073
Table 3: Comparison of diffusion methods synthetic HCOCO.
L EXPERIMENTAL RESULTS FOR IMAGE RESSOUTINOON
We showcase the versatility of our method by training the model described in Saharia et al. (2021) enhanced with our schedule leamingg mechanism, for the task of image super-resolution in ImageNet.
28

---

Published as a conference paper at ICLR 2024
l GT
CDDPM
CVDM-simple
CVDM
AJ3
Figure 14: Reconstructions by different diffusion methods on a synthetic HCOCO sample. The first column displays the defocused image from left right at distance d (La), with the ground truth (GT) directly below. The second, third, and fourth columns represent each a different method, with the reconstruction top and the error image the bottom.
Low Resolution
Reconstruction
GT
Figure 15: Super-resolution reconstructions ImageNet using CVDM (our method).
29

---

Published as a conference paper at ICLR 2024
For this task, we compare our method to SR3 (Saharia et al. (021) and Denoising Diffusion Restoration Models or DDRMs (Kawar et al., 2022). The images are sampled with T = 500 timesteps, for our method. For SR3 and DDRM, the metrics are taken from their respective works.
Table shows the results obtained by the three methods, using the peak signal-to-noise ratio (PSNR) and structural similarity index measure (SSIM) metries. The results obtained by our method are comparable to both SR3 and DDRM. Similarly to the other applications described in Section 4 our method achieves competitive results without requiring any fine-tuning of the schedule. Figure 13 shows three examples from the ImageNet dataset, including the low-resolution image, the high- resolution ground truth, and a reconstruction sampled with our method.
Method
PSNR
SSIM
CVDM
26.38
0.76
SR3
26.40
0.76
DDRM
26.55
0.74
Table 4: Comparison between methods for image super-resolution.
M UNCERTAINTY ESTIMATION AND THE SCHEDULE
Our method implicitly learns a conditional probability distribution, from which conditioned sampling can be performed. This means that we can get different reconstructions for given input, drawing attention to the question of how much uncertainty is there in the reconstruction. As a way of measuring uncertainty, several reconstructions can be sampled for a given input, and the element-wise (e.g.. pixel-wise) sample variance can be computed. Theoretically. we expect the schedule function / to be significantly linked to this sample variance. As described in Section 5 in continuous time, the reverse diffusion process can be characterized by stochastic differential equation with a diffusion coefficient of / [B(t,x), In that sense, higher values of / make the reverse process more diffusive and leads to more randomness in the reconstructions.
We illustrate these ideas using samples from both the BioSR and ImageNet datasets, corresponding respectively to the problems of super-resolution microscopy and image super-resolution. The top half of Figure| 16 shows widefield microscopy image from the BioSR dataset, along with the ground truth image and five sample reconstructions. The bottom half shows a low-resolution image from the ImageNet dataset, along with the ground truth super-resolution image and five sample reconstructions.
For each reconstruction, we include the absolute error with respect to the ground truth image, which can be seen at the bottom of the respective half. Also at the bottom, we include the sample variance over the five reconstructions for each pixel, besides an image showcasing the pixel-wise magnitude of the 33 schedule function (as integral with respect to l). Finally, we highlight zones of high sample variance, so that it is possible to visually appreciate how the reconstructions differ from each other.
From Figure 16 it is easy to visualize the ideas descrihed at the beginning of this Section. The sample variance is reflected in the reconstructions, which clearly vary in some of the details. There is also a clear relation between the magnitude of /3 and the sample variance, as expected theoretically. This is interesting, as it suggests that 3 could be used to assess uncertainty in the absence of a readily available ground truth image, which would be the case most real-world applications. We can also observe that the regions with high sample variance tend to exhibit a higher absolute eтror in the reconstruction. that sense, the regions with larger uncertainty (as measured by the sample variance or the magnitude of (3) correspond to details that are truly more difficult to reconstruct correctly.
In general, this highlights the importance of robust uncertainty estimation, even for reconstructions that look generally correct. For applications such microscopy, single pixel may contain relevant information, such as evidence of cell interaction. Something similar happens with MRI, where a few pixels can determine, for instance, the difference between the presence absence of tumor. Therefore, hallucinations in the reconstruction method can be easily misleading, which is in fact one of the challenges for the adoption of accelerated MRI in clinical environments (Bhadra et al. 2021
30

---

Published as a conference paper at ICLR 2024
Widefleld
CT
Sample
Sample 2
Semple
Sample 4
Sample
Sample Variance
Sample
Sample
Sample Variance
Schedule
Figure 16: Input and ground truth images, along with five reconstructions, for samples BioSR and ImageNet. For every reconstruction, absolute eтror with respect to the ground truth is included, and regions of high uncertainty (i.e., sample variance) are enlarged so that differences the details can be appreciated. Pixel-wise sample variance over the five reconstructions is also shown, along with the pixel-wise intensity the schedule function.
31

---

Published as a conference paper at ICLR 2024
Stucture Bacaground
Spuccuuee Bajckroound
0.8
0.6
8000
Comma
0.03
0.90
0.0
25
50
75
1.00
125
150
175
300
25
50
75
1000
125
150
375
200
Timestops
Timestops
(a)
0.05
Staucture Background
1.0
Structure Background
0.04
0.6
Comma
0.4
0.02
0.00
0.0
50
100
150
175
200
1.00
150
175
Timestops
Timesteps
(C)
(d)
Figure 17: Schedule for the represented images from the BioSR dataset. The graph shows the average of the pixels in the respective region (structure and background). (a) Schedule function B for an ER image. (b) Schedule function y for the same ER image. (c) Schedule function B for an MT image. (d) Schedule function y for the same MT image.
32


Reviews:
[{"review_id": "qKoLP7CQw0", "reviewer": "Reviewer_qRoi", "summary": "In this paper the authors propose to learn the schedule of the forward process for conditional diffusion models, building up on the work of   [1]. This is an interesting topic since fine-tuning the schedule can be quite time consuming. The main contribution of this paper is to make the schedule dependent on the conditioning variable and to devise an appropriate learning procedure. \n\n\n[1] Kingma, Diederik, et al. \"Variational diffusion models.\" Advances in neural information processing systems 34 (2021): 21696-21707.", "questions": "I have no further questions.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The paper is well written and quite easy to understand, which is great. The numerical experiments/applications are interesting and show that the proposed method provides comparable or superior performance to existing methods that fine-tune the schedule.", "weaknesses": "I think that the contribution of this paper is quite marginal and is not suited for ICLR. The main difference with the methodology developed in [1] is the replacement of the unstable SNR term. While this yields good results in practice, I am not sure if this is enough novelty for a conference like ICLR.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "o3tj1BEl12", "reviewer": "Reviewer_tcMU", "summary": "This paper presents a conditional extension to the prior work of Variational Diffusion Models, to allow learning the variance schedule which eliminates the need to fine-tune this choice. This goal is achieved by incorporating the conditioning information as an additional input to the diffusion denoising network and as the input to one of the component in the decomposition of the learnable variance schedule. Specifically,  the paper discusses a regularized learning approach that keeps the scale of the SNR curvatures with respect to time steps to be low. Empirically, the method shows promising results on super-resolution microscopy tasks and quantitative phase imaging tasks.", "questions": "1. In abstract and intro, what does \"causal factors\" mean?\n\n2. The condition $A(x) = y$ is brought up in the introducation, but is not used in the method section. Does this condition specification matter?\n\n3. Missing a background of Variational Diffusion Models (VDM) before starting the method section. And I would appreciate a more straightforward discussion to how this version extends beyond and distinguishes from VDM. \n\n4. I find the notation $SNR''$ on page 5 not clear enough; it might be worth out writing $SNR''(t,x)$ when it is first introduced. \n\n5. Missing (the reference to) the details on how the base model architecture and other training details for the competing methods. This question is important to understand whether the empirical compairison is a fair one.\n\n6. Missing a description on competing methods, i.e. DFCAN, CDDPM, and  a discussion on how is the proposed method different from them. \n\n7. For the CDDPM, did you fine-tune the hyperparameters, namely the variance schedule?\n\n8. Limited methods used for comparison. There are definitely other qualified competing methods in the diffusion space, e.g. [2], [3], [4]\n\n9. Table1: Maybe I missed this point but what does the \"resolution\" metric mean?\n\n10. Table 1: why there is only one underlined result?\n\n11. Table 1: Why some rows do not have underlines and bolded results?\n\n12. Table 1: How do you define statistically significant? Do you repeat the experiments for multiple times and compute the average results with standard errors? If so, please describe this detail in table or the text. \n\n13. Table 2a: I would suggest to add the \"synthetic\" description to the HCOCO dataset in the caption. \n\n14. Section 4.2.2, what is the \"US-TIE method\"?\n\n15. Figure 3a: maybe I missed this but what do the structure and background in the legend box mean? Are they fixed values instead of learned? \n\n16. Figure 3a: Where does the learned uncertainty reflect in the figure, or it is not? I was thinking of a figure like Figure 4a in [1]. And how do you draw the conclusion on \"We briefly point out that our uncertainty estimation is consistent with the values of β as described above\" if there is no comparison between the uncertainty estimation and the values of $\\beta$? \n\n17. I wonder how the uncertainty quantification around the reconstruction can be useful?\n\n18. Appendix A, second paragraph, first sentence, should be \"Ay=x\" instead of \"Ax=y\"? \n\n[1] Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David J. Fleet, and Mohammad Norouzi. Image Super-Resolution via Iterative Refinement, June 2021. URL http://arxiv.org/abs/ 2104.07636. arXiv:2104.07636 [cs, eess].\n[2] Denoising Diffusion Restoration Models. Bahjat Kawar, Michael Elad, Stefano Ermon, Jiaming Song. https://arxiv.org/abs/2201.11793\n[3] Diffusion Models Beat GANs on Image Synthesis. Prafulla Dhariwal, Alex Nichol. https://arxiv.org/abs/2105.05233\n[4] Pseudoinverse-Guided Diffusion Models for Inverse Problems. Jiaming Song, Arash Vahdat, Morteza Mardani, Jan Kautz. https://openreview.net/forum?id=9_gsMA8MRKQ", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "This paper's primary strength lies in its practical extension of variational diffusion models (VDM) to the conditioning case, including several technical improvements such as the incorporation of a regularization term on the signal-to-noise ratio. \n\nIn addition, by adopting the VDM framework, the paper eliminates the need of prior works that fine tunes the variance schedule. \n\nAnother strength of the paper is on the experiment study which includes two practical downstream benchmarks assessed with meaningful metrics. It also pays attention on the uncertainty quantification, which is rarely highlighted in prior works.", "weaknesses": "The biggest weakness is the paper's technical novelty. The proposed approach seems like a straightforward conditional extension to the of variational diffusion models. In addition, the paradigm of turning an unconditional model to a conditional version has been largely explored and established, e.g.  [1]. I can see the decomposition of the learnable variance schedule, and the regularized learning are novel. In Besides that, can the authors comment on the technical non-triviality of this extension?  \n\nAnother weakness is lack of ablation study. For example, how does the method perform without using the regularized learning approach in Section 3.4? \n\nFinally, there is a room for improvement in terms of clarity, especially in the experiment section. See my detailed comments in the below Questions section.\n\n\n[1] Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David J. Fleet, and Mohammad Norouzi. Image Super-Resolution via Iterative Refinement, June 2021. URL http://arxiv.org/abs/ 2104.07636. arXiv:2104.07636 [cs, eess].", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uXHiZFBOcK", "reviewer": "Reviewer_UQRa", "summary": "The authors aim to solve inverse problems, i.e. recovering an underlying true object from imperfect measurements. The authors propose a conditional variational diffusion model (CVDM) to learn the conditional distribution of the true object given the measurement while being able to report uncertainty unlike typical supervised deep learning methods for inverse problems. CVDMs extend *unconditional* variational diffusion models (VDM; Kingma et al., 2023) to the *conditional* case in order to solve inverse problems and avoid hand-tuning the variance schedule of the diffusion process.\n\nThe authors demonstrate their model by:\n1. showing competitive performance on the BioSR dataset for super-resolution in microscopy\n2. showing superior performance on quantitative phase imaging (phase retrieval) on both the HCOCO dataset and their own clinical dataset.", "questions": "In addition to the previous concerns, I would appreciate if the authors could clarify the following:\n\n1. Does the pattern of schedules for structured versus background regions shown in Figure 2 hold true across the whole sample and for all samples?\n1. Are there significant differences in the inference time of this model versus the methods being compared against?\n1. Is the \"US\" in \"US-TIE\" for universal solution? Only TIE is specified in the text.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The authors present an extension of VDMs to the conditional case, which has not yet been demonstrated and would indeed eliminate some hand-tuning when training conditional diffusion models. This is also well motivated by an interesting application to inverse problems in optics where uncertainty estimates would be very useful. The authors demonstrate that their learned schedules correspond well to more structured regions of their input on the super-resolution task.", "weaknesses": "While the authors motivate their work well, ~I have some concerns with the experiments performed and their presentation. I would be willing to raise my score if the following concerns are addressed.~\n\n**Update**: The authors have addressed the major concerns, so I have updated my score accordingly.\n\n**Major concerns**:\n\n**Uncertainty results**: One of the main motivations of the paper is the ability to report uncertainty in order to point out artifacts in the solutions. In Figure 2, the authors show error images between the reconstructions and the ground truth. Do the uncertainty estimates correspond well to the regions of the image with more error?\n\n**Performance on other datasets**: The performance is not significantly better than other methods for the BioSR super-resolution dataset. If one of the major goals is to demonstrate the flexibility and robustness of this approach, it would be nice to see comparisons against more conditional generative model datasets, especially those in the CDDPM paper (though I understand that there is a focus on inverse problems).\n\n**Flexibility of super-resolution results**: When applied to super-resolution, this method does not require information about the PSF. This means that the model is learning the PSF implicitly from the data, but also means that the model would need to be retrained in order to get accurate results on data from a new PSF, which might be impractical if too much data is required for training. Again, the flexibility and robustness of this method is highlighted, so it would be nice to see some experiments or discussion addressing this, e.g. comparison against other blind deconvolution methods on datasets with different PSFs.\n\n**Estimating pixelwise schedule**: The schedules demonstrated for structured versus background regions appear to be linear with different slopes for the different regions. Would other/baseline methods perform better using an estimated linear pixelwise schedule with a higher slope on \"structured\" regions?\n\n**Minor concerns**:\n\n**Fourier features**: The original VDM paper showed that a significant amount of the performance was due to the use of Fourier features. It would be helpful to know whether anything like that is the case here.\n\n**Differences from VDM**: It would be helpful to have an explicit list or section summarizing the differences/extensions from VDM.\n\n**Figure design**: Many of the figures are very small and laid out in a way that makes comparing images and reading text difficult, e.g. in Figure 3 it is difficult to see the sections from which the schedules are being estimated.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PPNXenCdnJ", "reviewer": "Reviewer_jHYB", "summary": "This papers proposes a learned scheduling approach for training conditioned diffusion model for solving imaging inverse problems. In effect, the paper takes the variational diffusion model developed by Kingma et al. and replaces its learned scheduling algorithm with one that is conditioned on observations. The paper further extends this approach by learning per-pixel variances. The proposed method is applied to super-resolution microscopy and quantitative phase imaging; it outperforms existing diffusion-based approaches in both contexts.", "questions": "## Minor comment\nThe typical convention in imaging inverse problems is y=A(x), rather than x=A(y).", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The proposed method is effective and, though relatively straightforward, novel.\n\nValidation on real-world data is highly valuable and effectively demonstrates the utility of the proposed method.\n\nThe proposed method is general-purpose.", "weaknesses": "The paper could do a clearer job differentating itself from VDM. A bullet-pointed list of contributions would have been appreciated.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "diXwQDxQP9", "reviewer": "Reviewer_gCsZ", "summary": "The paper presents an extension of variational diffusion models to the conditional setting. In particular, the paper makes the following novel contributions:\n1. An element-wise learned variance schedule is used.\n2. The learned variance schedule is factorized wrt. time and the conditioning variable.\n3. The continuous time shedule is formulated in a different way than in Kingma et al. and new loss terms enforcing constraints is introduced.\nThe method is demonstrated on two image processing problems.", "questions": "The introduction begins with a relatively high level discussion of inverse problems, which in my view is a bit removed from the specific contributions of the paper. Maybe it could be an idea to motivate the paper more directly by mentioning specific inverse problems that the method is suitable for, such as the super resolution microscopy problem?\n\n\"To achieve this, we adopt the framework proposed by Saharia et al. (2021), which focuses on training with the statistics of the noise at each timestep rather than direct timesteps, thus allowing flexible use of the model during inference.\" Could you clarify what this means?\t\n\nIn the final loss, you drop the term SNR'/SNR. A large portion of the paper is dedicated to the design of the variance schedule mechanism. Would it be possible to include some more discussion / insights into why this term can reasonably be dropped? I realize this if done in many other papers as well, but usually without much discussion.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "Developing variational diffusion models and making them practically useful is an important and timely problem.\n\nThe proposed methods appear sound and reasonable, and seem to work well in practice.\n\nThe paper includes supplementary material where ample technical details are given, allowing the reader to follow each step of the derivations.\n\nThe paper presents several novel contributions beyond the existing literature.", "weaknesses": "In the abstract, learning the variance schedule is stated as a main contribution, however it is not clear from the abstract in which way the proposed method differs from existing methods that learn the variance schedule.\n\nIn general, it is not clearly stated what the contributions are beyond previous work, particularly Kingma et al. 2023. It would be a strong improvement to have a clear list of technical contributions in the beginning of the paper, allowing the reader to have an overview from the beginning.\n\nBefore eq. 1 it is mentioned that a continuous time diffusion is used, however, the manuscripts proceeds with discrete time steps. Since in the end, the continuous time formulation is used, I wonder if the presentation could be more direct, building on the continuous time formulation in Kingma et al. 2023?\n\nThe technical novelty of the paper is fairly limited, but while the novel contributions might be minor, it could still be practically important. However, it is not absolutely clear from the paper, to what extent these contributions are important for performance. An ablation study or a direct comparison with Kingma et al. 2023 could have illuminated this more direcly.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)