DEFT: Efficient Fine-tuning of Diffusion Models by Learning the Generalised $h$-transform
Authors: Alexander Denker
Francisco Vargas
Shreyas Padhy
Kieran Didi
Simon V Mathis
Riccardo Barbano
Vincent Dutordoir
Emile Mathieu
Urszula Julia Komorowska
Pietro Lio
Venue: NeurIPS 2024 poster (2024)
Keywords: Inverse problems
Generative Modelling
Diffusion Models
Conditional Generative Modelling
Diffusion model guidance
Primary area: generative_models
TL;DR: Novel and highly efficient fine-tuning algorithm for diffusion models motivated by theoretical formulations of conditional sampling in SDEs

Abstract: Generative modelling paradigms based on denoising diffusion processes have emerged as a leading candidate for conditional sampling in inverse problems. 
In many real-world applications, we often have access to large, expensively trained unconditional diffusion models, which we aim to exploit for improving conditional sampling.
Most recent approaches are motivated heuristically and lack a unifying framework, obscuring connections between them. Further, they often suffer from issues such as being very sensitive to hyperparameters, being expensive to train or needing access to weights hidden behind a closed API. In this work, we unify conditional training and sampling using the mathematically well-understood Doob's h-transform. This new perspective allows us to unify many existing methods under a common umbrella. Under this framework, we propose DEFT (Doob's h-transform Efficient FineTuning), a new approach for conditional generation that simply fine-tunes a very small network to quickly learn the conditional $h$-transform, while keeping the larger unconditional network unchanged. DEFT is much faster than existing baselines while achieving state-of-the-art performance across a variety of linear and non-linear benchmarks. On image reconstruction tasks, we achieve speedups of up to 1.6$\times$, while having the best perceptual quality on natural images and reconstruction performance on medical images. Further, we also provide initial experiments on protein motif scaffolding and outperform reconstruction guidance methods.

DEFT: Efficient Fine-Tuning of Diffusion Models by Learning the Generalised h-transform
Alexander Denker" University College London a. d.n.eeruull... ac
Francisco Vargas University of Cambridge fav250cam ac uk
Shreyas Padhy University of Cambridge sp20580cam. ac uk
Kieran Didi* University of Cambridge ked488cam ac uk
Simon Mathis University of Cambridge svm340cam. ac uk
Vincent Dutordoir University of Cambridge vd3090cam.. ac. uk
Riccardo Barbano Atinary Technologies rbarbamo@atinnayy..o com
Emile Mathieu University of Cambridge ebm320cam. ac uk
Urszula Julia Komorowska University of Cambridge ujk21 @cam. ac uk
Pietro Lio University of Cambridge p12190cam. ac uk
Abstract
Generative modelling paradigms based on denoising diffusion processes have emerged leading candidate for conditional sampling in inverse problems. In many real-world applications, we often have access to large, expensively trained unconditional diffusion models, which we aim to exploit for improving conditional sampling- Most recent approaches are motivated heuristically and lack unify- ing framework, obscuring connections between them. Further, they often suffer from issues such as being very sensitive to hyperparameters, being expensive to train needing access to weights hidden behind closed API. In this work, we unify conditional training and sampling using the mathematically well- understood Doob's h-transform. This perspective allows to unify many existing meth- ods under a common umbrella. Under this framework, we propose DEFT (Doob' 's h-transform Efficient FineTTuning), simply fine-tunes a very small network to quickly learn the conditional h-transformm while keeping the larger unconditional network unchanged. DEFT is much faster than existing baselines while achieving state-of-the-ar performance across a variety of linear and non-linear benchmarks. On image reconstruction tasks, achieve speedups of up to 1.6x, while having the best perceptual quality on natural images and reconstruction performance medical images. Further, also provide initial experiments on protein motif scaffolding and outperform reconstruction guidance methods.
approach for conditional generation that
Introduction
Denoising diffusion models are powerful class of generative models where noise is gradually added data samples until they converge pure noise. The time reversal of this noising process then
cqual contributions
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

DEFT Reverse Diffusion
$1
z-0
27
TT
h
h2
NN So, V2o log ((yl(,)) t)
y
(t logp(u/ão)
Figure 1: DEFT reverse diffusion setup. The pre-trained unconditional diffusion model of and the fine-tuned h-transform lyg are combined at every sampling step. We propose a special network to parametrise the h-transform including the guidance term V In p(y) (o)) part of the architecture. Here deg denotes the unconditional denoised estimate given sy (xx). During training, only need fine-tune he (usually 4-9% the size of st) using a small dataset of paired measurements, keeping s fixed. During sampling, we do not need to backpropagate through either model, resulting in speed-ups during evaluation.
allows noise be transformed into samples. This process has been widely successful in generating high-quality images [28] and has more recently shown promise in designing protein backbones that have been validated in experimental protein design workflows [77]. Recently, there has been much interest in conditioning the time reversal process, in order to generate samples that are subject observed condition. Conditional sampling requires the posterior score Va In pt(2) Y y). given some observation As diffusion models typically approximate the score of the underlying distribution, i.e., of (ar) R V Ie pe a pre-trained diffusion model can be leveraged using Bayes' theorem
V (x 3) = (x) + V. (Y = y(z),
(I)
to approximate the posterior score. The time-dependent likelihood Va In (Y y) often termed guidance due to its interpretation to guide the reverse process the conditioned imputs, and is unfortunately analytically intractable. To tackle this problem, several approximations for the guidance have been proposed; see, for example, [12, 22, 29, 56, 65, 69] and further discussion in Appendix B. Instead of relying on the decomposition (1), another line of work aims to learn the posterior score directly [5. 27], which requires expensive training for new conditional sampling tasks, and access to large amounts of paired data points.
In the setting of conditional generation with diffusion models, our primary goal is to leverage large pre-trained foundation models which are prevalent in applications, but which typical front-end users are not able to backpropagate through, making approaches like [12, 22] infeasible. This might be due to their prohibitive computation times or because they lie behind an API preventing the usage of autodiff frameworks.
In this work, we propose unified framework for conditional generation using Doob's h-transform,
well-known result in the stochastic differential equations (SDE) literature [14, 55, 61, 78). Under this framework. we propose DEFT (Doob's h-transform Efficient FineTinnig)) an algorithm that estimates the time-dependent likelihood directly from data, i.e., h- = V In Pa (Y = y y)z), while being able leverage an existing pre-trained unconditional model. We learn the guidance term (h-trransfom) efficiently using 1) smaller networks, and 2) small training dataset of paired data points and corresponding observations. Furthermore, through connections to stochastic control, we propose a novel network architecture for purpose fine-tuning, which, in conjunction with

---

proposed loss, achieves competitive results across series of inverse problems in imaging and protein design, while having a much lower computational cost.
2 Conditioning diffusions via the h-transform
In this section, we explore the formal mechanism to condition the boundary points of an SDE mathematically, and connect to existing methodologies for conditioning diffusions in generative modelling. For more rigorous background to denoising diffusion models, see Appendix A. Let us first recap the score-based generative modelling framework of [68]; we start with a forward SDE, which progressively transforms the target distribution Po (e.g. Po = Pdata)
dx, = f((X, dt + a,dWe, Xo ~ Po.
(2)
with drift and diffusion o₁- Under some regular assumptions, there exists corresponding reverse SDE with corresponding drift by [2]. that allows us to take samples from Pr (typically NOO,I) and denoise them generate samples from Po.
dx, (f.(X.) oiV x, In (x.) di + o. Xr Pr,
where the time flows backwards, and b, fa(X,, of x, Inp, The goal of conditional sampling is to condition the reverse SDE on a particular observation, i.e., to produce samples that satisfy constraints. For example, might want to use (3) generate samples where we already know some dimensions of the sample (e.g. knowing some pixels of the image a-priori in image inpainting). Doob's h-transform [55, 14] provides formal mechanism for conditioning SDE to hit an event at given time. We will show that existing methods for conditional generative modelling arise as approximate instances of this proposed framework. Formally, we have:
Proposition 2.1. (Doob's 's h-transform [55]) Consider the reverse SDE in Eqn. (3). The conditioned process X, Xn E B is solution of
dH, (b. (b, of
H. In (Xo H.H dt + W.W HT Pr.
(4)
with a backward drift (H. = (. - of H. (H.),, such that Law HH))H))
=
(2) 20 E B) and P(X, E B)=1.
Note, that will refer to the conditional process with H, and the unconditional process with X.- Doob's hi-transform shows that by conditioning a diffusion process to hit a particular event X, E B at boundary time, the resulting conditional process is itself an SDE with additional drift term (shown in the blue box above). Furthermore, the resulting SDE will hit the specified event within a finite time T. The function htt, H₂) A Poje E B H,) is referred to as the h-transform [55, 14]. See also Appendix C.3 for a discussion about the connection reconstruction guidance methods.
Rather than conditioning an SDE on a deterministic event, one is often interested in posterior arising from noisy observations (e.g. noisy inverse problems)
noisy(A(Xo)), Xo ~ Pdata:
where A is a forward operator, "noisy" describes noise process and unlike the classical h-transform, we are not enforcing a deterministic condition such as A A(Xo) Y. We typically assume we can evaluate and sample from the likelihood p(y)X = xn). Our goal is to sample from the posterior p(æoo Y y) = p(y (ro Paaca (zo))p(y)) Sampling from the posterior p(æa/Y Y y) can be achieved by generalisation of the h-transform that build on results in [75], given as follows:
Proposition 2.2. (Generalised h-transform) Given the following backwards SDE with marginals pt
dx, (x.) dt + or dW
XT ~ Pr.
(6)
then it follows that the backward SDE
Hr ~ Q/: [p(2pp [o)ll)]] = (
= H) ) di + o, dW., dH, = (H.,)) - - WH. In Poje (
(7)
3

---

satisfies Law (Ho) p(au Y = y) with Pole! Y = yl / p Y = y y (2o)) dao- We recover guidance based diffusions via (H,) = (. o. V H. In pal H.).
(20)) d is the transition operator of the forward process. Note, that the initial distribution 2/11 [p( y)] the controlled SDE differs from the unconditional SDE. However, in Proposition G.2 show that for the VP-SDE the difference between them gets exponentially small for increasing T. To summarise, the above result gives a generalisation of the h-transform that allows sampling from posteriors; notice that it recovers the traditional h-transform in the no-noise setting. Whilst this more general formulation of the h--rrasformm has been explored in unconditional generative modelling [78), this is the first work cast conditional generative modelling in this light. We refer to the term in blue as the generalised h--rrassfrmm henceforth.
Here Qf[ [)(x))]
Proposition 2.2 provides theoretical backing to methodologies such as DPS [12] HIGDM [65], in which the reverse SDE (7) is used to solve noisy inverse problems. For a careful derivation of Proposition 2.2 see Appendix D. While prior works have explored using Bayes' rule to decompose the conditional score, we provide rigorous arguments for intermediate steps, and carefully formalise the connection between conditional generative modelling and the h-transform, providing concise result. This framework is flexible enough to also encompass prior work on conditional score matching, see e.g. 27], and the discussion Appendix H.
3 Learning the generalised h-transform
Prior works either learn the posterior score from scratch, see e.g. [5, 27], or use approximations to the generalised h hi-rranffomm see e-g. [12, 32]. Instead, we propose a method to learn the generalised h-transform. We refer to this process as fine-tuning, as the pre-trained unconditional network remains unchanged and only the approximation given in the following theorem, where we give several representations of the generalised le-transform.
the generalised h-transform is learned. Our main result is
Theorem 3.1. (Representations of conditional SDE sampling) For given y noisy(A(2r)), let Q be the path measure of the conditional SDE
dH, = (J.(H.) - of Inp(( + h₂(H.)) dt + a dW
(8)
where HT ~ of (p(o ly). The generalised h--raasforms admits the following representations:
1) The path measure induced by the h-transformed SDE satisfies dQ' dP dp by) is the path measure of the unconditioned SDE and Po is it's time 0 marginal.
where P
dP
2) The h-transform admits denoising score matching representation
h = arg min (fx) N
[1n((H)) (e Inp.( -VH. In (H,
Cy (he)=
Xa Xamp((mll)) t~U(O H. Pajoo
3) The h-transform admits the following stochastic control formulation
{
1 ( (H)))( - Peto-qullppee H
hg arg min R.EH
(%c(lh)= = Eq [
where Q the path measure for the conditional SDE being controlled.
4) The path measure induced by the h-tranformed SDE solves the a Schrödinger bridge problem with boundary conditions Qo = Q [p( R NO,, T), Qr p(æn y) and with the unconditional process IP as its reference.
Here, 4) and 1) follow directly from [7, 73]. For the proof 2) see Appendix D.2 and for 3) see Appendix G. Under appropriate conditions the likelihood, the space of admissible controls H

---

can be taken to be the set of C1 --veeco fields with linear growth in space; see [48]. In the following sections, we will discuss the representations in and 3) in more detail.
3.1 DEFT: Fine-tuning by score matching
The score matching objective in Theorem 3.1 2) offers a simulation-free loss function to estimate the generalised h h-transform. While the theorem's ' formulation focuses on learning the h-transform for specific measurement y, this loss function can naturally be extended and amortized over the full range of measurements, i.c.,
min B (h)],
REM
to obtain (z. 3) = V (z)lax)). Further, for settings where the operator may vary, we can additionally amortise over the forward operator A p and learn hī( y. A) V In pr (y) x A). We exploit this to amortise over inpainting masks, see Section 4.1, and motif scaffolding, see Section 4.3. For the DDPM [28] discretisation of the SDE and pre-trained epsilon matching model el the fine-tuning objective (9) reduces to
min E Ex )e [ Y) + +
o0 (H.) - -el33
(10)
with H, = Va + VI- (Xo, Y) p( y). NO I), where he represents the neural network used approximate the generalised h-transform. Note that the loss function (10) only requires evaluation of the pre-trained model, without needing to backpropagate through the weights 0", which is often quite expensive and sometimes impossible in closed APIs. Training under the DDPM discretisation can be performed according to Algorithm 5. Sampling with DEFT is further explained in Algorithm 6, and pictorially represented in Figure As an additional insight into the behaviour of the hi-transform that makes it more flexible and capable of modelling non-linear tasks than standard reconstruction guidance methods, we show that the l--raasformm can be interpreted as correction term for the Tweedie estimate [20]. We can express the conditional Tweedie's 's estimate as
E [ at, y) Ri Z. 8)
Dt vI Ott (nf (ar, + Ee
(11)
v1 -I- āi hig ( y),
Voa
Vā,
where do (x;)) is the unconditional Tweedic estimate. Equation (11) highlights that the h-transform can also be interpreted correction factor to the unconditional denoised estimate, similar to [52. 80).
3.2 Connections to variational inference and stochastic control
A limitation to the fine-tuning objective with DEFT is that it requires a small dataset of paired datapoints and measurements. In this section, we propose an alternative approach by expressing the solution to the conditional sampling problem highlighted in Theorem 3.1 3). This allows us to learn the h-transform by optimising a variational inference-type problem. Importantly, this stochastic control objective only requires the availability of a single noisy observation y instead of a paired fine-tuning dataset. Further, the stochastic control objective can even be used in other conditional sampling tasks, for example in reward tilted distributions, i.e. where the goal is to sample from ((x) (#)) Paata (a:). Here er()) serves the same purpose as the likelihood, but there is no explicit measurement y [18].
stochastic optimal control objective, which is
However, the stochastic control objective is not directly applicable for high high-dimensional training, as the complete chain (He) must be kept in memory and backpropagated through or adjoint methods have to be used [36]. In Appendix G.3, we discuss several alternatives and present experiments for scaling up the above objective, e.g., methods like VarGrad [53] and Trajectory Balance [42]. VarGrad allows to detach the trajectory from the gradient computation, drastically reducing the memory footprint. We discuss concurrent work G. and G.2. Further, we show initial experiments for conditional sampling in G.4. The stochastic control objective serves as a conceptual bridge between sampling from unnormalised densities using diffusion models [74, 75, 81] and conditional score-based generative modelling.
5

---

3.3 Maeelhood-informme inductive bias
If the likelihood is differentiable, we can impose an inductive bias on the h-transform approximation. Specifically, the generalized h h-transform can be expressed as an expectation, and we can apply the DPS approximation [12] as follows
V an In Pojo (4)) Vz, pxorpp (yy
~ Va, In p B
z: In In (z,)),
where we use Tweedie's 's estimate based on the pre-trained unconditional diffusion model in the last step. The DPS approximation has been validated in many different conditional sampling tasks, so it would make for good initialisation of the learned h-transform. However, the DPS approximation requires the Jacobian of the unconditional model, which is expensive to compute and known to be poorly conditioned. Further, in applications where we only have access to the forward pass of the unconditional model, the Jacobian is infeasible to compute. Similar to [50]. we found that omitting this term still leads to an expressive architecture, while greatly reducing the computational cost. Thus, propose the following network architecture
he (z) y) NN (z,, (x)), V In ln (t) In (z2())
(12)
to parametrise the h-transform, where the last layer of NNÝ is initialised with 0 and NN; is initialised output 1. This initialisation provides computationally efficient approximation to the h-transform, which still guides the sampling.
This type of network architecture has been proposed within the sampling community to apply diffusion models to normalising constant estimation [49, 74, 81]. The theoretical connection to stochastic control Section 3.2, motivates us further to adapt the architectures from the sampling field to the conditional generative modelling setting. We ablate the different components of our proposed architecture in Appendix F.1 and find that the additional components greatly improve performance empirically.
4 Experiments
We evaluate the DEFT framework from Section 3.1 on both linear and non-linew natural and medical image reconstruction tasks. as well as the motif scaffolding problem in protein design. Further, in Appendix H.2 we provide a comparison FLOWERRS [47] image dataset. We provide our code https: //github.. com/ alexdenker/DEFT.
the conditional training framework with DEFT on the
4.1 Image reconstruction
We test wide variety of both linear and non-linear image reconstruction tasks the 256 256px ImageNet dataset [58]. We make use of tra-trained unconditional diffusion model with ~ 500M parameters [16]². We perform all our evaluations on a 11 subset of the validation set For all inverse problems under consideration, the h-transform was trained on separate 1k subset of the validation set. For linear inverse problems, we compare against IIGDM [65], DDRM [32], DPS [12] and RED-diff [44]. Additionally, we evaluate FSB [39]. The performance of SSSB can be seen as an upper-bound to DEFT, it is a conditional diffusion trained on the complete ImageNet dataset. For non-linear tasks, we only compare against DPS and RED-dit as both IIGDM and DDRM are not directly applicable to non-linear forward operators. For DEFT we make use of the DDIM sampling scheme with 100 time steps [64]. For the comparison methods used the same hyperparameters as [44] without further tuning, including the number of sampling steps (1000 for DPS and RED-Diff, 20 for DDRM and 100 for IIGDM).
We compute PSNR and SSIM, which are commonly used distortion measures, along with perceptual metrics such as Learned Perceptual Image Patch Similarity (LPIPS) [83], Kernel Inception Distance (KID) [8], and top-1 classifier accuracy of a pre-trained ResNet50 model [26]. There a well-known tradeoff between optimising distortion metrics versus perceptual quality, and depending on the task, one may wish for better performance along one axis at the cost of the other. For natural image tasks
"Checkpoints are available at https: //tgtgggt.t//t/////// com/ comm/ppnai//////ommmoooo /gudddd diffusion hhttps://bii wwt.btttttt/eeeeeeeeeeeeeeeeeeee www.lllllllllaae -ixxppx

---

Ground truth
DPS
TTGDM
DDRM
RED-diff
DEFT
Figure 2: Results for inpainting. We show the ground truth with the inpainting mask superimposed. Table 1: Results on inpainting and 4x super-resolution. Best values are shown bold, second best values underlined. We report both the total time to sample 1k images, and the time per sample in seconds. The time to sample includes the training time for DEFT. These tasks aim to generate "natural" "nookia images and therefore perceptual similarity metries (KID, LPIPS and top-1) 1) more relevant. SSB (grey column) can be considered an upper bound on performance.
Inpainting DDRM
Super-resolation
DPS
IIGM
RED-difT
DEFT
1'SB
DP5
IIGDM
DDRM
RED-diff
DEFT
F'SB
PSNR (f)
21.27
20.30
20.72
23.29
22.18
26
21.83
25.25
25.32
25 95
21.92
23.95
SSIM (1)
0.67
0.86
0.73
0.75
KID (1)
15.28
4.50
2.50
0.86
0.29
0.238
10.01
10.9
14.0
10.0
1.78
0.004
LPIPS (L)
0.14
0.068
0.25
0.11
top-1 (1)
58.2
67.8
68.6
72.0
71.7
74.5
71.5
71.02
63.9
66.7
71.9
71.6
N/A
7.86
Time per sample (s) (U)
100.6
10.2
1.22
28.3
4.36
N/A
100.6
10.2
1.22
28.3
4.36
NIA
involving in-painting and super-resolution, is common prefer "natural"-lookiin images, which score better on perceptual similarity, whereas for tasks involving (medical) image reconstruction standard to prefer lower distortion metric [9]. Further, we calculate the total time (including training for DEFT) for evaluation 1k validation images in the "Time (hrs)" row. Furthermore, we report the effective time taken to sample single image in the "Time per sample (s)" row. This time is calculated by fitting the largest batch size of validation images that fit on single A100 GPU and dividing the time taken for the batch by the batch size.
Inpainting First, we evaluate DEFT on the linear inverse problem of image inpainting. make use of the inpainting masks for the 1k subset used by [ which includes masks that obscure 20%-30 the image. Results are shown in Table 1, including the computational time for sampling all 1000 validation images. For DEFT, this computational time additionally includes the 3.9 hrs of training time of the h-transform additionally with the 1.2 hrs of evaluation. Even with the added training time, reduce the overall computational time for DEFT, compared to DPS and RED- visual comparison is provided in Figure 2. Further, in Figure in the Appendix, we show the diversity of samples using different initial seeds. Even though IIGDM and DDRM are faster methods, they perform significantly worse, and are only applicable for linear inverse problems. Inpainting is task that prefers Inatrrll"-lookin image samples, and DEFT outperforms all other methods on perceptual metrics such as LPIPS and KID, being a close second on top-1 accuracy.
Super-resolution For another linear inverse problem, we evaluate 4x noiseless super-resolution. Here, the forward operator is given by bicubic downsampling. The results are presented Table I. While DEFT has lower PSNR compared to the baseline methods, see significant improvement on perceptual quality metrics (KID, LPIPS, and top-1 accuracy). We show visual results in Figure 9.
High dynamic range For the first non-linear tasks, make use of the high dynamic range (HDR) task described [44]. Here, the forward operator is given by A(æ=) clip(2x:; -1, 1), where denotes the RGB image scaled to the range [-1,1). The results are presented in Table 2. We observe

---

1
Ground truth
Measurements
DPS
RED-diff
DEFT
Results for non-linear deblurring. We show both the ground truth, the measurements and samples for DPS, RED-diff and DEFT. DEFT is able to reconstruct high-quality images.
Figure 3:
Table 2: Results on different non-linear image reconstruction tasks. Best values are shown in bold, second best values are underlined.
HDR RED-diff
Phase retrieval RED-diff
Non-linear Deblurring RED-diff
DPS
DEFT
DPS
DEFT
DPS
DEFT
PSNR (t)
7.94
28.51
9.99
10.53
13.03
17.57
21.21
25.16
SSIM (1)
0.21
0.79
0.89
0.12
0.17
0.32
0.39
0.53
0.79
KID (J)
272.5
1.2
0.10
93.2
114.0
80.89
12.89
66.8
0.34
LPIPS ())
0.72
0.1
0.04
0.66
0.6
0.52
0.42
0.42
0.09
top-1 (T)
4.0
68.5
74.0
1.5
7.2
13.1
30.2
23.5
69.9
Time (hrs) (L)
30.7
7.9
5.2
30.7
7.9
5.2
30.9
8.1
5.2
Time sample (8) (4)
100.4
28.3
4.4
100.6
28.4
4.4
101.2
30.4
4.6
that DPS struggles with this specific non-linear tasks, while DEFT achieves good results. We show visual comparison in the Appendix, see Figure 10.
Phase retrieval The goal in phase retrieval is to recover the image from intensity measurements only, i.e., the forward operator given by A(xx) = [Fx),, with F as the Fourier transform. We study the same 2x oversampling setting as in [12, 44]. Phase retrieval is challenging non-linear inverse problem, as the forward operator is invariant to translations, global phase shifts and complex conjugation. In Figure 7 in the appendix, we show samples for different initial seeds and observe wide variety of image quality. We also observe this behaviour for the baseline methods. However, DEFT is able to achieve better performance compared to RED-diff and DPS, also Table 2. However, there is room for further improvement to achieve good reconstructions on a consistent basis.
Non-linear deblurring The non-linear deblurring task was originally proposed by (12). Here, forward operator is defined by trained neural network [70], resulting in highly non-linear blur. Quantitative results are presented in Table This non-linear reconstruction task was also evaluated for RED-diff in [44]. However, leads to a nearly trivial reconstruction task. In Appendix F.2, we show results with the code from [44], while Table 2 shows the results with our implementation of the forward operator. Further, in Figure we provide a visual comparison, where DEFT able to recover the ground truth quite faithfully.
found that the forward operator of the original implementation"
Size of fine-tuning dataset As DEFT requires a dataset for fine-tuning we ablate the number of training samples. We trained DEFT subset of 10, 100 and 200 ImageNet images for Inpainting. We see improvements of all metrics, when training on a larger dataset. The results are presented in Table 3. For the KID, outperform RED-diff (KID: 0.86) when trained only 200 images. However, even with 10 images, we perform quite competitively, showcasing that our method is very sample-efficient when comes learning conditional transform.
Ablation:
'https://githan http:/ihhh..co//llaa//saa/aa// diff http://awwa....ddff/eeeaeaaaaea

---

Table 3: Varying the size of the fine-tuning dataset for DEFT for Inpainting on ImageNet.
DEFT on ImageNet for Inpainting
Number of images
10
100
200
1000
PSNR (1)
20.87
20.99
22.11
22.18
SSIM (1)
0.83
0.84
0.847
01.85
KID (4)
1.85
0.978
0.401
0.29
LPIPS (4)
0.123
0.112
0.096
0.09
top-1 (1)
68.8
09.6
70.6
71.7
PSNR: 34.77
PSNR: 35.35
PSNR: 36.6
Ground truth
DPS
RED-diff
DEFT
Figure 4: Reconstructions for computed tomography on LoDoPab-CT
4.2
Computed tomography
We are evaluating DEFT both on the 2016 American Association of Physicists in Medicine (AAPM) grand challenge dataset [45], and the LoDoPab-OC dataset [35], for details see Appendix E. For the un- conditional models we make use of the at- tention U-Net architecture [16]. For the model trained AAPM, we exactly the same architecture ( 374 params.) as in [11], while for LoDoPab-CT we smaller model (== 133M params.). For the forward operator, we use a parallel-beam radon transform with 60 angles and add Gaussian noise with Ø, = 1.0, which corresponds to approx. 3.5% relative noise. We compare against DPS [12] and RED-diff [44], where the parameters were obtained using grid search on a subset of the validation set to maximise the PSNR. In Table 4 we present PŠNR and SSIM, in addition to the sampling time, and provide visual comparison Figure 4. For both datasets, we choose the same DEFT architecture with about 23M parameters. In the Appendix F, we perform an ablation regarding the parametrisation of DEFT, see Table 6. In particular, these results show the necessity of providing the unconditional Tweedie estimate a as input to the h-transform in (12). We observe almost &dB difference PSNR for models without the Tweedie estimate and the log-likelihood term
Table 4: Results for CT on AAPM and LODOPAB- CT and sampling time per image on a single GeForce RTX 3090. Best values are shown in bold, second best values are underlined. For DEFT we use 100 DDIM steps, while RED-diff and DPS use 1000 time steps.
AAPM RED-diff
LoDoPab-CT RED-diff
DPS
DEFT
DPS
DEFT
34.85
31.95
SSIM
0.845
0.865
0.887
0.846
0.849
0.876
Time
208.9
83.4
16.3
156.8
20.1
13.8
4.3 Conditional protein design: motif scaffolding
We evaluate DEFT on the contiguous motifs of the RFDiffusion benchmark [77]. In this motif scaffolding task, we sample protein C. atom coordinates x E Rd such that the generated backbone contains a targeted motif, i.c. subset of Ca coordinates E R", similar to an image outpainting task. The forward operator is therefore given by A(x) = Ax, where A {0, 1 1 denotes masking matrix which only selects the n observed C₂ coordinates.
We leverage the pretrained Genie diffusion model which is an unconditional model for protein backbone generation [37]. To apply DEFT to it, we use a downsized version of the unconditional base model as our h--tranffrrm model which only uses 200k instead of the original 4.1M parameters. To adopt this model to the DEFT algorithm, we modify the SE(3)-invariant encoder by adding additional conditional pair feature networks for the motif coordinates as well as the unconditional Tweedie estimate d.p similar to the setting the previous experiments. As per Section 3.3, we add a time-dependent likelihood approximation term to the h-traasfform network. We train the t-transform

---

network on the same SCOPe dataset as in [17]. More details on the training details can be found in App. E.2. We compare DEFT against DPS [12] and a previously published version of Genie that was trained in an amortised fashion [17]. The guidance parameter of DPS was tuned over 5 different experiment runs. The amortised model serves here as an upper limit of how well DEFT can perform with Genie as base model.
The overall in-silico success, defined by scRMSD < 2Á and motifRMSD < 1Å, is provided in Figure 5. In the Appendix, we provide a detailed breakdown of these results, see Figure 13. Further, in Figure 14 and Figure 15 we provide a comparison of the task IYCR for the different methods. We observe that DEFT outperforms DPS, solving 10 out of the 12 tasks compared to only 5 tasks for DPS. While it has lower success rates than the amortised model, it still solves all but two tasks in that benchmark with only 9% of the parameter count and significantly shorter training time compared to the amortised model (800 epochs for DEFT vs 2100 epochs for amortised). The low performance of DPS indicates that the base Genie model is limiting the performance here and may partly explain the performance difference between DEFT and amortised training. Exploring DEFT with a more capable base model is therefore another promising avenue for research. Excitingly, the lower training time and data requirements of DEFT enable fine-tuning a model on specific protein families for particular applications, task that is left for future work.
Amortised
9%
60-
DPS
DEFT 4%
60-
In-iico
METRIC
DPS
DEFT
AMORTISED
Success (1)
1.3
24.5
sRRMSD <2t(()
41.3
28.9
42.
STAM
204T
IYCR
PHRK 60
42yp
STRY short
STRV ner
GEGR short
d86R. med
402 short
a med
mRMSD 1 A(t)
4.1
24.0
45.8
THRE
Table 5: RFDIFF benchmark met- rics (averaged over the 11 tar- gets, 100 samples each). Success: scRMSD < 2', motifRMSD 1Å. Details in Sec. 4.3.
Figure 5: Comparison of DPS, DEFT and amortised
training for motif scaffolding for 12 contiguous targets.
4% and 9% are the relative sizes of the h-transform compared to the unconditional model.
5 Conclusion
We presented a unified mathematical framework, based on Doob's h--rransfrmm to better understand and classify different conditional diffusion methods. Under this framework, we proposed DEFT, a novel parameter-efficient conditional fine-tuning method that does not require backpropagation through large pre-trained score networks, resulting efficient sampling. We evaluated DEFT on several image reconstruction tasks and showed that it reliably outperformed standard methods, both in time, reconstruction quality and perceptual similarity metrics. While DEFT requires additional training on small dataset of paired measurements, we find that it is still faster than many existing baselines due to being able to use fewer sampling steps during evaluation, and not needing to backpropagate during evaluation.
Linitations and future work The DEFT framework uses (small) fine-tuning dataset, in contrast to zero-shot conditional sampling approaches, e.g., DPS [12] or TIGDM [65]. Fine-tuning on small datasets may have the risk of overfitting to biases inherent in the data. In contrast to zero- shot conditional sampling, DEFT assumes no knowledge of the forward operator. However, the forward operator can be incorporated as an inductive bias within the network architecture to improve performance. We also proposed zero-shot approach through the optimal control loss in Section 3.2, which only needs a single observation y to learn the h-transform. Though we show promising results scaling this approach to the MNIST dataset in Appendix H, the computational burden of simulating the full SDE at each iteration is still high, which might make this optimal control loss infeasible for high-dimensional data. However, there is promising recent work on partial trajectory optimisation [79], which may reduce the computational burden of the stochastic control objective, making competitive with existing methods.
10

---

Acknowledgements
Alexander Denker acknowledges support by the EPSRC programme grant EP/V026259/1. Shreyas Padhy is funded by the University of Cambridge Harding Distinguished Postgraduate Scholars Programme.
References
[1] Michael S Albergo, Nicholas M Boffi, and Eric Vanden- Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797, 2023.
[2] Brian D.O, Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):31-332, 1982.
[3] Simon Arridge, Peter Maass, Ozan Öktem, and Carola-Bibiane Schönlieb. Solving inverse problems using data-driven models. Acta Numerica, 28:1-174, 2019.
[4] Dominique Bakry, Ivan Gentil, Michel Ledoux, et al. Analysis and geometry of Markov diffusion operators, volume 103. Springer, 2014.
[5] Georgios Batzolis, Jan Stanezuk, Carola- Billa-BBbiane Schönlieb, and Christian Etmann. Conditional image generation with score-based diffusion models. arXiv preprint arXiv:2111.1366 2021.
[6] Julius Berner, Lorenz Richter, and Karen Ullrich. An optimal control perspective on diffusion- based generative modeling. In NeurIPS 2022 Workshop on Score-Basee -aased Methods, 2022.
[7] Espen Bernton, Jeremy Heng, Arnaud Doucet, and Pierre Jacob. Schrödinger bridge samplers. arXiv preprint, 2019.
[8] Mikolaj Binkowski, Danica J Sutherland, Michael Arbel, and Arthur Gretton. Demystifying mmd gans, arXiv preprint arXiv:1801.01444 2018.
[9] Yochai Blau and Tomer Michaeli. The perception-distortion tradcoff. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 6228-6237, 2018.
[10] Valentin De Bortoli. Convergence of denoising diffusion models under the manifold hypothesis. Transactions on Machine Learning Research, 2022. ISSN 2835-8856.
[11] Hyungjin Chung, Byeongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving diffusion models for inverse problems using manifold constraints. In NeurlPS, 2022.
[12] Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Mare Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In ICLR, 2023.
[13] Kevin Clark, Paul Vicol, Kevin Swersky, and David J. Fleet. Directly fine-tuning diffusion models on differentiable rewards. In The Twelfth International Conference on Learning Repre- sentations, 2024. URL https:
[14] Valentin De Bortoli, Arnaud Doucet, Jeremy Heng. and James Thomton. Simulating diffusion bridges with score matching. arXiv preprint 1rXiv:2111.0722 2021.
[15] Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion Schrödinger bridge with applications score-based generative modeling. NeurIPS, 2021.
[16] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. NeurIPS, 2021.
[17] Kieran Didi, Francisco Vargas, Simon V Mathis, Vincent Dutordoir, Emile Mathicu, Urszula J Komorowska, and Pietro Lio. A framework for conditional diffusion modelling with applications in motif scaffolding for protein design. arXiv preprint (tt.iiv223220022360 2023.
[18] Carles Domingo-Enrinn Michal Drozdzal, Brian Karrer, and Ricky TQ Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. arXiv preprint arXiv:2409.000 08866, 2024.

---

[19] Vincent Dutordoir, Alan Saul, Zoubin Ghahramani, and Fergus Simpson. Neural diffusion processes. In ICML, pages 8990-9012. PMLR, 2023.
[20] Bradley Efron. Tweedie's formula and selection bias. Journal of the American Statistical Association, 106(496):600216611 2011.
[21] James R Fienup. Phase retrieval algorithms: a comparison. Applied optics, 21((1))2758-2769 1982.
[22] Mare Anton Finzi, Amudhyan Boral, Andrew Gordon Wilson, Fei Sha, and Leonardo Zepeda- Núñez. User-defined event sampling and uncertainty quantification in diffusion models for physical dynamical systems. /CML, pages 10136-10152. PMLR, 2023.
[23] Wendell H Fleming and Raymond W Rishel. Deterministic and stochastic optimal control, volume 1. Springer Science & Business Media, 2012.
[24] Xizewen Han, Huangjie Zheng, and Mingyuan Zhou. Card: Classification and regression diffusion models. NeurlPS, 35:18100-18115, 2022.
[25] Andreas Hauptmann, Jonas Adler, Simon Arridge, and Ozan Öktem. Multi-scale learned iterative reconstruction. IEEE transactions on computational imaging, 6:843-856 2020.
[26] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[27] Jonathan Ho and Tim Salimans, Classifier-free diffusion guidance. In NeurIPS 2021 Workshop on Deep Generative Models and Downstream Applications, 2021.
[28] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. NeurIPS, 33:6840-651,, 2020.
[29] Ajil Jalal, Marius Arvinte, Giannis Daras, Eric Price, Alexandros G Dimakis, and Jon Tamir. Robust compressed sensing mri with deep generative priors. NeurIPS, 34:14938-14954 2021.
[30] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool Russ Bates, Augustin Židek. Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. Nature, (996(773))583-5888 2021.
[31] Hilbert J Kappen. Linear theory for control of nonlinear stochastic systems. Physical review letters, 25020), 2005.
[32] Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. In ICLR Workshop on Deep Generative Models for Highly Structured Data, 2022.
[33] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. ICLR 2015, 2015.
[34] Yann LeCun and Corinna Cortes. MNIST handwritten digit database. 2010. URL http: //yann. lecun. com/ exdb/mnist/.
[35] Johannes Leuschner, Maximilian Schmidt, Daniel Otero Baguer, and Peter Maass. Lodopab-ctt ct, benchmark dataset for low-dose computed tomography reconstruction. Scientific Data, 8(1): 109, 2021.
[36] Xuechen Li, Ting- Kam Leonard Wong, Ricky T. Q. Chen, and David K. Duvenaud. Scalable gradients and variational inference for stochastic differential equations. In Symposium on Advances in Approximate Bayesian Inference, pages 1-28. PMLR, 2020.
[37] Yeqing Lin and Mohammed Alquraishi. Generating novel, designable, and diverse protein structures by equivariantly diffusing oriented residue clouds. In Proceedings of the 40th Inter- national Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 20978 -20022 PMLR, 23-29 Jul 2023.
12

---

[38] Yaron Lipman, Ricky TQ Chen, Heli Ben- Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
[39] Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos A. Theodorou, Weili Nie, and Anima Anandkumar. 12sb: image-to-image schrödinger bridge. In Proceedings of the 40th International Conference on Machine Learning, ICML'23, JMLR.org. 2023.
[40] Xingchao Liu and Lemeng Wu. Learning diffusion bridges on constrained domains. In ICLR, 2023.
[41] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.033: 03003, 2022.
[42] Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment in gflownets. NeurIPS, 35:5955-5966 2022.
[43] Nikolay Malkin, Salem Lahlou, Tristan Deleu, Xu Ji, Edward Hu, Katie Everett, Dinghuai Zhang, and Yoshua Bengio. Gflownets and variational inference. arXiv preprint arXiv:2210.00580, 2022.
[44] Morteza Mardani, Jiaming Song, Jan Kautz, and Arash Vahdat. A variational perspective on solving inverse problems with diffusion models. In ICLR, 2024.
[45] Cynthia H McCollough, Adam C Bartley, Rickey E Carter, Baiyu Chen, Tammy A Drees, Phillip Edwards, David R Holmes III, Alice E Huang, Farhana Khan, Shuai Leng, et al. Low-dose ct for the detection and classification of metastatic liver lesions: results of the 2016 low dose et grand challenge. Medical physics, 44(10):e339-e35) 2017.
[46] Xiangming Meng and Yoshiyuki Kabashima. Diffusion model based posterior sampling for noisy linear inverse problems. arXiv preprint arXiv:2211..2344, 2022.
[47] Maria-Flena Nilsback and Andrew Zisserman. Automated flower classification over a large number of classes. In Indian Conference on Computer Vision, Graphics and Image Processing,
Dec 2008.
[48] Nikolas Nüsken and Lorenz Richter. Solving high- dighedimensional Hamilton-Jacobi -Blllmn PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Partial Differential Equations and Applications, 2(4): 2021.
[49] Angus Phillips, Hai-Dang Dau, Michael John Hutchinson, Valentin De Bortoli, George Deligiannidis, and Arnau Doucet. Particle denoising diffusion sampler. arXiv preprint 20Xiv:2402.06320. 2024.
[50] Ben Poole, Ajay Jain, Jonathan T. Barron, and Ben Mildenhall. Dreamfusion: Text-to-3d using 2d diffusion. In ICLR, 2023.
[51] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. Marina Meila and Tong Zhang. editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 1821-8831. PMLR, 18-24 Jul 2021.
[52] Sriram Ravula, Brett Levac, Ajil Jalal, Jon Tamir, and Alex Dimakis. Optimizing sampling patterns for compressed sensing MRI with diffusion generative models. In NeurIPS 2023 Workshop on Deep Learning and Inverse Problems, 2023.
[53] Lorenz Richter, Ayman Boustati, Nikolas Nüsken, Francisco Ruiz, and Omer Deniz Akyildiz. Vargrad: low-variance gradient estimator for variational inference. NeurlPS, 1311381-114992 2020.
[54] Lorenz Richter, Julius Berner, and Guan-Horng Liu. Improved sampling via learned diffusions. arXiv preprint arXiv:2307.00 01 198, 2023.
[55] Chris G Rogers and David Williams. Diffusiins, Markov processes and martingales: Volume
2, Itō calculus, volume 2. Cambridge university press, 2000.
13

---

[56] Litu Rout, Negin Raoof, Giannis Daras, Constantine Caramanis, Alex Dimakis, and Sanjay Shakkottai. Solving linear inverse problems provably via posterior sampling with latent diffusion models. NewrIPS, 36, 2024.
[57] Francois Rozet and Gilles Louppe. Score-based data assimilation. arXiv preprini arXiv:2306. 10574, 2023.
[58] Olga Russakovsky. Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. IJCV, 115:211-223 2015.
[59] Chitwan Saharia, William Chan, Huiwen Chang, Chris Lee, Jonathan Ho, Tim Salimans, David Fleet, and Mohammad Norouzi. Palette: Image-to-image diffusion models. In ACM SIGGRAPH 2022 conference proceedings, pages 1-10, 2022.
[60] Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David Fleet. and Mohammad Norouzi. Image super-resolution via iterative refinement. TPAMI, 45(4):4713-477 2022.
[61] Simo Särkkä and Arno Solin. Applied stochastic differential equations, volume 10. Cambridge University Press, 2019.
[62] Mathis Simon V. Julia Komorowska Urszula, Jamnik Mateja, and Lio Pietro. Normal mode diffu- sion: Towards dynomics-informed protein design. The 2023 ICML Workshop on Computational Biology. Baltimore, Maryland, USA, 2023. C, 2023.
[63] Vignesh Ram Somnath, Matteo Pariset, Ya-Ping Hsieh, Maria Rodriguez Martinez, An- dreas Krause, and Charlotte Bunne. Aligned diffusion schrödinger bridges. arXiv preprint S30.:22302.11435 2023.
[64] Jiaming Song. Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In ICLR, 2021.
[65] Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Pseudoinverse-guiee diffusion models for inverse problems. In ICLR, 2022.
[66] Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. Advances neural information processing systems, 33:12438- 12448, 2020.
[67] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In ICLR, 2021.
[68] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In ICLR, 2021.
[69] Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medical imaging with score-based generative models. In ICLR, 2022.
[70] Phong Tran, Anh Tuan Tran, Quynh Phung. and Minh Hoai. Explore image deblurring via encoded blur kernel space. In Proceedings of the IEEF/CVF CVPR, pages 11956-11965, 2021.
[71] Belinda Tzen and Maxim Raginsky. Theoretical guarantees for sampling and inference in generative models with latent diffusions. In Conference on Learning Theory, pages 3084-31 14. PMLR, 2019.
[72] Masatoshi Uchara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathaniel Lee Diamant, Alex M Tseng, Tommaso Biancalani, and Sergey Levine. Fine- tuning of continuous-time diffusion models as entropy-regullrizzed control. arXiv preprint 240. 2024.
[73] Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving Schrödinger bridges via maximum likelihood. Entropy, 23(9):1134, 2021.
14

---

[74] Francisco Vargas, Will Sussman Grathwohl, and Arnaud Doucet. Denoising diffusion samplers. In ICLR, 2023.
[75] Francisco Vargas, Andrius Ovsianas, David Fernandes, Mark Girolami, Neil Lawrence, and Nikolas Nüsken. Bayesian learning via neural Schrödinger-Föllmer flows. Statistics and Computing, 33(1):3, 2023.
[76] Francisco Vargas, Shreyas Padhy, Denis Blessing, and Nikolas Nüsken. Transport meets variational inference: Controlled monte carlo diffusions. In ICLR, 2024.
[77] Joseph L Watson, David Juergens, Nathaniel Bennett, Brian Trippe, Jason Yim, Helen E Eisenach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, et al. De novo design of protein structure and function with rfdiffusion. Nature, pages 1-3, 2023,
[78] Mao Ye, Lemeng Wu, and Qiang Liu. First hitting diffusion models for generating manifold, graph and categorical data. In NewrIPS, 2022.
[79] Dinghuai Zhang, Ricky Tian Qi Chen, Cheng- Hao Liu, Aaron Courville, and Yoshua Bengio. Diffusion generative flow samplers: Improving learning signals through partial trajectory optimization. arXiv preprint arXiv:2310.02679, 2023.
[80] Lvmin Zhang, Anyi Rao, and Maneesh Agrawala. Adding conditional control text-to-image diffusion models. In Proceedings of the IEEE/CVF ICCV, pages 3836-3847, 2023.
[81] Qinsheng Zhang and Yongxin Chen. Path integral sampler: A stochastic control approach for sampling. ICLR, 2022. URL httpps://openreevee http:////pppareeeee..eeeeeemrrrp httt://orrurttttttuuoaaaaarrrrr
[82] Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. arXiv preprint arXiv:2204. 13902, 2022.
[83] Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unrea- sonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 586-595, 2018.
[84] Linqi Zhou, Aaron Lou, Samar Khanna, and Stefano Ermon. Denoising diffusion bridge models. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreviex wwtp:/oruu///eee wtt./oorrmmd---Pkeilrll FKkkaaayyoo...
15

---

A Background on diffusion formulations
A.1 Recap- continuous and discrete diffusion formulations
The discretised DDPM versions with various discrete time schedules amount the time-dependent
OU process
3(1) X,dt + Bt)) dW
dx.
(13)
where choosing different time schedules amounts to choosing different functions B(t). This process gives rise to the following transition probabilities
p( f200 0) ( (x)
(14)
N ThE 1o 2i [ B(1) 10-- (0'=)so)daddl
(15)
N (agee
So 3
ds
(1- -o a(s)aa) r),
(16)
see also Appendix B in [67]. With a(t) = e B(s)ds, this is the familiar form ([28)):
paa, fl 0) Pejo ax)) = N ( a(t), (1-((t)))))
(17)
with (t) time-dependent and we can therefore choose different functional forms for the noise schedule by either choosing the transition parameters B(t) or the cumulative parameters att).
If define the noise schedule in terms of (()), the time-dependent OU process is immediately apparent (see (13)). If we define the noise schedule terms of (t),, the mean and variance of the corresponding OU process can simply be obtained from
d B(1) dt [())
(18)
A.2 Score, noise and mean diffusion formulations
The score-based model used for generation inference time can be parametrised model different quantities. The three most common one are the score, the noise and the mean. Using the score based SDE formulation we parametrise the network as the score that is the network approximates the quantity:
Ve In pi(x)) (x.)
(19)
Moving on the DDPM formulation one typically trains noise prediction network instead which is proportional to the score
(xx)) VI- a(l)
e (z)).
(20)
This formulation typically preferable for training as it is known to learn a less stiff vector field [82]. Finally in its most naive form DDPM also admits mean matching formulation
1- ( VI a(1)(
(x f)
(21)
Voi
whilst not the ideal parametrisation for direct training, it is useful expression/macro, for expressing the sampling updates more succinetly.
In this work we parametrise ef directly with our novel architecture for finetuning, using noise matching objective in DDPM, however we allude and use the above parametrisations across our propositions and novel architectures. See Algorithm I for training and Algorithm 2 for sampling.
16

---

B Related Work Discussion
common distinction to all the works we will discuss in this section is that they all either train the conditional network from scratch or they initialise the conditional network with a pretrained score and fully train a conditional network. This is in stark contrast DEFT which completely freezes the unconditional score and trains a highly efficient network to learn the h-transform which ranges from 4 17% in total parameter size of the pretrained unconditional score network.
Classifier free guidance Methodologies such as classifier free guidance [27] not model the forward operator explicitly. As result, if these methods are applied to settings such as motif- scaffolding or image out-painting (where the conditioning is on a subset of the random variable), these methodologies would only denoise the scaffolding and the missing image patches. This is different to our approach which adds noise to both motif and scaffolding and then proceeds to denoise both jointly part of the same space. In way, one view RFDiffusion's 's conditional training as an application of classifier-ff free guidance to this subset conditioning setting.
Image 2 Image Schrödinger Bridges (I2SB [39]) I2SB and more generally aligned Schrödinger Bridges [63] are a recently proposed class of conditional generative models based on ideas from Schrödinger bridges. The premise of these methods is that they aim to learn an interpolating diffusion between a clean data sample and a altered or corrupted data sample. This is contrast to our framework: we consider an unconditioned SDE and condition it to hit an event at particular time, thus learning an interpolating distribution between noise and an un-corrupted target distribution. This results in several algorithmic differences:
At its core, 12SB treat spectively; thus, at sampling time, Y is provided to the learned SDE which generates approximate samples from Law (Xu). However, in our approach, the source distribution is (0, I) and we pass Y to the score network to then obtain approximate samples from Law (Xo).
A(xo) + 7) and Xo as source and target distributions re-
The score network in 12SB is a function only of X; and not Y A(Xx)) + 17. This means that in 12SB, the network is parametrised as (X,), whilst in our setting we parametrise as (((, + 77 A). In the case of completion tasks like motif-scaffolding or image out-painting, our paramerisation looks something like ef' (X,, X, Ximakkk mask). This makes the task much easier for the network indicating which parts of the image are conditioned and which are not. In I2SB, the network must learn this on its own. Furthermore, parametrisation is essential to allow recovering the true conditional score.
effectively provide it with binary variable
show in Prop. H.1, adding this to the network
The training procedure in 12SB uses the diffusion bridge p(æ, | y) to add noise both the source and target distributions, whilst our forward process is given by the transition density of OU process p(2) and is identical to standard DDPM/VP-SDE [28, 68] noise adding procedures.
Finally and most importantly 12SB does full fine-tuning firstly initialising with a pretrained score and training all parameters of this large pretrained network to learn an unconditional network, this requires longer training times and significantly larger networks as they must be at least the same size as the unconditional score network.
To summarise: whilst both methodologies employ similar mathematical methodologies (e.g. Dif- fusion Bridges [14]), their ideations and resulting methods are fundamentally different: on one side, [39] learns an interpolating distribution between the unconditioned p(xe) and conditioned p(y ( samples. On the other, we learn a denoising procedure that directly samples from the poste- rior p( y): via this, we derive and explain most popular approaches for conditioning denoising diffusion models as part of our framework.
It's important to highlight that another akin approach to I2SB, also based on the h-transform but leveraging VP-SDes, was recently proposed in [84].
CDE [5] Similar to classifier free guidance. CDE [5] trains a conditional network from scratch without leveraging pretrained unconditional score. For more details on CDE please see our detailed discussion in Appendix H.
17

---

First Hitting Diffusions A line of generative modelling methods proposed in [40, 78] utilise the h-transform for unconditional generative modelling in the following settings:
Hitting the target distribution Pdata in finite amount of time [0, 7' via time reversing an h-transformed VP-SDE conditioned hit at time T.
Constraining diffusion process time T lie in subset of the reals Q RR4.
Whilst the aforementioned work uses a similar methodology and theory the focus is more in line with unconditional generative modelling rather than our setting which seeks sample from the posterior arising in an inverse problem / conditional generative modelling.
RFDiffusion As highlighted in Alg. 3 and in contrast AMORTISED TRAINING, RFDiffusion [77] does not noise the motif coordinates x!M! with the forward OU-Process, instead directly aims to
sample from p X [MM]] X (xx) and estimate this score while keeping the motif fixed. We can relate
this approach to our amortised learning of Doob's h-transform, by noting that RF diffusion can be understood as learning the marginal conditional score
p(a!
[MM) ((11)) =
p( (() dx. [M]
(22)
This can be viewed RFDiffusion estimating marginal counterpart of our amortised h-transform approach. See Algs. 3 and for more details on how these approaches differ in a pseudo-code implementation.
C Doob's h-transform
C.1 Doob's h-transform intuition
Doob's transform provides a formal mechanism for conditioning stochastic differential equation (SDE) to hit an event at a given time. The l-transform drift decomposes into two terms via Bayes rule, conditional and prior score
HH In Pop X. E B H.) = HH, In Prio (H, X B)
VH (HH,),
(23)
whereby the conditional score ensures that the event is hit at the specified boundary time, while the prior score ensures it is time-reversal of the correct forward process [14]. Doob's h-transform adds new drift to the SDE which amounts to two terms (via Bayes Theorem), a conditional and an unconditional score
V In Po E B = V In P ( E B) Vln
(24)
Interestingly. these two terms provide for unique intuition: the Doob's 's transform SDE is the time reversal of the forward SDE corresponding to (3), that is the time reversal of the forward SDE
dx; b, (X,) dt + o.d t
Xo
Po X E B).
(25)
coincides with the Doob transformed SDE (4) [14]. Thus following series of steps:
view Doob's transform as the
1. Time reverse the SDE we want to condition ((4) to (25)).
2. Impose the condition via ancestral sampling from the conditioned distributiion/postrrior
3. Time reverse once more to in the same time direction as started.
C.2 Example: Truncated normal
Here for illustrative purposes we frame the problem of sampling from a truncated normal distribution simulating SDE that is given by Doob's 's h-transform.
Let's remind that a 1d truncated normal distribution had density p(zla, b) x ( (2) (alp,, o²). Now, let's assume a data distribution po(x) N(u, o²) which noised with an OU pro- NF (o(((z)))) is Gaussian, and so is
cess (13). Thus we have that p(zolæz))
18

---

p(a,) Mireiee Let's add the constraint that the process hit at time = 0 the event Xo E b).
H 2
+ H. ( VH, In ( (a, H.) dt + /B(t) aw.,
dH; = B(t)
(26)
We have that the h-transform is given by
htt, H.) (Xo E (a, H,) = / Exe(am)) (a (2 (Zo)Poje (2ol H, dæo
/ Lec(a,)) (o (H)), Hoje (Hi) )
Hi (H--Aoj((H.)
1
Boje H()
(27)
0 (H+) Φb (H,)
(-Pore(H..)
(H.)(HH) (H.)
+
(H,)
6o1 (H,)
where o(E)
V2x exp (-15²) is the pdf of a standard normal distribution, 2({)
+ (1+ee((E// erf(e//(2) its cumulative function. The corrective drift term due to the h-transform can then be computed via autograd. The unconditional score term can be computed in closed form.
C.3 Doob's i--rranffrrm classical noiseless setting
We now consider events of the form Xo B which are described by an equality constraint A(Xo) y with a known measurement (or forward) operator and y an observation, which a common setup in inverse problems such as inpainting super-resolution.
Corollary C.I. Consider the reverse SDE (3), then it follows that
dH, = ([b((H))) - o H. In Paj((A()))) = Ht)) dt WWW
(28)
with (H.) f. HH) - V In (H. satisfies Law HH, = Law (X, X A(Xo)
y)
thus Law (Ho) Law (Xo)A(X,)) =
Sampling from (28) directly provides samples æ ~ Pdats which also satisfy A(x) y. Crucially, this SDE is guaranteed to hit the conditioning in finite time, unlike prior cuullibruummmotivaen
approaches [12, 19, 22, 24, 46, 65].
Reconstruction guidance To get better insight into the challenge of sampling via Doob's h- transform in (28) let us re-express the h-transform as
(ox(A(Xo) = Hi) = L LA(m))) Poj ( H, H.)d.
(29)
where in the case of denoising diffusion models ( is the transition density of the reverse SDE (3). In practice, one does not have access to this transition density - i.e. we can sample from this distribution, but we cannot easily get its value certain point. This makes it difficult approximate the integral. To alleviate this, strand of recent works [12, 22, 57, 65] have proposed to apply Gaussian approximation of and the pre-trained score network. This line of work referred reconstruction guidance. We yl.). they do make the
A M(oo E[X,] X. = +I, Tv) leveraging Tweedie's formula
note that whilst proposing to approximate the quantity (A)o(A((X)) connection to Doob's 's transform and thus are unable to provide guarantees on the conditional sampling that Cor. C.I provides. Overall, the Gaussian- approximations of Doob's h-transform lead to reconstruction guidance-based approaches
dH. (H.. ofV AE X, H.Jll.) di a,dW, c, X Pr.
(30)
where T, acts as guidance scale [57, 62], and A a matrix if is linear otherwise A d X, H.)).
19

---

D Proofs
D.1 Proof of Proposition 2.2
Proof. Starting from the unconditioned reverse SDE
dX;- = b. (x dW Xy Pr Qf (p(az)),
(31)
we consider its reversal, the forward SDE, but we change its initial from distribution p(ar)) to the target posterior p(zoly), ie.,
dH, fiHH) dt of dW., Ho ply 20 panta(
(32)
ply)
where (H = f( of Inpi (H.)
Now let us use Ptly (a: = p(ac, æo (pppplp)))))) to denote the marginal of the above SDE and as before Pt to denote the marginal of the reference starting the data distribution. Then it follows that
/
Pejo(ar lano)
Ptly pa())p(1)-
p(y (D)) Paata (
(33)
Pe(æ))
-((())(())1 f
Poj (wo (a)) ply an dood
(34)
and thus the score of the reference starting at the posterior is given by:
V In Paly (aly) V pe + Va In /
Poje (2 ( (pp((()))da (
(35)
Now that we have the score of the SDE in Equation 32 we can reverse it yet another time to obtain the conditional backwards SDE:
Hr (/1o(a)))))) (y)) = pro d
dH, (fi(H)) of (V In pe (H)) VH, Pn = H())) dt or dWo,
dH, (H..H. V In P y H.)) dt + ar aw.,
(36)
Where it very important to notice that the backward SDE starts a the terminal distribution of the conditional forward SDE Of [p()oly)] and not 1PT = Q [p(x)]], for a VP-SDE this happens to approximately be (0, I) in both cases (i.e. Of (p(zolv) x Qf [(()))] se NO, I)) however more generally it is not and one needs to be careful.
Note that this remark highlights that the score used in DPS [12] (i.e. V In Ptjy (aly)) is in fact the score of an OU process starting at p(æuly) y) notice the cancellation going from Equations (33) to (34) was only possible since the prior in our target posterior is the initial distribution for the forward SDE (in their case an OU-process) with marginal Pt, these considerations are subtle yet important and omitted in prior works. Whilst this is akin the relationship motivated in DPS as V. In Ptly (mly) = V, In pn (2x) + Va In (y ar), DPS fails to convey that this is in fact the score of a VP-SDE with the posterior p( yy) as its initial distribution.
D.2 Proof of Theorem 3.1 2)
Proof. The idea of the proof is similar to Proposition H.1. Here, we directly proof the amortised variation from Equation (9), First, we state the well known score matching identity for the posterior score:
Va, In (x = / Ve, In Pr / (ajo ( (z,, (2, yd
Via the mean squared error property of the conditional expectation we get the minimiser as V In (H, VH. In.. H. =
E
20

---

Then
/ V (n - V (app(()))Pp (F) HP = x, Y = =yddoo
h5 (a. y)
oe H: = x. Y y)dzo Vx (np)
= Ve ( Ve V (z)lo)
E Experimental details
E.1 Image Experiments
In the image experiment, we use the DDPM [28] formulation for the diffusion model with N = 1000 steps, linear 3-schedule with 8o 10 and 3N 10 2, In all our imaging experiments the h-transform was implemented according to the parametrisation in Section 3.3 using an attention U-Net [16] for NNT. The network N in the residual pathway was implemented as small three layer fully connected network with SiL.U activation functions, predicting a single scalar value. The final layer in NN was initialised with zeros. All weights NN; where initialised with zeros, expect for the bias in the last layer, which was initialised to 0.01. We found that this initialisation was stable over all imaging tasks, i.e., close to the unconditional model. All tasks on ImageNet were noiseless, i.e., the log- likelihood p( y would be Delta function. However, for the guidance term Va, In p(y kn) we always approximated the likeliiod using a Gaussian with σy 1.0.
Computed Tomography The 2016 American Association of Physicists in Medicine (AAPM) grand challenge dataset [45] contains CT scans of 10 patients. We only make use of the abdomen scans. We the scans of 9 patients to train both the unconditional model and the h--rrasform. The remaining patient, i.e., id L035 with 87 slices, is used for testing only. We used a bilinear interpolation to resize the images 256 256px. The unconditional model follows the same architecture as in [11] with about 374M parameters. The model was trained for 300 epochs using the Adam optimiser [33]. We used exponential moving average on the weights as [66] with a parameter of 0.999.
The LoDoPab-CT dataset [35] is a collection of human chest CT scans. We resize the images to 256 x 256px using a bilinear interpolation. The training set contains 35820 images and was used to train the unconditional diffusion model. The validation set contains 3522 images and was used to train the h--transform model and for hyperparameter search for DPS and RED-diff. Finally, take 178 images (every 20th) from the test set to compute the final results. The unconditional model an attention U-Net [16] with about 133M parameters. The unconditional model was trained for 75 epochs, corresponding to about 671625 gradient steps using the Adam optimiser [33]. We used exponential moving average on the weights in [66] with parameter of 0.999.
For both datasets we used the same parametrisation for the h-transform, with an attention U-Net [16] and residual path as in Section 3.3. The fi-transform has about 23M parameters. Thus, the h--ransform is about 6% and 17% the size of the unconditional model for AAPM and LoDoPab-CT, trained set of 3522 images, while for AAPM
respectively. For LoDoPab-CTi -T, the h-transform we trained the h--rrasfform on the same dataset as the unconditional model. The forward operator A is linear and we added Gaussian noise with standard deviation Bg 1.0. Thus, the gradient of the log likelihood reduces to
1 A A~ - y),
Eo In p(y)@)))
20
where A" denotes the backprojection. Here, instead of the backprojection, we made use of the filtered backprojection, which is common technique for computed tomography [25].
Inpainting We made use of the pre-trained ImageNet dataset. For this task, we also amortised over the different sampling masks. We implemented the l--tranfform using an attention U-Net [16] as the base network NN, with an initial convolution with 13 channels, i.e., the noisy image 2,, the denoised estimate Eo. the observation y (where the missing pixels filled with zeros), the mask M and the cheap
21

---

guidance term Va (n))) The pre-trained unconditional score model has 550M parameters, and we use fine-tuning network with 23M parameters (4.2% the size). The fine-tuning network was trained for 200 epochs using batch size of 16, and the Adam optimiser with a learning rate of 5e-1 and annealing.
Super-resolution We made use of the pre-trained ImageNet dataset. We implemented the h- transform using an attention U-Net [16] as the base network NN, with an initial convolution with 9 channels, i.e., the noisy image 24, the denoised estimate &g. the observation y bilinear upsampled to the original size and the cheap guidance term Væ p ( The pre-trained unconditional score model has 550M parameters, and we use a fine-tuning network with 23M parameters (4.2% of the size). The fine-tuning network was trained for 200 epochs using a batch size of 16, and the Adam optimiser with a learning rate 54 and annealing.
HDR For HDR the forward operator is given by A(x) clip(2x: 1) for the RGB image scaled to [-,, This leads to a reduction of high intensity values. We approximated the cheap guidance term
do In Inp(y 20) = 0.5(A(2)) yy,
and found this to achieve good results. The pre-trained unconditional score model has 550M parameters, and we use a fine-tuning network with 23M parameters (4.2% of the size). The fine- tuning network was trained for 200 epochs using a batch size of 16, and the Adam optimiser with learning rate of 5e-1 and annealing.
Phase retrieval The h-traasformm was again implemented as an attention U-Net [16] as the base network NN1. The observations in phase retrieval corresponds to the magnitude values of the Fourier transform. In our initial experiments, we feed these observation directly into the model. However, this model failed to create convincing reconstruction. Instead we used two rough reconstruction. For the first initial reconstruction, we used the phase of the unconditional Tweedie estimate &y and the magnitude of the observation to construct an image. For the second initial reconstruction, we ran 350 steps of an ER algorithm [21] with a random initialisation. Further, we used the cheap guidance term V byy A(éo)ll. calculated using torch autograd. The h-traasformm had about 23M parameter, i.e., 4% of the size of the unconditional model.
Non-linear Deblurring For non-linear deblurring we found that there was little improvement when using the cheap guidance term Va, y A( calculated using torch autograd. As the forward operator is defined by a trained neural network this additional autograd term adds to the computational expense. We found that the DEFT provided results of similar quality, with a reduced computational burden, by using A(A(#))) - y) as a rough approximation. The pre-trained unconditional score model has 550M parameters, and size). The fine-tuning network was trained for 200 epochs using a batch size of 16, and the Adam optimiser with learning rate of 5e and annealing.
fine-tuning network with 23M parameters (4.2% of the
E,2 Protein Motif Scaffolding
Diffusion process We use discrete-time DDPM [28] formulation for the diffusion model with N
1000 steps and cosine 3-schedule [16].
Noise model The denoising model Eg is adapted from the Genie diffusion model [37]. ). In Genie, the denoiser architecture consists of an SE(3)-invariant encoder and SE(3)-cquivariant decoder. While the network uses Frenet-Serret frames as intermediate representations, the diffusion process itself is defined in Euclidean space over the C., coordinates. Similar to AlphaFold2, the denoiser network consists of a single representation track that is initialised via a single feature network and a pair representation track that is initialised via pair feature network. These two representations further transformed via a pair transform network and are used in the decoder for noise prediction via IPA [30].
To evaluate unconditional sampling- based methods, we used pre-trained version of the unconditional Genie model.
22

---

To evaluate the AMORTISED approach (Alg. 4), we perform a minor modification to the unconditional Genie model as described in [17]: we add an additional conditional pair feature network that takes the motif frames as input with the ground truth coordinates for the motif and 0 as values for all other coordinates that are not part of the motif. The output of this motif-conditional pair feature network is concatenated with the output of the unconditional pair feature network to form an intermediate dimension of twice the channel size compared to the unconditional model before being linearly projected down to the channel size of the unconditional model. From then onward the output is processed by the remaining Genie components as in the unconditional model. The implementation is therefore similar to the image case, where the motif features presented additional input and the model learns to use these for reconstructing the motif. This minor alteration of the Genie architecture means that the amortised network has 4.162M parameters while the unconditional Genie networks have 4. 4.087M parameters (~ 1.8% fewer). In 80% of the training steps for the amortised model, we pass condition to the network. The other 20% contains empty mask consisting of only O's.
For the DEFT implementation, we follow similar way of feeding in the additional inputs via conditional pair feature networks, but as part of the downsized h-transform model as described in the main text.
Metrics We measure the performance of the methods across two axes: designability and success rate. To assess whether a particular protein scaffold designable, we run the same pipeline as [37], consisting of an inverse folding generated C. backhones with ProteinMPNN and then re-folding the designed sequences via ESMFold. The considered metrics and their corresponding thresholds are the following:
scTM 0.5: This refers to the TM-score between the structure that's been designed and the predicted structure based self-consistency previously described. The scTM-score ranges from 0 to 1. Higher scores indicate a higher likelihood that the input structure can be designed.
scRMSD < 2 À The scRMSD metric is akin to the scTM metrie. However, it uses the RMSD (Root Mean Square Deviation) to measure the difference between the designed and predicted structures, instead of the TM-score. This metric is more stringent than scTM as RMSD, being local metric, is more sensitive to minor structural variances.
PLDDT> 70 and PAE <:: Both scTM and scRMSD metrics depend on structure prediction method like AlphaFold2 or ESMFold to be reliable. Hence, additional confidence metrics such as PLDDT and pAE are employed to ascertain the reliability of the celf-connistency metrics.
In addition, we want to judge whether the motif scaffolding was successful or not. Therefore, similar to previous work by [77], we calculate the motifRMSD between the predicted design structure and the original input motif and judge samples with 1 Á motifRMSD
successful motif scaffold.
We follow previous work and call a sample a "success" if scRMSD < 2 and motifRMSD < A. Similar to previous work we call task "solved" if among 100 samples for this task least 1 sample is a success.
F
Additional Results
F1
Ablation of the DEFT parametrisation
The parametrisation of the h-trannformm is motivated by the sampling theory in Section 3.3. We evaluate different parametrisations of this choice for the CT experiment on LoDoPab-CT For all architecture choices, we used the same training setup. For quantitative results, see Table In particular, we see that the naive choice NN (x, A"y, t) only achieves a PSNR of 26. 96.62dB. Note, that we do not input the observations directly into the network, but first transform them using the backprojection, which is a standard technique for computed tomography reconstruction [3]. In CT the observations correspond to sinograms and have a different geometry to the images. In difference, if we add the unconditional Tweedie estimate 2o to the model architecture, we get an improvement to 34.04dB. This shows that beneficial to supply the h-traansfomm with the information of the unconditional diffusion model. Further, given our architecture, i.c., adding the additional information
23

---

Trained residual scaling (N)
LoDoPab-CT
AAPM
0.04
0.060
0.03
0.045
0.02
0.030
0.01
0.015
0.00
0,000
0
200
400
600
800
1000
200
400
600
800
1000
time steps
time steps
Figure 6: The trained residual scaling network NN; in the DEFT architecture (see Section 3.3) for computed tomography reconstruction on LoDoPab-CTT and AAPM.
Table 6: PSNR and SSIM for computed tomography on LoDoPab-CT with different parametrisation of the h-trrasffrmm
Parametrisation
PSNR
SSIM
NN az Lo. Va. In pn + (t Inp(y(aa)
35.81
0.876
NN,
(x, V In
+ NN; Ve. (np(y(
35.74
0.875
NN
x To. y)
34.04
0.851
NN
x t)
26.62
0.724
Method
PSNR (t)
SSIM (t)
LPIPS (L)
Time in hrs (L)
RED-diff
45.00
0.98
1.2e
7.9
DEFT
64.64
0.99
1.0e
5.2
Table 7: Results the non-linear blurring operator from [44], where both methods achieve very high reconstruction and perceptual clarity due to the operator resulting in a trivial forward operation, rather than non-linead blur. On this task, still find DEFT to outperform RED-diff, with less time taken overall.
of the log- ligeliidl term, achieves PSNR of 35.81dB. Further, show the learning residual scaling network Ç in Figure 6. We observe similar behaviour for all tasks,i.e., at the start of sampling / Re 1000 the scaling network assigned small weighting to the guidance part, which increases during sampling (1 0).
F.2 Non-linear Deblurring: Implementation from [44]
We find that the original implementation of the non-linear forward operator from the codebase provided in [44] results in an almost trivial reconstruction task, due to incorrect loading of weights for the neural network used for the non-linear deblurring. can achieve highly performant results on this trivial task, as shown in Table 7. For the non-linear deblurring tasks in the main text in Table 2, we load the weights for the neural network correctly, and much lower performance, corresponding to the difficulty of the non-linear task.
result, both RED-diff and DEFT
G Generalised h-transform and Stochastic Control
Thanks to our formal framework this section we develop a new VI objective for learning the conditional score in the noisy inverse problems setting. That is by minimising the following ELBO
24

---

Ground Truth
Sample #1
Sample #2
Sample #3
Sample #4
Figure 7: Example reconstructions for phase retrieval. Phase retrieval has many local minima, which fully satisfy the data consistency constraints (c.g., complex conjugate, global phase sign). We often see flipped (and perturbed) images as our reconstruction. some instances, even only one colour channel is flipped. This leads to strong diversity of samples.
Ground truth
Sample #1
Sample #2
Sample #3
Sample #4
Figure 8: Diversity of samples for inpainting. On the left, we show the inpainting mask as overlay over the ground truth.
with respect to an additional fine-tuning network, one can learn the conditional score
arg min Eg J. (HH(( EHo~ [In p(y) Ho)),
(37)
where H, follows the unconditioned score SDE with added control h. This objective provides way to learn the conditioned SDE from the unconditioned one, without making Gaussian approxima- tions. We formalise this connection in the following proposition.
Proposition G.1. The following stochastic control problem
1 {" of of))(()))))
ht arg min Eq
EHOO-GC [Inp(g) Ho)]
(38)
2
with
Hr Q4 (p(aa)))))
dH, = (H, - of H. Inp. h(H,)) dt + t dW,,
(39)
25

---

Measurements
DPS
IIGDM
DDRM
RED-diff
DEFT
Figure 9: Results for 4x super-resolution.
Ground truth
Measurements
RED-diff
DEFT
Figure 10: Results for non-linear HDR. Similar [44], found that DPS does converge good solution, returning often black or images not consistent with the ground truth. Given that the forward operator A(x) clip(2x) 1) has zero gradient for [x > 0.5 the guidance term is often not informative.
is minimised by the conditional score SDE in Equation (28), that is
(xxx)) V In Exo-Poiee [p Xo)] V (ggie(x))
(40)
Furthermore, ht solves an associated half- half-bridge problem [7] with the SDE in Eqn. (3) as its reference process and p(acly) as its source distribution.
Proof. The derivation for this objective is inspired by the sequential Bayesian learning scheme proposed in Lemma 1, Appendix B of [75].
Let IP denote the distribution for the forward SDE in Eqn. (25). Now consider the following variational problem termed half-bridge [7, 15, 73].
Q
arg min
DKL
(41)
Q:Qo=p(#o w)
where the constraint enforces that time we hit the target posterior p(æly) then via standard results in half bridges know that the above optimisation problem has an unconstrained formulation (e.g.
26

---

see [74]) that is dQ = dIP dp(wlv) in the posterior term:
Now following [75] we notice that we can cancel the Pdata prior
deo
dp( y)
dp(zlv) dP dp(y)a=
dF
dP
(42)
dPo
dpdata
dp(y)
and thus:
Q' arg min DRL. (Q - E [In p(yl Ho)]
(43)
with Qr = Law (XT) 2 N/ I) when Xo p(zly) Furthermore, we can parametrise P as
dX; - (fr(X,, oi (X:)) dt o. Xo Qf (2dara
(44)
and thus Q as
HT ~ Q (p(ppp<ll)) [y]
dH, (H, of(V H, In (H.) + (a()))) dt + or dW., then via Girsanov Theorem we can re-express the KL. term in Eqn. (43) as
(45)
2 / T
h" arg min Eq
o
En (npply Ho)]
(46)
Now noticing that Eqn. (46) is standard stochastic control problem [48, 31] can characterise its minimiser as (using Theorem 2.2 in [48] and the Hopf-Cole transform [23])
h (z) Va In Exo~Poi (- [P(y) Xo)],
(47)
and thus the SDE in Eqn. 45 with hg hits the target posterior p(y ar)) at time as is the minimiser of half-bridge posed in Eqn. (41).
Notice that in the case of a VP-SDE, the stochastic control objective reduces to:
''" [I(H))))
T Be
arg min EQ
EH p(y Hu)]
(48)
h
2
Discretisation of inverse problem objective Following [74] we will discretise the objective presented in Eqn. (48). Let consider the pre-trained score SDE with added tuning network:
Hr ~ (O.
dH, -3, (H, + 280 H.) + 2h H.))
(H,)) dt + 28, dW f,
(49)
now using an exponential-like discretisation [10] (Ideally we want to discretise in the same way we trained the model):
Hix ~ 1),
H (VI ax + vI ax) (sa- (H2,) + hes (Hi,)) + Varsk.
(50)
where og 1- exp (Iit- 3,ds) note will denote the distribution of the above discrete time chain as 9ø- Now if we follow the sketch in Proposition 3 of [74] the discretised objective then becomes:
K
X2 (ho H1, In p(y) Ho)
arg min E
(51)
Okk
where A 1 v1 For more stable/simple objective following [74] we can make the approximation A 1 - VI - ag ss /2/2 for small time steps. This leads to the following iteration (which is possibly more akin to the training update being used):
Hix ~ ~((
H = (vI (VI- + aa (sg. (Hw,) + hes (H+,))) + VouER.
(52)
and objective:
arg min E EHogo
Σ
ak [ha HI, Inp(y)
(53)
27

---

Discrete Time Intuition For further intuition, we will provide a discrete-time derivation as to how this objective arises. Consider the following discrete-time VP-SDE (i.e let. 11+1IR(h,,,, (ti,) (h /1 aghir as)) starting from the posterior:
p(herat)) p(y (tyyhho)Patta (20)
K
II
Pa+ P ( (h(e)
(54)
p(y)
Ph (hey p IPE12 (h (h
unand
unend
Now applying Bayes rule and [2, Section 5] Pk+ (teeess hea))) telescoping cancel the marginals have:
and
pa (her))
pumend (hr) (po))))) Ddaaa (ha)
p(h(, Pk
II R=1
Ph [K++ (hey (btu41)
unond
(55)
p(y) Pducahho)
Pk ((y)p(y) heo)
umend
K unend Pk|K+1
(tee
(56)
p
uncudd where Pk|K++1 (hr,) ( is the transition density of the unconditional score SDE and pr (h..) correspond to its marginals. Now we would like to learn a backwards process that matches the above process (reverses the VP-SDE starting from the posterior). We can so by minimising the KL:
Ih g6 (hee, (houts) In (ho)
DKL (a lpp) oc Ey
In
(57)
II Pk]K+1 (he, (heet))
where we can upproximate the score transition via:
Pk Pk-1jE (tay) ((re))) ) x (i autt + - V1 S8- (tes):: aus)
and parametrise the new conditional denoiser as
I| (he, ((r.)) = W(th,_,//V aghi, + +2(1 v1 04k) (sg. (hea) + hop(1))))) ass)
making these two substitutions will lead to the objective in Eqn. (53).
G.1 Related Work
Fine-tuning diffusion models via a optimal control perspective, e.g., devolved [6], has received lot of attention in recent years. In particular, in the context of fine-tuning with respect to a differentiable reward function, i.c., considering tilted posterior [18].
( pdata (xo)
π xxx))=
(58)
Z
where e=(o) serves an equivalent role to the likelihood p(y|zo) in our setting.
The DRaFT framework [13] proposes heuristic method to estimate the fi-transform by only opti- mising the reward function. We want to highly that concurrently [72] develop the same stochastic control formulation as we do, and arrive at the same insight that the optimal starting distribution is given by Pr Of (mr. however, they chose to learn this distribution which argue necessary for diffusion models due to the mixing property of the OU process leading to negligible error by approximating Qf [ ~ (0, I) with Gaussian.
G.2 Connection to [18]- Value Function Bias
In [18], it is argued that minimising the stochastic control objective does not lead hitting the posterior p(zol! y) at time = 0 due to bias introduced by the value function. We can apply Proposition G.I to the tilted posterior (xo) from (58) which yields the following objective:
1
h" arg min Eg h
o 2
(II(HH)))
-H (Ho))
(59)
2
J
28

---

with
Hr ~ Q [p( [y]
dH, (Fr oi(V H. In (H..) + he (H.)) dt a o dW.,
(60)
then by Theorem 2.1 in [71] the optimal transition density of the controlled process is given by (8 S I and let h(ma)) = In (y)e(y(az)))
(xly) (ice,))
Pajo by)
(61)
where the log of hi-transform h (note we have used à to denote the log of the h-transform) coincides with the negative of the value function in [18, 71J. Then for 0 and = T this induces the following joint distribution:
( = (1)) klu. n
Paje ( (xlv))r(())
(62)
where [18] argue that in general the term e h(y.7) induces a bias such that when we marginalise out y Po (a)) ) not the tilted distribution, more precisely:
/ e -h...TT sef Pojr (2: (xll)rr(())d +
e e0(0) (a)ata())
pū (x) = rn z
(63)
However, let's look more closely as why this the case; first let us re-express the h-transform at time 7 as ratio of densities:
My.T)
dı (vo)puj (2o yr)dyo
(64)
dpdata
e prlMolpaataa( (Vo)
Zpaata(B1)) por 10rlyy ddyo
ref
(65)
/ ((
(66)
prc (yr)
Q (T](T))T)
(67)
Pf (yr)
substituting back into (62) and marginalizing we have:
/ ref (al
zef PT (3))
(a = 1n
In z
dy
(68)
Of [T]
now notice pr( is the distribution we simulate our stochastic control from which in our case we have chosen to be pr = of [T] making the cancellation
(x(x) = (()) /
ref (xle))) (y))) (y)dy- = p0(0)
(69)
Z
Leading to the following remark
Remark 1. The choice of setting Pr Of [x] leads removing the value function bias [18] in Equation 63 Note the authors of [18] pursue a different avenue for removing this bias by altering the noise of the controlled SDE.
Proposition G.2. (Value function bias when approximating Qfil=] NNO,T) Π)) In practice, often not have access to Q/: [x] and thus we may make an approximation with some tractable distribution Pr, then bound the value function bias as follows
r(w) ( (d)
pō(x)
pT -ar--2tin TV'
(70)
TV
then in the VP-SDE with a time homogenous Be chose pr (0, we can obtain tight bound,
33 based diffusion model (for simplicity), where
e(xx) (paaaa (2+)
(p(ar)
BB
(71)
PrV
for some constant C O, thus the value function bias [18] is exponentially small for score based diffusion models.
29

---

Proof. Applying Theorem 17 from [15] we have:
e²(0)) Pdata (x)
Po(x))
phthi [W(0, [)] Po
pf 10 [+)]
(72)
2
TV
- 1) 0- Ce BT
(73)
the final equality follows from the mixing properties of the OU process [4], where
Poithi [u](x)) = Po
(74)
Finally we want highlight that the benefits of Proposition G.2 can only be leveraged in score matching settings where we have a clear characterisation of the forward process and we are able to tractably characterise Pa however in settings such as flow matching [38, 41] and stochastic interpolants [1] where the forward process is not explicitly characterised we wither have to learn Pe like in [72] or use memoryless noise schedule as proposed in [18].
G.3 Scaling up the Control Objective
Naively trying to minimise Eqn. (37) is demanding, as the full chain has to be kept in memory, which is infeasible for high dimensional problems. To alleviate this problem, one could make use of the stochastic adjoint sensitivity method [36]. in which an adjoint SDE is solved to estimate the gradients of the stochastic control loss in Theorem 3.1 3). This method has the advantage of constant memory cost. However, the computational cost increases as both the reverse SDE and the adjoint SDE must be simulated. Instead, discuss two alternative approaches to reduce the memory requirements.
VarGrad
We can make use of a VarGrad [53, 54] type loss to reduce the memory requirements. In contrast to the KL loss of Eqn. (37), then the VarGrad loss is given by:
- ( dp
27
D Dingaar 12 W)
O:7 - E
9+ O:T
(75)
dP
where Q and are defined in the proof of Proposition G.I. i.e., Q is given by the conditional SDE and IP is given by the unconditional SDE. The RND in Eqn. (77) is evaluated at the trajectory of reference process W- Law(H given by
Hr ~ 0f (p(2o))))
dH,- (H, of(VH, In (H.) 96 (H.))) dt o dw.
(76)
The Ridon-Nikodym derivative (RND) in Eqn (75) can be evaluated as (Using the RND for time reverse SDEs see Equation 64 in [76]):
dQ dP H
L' of (HP + L - Inp(yl=))
= 2
+ J.
T
o (HW
(77)
for the reference process W. The core advantage of VarGrad that we can choose this reference process. In particular, the choice gt = stop_grad haap_grad(M)) (h+) gives us way to detach the trajectories, saving us from having to score all gradients memory.
Trajectory Balance An alternative to the VarGrad loss in Eqn. (75) is the following trajectory balance [42, 43] loss
( dQ dlp
01 (T) - k 2
}
2 (Q.(P.P = E E w
(78)
where a learnable parameter and W the same reference process as above. We again choose Dr stop_grrrrrrr stop_grad(hh) This loss is also motivated by valid divergence, see [48]. In difference to the VarGrad loss (75), the inner expectation is exchanged with k: which approximates running mean. In practice, optimise & and and ht at the same time.
30

---

Stoch. optimal control loss
PSNR
10a
EM-backprop
20.0
VarGrad
x 10
Trajectory Balance
17.5
102
15.0
3 x 10?
12.5
2 x 102
10.0
200
300
200
300
Gradient steps
Gradient steps
Figure 11: Left: Tracking the stochastic optimal control loss (37) for the three methods. Right: Mean PSNR of samples.
Trajectory Subsampling We observed that backpropagating gradients for only a random subset of discrete timesteps in the RND can be an effective strategy. This subsampling reduces memory costs the expense of increased variance and the introduction of small bias. Nevertheless, the reduced memory usage per example enables larger batch sizes, which can mitigate these effects. Notably, we found that backpropagating only 20% of the timesteps still achieves comparable performance.
G.4 Stochastic Optimal Control Experiments
We provide some initial proof of concept experiments on the MNIST dataset [34] of handwritten digits. In particular, we make use of a parallel beam Radon transform with 5 angles as the forward operator and perturb the observations with 10% additive Gaussian noise. All SDEs are discretised using an Fuler- Marlyma scheme and the integrals are estimated using simple quadrature rules. We use a -qnuiddiisttnt time grid according to square root function, i.e., let 0 to S S f
be an equidistant grid of [0, 1] for time points. We then use th f ........... &x as the time points for evaluating the SDEs. This gives us a finer discretisation closer to f O. The unconditional MNIST model based on the attention U-Net architecture [16] with about 3M parameters. The l--transform is implemented using the same DEFT parametrisation in Section 3.3 with about 70 000 parameters. compare EM-backprop, i.e., directly backpropagting through
discrete SDE solver (see e.g. [36]), with VarGrad and TrajectoryBalance from Section G.3. For EM backprop size of and use 60-80 time steps. Instead, for VarGrad and TrajectoryBalance we were able to use batch size of 26 and 80-140 time steps. For training we used single GeForee RTX 3090 and the training time took about 1h. The results are presented in Figure 11, where the loss trajectory of Eqn. (37) and the mean PSNR of samples is shown. Both VarGrad and frajectoryBalance are able to minimise the stochastic optimal control objective to the same extend as EM- backpoc Further, in Figure 12 we provide an overview of the TrajectoryBalang training. Here, we observe that k is working as a estimator of the mean RND with lower variance.
batch
H Amortised Conditional Training
In this section, we discuss an objective for learning the full conditional score training time in an amortised fashion instead of enforcing the constraint during inference time as before in reconstruction guidance approaches. This objective is akin to CDE [5] with the difference that we propose amortising over the the forward operator, for example in image inpainting or motif-scaffolding
Note that since Po (Y = X, x) = Palo y)po (Y v)/pe(X, x), we can re-express the Doob's transformed SDE of reversed OU process
HH. -3: H. 2V H. In (H. vi) di + 2 dw,, HT Law (XT).
31

---

Trajectory Balance loss
Mean RND and
107
400
mean RND
350
k
10a
300
250
10a
200
150
0
100
200
300
0
100
200
300
Gradient steps
Gradient steps
Figure 12: Training using TrajectoryBalance. Left: The trajectory balance objective loss (78)
over training steps. Right: The mean RND and the trained k. We see that follows the mean RND. However, has smaller variance.
Proposition H.1. The minimiser of
arg min E h
(Xo,Y )mppaa, Y) (~,((, (P (m))
[1h(t, Hr. y) VH Pn ( (o)1
(79)
is given by the conditional score S7 (xx V Iu P =y)-
Proof.
Via the mean squared error property of the conditional expectation the minimiser is given by:
(x y) = H.. Inpro(X Xo) = H, a
(80)
Then:
(x, v) / Va In Pojo (x) DD Poje (wo H, x, Y y)dx
V (x x0) Prjo (He = x oy. = y)p(2ol Y = y. das Pejo(ar ar ara)
pl Hi xaYY = y)
1 (0))
Pejo (x) Dejo (H, x x/zo)p(<00| yydao
HHH x Y y)
1
V P:jo 220022/2))p((0011 yddXo
p(H: = x = y)
V V x/Y = =
(81)
7(,; = x Y = =
= V. Lnpp(XX x = y)
where that Pnjo (H: Zb Y y) prpo HH x/20) H, is independent from Y given X.-
As with DEFT, for settings where A varies like in image completion we sample A randomly and amortise it over learned h-trransform, i.c. estimating hū (æ, A).
We refer this approach as amortised learning for conditional sampling, since practically the neural network approximating the (conditional) score is amortised over AA and instead of learning separate network for each condition. This approach is also reminiscent of 'classifier free guid- ance' [27] where the score network is amortised over some auxiliary variable (e.g. as in text-to-image models [51]), or of RFDiffusion [77] where proteins are designed given specific subset motif, or similar to [5]. See also Appendix B for discussion of related conditional training methodologies.
32

---

Note that conditional amortised learning is different to "classifier free guidance' as A is assumed to be known (e.g. an inpainting mask). Also note that due to its formulation, classifier guidance would be unable to noise a subset of X (the motif) as we do and would instead be more akin to RFDiffusion.
H.1 Relationship to Conditional denoising estimator (CDE)
Conditional denoising estimator (CDE) [5] is the adaptation of [27, 60] to inverse problem- like settings, deriving a variation of classifier-free guidance to a measurement model styled scenario. Whilst they do not focus the measurement model, they estimate very similar quantity as our Proposition 2.5
JCDE (x. y) V Pnjoaa/ = y)
(82)
In contrast to to the amortised conditional training:
pamortisod (2r y, A) V In l y, A A)
(83)
when explicitly considering the distribution over the measurement model, one can see that the quantities are related one another via marginalizing the measurement model pA. This introduces several practical and conceptual differences:
If we consider in/out painting as example, the score network estimating JCDE is not explicitly aware of where in the image the missing pixels are. As a result, it must perform inference over A (effectively marginalizing it) in order know where complete the image. This is clearly a much harder task for a single network to learn than conditioning on A where we provide this information.
Viewed under the lens of the h-transform, fCDE can be viewed as amortising the event A(Xo) = y for random A. It therefore falls under the soft constraint settings since A(Xu) ) Xx is not delta. Our quantity famortised is amortising over A(X,) = y for deterministic A and is therefore part of the more classical hard constraint domain of Doobs transform. We believe amortising over these simpler deterministic events can offer an advantage in making the problem easier to learn.
H.2 Comparison to /-traasform fine-tuning
We compare the amortised training framework against our h-transform fine-tuning on the FLOWERS dataset. The preprocessing procedure consisted of centrally cropping the image to size 64 x 64, and rescaling to pixel values [-1, 1]. The dataset is split into three parts containing 6149, 1020 and 1020 images each. We use the first part to train the unconditional and amortised model. The second part is used for the h-transform fine-tuning. The third part is used for evaluation.
For this experiment, we choose both an inpainting and an outpainting task. For the inpainting task a random 18 18px patch from the image removed. In difference. for outpainting only random 18 x 18px patch remains and the rest of the image is removed. Thus, the outpainting tasks tests better the generational capabilities of our framework. For the unconditional and amortised model, we use a standard attention U-Net [16] in the discrete DDPM framework. Both the unconditional and the amortised model have about 24M parameters. There is minor difference due the fact that the unconditional network has 3 input channels and the amortised model has 7 input channels, i.e., the noisy image, the observations and the mask. The h-transform implemented according the parametrisation in Section 3.3 with an attention U-Net [16] for NNŞ. In total, the f--transform has about 4M parameters. i.e., about 18% of the size of the amortised model. We evaluate three different settings for the amortised model:
AMORTISED (20x, FULL DATA): trained on the full training dataset for 1200 epochs,
AMORTISED (2x): trained on the fine-tuning dataset for 700 epochs,
AMORTISED (1x): trained on the fine-tuning dataset with the same computational budget as DEFT (300 epochs).
We trained all models on single GeForce RTX 3090. Training time for AMORTISED (20x, FULL DATA) and the unconditional model was about 50h. The training time for AMORTISED (2x) was 4.5h, while the fine-tuning and AMORTISED (1x) took about 2.5h.
33

---

Table 8: Comparing the full amortised training with our conditional fine-tume fine-tuning objective on the Flowers dataset for impainting, outpainting and blur.
INPAINTING
OUTPAINTING SSIM
BLUR SSIM
PSNR
SSIM
KID
PSNR
KID
PSNR
KID
AMORTISED (20x, FULL DATA)
27.81
0.936
0.000057
12.14
0.221
0.028
23.02
0.734
0.0196
AMORTISED (2x)
25.22
0.912
0.0016
10.84
0.192
0.0821
22.66
0.716
0.0289
AMORTISED (IX)
16.28
0.806
0.006
9.985
0.173
0.1
21.75
0.705
0.0374
DPS [12]
26.29
0.897
0.0036
11.88
0.215
0.0389
22.51
0.683
0.0624
DEFT
26.18
0.916
0.0019
11.18
0.160
0.11
23.16
0.709
0.0529
Ninds Mdm2
30X7
BSV protein
1.0
1.0
(ir 41M42
Surceae 12%
Sundtarr 46%
Succemo: 2%
-0
06
022
10
0.0
6/
2.5
5.0
7.5
10.0
12.5
STPK BSV Pepmteii Sitn V
STRV med De nown designed protein 44500044344
STRV short- De nove designed protein
Bacgmaiee 35
4.8
Dulmatia 4%
-08
Beeignahin s5
08
s.8
-04
004
5.2
02
02
12
5.0
10.0
12.5
15.0
255
2.0
SEGR med Ferredotia Protein
GB6R short Ferredosia Proteis
6EXZ med RNA export factor Semeers
2.0
Desgnabie: 3%
Daignable 1%
-08
Socmaa: 12%
ROERMMD
2
2.4
5.0
16
>CRM5D
short- export factor
7MEX 60
inhibitor
Barnase ribonuckase inhibitoc
1.0
(oddyaabi 20%
(allgaaa 58
Breignable 1950
Seecesss 22%
Scncees 1%
Sutceer 0%
30100005
06
022
16
10
15 5.0 25
5.0
7.5
10.0
12.5
KIRMSD
Figure 13: Full results for DEFT (9% model). For each task, we show the full scatter plot of scRMSD and Mot for all 100 samples. The colour indicates the PLDDT confidence score of the re-folded structure with ESMFold. Samples with PLDDT 2 0.7 are outlined.
Results are presented in Table 8. With same computational budget, DEFT outperforms the amortised model on all tasks. Training the amortised model with a larger computational budget, recovers a similar performance to DEFT. Finally, in the scenario of full access to the complete dataset and large computational budget the amortised model is able to outperform both DEFT and DPS.
34

---

Becigallo
10%
Dosignable
(Scccsss 29%
Success: 12%
04
0.6
6530
02
2.5
5.5
7.5
10.0
12.5
2.5
5.0
7.5
10.0
12.5
WIR06DD
<CRMSD
Amortised
DPS
DEFT 9%
Figure 14: Comparison the amortised model, DPS and DEFT (9%) on the task IYCR. We see the general trend for DPS that for low guidance scales the samples have high designability but do not adhere to the motif constraint, while for higher guidance scales they adhere to the motif constraint but have low designability.
Amortised
DPS
DEFT 9%
Comparison of samples from the amortised model, DPS and DEFT (9%) the tasks 6EXZ med. One can see that while the amortised and DEFT samples incorporate the motif into a realistic backbone, this is not the case for DPS. We generally observed that small guidance scales DPS produced realistic backhones without the desired motif and at high guidance scales it placed the motif into unrealistic backbone.
Figure 15:
35

---

I Algorithms
In this section, we reformulate multiple algorithms from the literature under our common framework as reference for practitioners. In these algorithms, we use the following conventions: our dataset drawn from the law Paua. but we can only sample from the simpler law P sumpling at inference time, which is often chosen as multivariate standard normal P sampling =0, I). Therefore, we construct a forward noising process Pasta + Pammlinn that is parametrised via the noise schedule 3 B(t), a; a(t) and try to learn the reverse denoising process P Psampling - Paata Due this notion of "forward", and to keep consistency with the literature on denoising diffusion models, we explicate the nomenclature PaZaa Po and sampling Pr.
There an additional law Puoise that is sometimes confused with Psampling since in practice both are often chosen as (N0, I), but they are two distinct laws that could in principle be different. P Pmoiee is the law from which the noise added during the forward noising process as well the during the reverse diffusion process is drawn from.
Algorithm
Unconditional training of denoising diffusion models [28]
Require: Dataset drawn from law Pdara Po
Dataset law Pdnta
Require:
Noise schedule B, = (t), ar = at),, parametrising process Panta P sampling
Require: Untrained noise predictor function e, (a) with parameters 0
1: repeat
2: To ~ Po = Paata
3:
Uniform {I, T})
4:
Forward noise sample, at (o))
5:
e, Phoise
P Often Brownian motion, Proise N()
6:
0z f Vā,a0 + /1 a,E,
7:
Estimate noise of noised sample
8:
Eo + (x,)
9:
Take gradient descent step VoI(e,,
b Typically, loss L(Duue ppred) =
D Fpred
10:
until converged epoch reached
Algorithm 2 Unconditional sampling with denoising diffusion models [28]
Require: Unconditionally trained noise predictor (x,, Require: Noise schedule B, (()t), āg a(t), parametrising process Paata P Psampling
1: Sample starting point 27
2: 2T Pr
Saaplling
D Often Pr I
3: > Iteratively denoise for steps
4: for in (T, - 1,..., .., 1) do 5: Predict noise with learned network
6: éo + (x))
7:
Denoise sample with learned reverse process #t-1 (P=-1j((x)) Perform reverse drift
8:
3t
9:
It-1 V1 3
VI
Öe
10:
Perform reverse diffusion, which often Brownian motion in R", i.e. Proise ~\
11:
E, Puoise if > 1 else Et 0
12:
Ft-1 f- 0t-1 + 0,E,
A common choice or = 3(t)
13:
return Z0g
36

---

Algorithm 3 RFDiffusion conditional training [77]
Require:
Dataset drawn from Pdata
Dataset law Pdata
Require:
Noise schedule 3, = _()), â₂ = a()), parametrising process Paata - P Pampling
Require:
Untrained conditional noise predictor function fM), with parameters
1: repeat
2: 3o Po Pdata
3: ~~ (iniform(11....//
4:
[M U u t æ0
Randomly partition data point into motif and rest
5:
Forward noise the non-motif rest via sampling from (oja(
6:
e, Puoise
7:
x AMM
t Varz
[MM]
+ VT بie;" a (M)
8:
Combine unnoised motif with noised rest and timestep of motif part to 0
9:
Tt - 20
[M]
U x
10:
M t 0
éo f- fa M)
Estimate noise of sample with noised rest
12:
Take gradient descent step on
VoL(E, Eo)
Typically, [( Tpred)
Firee Tpred
13:
until converged max epoch reached
Algorithm Amortised training i.e. Doob's h-trrasffomm conditional training for motif-scaffolding
Require: Dataset drawn from Panta
Dataset law Pana
Require: Noise schedule B, 3(1)), a. ā(l), parametrising process Pastn P ampling
Require:
Untrained amortised noise predictor function fo(x, f, x(M M) with parameters 0
I: repeat
2: Xa PP Pasta
3: t ~ Uniform f1. 77)
4:
MMM J Xo
Randomly partition data point into motif and rest
5:
Forward noise full sample via sampling from Pojy (xo)
6:
E, Puoise
7:
X, + VazXo + -I- a
8:
Estimate noise of sample with original motif as additional input
éa fo(x,, I x [M M)
9:
10:
Take gradient descent step on
VoL(E,
Typically, L [(Fuee Dpred) |uuo Fpred
11: until converged or max epoch reached
37

---

Algorithm 5
h-transform fine-tuning (new)
Require: Dataset drawn from Pdata
Dataset law Pdata
Require:
Noise schedule 3; = _((),, ar = a(t), parametrising process Paata P sampling
Require:
Trained noise predictor function (2) with parameters
Require:
Untrained h--transform hy (ax, aoo y) with parameters Φ
1: repeat
2: 2o Po Pdata
3: ~~ Uniform ((., T})
4:
y~u p(y)
Simulate observations
5:
Forward noise full sample via sampling from (xo)
6:
Et Pnoise
7:
Ft Vango + -ng
8:
éo t (xx)
Estimate noise of sample with pretrained model
9:
Tu (It -1-
10:
Eo t (x t0. y)
Estimate noise of sample with h-transform
11:
Take gradient descent step w.r.t. Φ on
Vol(E, éo + ε②)
Typically, MXDiue ppeed) = D Fpred
12:
until converged max epoch reached
Algorithm 6 h-transform DDIM sampling (new)
Require: Trained h-trransfrmm he (x, 2o. y) with parameters á
Require: Unconditionally trained noise predictor (x,)
Require: Noise schedule B, 3(t)) â. a(t), parametrising process Paata - P sampling
Require: Schedule a; (()
Require: Observation y
1: b> Sample starting point æy
2: ØT Pr P Psampling
Often Py -((11
3: Iteratively denoise for steps
4: for in (T,T- T 1,.... .., 1)
do
5: Predict unconditional noise with learned network
6: ég + (x,)
$o < F1 VI - E,E9
7:
Vas
8:
Eo h1 (x to y)
9:
Estimate posterior noise
t ED + E6
11:
e, ~ P Puoise i
1 else E, f- 0
VI- vT-
12:
Fz va--
+ vI ā, - 0.0 + 0
Vāl
13:
return Z0
38

---

Algorithm 7 RFDiffusion conditional sampling [77]
Require:
Conditionally trained noise predictor fx f M)
Require: Target motiffcontext x
Require: Noise schedule 3 = (()), at = a(t), parametrising process Paata - P sampling 1: Sample starting point XT
2: Xy ~ Pr = P sampling
3: I- Iteratively denoise for I steps
Often Pf - (00,1)
4: for in (T, 1,, ,,...,,] 1) do
5:
Overwrite motif variables with target motif and reset their time parameter
Note: Original RFDiffusion zero-centers and X [M] individually for equivariance.
6:
x [M] f X M t- 0
[M]
7:
Set noisy motif to unnoised motif
8:
Set timesteps for motif
9:
ég (x,, M)
Predict noise with learned network
10:
Denoise sample with learned reverse process X,-1 ~ (xx))
11:
Perform reverse drift
31
12:
X
V1 āt En
V1 B,
13:
Perform reverse diffusion, which often Brownian motion in IR,, i... Pooise = =N(0,1)
14:
E, Puoise ( > 1 else e 0
15:
X f- Xt-1 + J,Et
A common choice or = 8(t)
16: return Xo
Algorithm 8 Replacement conditional sampling for motif-scaffolding
Require: Unconditionally trained noise predictor ei (x;,)
Require: Noise schedule Br (()), 0r = a()), parametrising process Pasta - P Pampling
Require: Target motif æo
JM[
1: b Sample starting point øT
2: 2T Pr
P Pampling
3: Iteratively denoise for T steps
Often Pr =N
4: for in (T, .,...,,1) 1) do
5: Predict noise with learned network
6: éa + (x,)
7:
Denoise sample with learned reverse process #t-1 ~ (2-1)(((2)) Perform reverse drift
8:
3r
9:
Br Ct
Eo
VI
Perform reverse diffusion, which is often Brownian motion in RR" i.c. Pooise (o E, ~ Pmoise if > 1 else E, +- 0
11:
x;-1 f- 8t + o,E, 0,E
A common choice a, 8())
Forward noise the target motif _/-1 ~ Poje-1 (M
13:
n M
Puoise if else 77-1 +- 0
[M]
15:
x + Vār 120 + VI 0;-177-1
U x [M]
16:
Φ f- 1-1
Insert noised motif into current sample
17:
return X9
39

---

Algorithm 9 Reconstruction Guidance (i.e. Moment Matching (MM) Approximation to h-transform, DPS [12]) for general inverse problems y ~ noise(A(a)))
Require:
Unconditionally trained noise predictor (azz) observation y.
Require: Noise schedule 5 = 3(t), ar = @(t), parameterising process Pauta -> P sammlling
Require: Guidance scale (schedule) 7 (t)
Require: Conditioning loss 1( (/rred y) e.g, Gaussian MM (3ppod y)
l| yl1
1: Sample starting point TT
2: 2r ~ Pr = P sampling
& Often Pr -N
3: & Iteratively denoise and condition for T steps
4: for in (1.., do
5:
Ég (z,)
Predict noise with learned network
6:
Estimate current denoised estimate via Tweedie's ' formula (2t, Eg) + Var (a+ - Vāā
7:
b c.f. also eq. 15 in (28)
8:
Perform gradient descent step towards data consistency æz 82 74 V (..(((())) y)
9:
Requires backprop through ef via e.g. La loss
Denoise sample with learned reverse process #1-1 (20))
11:
It-1 + -1- - (x, 3, [1 -
>> Perform reverse drift
12:
Perform reverse diffusion, which is often Brownian motion in IR", L.e. Posise = (0
13:
E, Puoise f > 1 else Et 0
14:
1t-1 f- -t-1 + orEt
D- A common choice or = 3())
15:
return 250
Algorithm 10
Reconstruction Guidance (i.e. Moment Matching (MM) Approximation to /--
transform, DPS (12]) for motif scaffolding
Require: Unconditionally trained noise predictor (x;) target motif/context 2g
[M]]
Require: Noise schedule 3; = 3()), ā = a(),, parameterising process Paata - PSsamllnng
Require: Guidance scale (schedule) 7 (t)
Require: Conditioning loss 1( (Zimee pred)) e.g, Gaussian MM [[eme, Fpred) = Drrre
2 pred
1: Sample starting point 2T
2: FT ~ Pr = P Saapllngg
b Often Pr Pr=N((0,1)
3: & Iteratively denoise and condition for T steps
4: for in (T, (1, 1) do
5: ég (xx))
Predict noise with learned network
6:
Estimate current denoised estimate via Tweedie's 's formula (z, Eo) + vay (z: - VI ā(Ep)
7:
b c.f. also eq. 15 in [28)
8:
Perform gradient descent step towards condition on motif dimensions M [M]
9:
(x, Ed)) Requires backprop through ef via e.g. L2 loss
10:
Denoise sample with learned reverse process #t-1 (1)))(
11:
1 + (1 - (.) (x. 3, - -1/2gg)
> Perform reverse drift
12:
Perform reverse diffusion, which is often Brownian motion in RR" L.. Pooise = (0,
13:
E, Puuise if > 1 else e 0
14:
Et-1 f- It-1 + 0g5₁
A common choice a, = 8(t)
15:
return Zg
40

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We propose new framework for conditional sampling from diffusion models, referred to as DEFT in the paper. This is stated both in the abstract and introduction, see 1.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
Limitation of DEFT are discussed special paragraph in Section 5.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
41

---

The assumptions for the proposition are always stated in the corresponding block. Proofs for statements are provided in Appendix D, with an exception for the optimal control loss. The derivation for this loss function is given in Appendix G.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We clearly explain the settings and choices of our experiments either directly in the main paper, see Section 4, or in Appendix E.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
42

---

reproducible, they Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: We provide anonymised code for our experiments here, as well details on how use the datasets we refer to. All datasets we use are publicly available and open-source.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
Atollic/guides/CdeStttiiis ssiiiiocccyy) for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / //iips.c//ppii
Soaeeeeee/ooSSSumiieeeee cesssooooolll onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permirted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Full training details provided in the released code base. Further, the setup is explained both the experimental section, see Section 4, or in Appendix E.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Due to the computational expense needed for sampling, providing error bars for conditional sampling results is not standard.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
43

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources Question:
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification:
We provide training times and experimental setup in the relevant sections.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: https://meee /ttee://eeeuueeeeeeeeeeeee cttpp//////c/////ccccccccceaaaaa
Answer: [Yes]
Justification: The authors have reviewed the NeurIPS ethics guidelines and conducted the research according to these guidelines.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
Justification: As DEFT is generative model, it suffers from the sample issues as other generative approaches. The model can easily reproduce biases inherent in the training data. Further, due to the ability for conditional sampling, i.e., drawing samples according to specific constraints, the method can be used for "deep fakes" or nisinformation. However, these impacts do not only apply to DEFT, but to other conditional sampling method well. We added sentence to the conclusion.
44

---

Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considcrations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning conrectly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: We only pre-trained diffusion models, which are already publicly available in open-source code bases or publicly available datasets.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-us, use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Where applicable, we provide references to the code bases, datasets and implementation used.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
45

---

The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. cow com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We provide framework for conditional sampling from diffuusin models. The codebase is provided as an anonymized Github repository. We re-use datasets and models with open licenses.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsoureing or research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing research with human subjects
Guidelines:
46

---

The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
47


Reviews:
[{"review_id": "6bfNLQLQjm", "reviewer": "Reviewer_no8y", "summary": "The paper proposes a new framework for fine-tuning unconditional diffusion models for conditional generation based on Doob’s $h$-transform. By utilizing a small set of observations and ground truth samples, the algorithm can learn the conditional $h$-transform that is used to", "questions": "- How do the $h$-transform networks perform on out-of-distribution samples? Given that they still contain a sizeable number of parameters how important is the number of samples used to train them in the final image quality?", "limitations": "The limitations have been addressed.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The work is a novel approach to fine-tuning pre-trained unconditional diffusion models for conditional generation. The authors provide an extensive formulation of sampling from the posterior given a diffusion prior using Doob's $h$-transform. The proposed method has the potential to unify existing diffusion posterior sampling methods under a common framework.\n- The work can be extended to non-linear inverse problems, which is a significant limitation of many existing diffusion posterior sampling methodologies.\n- By learning the transform, the proposed algorithm significantly speeds up the conditional inference in comparison to existing methods. Many of the previous approaches also required backpropagating through the denoising network during inference, making their usage impractical in many applications.", "weaknesses": "- The method requires a non-negligible dataset of observations and samples to train the $h$-transform on. Although there is a significant speed advantage during inference, access to this dataset is not guaranteed for every task and there could be issues with generalization. I.e. for training a generic non-linear deblurring operator, the network that parametrizes the $h$-transform has to be trained on a diverse enough set of observations and images. This is an important limitation that some other methods (such as [9]) do not suffer from. Even in the stochastic optimal control case, which is training-free, using VarGrad or Trajectory Balance requires significant computational resources that can exceed the requirements of previous posterior inference approaches.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable comments. We ran additional experiments, the complete results are available in Table 1 in the 1-page PDF.\n\n### Q1: How does DEFT perform on OOD samples?\n**A1**: We evaluated DEFT, trained on ImageNet (Section 4.1), on a subset of 200 images of the ImageNet-O [1] dataset for both inpainting and HDR. The ImageNet-O dataset contains unique images that do not belong to any of the classes present in the original ImageNet dataset and is considered OOD. The results for inpainting are similar to the the evaluations of the main paper, see also Table 1 and Figure 3 in the rebuttal pdf. For a nonlinear inverse problem such as HDR, we still significantly outperform our best baseline RED-Diff, showing that our method is robust on images outside the training and finetuning distribution, and it has the potential to learn an agnostic solution to an inverse problem.\n### Q2: Dependence on number of training samples and image quality?\n**A2**: As DEFT requires a dataset for fine-tuning, we ablate the number of training samples that can still result in good performance. We trained DEFT on a subset of 10, 100 and 200 ImageNet images for Inpainting. We see improvements of all metrics, when training on a larger dataset. For the KID, we can outperform RED-diff (KID: 0.86) even when trained on only 200 images.\nSee also Figure 2 in the rebuttal pdf for an example reconstruction. Here, we can see that with increasing number of training samples, the resulting image looks more realistic. However, even with 10 images, we perform quite competitively, showcasing that our method is very sample-efficient when it comes to learning a conditional transform.\n\nSee the following table for the results on inpainting on ImageNet. Here, we see that even with 10 samples we can outperform DPS w.r.t. to the KID and with 200 samples we outperform RED-diff.\n\n| DEFT, train on | 10 | 100 | 200 | 1000 (original) | | RED-diff | DPS |\n|----------------|----|-----|-----|-----------------|-|--------|-----|\n| KID        \t| 1.85   | 0.978\t| 0.401\t| 0.29   | | 0.86 |15.2   |\n\n\nFurther, we view the requirement of a fine-tuning dataset as a **feature**, not only as a limitation. While inference-time, training free methods (e.g. FreeDoM, DPS or RED-diff) have a ceiling on their performance when applied to new tasks or datasets, DEFT leverages fine-tuning to achieve high efficiency and performance even with a small dataset. We see that even with only 10 images, we can achieve a PSNR of 20.87 compared to 21.27 for DEFT. This efficiency is mostly due to the initialisation and network parameterization in Eq (13), where DEFT is initialized to mimic a cheap guidance term, similar to what is proposed in DPS or MPGD.\n\n### Q3: Computational expense of VarGrad/Trajectory Balance for the online fine-tuning loss (Section 3.2)\n\n**A3**: We acknowledge that the online fine-tuning in its current form comes with a high computational cost. We included the online objective first and foremost to highlight the interesting connection of stochastic optimal control and conditional sampling.\nAfter the Neurips submission deadline Venkatraman et al. [2] published a fine-tuning approach for text-to-image diffusion models based on trajectory balance with a similar architecture (cf. eq. (11) in [2] with eq. (13) in our work). They were able to scale trajectory balance to higher dimensional problems by using off-policy training, i.e., re-using previous samples, and stochastic subsampling, i.e., calculating the gradient only for a randomly sampled subset of the trajectory. Similar tricks can be used for our online fine-tuning objective, to reduce the computational cost and scale it to high-dimensional settings. We believe that our proposed online objective can lead to interesting future work to scale up, similar to [2].\n\n[1] Hendrycks et al., Natural Adversarial Examples, IEEE CVPR 2021.\n\n[2] Venkatraman et al., Amortizing intractable inference in diffusion models for vision, language, and control, arXiv preprint (2024)", "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": "r5ZzohRmHe", "reviewer": "Reviewer_qBFS", "summary": "The paper tackles the problem of utilization of generative modelling to solve inverse problems. The main highlight is that the authors developed a technique that can solve inverse problems without the need for backward pass through the generative model. Hence enabling deriving the prior knowledge from even closed source models for solving inverse problems. Experiments are performed across multiple datasets to show the results. The results show that DEFT achieve a speed up and performance boost over some methods.", "questions": "1. Could the authors give an analysis of the benefits of the method over [1,2] referred in weakness\n2. I would also like to see the computational overhead in terms of memory and time involved in the fine-tuning process.", "limitations": "Yes, The limitations and potential negative impacts section looks reasonable to me.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The paper introduces a new novel method for conditional generation thorough efficient fine-tuning of a small network using Doob's h-transform\n2. The proposed method enables learning a solution for inverse problems without backpropagation through the diffusion network, hence enabling learning solution for inverse problems even from closed-source models since it doesn't require a backpropagation operation\n3. By bypassing the backpropogation through the Diffusion U-Net, the method achieves a speed up over existing methods.\n4. The paper is well written and extensive experiments are performed across multiple tasks to validate the effectiveness of the method.", "weaknesses": "1. Although the paper claims that existing baselines require backpropagation through U-Net for solving inverse problems (Ln 37-38). This is not always the case. I think some relevant baselines like \n\n[1] Manifold Preserving Guided Diffusion\n[2] FreeDoM: Training-Free Energy-Guided Conditional Diffusion Model\n\nhave missed the author's notice. These methods do not require any backpropagation through the UNet and hence are faster and perform better.\n\n2. The baselines compared in the paper are very old and refer to works more than 1 years ago.\n3. An analysis of the computational overhead caused due to the fine-tuning process is missing", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for pointing us to further baselines beyond the ones we provide that would provide an interesting axis of comparison to our method.\n\n### Q1: Comparison against FreeDoM and MPGD\n\n**MPGD** [1]: MPGD proposes three variants: MPGD w/o projection, MPGD-AE and MPGD-Z. MPGD-Z requires a latent diffusion model and is not applicable to pixel-based diffusion. MPGD-AE requires to train an additional auto-encoder and thus, similar to DEFT, requires an additional dataset and training time. Unfortunately, the codebase for MPGD does not provide a pre-trained autoencoder, and we were unable to replicate the MPGD-AE variant to match the paper. We provide additional experiments of MPGD w/o projection on ImageNet for inpainting, super-resolution and HDR. In all these settings, DEFT is able to outperform MPGD w/o projection, see the following Table were we present the KID. DEFT also outperforms MPGD w/o projection on the other metrics, see the one page PDF.\n\nFor the KID, we can see that DEFT outperforms MPGD w\\o projection on both the linear (inpainting, super-resolution) as well as the non-linear task (hdr).\n\n| KID (↓)       \t| Inpainting | Super-resolution | HDR   |\n|----------------|------------|------------------|-------|\n| MPGD w/o proj. | 3.02   \t| 3.693        \t| 3.571 |\n| DEFT       \t| 0.29   \t| 1.78         \t| 0.10  |\n\nMPGD-AE defines an interesting conditional sampling approach. The use of an additional auto-encoder could also be used for DEFT to learn a conditional update in the latent space. This could possibly speed up the training. However, MPGD-AE requires training the auto-encoder, and thus an additional dataset is required, which is often much larger than our fine-tuning dataset, as an autoencoder such as VQVAE or VQGAN typically need a large amount of samples to train from scratch.\nNote, that the conditional update step of MPGD w/o projection mimics the initialization of DEFT (second term in Equation 13 of the manuscript), which can be derived by omitting the Jacobian of the unconditional diffusion model from the DPS [3] update.\n\n**FreeDoM** [2]: FreeDoM defines a distance-measuring function D. In the context of inverse problems and image reconstruction this is chosen as the negative log-likelihood (see also Section 3.3 and Appendix B in FreeDoM). In this case, the FreeDoM sampling scheme reduces to DPS [3], compare line 7 in Alg. 1 in FreeDoM against line 7 in Alg. 1 of DPS . Further, FreeDoM requires backpropagation through the trained score model, thus it would have the same computational cost as DPS. As we are already presenting comparisons against DPS as our reconstruction guidance baseline, we do not benchmark against FreeDoM.\n\n\n\n\n### Q2: The baselines compared in the paper are very old and refer to works more than 1 years ago.\n **A2:** We respectfully disagree on this point with the reviewer. While we compare against relevant, yet older, methods such as DPS (arxived Sep ‘22), we also compare against **RED-diff** (published at ICLR 2024 in May ‘24 with a corresponding complete arxiv on Sept. ‘23). Given that our submission to NeurIPS (May ‘24) falls within a year of this timeline, we believe RED-diff does not qualify as *very old*. At most, it is a year old, but in fact, the complete version (with the experiments we compare to) has only been available for about **nine months**.\n\nHowever, we thank the reviewer for pointing us to more recent methods that we could compare against, and we appreciate the references they have provided. We thank the reviewer for suggesting MPGD, (arxived Nov’23) which we now compare against in our rebuttal experiments (c.f. point above). We are also happy to run further experiments against other baselines that were released in more recent months if the reviewer has any suggestions. Based on suggestions by other reviewers, we have also added comparisons against further baselines such as Controlnet [5] and I2SB [4], providing a more thorough ablation (see general response).\n### Q3: Computational overhead in terms of memory and time\n**A3**: In our image reconstruction experiments, we report the combined sampling and training time of DEFT in Table 1 and Table 2 of the submission. However, we only give the total time. However in the text for each experiment, we also have mentioned the split between finetuning and sampling time -\n> For DEFT, this computational time additionally includes the 3.9 hrs of training time of the h-transform additionally with the 1.2 hrs of evaluation).\n\nWe will endeavor to make these numbers more clear in the text and the table.\nThe size of the DEFT model is about 5-10% (depending on the task) the size of the pre-trained unconditional model. This means that during inference, the DEFT model incurs about an additional 5-10% memory overhead, whereas baselines such as DPS incur an additional memory overhead due to needing to backpropagate through the pretrained unconditional score model. In particular, the memory cost of DPS is O(sum_i^L  layer-width_i), where L is the number of layers on the unconditional model, as all activations have to be kept in memory for the backward pass. In contrast the memory cost of DEFT during inference is just O(max layer width) as only the forward pass is needed.\n\n[1] He et al., Manifold preserving guided diffusion. ICLR 2024\n\n[2] Yu et. al., Freedom: Training-free energy-guided conditional diffusion model. IEEE CVPR 2023\n\n[3] Chung et al., Diffusion Posterior Sampling for General Noisy Inverse Problems. ICLR 2023\n\n[4] Liu et al., I2SB: Image-to-Image Schrödinger Bridge, ICML 2023.\n\n[5] Zhang et al., Adding Conditional Control to Text-to-Image Diffusion Models. IEEE CVPR 2023", "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": "7WjEFZmYbm", "reviewer": "Reviewer_ESdx", "summary": "The authors propose a novel conditional diffusion sampling strategy for solving inverse problems. Previous conditional diffusion-based inverse solvers are heuristically motivated, lack a unifying framework, and suffer from sensitivity to hyperparameters and heave computation of the Jacobian of the trained score network. The proposed method involves fine-tuning of a small network to quickly learn the conditional $h$-transform, enabling conditional sampling without altering the large unconditional network. The authors demonstrated the efficiency of their method by achieving SOTA performance across various benchmarks with faster inference times.", "questions": "See the weaknesses part", "limitations": "Nothing to mention", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "### Theoretical Foundation\nThis paper is built on solid mathematical theory, specifically “Doob’s $h$-transform,” which enhances our understanding to diffusion models and paves the way for future improvements in their applications.\n\n### Comprehensive Experiments\nThe authors conducted extensive experiments on their method, covering both linear and nonlinear inverse problems across domains such as natural images, medical images, and conditional protein design. They provide clear evaluation metrics and inference time, supporting their claims about the method’s efficiency.\n\nFurthermore, the network architecture is derived from theoretical principles, enabling systematic improvements. Additionally, they propose an extra loss function for fine-tuning that is applicable without paired data, potentially inspiring future research.", "weaknesses": "The proposed claims are poorly presented or difficult to understand. It would be beneficial to supplement the contents in the paper by addressing the following questions.\n1. How does Doob’s h transform unify existing methods for diffusion-based inverse solvers?\n2. What is the difference between $X_t$ and $H_t$?\n3. I do not understand the rationale behind the DEFT network parametrization in Section 3.3. How is equation (13) derived from Doob’s $h$-transform?\n4. In Line 281, it is stated that “DEFT assumes no knowledge of the forward operator.” How is this possible? To my understanding, the forward operator is needed to compute $\\ln p(y|X_t)$. \n\nErrata bold 5.2 for Time(hrs) of DEFT for super-resolution", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for this valuable feedback on the presentation - this is very helpful for clarifying our work. We have made updates to the manuscript to address the points you suggest for clarification. The error in Table 1 will be fixed in the camera-ready version. In detail:\n### Q1: How does Doob’s h transform unify existing methods for diffusion-based inverse solvers?\nA1: Doob’s h-transform is the formal approach to conditioning an SDE and well established in SDE literature, see for example [1]. Yet, Doob’s transform has not been discussed in or been connected to the conditional generative modeling literature. By spelling out the mathematics explicitly, we can identify previous work as special case approximations to the Doob’s h-transform term, thereby providing an underlying framework in conditional generative modeling. Specifically, we derive previous methods as Doob’s h-transform approximations: 1. Reconstruction Guidance (DPS, FreeDoM [3], MPGD [4] etc), 2. CDE (amortized, conditional training) and 3. Classifier Guidance [2].\n### Q2: What is the difference between $X_t$ and $H_t$?\nA2: We denote with $X_t$ the unconditional diffusion process and with $H_t$ the conditional diffusion process. We will make sure to clarify this in the manuscript when $H_t$ is first introduced.\n### Q3: How is equation (13) derived from Doob’s ℎ-transform?\nA3: The network architecture in equation (13) defines an inductive bias for DEFT. As discussed, we can derive reconstruction guidance (i.e., DPS) as a special case of the h-transform (see Equations 30-31 Appendix C3 of the manuscript). As the conditional update in DPS has a high computational cost, we omit the Jacobian of the unconditional diffusion model. This then serves as an inductive bias for the architecture of DEFT and is an important ingredient for the approach. We will make this clearer in the camera-ready, ensuring we talk more about the intuition behind using the specific network architecture we recommend.\nPlease also refer to Table 5 in the appendix. In this ablation we test a version of DEFT without this inductive bias, leading to a worse PSNR/SSIM for the low-dose computed tomography experiments, showing that this inductive bias improves performance.\n### Q4: In Line 281, it is stated that “DEFT assumes no knowledge of the forward operator.” How is this possible?\nThank you for pointing out this imprecision on our part. Indeed, the forward operator is necessary to compute $p(y|x_t)$ and also to use the parameterization in Eqn. (13). What we meant to emphasize was that the training objective in Eqn. (8) does not require the likelihood.  Thus, all we require for the DEFT objective is to be able to evaluate the forward operator such that we can obtain the measurement $y$ or alternatively, access to a paired data set of images and corrupted measurements. For the DEFT architecture in Eqn. (13), we use $\\nabla_{x_0} || y- A(\\hat{x}_0(x_t)) ||^2$ in most cases including cases where there's not necessarily an explicit likelihood available. This is again an inductive bias that helps guide our network in early iterations but does not require an explicit form for $p(y|x_0)$, although it does assume the forward operator can be differentiated. In cases where it can’t, one can resort to approaches such as ΠGDM [5]. To summarize, the DEFT objective **does not require explicit assumptions of the forward operator**, other than some very weak / regular assumptions for the existence of the score. We have adjusted the manuscript to reflect this.\nAlso, here we want to point to the ablation in Table 5 in the Appendix, where we show that the architecture defined in Eqn. (13) leads to a boost in performance. However, without this inductive bias, DEFT can still work.\n\n[1] Särkkä and Solin, Applied Stochastic Differential Equation, Cambridge University Press. (2019)\n\n[2]  Dhariwal and Nichol, Diffusion models beat gans on image synthesis, Neurips 2021.\n\n[3] Yu et. al., Freedom: Training-free energy-guided conditional diffusion model. IEEE CVPR 2023\n\n[4] He et al., Manifold preserving guided diffusion. ICLR 2024\n\n[5] Song et al., Pseudoinverse-Guided Diffusion Models for Inverse Problems, ICLR 2023.", "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": "tifymnAR2d", "reviewer": "Reviewer_ZJuc", "summary": "The authors propose a method for conditional generation using diffusion models, named DEFT. The idea is to combine a fixed, pre-trained unconditional model with an additionally learned conditional correction term to generate conditionally. The authors provide extensive experiments to demonstrate the effectiveness of DEFT.", "questions": "Honestly speaking, I have only one question that may affect my decision. Do you claim to be the first to come up with the idea of learning a small conditional corrector to make an unconditional model conditional?\n\nI’m not familiar with all relevant research on conditional generation. However, if so, I believe your method might be very useful for others, and the paper must be accepted. If not, I would like to see a detailed discussion of DEFT’s contributions compared to prior works.", "limitations": "The authors have adequately addressed the limitations.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is well-written and easy to follow.\n- The proposed method is simple, easy to understand, and should be easy to implement, making it potentially very useful for practical applications. Additionally, DEFT only requires the inference of a pretrained model, eliminating the need for fine-tuning or differentiation.\n- The authors provide theoretical motivation for their method through the Doob h-transform.\n- The authors conduct extensive experiments across different domains and provide an ablation study.\n- The code is provided.", "weaknesses": "- Compared to some other approaches like DPS, DEFT requires additional training.\n- Some image-to-image generative models (such as I2SB and DDBM, which the authors discuss in the appendix) demonstrate great performance in conditional generation. Unlike DEFT, they require fully training a generative model. Despite this difference, for completeness of comparison, it would be good to include both: resulting performance and training budget for these models and DEFT.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback, we will incorporate the following discussion in a revised version of the manuscript.\n\n### Q1: Do you claim to be the first to come up with the idea of learning a small conditional corrector to make an unconditional model conditional?\n**A1**: Thank you for your question. We appreciate the opportunity to clarify the novelty of our work.\nThe idea of learning a conditional corrector to make an unconditional model conditional has indeed been explored in earlier works. In [1] Dhariwal and Nichol, propose classifier guidance and utilize a pre-trained classifier to enable conditional sampling of an unconditional diffusion model. In a different work, ControlNet [3] proposes learning a conditional corrector based on an additional dataset for text-to-image diffusion models. However, our approach introduces several novel strategies and a mathematical foundation for these fine-tuning approaches.\nAs far as we know, DEFT is the first fine-tuning approach that learns a purely **additive** corrector ($\\epsilon_\\text{new} = \\epsilon_\\text{old} + \\epsilon_\\text{corrector}$), allowing for a small  (in model size) corrector term applicable to a wide range of inverse problems. This includes scenarios where measurements $y$ are either real-valued or discrete (whilst [1] is strictly limited to the discrete setting).\nFurther, DEFT is **model agnostic** and makes no assumptions on the specific implementation of the pre-trained unconditional diffusion model. This is in contrast to ControlNet, which employs a non-linear composition of the pre-trained and fine-tuned network, which requires knowledge (and access) of the underlying architecture. In DEFT, we only need the availability of evaluating the pre-trained unconditional diffusion models, making our approach applicable even in cases where the pre-trained model is hidden behind an API. Further, DEFT **only adds 3-10% of additional parameters**, making it more parameter-efficient and making it possible to store multiple fine-tuning networks for different tasks. In contrast ControlNet adds about 30-40% additional parameters. \n### Q2: DEFT requires additional training\n**A2**: We acknowledge that DEFT requires additional training, which in turn requires an additional dataset and an additional training budget. This is a common characteristic of many fine-tuning methods. However, as we discussed in the overall response we see that even when training with 100 or 200 images, we can get quite good results for inpainting on ImageNet. Further, after the initial training phase, sampling is more efficient compared to approaches such as DPS or RED-diff. Please also see our discussion in the overall response. \n\n### Q3: Comparison against conditional trained image-to-image generative models\n**A3**: We were able to compare DEFT against I2SB [3]. Here, we used the pre-trained checkpoint on their github for `sr4x-pool` and `inpaint-freeform2030`, which directly correspond to our super-resolution and inpainting settings on ImageNet. The results can be found in the 1-page PDF. On inpainting, I2SB outperforms DEFT on all metrics, for example top-1 71.7% vs 74.5% or PSNR 22.18dB vs 23.26dB. \n\nIn the following tables we present the results for inpainting and super-resolution on ImageNet. We see that the results are comparable even though I2SB was trained on the complete ImageNet training set. On super-resolution we are even able to outperform I2SB on PSNR and SSIM.\n\n| Inpainting | PSNR (↑)  | SSIM (↑) | LPIPS (↓) | top-1 (↑) | KID (↓)  |\n|------------|-------|------|-------|-------|-------|\n| DEFT   \t| 22.18 | 0.85 | 0.09  | 71.7  | 0.29  |\n| I2SB   \t| 23.26 | 0.86 | 0.068 | 74.5  | 0.238 |\n\n| Super-resolution | PSNR (↑) | SSIM (↑) | LPIPS (↓) | top-1 (↑) | KID (↓)   |\n|------------------|-------|------|-------|-------|-------|\n| DEFT         \t| 24.92 | 0.71 | 0.12  | 71.9  | 1.78  |\n| I2SB         \t| 23.95 | 0.64 | 0.11  | 71.6  | 0.004 |\n\n\n\nThe conditional trained generative models can be seen as an upper limit on the image quality of fine-tuned models, as they are generally trained on a large dataset. For example, the I2SB models were trained on 1M gradient steps on the full ImageNet training dataset. While DEFT was only trained on a subset of 1000 images.  \nIn addition, we present a comparison with conditional diffusion models in Appendix H.2 (see Table 7) using the Flowers dataset. Our results show that, with identical training budgets and dataset sizes, DEFT surpasses the conditional model [4] on several image restoration tasks. However, when provided with a substantially larger dataset (6 times the size) and a significantly greater compute budget (20 times more), the conditional diffusion model outperforms DEFT. Note that in many situations (such as medical imaging) the available fine-tuning dataset may be too small to effectively train a conditional diffusion model from scratch.\n\n[1] Dhariwal and Nichol, Diffusion models beat gans on image synthesis, Neurips 2021.\n\n[2] Zhang et al., Adding conditional control to text-to-image diffusion models, IEEE CVPR 2023.\n\n[3] Liu et al., I2SB: Image-to-Image Schrödinger Bridge, ICML 2023.\n\n[4] Batzolis et. al, Conditional Image Generation with Score-Based Diffusion Models, arXiv preprint 2021.", "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:
We thank the reviewers for their valuable and thorough feedback, we will revise the paper accordingly. Below, we will address the main points and describe improvements we made to our submission.
#### Strengths and contributions
We appreciate the recognition of our work's **novelty** (reviewers `no8y, qBFS, ESdx`) in fine-tuning pre-trained unconditional diffusion models for conditional generation using Doob’s ℎ-transform. The **solid theoretical foundation** (reviewers`no8y, qBFS, ESdx, ZJuc`), **comprehensive experiments** (reviewers `qBFS, ESdx, ZJuc`), and **speed improvements during inference** (reviewers `no8y, qBFS`) were also noted by the reviewers.
 
#### Addressing weaknesses and reviewer questions
Please find detailed point-to-point responses to each reviewer in the corresponding thread. Here, we summarize the main points from these discussions:


1. **DEFT requires a fine-tuning dataset**: We acknowledge that DEFT requires a fine-tuning dataset for training the h-transform. In applications where only a handful of fine-tuning examples are available this may limit performance. Yet, our additional experiments show that even in low data settings (n=100,200) DEFT is able to produce similar results to DPS or RED-diff. For many other applications for which fine-tuning datasets are available (conditional protein design, medical imaging) the fine-tuning phase of DEFT allows the model to **benefit from this data** (as compared to inference-time-only, training-free strategies such as DPS) leading to improved performance, as demonstrated in the experiments in our paper, with comparable or faster total evaluation time on the eval dataset even when taking the fine-tuning time into account (as the faster DEFT inference makes up for the computational overhead of fine-tuning). 
In the experiments on conditional protein design, we saw that reconstruction guidance methods such as DPS often fail on giving convincing results, whereas already a small fine-tuning dataset can improve the performance for DEFT.


2. **DEFT fine-tuning has a computational overhead**: Reviewers rightly point out that DEFT requires an initial training phase. This is common across many fine-tuning methods. However, the overall computational cost, combining both training and inference phases, may vary depending on the specific training budget and application.
We demonstrate that DEFT provides a fast and efficient sampling during inference time. This speed-up during inference helps to offset the initial training time. To illustrate this, our experiments include the total sampling time, which encompasses both training and inference phases (see Table 1 and Table 2 in the manuscript). These results show that DEFT is comparable to or faster than other baseline methods, such as DPS and RED-Diff, when the goal is to sample 1000 images. In the revised version, we include a detailed breakdown of the computation time, not just the total time in the tables.

3. **Additional experiments for rebuttal**:
According to the comments of reviewers, we performed a variety of experiments on the ImageNet dataset. We evaluated:
- *MPGD [2] on inpainting, hdr and super-res* as a new training-free baselines, which does not require to backpropagate through the diffusion model. Our results show that DEFT is able to outperfrom MPGD.
- *ControlNet [4] on inpainting, hdr and super-res* as a comparison to a different fine-tuning method. Our results show that DEFT is able to outperform ControlNet on these image reconstruction tasks.
- *I2SB [3] on inpainting and super-res* as an example of a fully-trained conditional diffusion model. Here, our results show that I2SB outperform DEFT on some tasks. However, I2SB was trained on the full ImageNet training set with a big compute budget, whereas DEFT was trained on a subset of 1000 images.
- *DEFT on inpainting, trained on a subset of 10, 100 and 200 images* to study the effect of the size of the training set. Here, we see that a larger training set, results in better image quality. However, already with 100 images is DEFT comparable to training-free methods as RED-diff or DPS.
- *DEFT on inpainting, hdr for out-of-distribution data (ImageNet-O [1])* to study the generalisability. On ImageNet-O the quality metrics deteriorate. However, we are still competitive to RED-diff.

**A table with all new results and some examples can be found in the 1-page PDF**, notice we strongly outperform the other requested inference-time baselines and ControlNet whilst having comparable to slightly worse performance compared to fully conditional training methods like I2SB, with significantly smaller models, training time and datasets.


[1] Hendrycks et al., Natural Adversarial Examples, IEEE CVPR 2021.

[2] He et al., Manifold preserving guided diffusion. ICLR 2024

[3] Liu et al., I2SB: Image-to-Image Schrödinger Bridge, ICML 2023.

[4] Zhang et al., Adding Conditional Control to Text-to-Image Diffusion Models, IEEE CVPR 2023

Meta-review: This paper proposes a novel approach based upon the Doob's h-transform for turning an unconditional generative diffusion process into a conditional one. 

The reviewers all appreciate the paper's originality and empirical work with some still seeing the paper as borderline. The authors have gone out of their way to answer the reviewers also providing new benchmarking. Taking together it lends to convince that the paper deserves to be presented at Neurips and that given the popularity of the topic, the paper can have a substantial impact.

In this AC's opinion the developed methodology can be presented more concisely. So the authors are encouraged to work on this for the final version.
Decision: Accept (poster)