Repulsive Latent Score Distillation for Solving Inverse Problems

Thank you for your detailed review and constructive feedback. Please find the comments addressed below. All new experiments are in blue in Appendix D.1 to D.9 in the updated pdf.

Weaknesses

W1: Full RLSD model does not consistently achieve the best results across all tasks

We thank the reviewer for raising this excellent point. The added diversity gained from the repulsion between particles can indeed affect the average performance of the images. However, there are three important points to notice:

1. If instead of focusing on the average performance, we focus on the performance achieved by the best image among the ensemble of particles, then the conclusion differs in favor of the full RLSD; see Table 4. In particular, the best image generated by RLSD with repulsion ($\gamma = 50$) achieves better PSNR and LPIPS metrics than the best image generated when repulsion is ignored. This improvement can be explained as follows: while some modes obtained with RLSD might not be as good as NonRepuls-RLSD, the additional diversity obtained by repulsion can allow for the exploration of better modes.

Table 4 (Box inpainting):

2. Even if one is interested in a good average performance while boosting diversity, then one can further explore the hybrid approach mentioned in Section 5.1. More precisely, given N particles, we can have one or two that are not affected by repulsion, so that these will discover the same modes that the non-repulsive RLSD would. We can then use the remaining particles with repulsion to promote the discovery of (potentially better) modes. In this way, we can have an additional degree of freedom (the number of particles without repulsion) to interpolate between the non-repulsive and the full RLSD, enabling the user to find a sweet spot for their problem at hand.
3. Lastly, how much the gain in diversity affects the performance depends on intrinsic aspects of the problem at hand, which brings us to your next point.

W2: Strength of repulsion influence the landscape of the distribution

Indeed, this is true. For instance, in tasks such as phase retrieval (highly multimodal with discrete modes), how much strength is considered is fundamental to discovering different -- and good -- modes.

Furthermore, if the strength is too high, the solution might diverge, as illustrated in Limitations (Appendix D.9). We believe the hybrid approach mentioned in the previous point can provide a well-balanced trade-off to guarantee a good solution while allowing exploration to promote diversity in other samples.

W3: Additional ablation w.r.t. hyperparemeters

Besides the ablation of $\gamma$ and $L$ in Appendix D.8.2, we added one comparing the number of particles $N$ in the repulsion (Figures 20, 21 in the updated pdf). Given 4 particles, we considered 3 scenarios:

1. All particles independent ($N=0$),
2. Two particles interacting, two independent ($N=2$),
3. All particles interacting ($N=4$).

Additionally, we include an ablation for the unconstrained case in Appendix D.10.2.

Questions

Q1: Version of Stable Diffusion and comparison with pixel-based models

We tried both versions of Stable Diffusion (v1.5 and v2.1) for RLSD and PSLD, and in each case, we chose the version that entails the best performance.

• PSLD: We tried with Stable Diffusion v2.1, but we did not get good results even after fine-tuning. Hence, we kept v1.5.
• RLSD: Achieved better performance with v2.1 (although it still outperforms PSLD when using v1.5).

This analysis yields an important implementation remark to be made. PSLD's implementation does not use Hugging Face's diffuser library. Thus, we needed to adapt the code to run with Stable Diffusion v2.1. By contrast, our code leverages the diffuser library, making the code transparent to the version of Stable Diffusion and making the transition between different Stable Diffusion versions completely seamless.

Regarding the comparison with pixel-space methods, there is no official implementation for FFHQ trained on 512 x 512 resolutions. While we agree that this might not be 100% fair, we believe this is one of the strengths of our method: it generates better and higher-resolution images, while being as fast as the samplers working with images of resolution 256 x 256.

Q2: How metrics are calculated

As explained in the response to your W1, we averaged across all particles to get the results with the exception of phase retrieval, where the average is not useful. However, considering the best image can also provides important insights. In particular, when looking at Table 4, we observe that RLSD is the best one.

Thanks again for your review. We run more experiments to further improve the results and address your concerns.

Q1: Full RLSD model does not consistently achieve the best results across all tasks

We reran the HDR and Phase Retrieval experiments with improved fine-tuning. Also, we ran inpainting with free mask for the ablation. The results, especially for HDR, confirm the effectiveness of RLSD.

HDR. We provide below the results for HDR. Table 1 is mean among four particles and Table 2 is with max among the four ones.

Table 1 (HDR with mean among four particles)

Table 2 (HDR with max to choose the sample)

We observe that RLSD achieves higher PSNR and lower LPIPS when choosing the best sample among 4 particles, confirming its advantage in HDR reconstruction.

Inpainting Free Mask. In free-mask inpainting, RLSD ($\gamma = 25$) also outperforms NonRepuls-RLSD (see corresponding Table 4 in Ablation $\gamma$ below). These findings align with our results for inpainting with box masks that we shared before, further demonstrating RLSD's robustness across different mask configurations.

Limitations. Unlike HDR and inpainting, phase retrieval presents challenges where RLSD does not outperform NonRepuls-RLSD. In particular, neither approach achieves satisfactory results compared to the non-augmented case. We hypothesize that this is due to the limited information that provides the measurement compared to the other experiments, making the optimization w.r.t. $\mathbf{x}$ unstable.

Conclusion. RLSD demonstrates superior performance over NonRepuls-RLSD in HDR and inpainting (both box mask and free mask). These results emphasize RLSD's effectiveness in solving diverse and challenging inverse problems.

Q2: Ablation

We conducted an additional ablation study to analyze the impact of the coupling parameter $\tilde{\rho}$ and the repulsion term $\gamma$ on inpainting with free mask. In addition, a qualitative result that illustrates the effect of $\tilde{\rho}$ has been included in Appendix D.8.5.

Ablation $\tilde{\rho}$. The coupling parameter $\tilde{\rho}$ controls the strength of the prior during the optimization of $\mathbf{x}$. The results below indicate that increasing $\tilde{\rho}$ enhances diversity, but degrades performance metrics such as PSNR and LPIPS. Notice that $\tilde{\rho}=0.07$ is the one we used for the experiments in the main apper

Table 3 (Ablation $\tilde{\rho}$ for inpainting with free-mask - $\gamma = 0$)

Observations:

• A smaller $\tilde{\rho}$ balances prior influence and optimization, resulting in higher performance (PSNR and LPIPS).
• Increasing $\tilde{\rho}$ improves diversity but sacrifices reconstruction quality.

Ablation $\gamma$. The repulsion term $\gamma$ controls the strength of the repulsion effect during sampling. The results are below

Table 4 (Ablation $\gamma$ for inpainting with free-mask.)

Observations:

• For this experiment, increasing $\gamma$ improves diversity, with diminishing returns for values above 25. However, as mentioned in the limitations (Appendix D.9), too much repulsion can have a negative impact.
• RLSD with $\gamma = 25$ outperforms NonRepuls-RLSD in both diversity and performance metrics, highlighting the repulsion term's dual benefit. Thus, this configuration achieves the best balance, improving both diversity and performance compared to NonRepuls-RLSD

Summary of the ablation

• Increasing $\tilde{\rho}$ enhances diversity but degrades performance.
• The repulsion term significantly improves diversity and can also boost performance.
• In general, repulsion is more effective than increasing the coupling to improve diversity while maintaining good performance metrics.
• RLSD is robust with respect to the hyper parameters, as in all the cases the performance is still good.

Thank you for reviewing our paper. We address your concerns below, and added to Appendix D in the pdf.

Q1: Additional comparison with baselines.

Following your suggestion, we have included additional comparisons with $\Pi$GDM and FPS-SMC for half-face inpainting and super-resolution ($\times$ 8). We report these additional results in Appendix D.2-D.4 and also here. Notice that both new baseline methods can be applied to linear inverse problems, and in the case of $\Pi$GDM for some non-linear inverse problems like JPEG, but not to phase retrieval or HDR, limiting their application to challenging inverse problems.

Table 5 (half-face inpanting)

Table 6 (super-resolution)

Also, in response to Reviewer 2, we included ReSample (a recent method working with latent diffusion models) in all the experiments we have in the appendix.

Q2: How sensitive is RLSD to the initialization of particle positions?

This is a great point. We consider a random initialization (from a standard Gaussian distribution). While changing the initial point (by using other seeds or other distributions) renders different solutions, we saw that, on average, we get similar results in terms of diversity/quality.

We tried other initializations, such as $f^{-1}\mathbf{y}$ (when $f$ is linear and invertible), but this is counterproductive for promoting diversity as we tend to start close to a local minimum, and it takes more time to explore the search space. This is particularly relevant in phase retrieval.

Q3: Strategies to improve convergence to distinct modes

Below, we give details of two strategies that might help to improve convergence.

1. Adaptive $\gamma$
   • The simplest way is by increasing its value (strength of repulsion). While this might keep the particles far away from each other, it can also generate a bad trajectory. In other words, it might diverge.
   • A more sophisticated way is to define an adaptive $\gamma$. Doing this, one could enforce a strong repulsion at the beginning and gradually decrease it to converge to a good mode. While we did not include an experiment using an adaptive $\gamma$, the experiments in Appendix D.8 using repulsion in a truncated interval go in this direction.
2. Adaptive kernel
   • Use an adaptive kernel based on noise level. This resembles an annealing process, where the width of the RBF will depend on the noise level. By incorporating more information in the kernel's definition, the repulsion will help to improve convergence to distinct modes.
   • Furthermore, if one has access to a dataset of points, we could aim to learn a kernel that promotes diversity.

Q4: Similar performance when the number of particles is reduced

This is an excellent question. Diversity is attained from the interaction between particles. To achieve good performance with a small number of particles, we can explore the different degrees of freedom in the design of this repulsion force. In particular, we can focus on and modify three aspects:

1. Kernel function. In the paper, we consider the radial basis function (RBF) to be a width that depends on the distance between particles. A possible modification would be using the adaptive kernel mentioned in point 2 in response to your Q3, either by rescaling the width to account for the small number of particles or using a learned kernel.
2. Kernel domain. The domain where the repulsion is computed is critical for having a well-defined repulsion. In this context, we use latent diffusion models + DINO to avoid the curse of dimensionality. Thus, this allows for the achievement of diversity with a relatively small number of particles (four). If we want to achieve more diversity, other domains might be explored, such as LPIPS; however, our ablation in Appendix D.10.2 suggests that LPIPS does not render a good repulsion for a small number of particles.
3. Differentiable treatment. We can use a hybrid approach, where only a few particles are affected by the repulsion force. In this case, some particles will promote diversity and others will converge to the modes that would have been discovered without repulsion. This additional degree of freedom (the number of particles affected by repulsion) can be useful in maintaining high performance even for varying numbers of total particles. We expand on this approach in the response to W1 and W2 of Reviewer 3.

Thank you for your positive comment and for considering our response in your assessment of our manuscript!

Thank you for reviewing our paper. We provide answers to your three questions below; you can find also in Appendix D.1 to D.7 in blue in the updated pdf.

Q1: Comparison with ReSample

We used their implementation, which only includes an implementation with LDM trained on FFHQ. For completeness, we also adapted the code to handle Stable Diffusion. We report quantitative results for half-face inpainting, and motion deblurring and random inpainting in the response to Q2.

Table 1 (Half face inpainting):

We can summarize the following findings when comparing with ReSample:

1. Performance. RLSD performs better than ReSample across the different inverse problems; ReSample only outperforms in one metric (LPIPS) in one experiment (rand. inpainting). It is worth saying that we also tried Stable Diffusion as a diffusion prior for ReSample, but the results were worse than those with LDM.
2. Computational complexity. RLSD is computationally more efficient. Specifically, ReSample involves an inner optimization within the sampling. Consequently, when using LDM (256x256 resolution), it takes 4 minutes in one A100 GPU to generate 1 image. Moreover, when we tried Stable Diffusion, it took 16 minutes (considering the same hyperparameter configuration). We hypothesize that ReSample takes a long time with Stable Diffusion due to the optimization step, where it might struggle to converge at every step. While this might be improved, it demonstrates that the method needs more fine-tuning than ours. On the other hand, our method, using Stable Diffusion as prior, can generate 4 samples with just 200 steps in less than 4 minutes (without repulsion, 1 image per minute) and 5 minutes (with full repulsion).

It is worth noting that ReSample performs particularly poorly in half-face inpainting. We fine-tuned hyperparameters (step size and the hard consistency step), but the results in Table 1 are the best we could get. We have contacted the authors of [1], and they told us that, for this task, there is a step in their algorithm that is not used. However, they haven't made that script publicly available. This example demonstrates that our method generalizes well to different inverse problems, while ReSample implementations might change based on the task.

Q2: Breadth of Inverse problems considered

We added all the missing baselines for motion deblurring and random inpainting. We report these results in Table 2 and 3. Overall, we experimented with deblurring, inpainting, super resolution (linear), and phase retrieval, HDR (non linear), covering a wide range of inverse problems.

Table 2 (Motion Blurring):

Table 3 (Random inpainting (80%)):

We did not include these inverse problems in the main text due to space limitations. Instead, we prioritized those with

1. A high level of ambiguity where diversity is important
2. Multi-modal non-linear inverse problems (phase retrieval).

Q3: Computational efficiency.

RLSD involves two optimization steps due to the augmentation. The optimization w.r.t. $\mathbf{x}$ is cheap: in particular, when the operator $f(.)$ is linear, the loss function is quadratic. This allows acceleration by computing the gradient in a closed-form expression. On the other hand, the optimization w.r.t. $\mathbf{z}$ is more expensive: we need to backpropagate through the decoder and compute the repulsion term (when used). Despite this, RLSD is more efficient than its baselines, as explained below.

Compared to PSLD:

• RLSD does not need to compute the score Jacobian.
• RLSD can solve inverse problems with fewer steps (e.g., 200 steps for inpainting as shown in appendix) due to its plug-and-play formulation. In constrast, PSLD requires more steps due to its posterior sampling formulation.
• The repulsion term indeed increases the computational cost as we need to compute the gradients of the kernel w.r.t. particles; however, we can trade-off between number of particles and diversity/quality depending on the downstream task. Furthermore, as mentioned in Section 5.1 and in response to reviewer 3 (point W1), we can consider a hybrid approach, where only a subset of particles interact, reducing the computational cost.

Compared to ReSample, RLSD do not have an inner optimization (the hard consistency step in [1]). As explained in the answer to your Q1, this yields a marked difference in running time.

We thank the reviewer for the positive comments about our paper and the constructive feedback. Please find the comments addressed below, and in Appendix E and G in blue in the main paper.

Q1: Novelty

We thank the reviewer for raising this point. Indeed, repulsion and augmentation have been studied in the past (see [1], [2]). Also, augmentation for inverse problems (decoupling) can be related to plug-and-play methods (we added details on this in Appendix G). However, as far as we know, nobody has combined repulsion and augmentation in the context of variational inference. Moreover, repulsion in the context of diffusion models has been applied only in a few works (see Appendix F in the paper), but none for solving inverse problems with diffusion priors. See also Appendix F for our description of how the proposed repulsion can be related to copula models.

More concretely, our method is novel in the following aspects:

1. It proposes a variational approximation by leveraging both particle coupling (repulsion) and distribution augmentation. This is the first work to explicitly promote diversity for solving inverse problems.
2. It proposes to solve inverse problems with latent diffusion model as an expressive prior in a more efficient way than previous methods such as ReSample (see the response to Reviewer 2) and PSLD. This markedly reduces the latency in image generation (see response to Q1 from Reviewer 2)

Q2: Large-scale, high-dimensional data and practical limitations

Thanks for the comment. Regarding the limitations of our method, we expanded the related discussion in Appendix D.9. Concerning the dimensionality of the data, it is constrained by the existing pretrained diffusion priors (currently, stable diffusion provides resolutions up to 512x512). This is still considered a high dimensional data, compared with existing methods that work with resolutions up to 256x256; see e.g., [5, 7]. Also, note that since our method works in the latent space, it is scalable to the dimension as all the diffusion operates in a low-dimensional kernel space. A comparison between different domains where the kernel is defined is in Appendix D.10.2.

Notice that since diffusion and the diversity kernel works in the latest space, our method is scalable with respect to the dimension. A comparison between different domains where the kernel is defined is in Appendix D.10.2.

Q3: Comparison with other variational samplers

We added a new section in Appendix G (see the updated pdf for more details) with a comparison with previous variational samplers for solving inverse problems. To summarize, the new section includes:

1. A comparison with traditional variational samplers, such as [3], and more recent ones like [4].

2. We expanded on the discussion with two recent works that used diffusion priors within a variational framework [5] and [6].

   Compared to [5], our method has three main differences:

   1. While RED-Diff is formulated in the pixel domain, our formulation is in a latent space, allowing us to exploit the latent diffusion model
   2. We consider a multi-modal variational distribution with a repulsion term instead of a Gaussian
   3. We decouple the data and prior term to handle the challenges of latent inversion.

   Compared to [6], there are also a few differences.

   1. They also consider pixel-based diffusion models, which simplifies the formulation.
   2. They incorporate the diffusion prior by computing the log probabilities: this requires solving the underlying ODE and estimating a divergence, which is computationally expensive.
   3. They consider a normalizing flow as a variational distribution. While this differs from our particle-based variational distribution, our framework allows incorporating a normalizing flow as a variational distribution. Exploring this type of parametric distribution opens an interesting research direction, where we define an amortized distribution that can generate solutions in only a few steps and it relates to the modern one-step distillation; we leave this as future work.

[1] D'Angelo, Francesco, and Vincent Fortuin. "Repulsive deep ensembles are bayesian." NeurIPS 2021

[2] Wild, Veit David, et al. "A rigorous link between deep ensembles and (variational) Bayesian methods." NeurIPS 2024.

[3] Kobler, Erich, et al. "Variational networks: connecting variational methods and deep learning." GCPR, 2017.

[4] Tonolini, Francesco, et al. "Variational inference for computational imaging inverse problems." JMLR 2020.

[5] Mardani, Morteza, et al. "A Variational Perspective on Solving Inverse Problems with Diffusion Models." ICLR 2024.

[6] Feng, Berthy T., et al. "Score-based diffusion models as principled priors for inverse imaging." CVPR 2023.

[7] Chung, Hyungjin, et al. "Diffusion Posterior Sampling for General Noisy Inverse Problems." ICLR 2023.

We want to let you know that we ran new experiments that further showed the advantages of RLSD; you can find these results in the last response to reviewer 3.

Besides these experiments, please let us know if there is something else that we can include to improve your impression of the manuscript. After addressing your original comments and the addition of these experiments, if there is any reason why you still think that the paper is only “marginally above the acceptance threshold”, we would like the opportunity to improve our manuscript further.