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.
Paper
Similar papers
Peer review
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.
Strengths
- The paper is well-written and easy to follow. - 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. - The authors provide theoretical motivation for their method through the Doob h-transform. - The authors conduct extensive experiments across different domains and provide an ablation study. - The code is provided.
Weaknesses
- Compared to some other approaches like DPS, DEFT requires additional training. - 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.
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? I’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.
Rating
5
Confidence
3
Soundness
3
Presentation
3
Contribution
2
Limitations
The authors have adequately addressed the limitations.
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.
Strengths
### Theoretical Foundation This 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. ### Comprehensive Experiments The 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. Furthermore, 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. 1. How does Doob’s h transform unify existing methods for diffusion-based inverse solvers? 2. What is the difference between $X_t$ and $H_t$? 3. 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? 4. 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)$. Errata bold 5.2 for Time(hrs) of DEFT for super-resolution
Questions
See the weaknesses part
Rating
5
Confidence
4
Soundness
3
Presentation
3
Contribution
3
Limitations
Nothing to mention
Derivation of the DEFT architecture and inductive bias. Part I
Thank you for your response, please allow us to clarify the inductive bias and the motivation behind the DEFT network parameterization. We will provide a **short derivation** as well as **empirical evidence** motivating the architectural choice, we hope that this clarifies the reviewers concern and if not we would like to engage further and understand what is missing to make this more clear. In conditional sampling methods, the conditional score $\nabla \ln p_t(x_t | y)$ can be decomposed into the unconditional score, approximated with an unconditional score model $s_\theta(x_t, t) \approx \nabla_{x_t} \ln p_t(x_t)$, and a likelihood term, i.e., $$\nabla_{x_t} \ln p_t(x_t | y) = s_\theta(x_t, t) + \nabla_{x_t} \ln p_t(y | x_t) $$ Here, the likelihood term $\nabla_{x_t} \ln p_t(y | x_t)$ is the h-transform that we aim to learn with a neural network in this work. We can express this term as an expectation of the inverse problem likelihood with respect to the denoiser (as done in works such as DPS [1]): $$ \nabla_{x_t} \ln p_t(y | x_t) = \nabla_{x_t} \ln \mathbb{E}_{x_0 \sim p(x_t |x_0)}[p(y| x_0)] $$ As a next step, we can then make a MAP styled approximation to the posterior (more precisely, approximating a posterior with a mean point mass rather than MAP is known as a “Bayes Point Machine” [2]) i.e. $p(x_t |x_0) \approx \delta_{\mathbb{E}[x_0|x_t]}(x_t)$. This results in the following: $$ \nabla_{x_t} \ln p_t(y | x_t) \approx \nabla_{x_t} \ln p(y|\mathbb{E}[x_0|x_t]) \approx \nabla_{x_t} \ln p(y| \hat{x}_0(x_t))$$ Where $\mathbb{E}[x_0|x_t]$ can be estimated with Tweedies formula and the learned approximate score (i.e. $\hat{x}_0(x_t)$). If we now initialise the h-transform neural network with $\nabla_{x_t} \ln p(y| \hat{x_0})$, this is clearly a much better starting point than initialising it at $0$, as this term is an approximation of the h-transform and has been validated to perform well in these tasks. However, this approximate expression is prohibitive to train with as it requires the Jacobian $\partial_{x_t}\hat{x}_0(x_t)$ which backpropagates through the score network. To mitigate this, we follow works such as DreamFusion [3] which take the gradient with respect to $\hat{x}_0(x_t)$ rather than $x_t$. This step is completely heuristic, but it has been validated empirically. After this, we are left with: $$ \nabla_{x_t} \ln p_t(y | x_t) \approx \nabla_{\hat{x}_0} \ln p(y|\hat{x}_0) $$ As we have already motivated conceptually, this expression approximates the h-transform and thus it makes sense to incorporate this in our h-transform architecture as it provides a good warm start (also note in the MCMC/sampling community these style of gradient aided NN architectures have already demonstrated a lot of succes [4]): $$\text{NN}(x_t, y, t) = \text{NN2}(x_t, y, t) + \text{NN1}(t) \nabla_{\hat{x}_0} \ln p(y | \hat{x}_0) $$ Where prior to training, NN2 is initiatilised to $0$ and NN1 is initialised to 1 **such that at epoch 0 our network is initialised at this cheap approximate h-transform** $\text{NN}(x_t, y, t) = \nabla_{\hat{x}_0} \ln p(y | \hat{x}_0) $. Empirically we found that this initialisation gave much lower starting DEFT losses than without and non surprisingly it lead to converging faster as well as better results. In a way, the architecture $ \text{NN2}(x_t, y, t) + \text{NN1}(t) \nabla_{\hat{x}_0} \ln p(y | \hat{x}_0) $ achieves the best of both worlds: cheap like conditional sampling methods, accurate like conditional training methods. It starts off with the cheap guidance term in early epochs, thereby providing a good warm start to our objective. But in later epochs, it is able to learn a more accurate approximation (without MAP-like approximations or heuristics) to the h-transform with the term $ \text{NN2}(x_t, y, t) $. Note the term $\text{NN1}(t)$ is there to serve as a “guidance scale like” term. In methods like DPS [1] these terms are typically tuned on a small dataset. In contrast, we believe that the term $\text{NN1}(t)$ is particularly helpful in early iterations as it allows the guidance term to quickly become well tuned, whilst the NN2 network uses this warm start to more slowly learn the full h-transform. For empirical evidence of this hypothesis, see Figure 6 where NN1(t) is plotted and learns very expected guidance scales that increase along the diffusion trajectory as expected.
Derivation of the DEFT architecture and inductive bias. Part II
Please see the following table (can also be found in our appendix as Table 6) where we carefully ablate our architectural choice and prove its success. | Parametrisation | PSNR | SSIM | |------------------------------------------------------------------------------------------------------------------------------------------------------------|:-------:|:-------:| | $\text{NN2}(x, \hat{x}_0, \nabla \ln p(y \hat{x}_0), t) + \text{NN1}(t) \nabla \ln p(y \| \hat{x}_0)$ | $35.81$ | $0.876$ | | $\text{NN2}(x, \nabla \ln p(y \| \hat{x}_0), t) + \text{NN1}(t) \nabla \ln p(y \| \hat{x}_0)$ | $35.74$ | $0.875$ | | $\text{NN2}(x, \hat{x}_0, A^*y, t)$ | $34.04$ | $0.851$ | | $\text{NN2}(x, A^*y, t)$ | $26.62$ | $0.724$ | $\nabla = \nabla_{\hat{x}_0}$ in the table above. As you can see, the added gradient-guided inductive bias boosts the performance of the DEFT objective significantly, showing that it is a good architectural choice. We have now provided both thorough empirical and conceptual motivations to the reviewer for our architecture. We are happy to continue clarifying where needed. We will add a detailed derivation in the revised version of the manuscript. [1] Chung et al. 2022. Diffusion posterior sampling for general noisy inverse problems. arXiv preprint arXiv:2209.14687. [2] Herbrich et al.. Bayes point machines. Journal of Machine Learning Research, 2001. [3] Poole et al. DreamFusion: Text-to-3D using 2D Diffusion, ICLR 2023. [4] Zhang and Chen, 2021. Path integral sampler: a stochastic control approach for sampling. arXiv preprint arXiv:2111.15141.
Small typo correction to - "Derivation of the DEFT architecture and inductive bias. Part I"
Dear Reviewer ESdx, We have just made some very small typo corrections (the cheap initialisation gradient we use in DEFT is with respect to $\hat{x}_0$ and there was a typo in the part I response ) to the above derivation (Part I) detailing how we arrive at the additional term in our architecture from the h-transform. We hope that this makes our derivation / conceptual motivation of the architecture more clear and as before we thank you for your feedback and continued input.
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.
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 2. 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 3. By bypassing the backpropogation through the Diffusion U-Net, the method achieves a speed up over existing methods. 4. 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 [1] Manifold Preserving Guided Diffusion [2] FreeDoM: Training-Free Energy-Guided Conditional Diffusion Model have missed the author's notice. These methods do not require any backpropagation through the UNet and hence are faster and perform better. 2. The baselines compared in the paper are very old and refer to works more than 1 years ago. 3. An analysis of the computational overhead caused due to the fine-tuning process is missing
Questions
1. Could the authors give an analysis of the benefits of the method over [1,2] referred in weakness 2. I would also like to see the computational overhead in terms of memory and time involved in the fine-tuning process.
Rating
5
Confidence
4
Soundness
3
Presentation
4
Contribution
3
Limitations
Yes, The limitations and potential negative impacts section looks reasonable to me.
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
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. - The work can be extended to non-linear inverse problems, which is a significant limitation of many existing diffusion posterior sampling methodologies. - 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.
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?
Rating
7
Confidence
3
Soundness
3
Presentation
3
Contribution
3
Limitations
The limitations have been addressed.
Final Score
Thank you for the detailed responses. Considering the additional clarifications and results, I believe that this work could have significant impact and will be raising my score to reflect it.
I thank the authors for their comprehensive clarifications and additional comparisons. I will maintain my score. My final recommendation will depend on the entire discussion.
Thank you for your reponse !
We thank the reviewer for their prompt follow-up and helpful discussion. To aid the discussion further we would like to highlight that have conducted a more extensive literature survey, and we can conclude that our method is uniquely novel in its scalability and design of the conditional model, which is very different from existing techniques that suffer a litany of issues. In short to clarify and answer the reviewers' question we would like to highlight that: ``` We believe our approach is the first method that enables learning a small corrector network for real valued inverse problems. ``` Due to our corrector’s additive nature and inductive biases, it allows for parametrising **very small and efficient correctors**, unlike approaches such as ControlNet, which do not offer a significant enough parameter reduction. For example, we are able to achieve much better performance with as little as 3-10% of the original model’s size compared to ControlNet’s 40%. We would also like to point out that prior well-known approaches such as classifier guidance _do not apply_ when considering real valued inverse problems, which is the main task we tackle in the formulation we propose. We hope this shows how our approach is valuable to the community and that it is the first highly efficient learned corrector for general inverse problems, also highlighted by reviewer **no8y** `I believe that this work could have significant impact`. As before we thank the reviewer for their input and please let us know if there are any other questions which can help aid the dicussion. P.S. We have added all the additional experiments requested by the reviewer. Let us know if there are any additional comparisons we can perform in order to strengthen the presentation of the paper for the reviewer’s consideration.
Further clarification on reported total cost of DEFT (fine-tuning + sampling)
To further aid the discussion and address the computational overhead comment raised, we would like to highlight some key points of our experiments. In Tables 1-3 in our paper, we report the total wall clock time for DEFT, which includes the sum of both training time (fine-tuning time) and inference time. Importantly, even when accounting for the fine-tuning time, DEFT still **take less time** than many inference-only or training-free methods such as DPS and REDDiff. Additionally, we believe that the memory complexity analysis presented in our rebuttal further clarifies our approaches efficiency. Please let us know if there are any other points we can clarify or if you have additional questions.
Dear author, I thank you for the detailed rebuttal. After going through the rebuttal, my concerns regarding the theoretical novelty is clear. have decided to improve my rating to BA. I believe the comparison with freedom and mid is not fair since these methods are training free and the authors were not give satisfactory explanation/ comparisons.
Title: Choice of comparison methods
Dear reviewer, Thank you for your feedback. We appreciate your decision to improve your rating and value your continued input and discussion. We understand your concerns regarding the comparisons with training-free methods such as MPGD and we will add a note to the manuscript making this explicitly clear so that no unfair conclusions can be drawn. It is not our intention to perform any unfair comparisons. We really want to make this clear and for it not to damage the perception of our work. We still think there is value in the comparison to training-free, and hope that you see it the same way since you suggested comparing to the training-free methods MPGD and FreeDoM in your first response. You can interpret DEFT as fine-tuning of a training-free method (see also the response to reviewer ESdx about the motivation of the network parameterization). Thus this comparison is almost like an ablation, i.e., how is method X improved when we fine-tune it this way. Moreover, it is important to note that many conditional sampling methods, while training-free, require setting hyperparameters, such as the time-dependent strength $\lambda(t)$ for DPS or RED-diff. As these methods can be quite sensitive to the choice of $\lambda(t)$, in practice a small dataset is often necessary to appropriately tune these hyperparameters. In addition, we see the ability of DEFT to leverage a small dataset as **a positive feature of our method**. In many practical applications such as the medical imaging or protein design setting described in our paper, there are small datasets specific to the task available. Leveraging these for optimal performance is a strength of the DEFT method. We would also like to add further clarification pertaining to the fair element of comparison. Many of the training free methods such as DPS require backpropagating through the score network over many iterations; from a computational budget perspective this is not so different to training (especially in the small-scale fine-tuning that DEFT requires, i.e. small fine-tuning network and small dataset). As you can see in our total wall clock time comparisons, such training free methodologies actually result in longer GPU hours overall than our fine-tuning approaches. Therefore, we do believe this comparison to be very helpful as the validation compute time becomes **comparable / higher than our required finet-uning time**. Finally, we would like to emphasize that we did add some trained comparisons (ControlNet and I2SB as proposed by the reviewers). To the best of our knowledge there are not many other applicable trained methods. We did try out all the methods suggested by reviewers and we are happy to commit to adding more comparisons in the final version of the manuscript if you think there are other suitable methods to compare against.
I find some aspects of the DEFT network parameterization, including the rationale behind it and the inductive bias, to be unclear. However, the strengthened experimental results demonstrate significant value in this work. Therefore, I remain positive about the acceptance of this paper and will keep my score unchanged.
Final Rebuttal Response
Thank you to all of the reviewers for engaging with our rebuttal and the discussion period. We'd like to highlight a few of the comments that came up during the discussion period: 1. “I believe that this work could have significant impact.” - Reviewer no8y 2. “the strengthened experimental results demonstrate significant value in this work.” - Reviewer ESdx 3. “comprehensive clarifications and additional comparisons. I will maintain my score.” - Reviewer ZJuc 4. “my concerns regarding the theoretical novelty is clear. have decided to improve my rating” - Reviewer qBFS We are glad that we were able to run additional experiments and provide further clarifications to the reviewers to address their initial questions. Furthermore, we thank the reviewers for a productive rebuttal period and their continued engagement with our work. We believe the main remaining points raised by the reviewers are as follows, and we summarize our clarifications to them below - > “I believe the comparison with freedom and mid is not fair since these methods are training free and the authors were not give satisfactory explanation/ comparisons.” We have now provided several comparisons to methods that are not training free (such as I2SB and ControlNet), furthermore we have demonstrated how the compute budget (total GPU time) on many of these training free methods exceeds the highly efficient training done by DEFT, thus from a compute perspective there is a value in these comparisons. Finally comparisons to methods such as MPGD serve as an ablation to DEFT as this showcases how extra fine-tuning such as the one we do improves on MPGD. We will be super clear in the final version of the manuscript when mentioning which methods are trained/fine-tuned and which methods are not so that no unfair conclusions can be drawn. > “I find some aspects of the DEFT network parameterization, including the rationale behind it and the inductive bias, to be unclear. “ We provided a detailed explanation in response to reviewer ESdx carefully deriving (from the h-transform) the cheap “warm start” inductive bias in the DEFT architecture (i.e. $\mathrm{NN1}(t) \nabla_{\hat{x}_0} \ln p(y|\hat{x}_0)$). We then discuss how this term empirically leads to a * Lower loss / Better initialisation * Faster and better training (Converges in less epochs) * Quickly (in early epochs) tunes $\mathrm{NN1}(t)$ to a reasonable/ expected guidance scale (low near the noise, then smoothly increases as we reach the target/data) Intuitively we also discuss how this architecture brings a trade off between compute (cheap approximate MPGD styled warm start) and accuracy - learns a final network $\mathrm{NN2}(x,y,t)$ which approximates the h-transform without strong assumptions, yet leveraging the cheap warm start. Finally we provided the reviewer with an ablations table from the manuscript empirically demonstrating the success of this architecture. We hope these additional points address the remaining issues the reviewers might have had with our submission, and we thank the reviewers and the AC for an excellent NeurIPS reviewing experience!
Decision
Accept (poster)