Improving Diffusion-based Inverse Algorithms under Few-Step Constraint via Linear Extrapolation

We thank the reviewer for recognizing the lightweight advantage and general applicability of our method. Below, we provide point-by-point responses to the concerns and questions raised. All additional experimental results are reported in terms of PSNR ↑ / SSIM ↑ / LPIPS ↓.

1. "I am not sure how effective this method would actually be considering that the improvements are not great. I would suggest the author to put more qualitative results."

We sincerely thank the reviewer for the valuable comments and helpful suggestion. We would like to clarify that in many cases, LLE improves the PSNR of the original algorithm by close to or more than 1dB, which is generally regarded a considerable improvement in the context of image restoration. Moreover, we consistently observe improvements across all three metrics—PSNR, SSIM, and LPIPS—indicating that LLE not only reduces pixel-wise reconstruction error but also enhances perceptual quality.

Regarding the qualitative results, we appreciate the suggestion and will include more visualizations in the camera ready version. In the meantime, we refer the reviewer to Appendix I (pages 35–39) for additional qualitative examples beyond those presented in the main paper. We apologize for the inconvenience that, per this year's policy, visual results cannot be included during the rebuttal phase. Thank you again for this thoughtful feedback.

2. "This method assumes that there exists a reference dataset to perform the optimization, which I am not sure holds in practice. How were these reference points chosen for experiments?"

In all of our experiments, we use 999-step DDIM to sample 50 synthetic images from the diffusion model as the reference dataset. This design ensures that LLE remains practical in real-world scenarios: since diffusion-based inverse algorithms assume access to a pre-trained diffusion model, it is straightforward to generate the reference dataset using the same model without requiring any real images.

While this setup is mentioned at several points in the paper (e.g., line 52 on page 2, line 223 on page 7, and line 722 on page 25), we acknowledge that the current presentation may not be sufficiently clear or direct, and could lead to misunderstanding. We will revise the Experimental Setup section to explicitly describe the construction of the reference dataset used by LLE. We thank the reviewer again for this helpful question.

3. "I found the presentation of the methodology a bit convoluted. Perhaps there is a way to condense the main text to bring the algorithm box up to the main part."

We appreciate the reviewer's suggestion. We will move the algorithm box into Section 4 of the main text in the camera ready version to improve the overall flow and coherence. We will also refine the presentation of the methodology section to make the description clearer and more concise.

4. "Why are the results for DMPS in Table 1 so poor? I suggest the authors to look into this to provide a more fair comparison."

We thank the reviewer for pointing out this issue. We have double-checked our implementation of DMPS and confirmed that it is consistent with the original paper. We also verified our implementation under 100 and 1000 sampling steps, and the results closely match those reported in the original work. In fact, we observe that our implementation achieves comparable performance with the original paper using as few as 15 steps (see Table 18 and Table 19 on page 30), which supports the correctness of our implementation.

Upon further investigation into the hyperparameter settings of DMPS, we find the primary issue lies in the scaling parameter $\lambda$. While the original paper suggests setting $\lambda \geq 1$, we find this choice to be suboptimal in low-step regimes, especially when the number of steps is below 10. To address this, we conduct a step-by-step and task-by-task hyperparameter search for DMPS on the CelebA-HQ dataset with 3, 4, 5, and 7 sampling steps. The table below shows the optimal values of $\lambda$ we identified.

Using these revised hyperparameters, we re-evaluate both the original DMPS and the LLE-enhanced version on noisy tasks on CelebA-HQ. The updated results are shown in the table below. We believe these results provide a more fair and accurate comparison for DMPS. We will incorporate them into the camera ready version, along with revised results on noiseless tasks on CelebA-HQ and on FFHQ. We sincerely thank the reviewer for helping us improve the fairness and credibility of our experimental evaluation.

5. "How precisely do Equations (19–20) need to be optimized? Did the authors find that the loss associated with this problem is easy to optimize? In cases where the optimization falls into a local minimum, did the authors find that LLE still outperforms the original method?"

We thank the reviewer for the insightful question. We find that the optimization of the objective function (19–20) requires relatively few steps, is easy to optimize, and is not prone to getting stuck in local minima.

To better understand the impact of optimization iterations on LLE's performance, we conduct an additional experiment evaluating LLE trained for {25, 50, 100, 150, 200} iterations. The experiment is performed using 5-step DDNM on the FFHQ noiseless tasks. The results are summarized in the table below. We observe that LLE achieves comparable and strong performance when trained for iterations between 50 and 200. This suggests that LLE converges quickly with relatively few iterations and does not show obvious signs of overfitting.

Theoretically, this is reasonable since LLE learns interpolation coefficients by minimizing a weighted sum of MSE loss and LPIPS loss. When the LPIPS weight $\omega$ is zero, the problem reduces to a least squares problem, which is convex and free of local minima. Although the LPIPS loss is non-convex, our choice of a relatively small weight ($\omega = 0.1$) ensures that the overall loss landscape does not contain many undesirable local minima, making optimization easier in practice. We sincerely thank the reviewer for this insightful question and will include these additional experiments and discussions in the final version of the paper.

We would like to express our gratitude once again for the reviewer's comments and suggestions for our work. Should there be any further questions, we welcome continued discussion.

Dear Reviewer 7BP1,

Thank you once again for your helpful and constructive review.

As the Discussion phase is approaching its end in less than three days, we would like to kindly remind you to take a look at our rebuttal. We have added substantial experiments and clarifications in response to your valuable comments and concerns.

If there are any points we may have not fully addressed, we would be grateful if you could let us know.

Looking forward to your feedback.

Best regards,

The Authors.

We sincerely thank the reviewer for acknowledging the motivation and general applicability of our work. Below, we provide detailed responses to the concerns and questions raised. For clarity, we organize our response as follows:

(1) We first present additional experiments on ImageNet to further validate the generality of LLE in more challenging settings.

(2) We then report the performance of additional baselines on the FFHQ dataset as suggested by the reviewer, including five pixel-space methods, three latent-space methods, and one method based on consistency model.

All results are reported using PSNR ↑ / SSIM ↑ / LPIPS ↓ for consistency.

1. "Validation on a diverse benchmark like ImageNet is needed to prove general applicability."

We supplement the results of $\Pi$GDM, RED-diff, DPS, and ReSample on ImageNet with 4 steps and 10 steps in the following table. Following [12, 14] in the main paper, we sample 100 images from the ImageNet test set for evaluation. For the noisy tasks, we add Gaussian noise with $\sigma = 0.1$ with respect to data range [-1, 1] (same in the following experiments). We use the ADM pre-trained checkpoint ([4] in the main paper).

The results demonstrate that LLE consistently improves performance across all methods in the few-step setting on ImageNet. This confirms that LLE is effective and broadly applicable on more complex datasets.

In addition, we conduct a cross-domain experiment between FFHQ and ImageNet. Due to space limitations, we kindly refer the reviewer to our response to Reviewer enUs (Point 4) for detailed results. These results further supports LLE's out-of-domain robustness.

2. Including more baseline methods like LGD, DCDP, MPGD, FPS-SMC, and SGS-EDM.

We supplement experiments with LGD, DCDP, MPGD, FPS-SMC, and SGS-EDM on noisy tasks on FFHQ with 4 steps. The results are reported in the Pixel space part of the table below. The results show that LLE consistently improves the performance of all five baselines in the few-step regime. This further validates the generality and broad applicability of LLE.

3. "What are the specific technical challenges in applying the canonical form and LLE to Latent Diffusion Models?"

There are no difficulties in extending the canonical form and LLE from pixel space to latent space. The core theoretical analysis of linear extrapolation still holds. The observation function with LDMs is $\mathbf{y} = \mathcal{A}(\mathcal{D}(\mathbf{z}_0)) + \sigma \mathbf{n}$, where $\mathcal{D}(\cdot)$ denotes the decoder of the VAE, $\mathcal{A}(\cdot)$ is the observation operator, and $\mathbf{z}_0$ is the latent. This formulation corresponds to a nonlinear inverse problem with respect to $\mathbf{z}_0$, since the composed operator $\mathcal{A} \circ \mathcal{D}$ is highly nonlinear. Our theoretical motivation remains valid, as it does not rely on any specific assumptions about the form of the observation function.

To solve inverse problems with LDMs, the Corrector (or the inverse algorithm) must be capable of handling nonlinear observation functions. Moreover, as long as the inverse algorithm is applicable with LDM, LLE can still serve as an effective enhancement strategy.

To validate LLE with LDMs, we include results of Latent-DPS (introduced in [12] of the main paper), PSLD [R.1], and STSL [R.2] on the FFHQ dataset. The results are reported in the Latent space part of the table below. LLE consistently improves performance for these LDM baselines. The only adjustment required for training LLE in the latent space is a minor change to the loss in Equation (20), which we modify to $\mathcal{L}(\mathcal{D}(\tilde{\mathbf{z}}_{0, t_i}^{(n)}), \mathcal{D}(\mathbf{z}_0^{(n)}))$, since LPIPS is computed in the pixel space. These results further confirm the generality and effectiveness of LLE when applied to LDMs.

4. The conceptual differences and key trade-offs between the LLE method and consistency-based approaches like Consistency Posterior Sampling.

We would like to clarify that LLE and Consistency Posterior Sampling (referred to CPS below) [R.3] aim to improve performance from different but complementary perspectives. Specifically, LLE enhances the efficiency by learning a better solver over the original trajectory, while CPS leverages the strong single-/few-step performance of consistency models to improve reconstruction quality.

Moreover, we note that CPS also fits into our proposed canonical form, and LLE can be applied to further enhance CPS's performance in the few-step regime. More specifically, CPS uses a consistency model as the Sampler, rather than a traditional diffusion model, but it still generates a sequence of $\mathbf{x}_0$ estimates, forming a valid trajectory. This trajectory can be directly used with LLE's linear extrapolation framework.

We supplement experiments with CPS on the ImageNet-64 dataset. The results are shown in the CM part of the table below. We use the openai/diffusers-cd_imagenet64_l2 checkpoint available on Hugging Face, and test on 100 randomly sampled images from the ImageNet test set.

LLE consistently improves CPS in the few-step setting, demonstrating that LLE and CPS are not exclusive, but rather complementary approaches. This highlights LLE's broad applicability and its ability to boost performance when applied to advanced sampling frameworks with consistency models.

5. Ethical Concern

It was unexpected for us to receive an Ethical Concern flag, as all datasets and algorithms used in our work are publicly available and widely adopted. We'd appreciate it if the reviewer could elaborate on the specific ethical concerns in detail.

We would like to express our gratitude once again for the reviewer's comments and suggestions for our work. All additional experiments and comparisons, including those on ImageNet and nine baseline methods on FFHQ, will be included in the camera ready version to further support the generality and validity of LLE. Should there be any further questions, we welcome continued discussion.

[R.1] Rout, Litu, et al. "Solving linear inverse problems provably via posterior sampling with latent diffusion models." NeurIPS, 2023.

[R.2] Rout, Litu, et al. "Beyond first-order tweedie: Solving inverse problems using latent diffusion." CVPR, 2024.

[R.3] Xu, Tongda, et al. "Consistency model is an effective posterior sample approximation for diffusion inverse solvers." ICLR, 2025.

Dear Reviewer bt2H,

Thank you once again for your helpful and constructive review.

As the Discussion phase is approaching its end in less than three days, we would like to kindly remind you to take a look at our rebuttal. We have added substantial experiments and clarifications in response to your valuable comments and concerns.

If there are any points we may have not fully addressed, we would be grateful if you could let us know.

Looking forward to your feedback.

Best regards,

The Authors.

We sincerely thank the reviewer for recognizing the novelty and generality of our proposed method. Below, we provide point-by-point responses to the raised concerns and questions. All additional experimental results are reported using the metrics PSNR ↑ / SSIM ↑ / LPIPS ↓.

1. "Ambiguity in the justification for being a linear combination of previous estimate, particularly in Equation (10), it is unclear why it holds and a reference would great."

Equation (10) employs a forward difference approximation of high-order derivatives. The related derivation can be found in [R.1] and in book [R.2] (Chapter 6.5). We will add these citations in the revised version to make the presentation clearer. We sincerely thank the reviewer for the helpful suggestion.

2. "The matrix of weights in Figure 6-8 is dominated by the diagonal or first diagonal elements, while the rest are almost zero."

Yes, and this is consistent with our expectations. According to Equation (10), elements farther from the diagonal correspond to higher-order derivatives of the sampling trajectory. Since most sampling trajectories are expected to be smooth, the high-order derivatives decay rapidly to zero, which results in the weight matrix being dominated by the diagonal and first off-diagonal elements. In a few cases, we observed non-zero elements farther from the diagonal, indicating possible non-smooth regions in the trajectory with larger high-order derivatives. We will include this discussion in Appendix C.1. We thank the reviewer for this insightful comment.

3. "The choice and size of the ground-truth datasets are not adequately addressed."

Regarding the train set for LLE, we used 999-step DDIM to sample 50 synthetic images from the diffusion model as the reference set, as stated in the paper (line 52 on page 2, line 223 on page 7, and line 722 on page 25). Using only synthetic images makes LLE more practical in real world applications, since diffusion-based inverse algorithms assume access to a pre-trained diffusion model, thus making it straightforward to construct the reference set required by LLE.

Regarding the test set in our experiments, we follow common practice (e.g., references [5, 6, 12, 14] in the main paper) to sample test examples, as described in the experiment section. Specifically, we use 1000 test samples for CelebA-HQ and 100 samples for FFHQ (and ImageNet in the additional experiments).

We will further add a paragraph at the beginning of the Experiment section to clarify these settings. We sincerely thank the reviewer for helping us improve the clarity of our paper.

4. A more relevant discussion of cross-domain performance on Imagenet

Here we supplement cross-domain experiments of LLE between the FFHQ and ImageNet datasets. We test DDNM with 5 steps, and the results are shown in the table below ("–" denotes results without using LLE). LLE demonstrates generalization ability across FFHQ and ImageNet. The coefficients trained on one dataset can still improve performance on the other in most cases, and sometimes even outperform those trained in-domain. This further indicates the out-of-domain robustness of LLE, which is critical for ensuring its general applicability.

Moreover, we have conducted broader evaluations of LLE with multiple inverse algorithms on the ImageNet test set; due to space limitations, we kindly refer the reviewer to our response to Reviewer bt2H (Point 1) for details.

We will include these new experimental results in the camera ready version. We sincerely thank the reviewer for the valuable suggestion.

5. Training time required to obtain the extrapolation coefficients and the memory requirements during inference.

Regarding the training time, we reported the training time of LLE with different inverse algorithms in Table 6 of Appendix C.4 (page 19). Training LLE takes only 2 to 20 minutes, thanks to its lightweight design and the use of a small reference set.

Regarding the memory footprint, here we further report the GPU memory usage (in MB) of DDNM and DPS with and without LLE in the table below. We measure the CUDA peak allocated memory, and we account only for the memory used by model parameters, data, and additional storage required by the algorithm. The results show that LLE introduces only a few to several dozen MB of additional memory overhead, even under the 15-step setting. This is consistent with our expectations: during inference, LLE only needs to store a few 256$\times$256 RGB images from previous steps, typically around 0.7 MB per image, resulting in at most a few dozen MB of additional cost. This is negligible compared to model weights or other algorithm-specific overhead (e.g., back-propagation in DPS).

In summary, both the training and inference overhead of LLE are minimal, which further supports the lightweight advantage of LLE. We will include these additional results and discussions in the camera ready version. We sincerely thank the reviewer for the helpful suggestion.

6. Missing logarithm in Equation (2).

We thank the reviewer for pointing out the error, and we will correct it accordingly. The correct form is: $\epsilon_{\theta}(\mathbf{x}_t, t)=-\sqrt{1-\overline{\alpha}_t} \nabla\_{\mathbf{x}\_t}\log q_t(\mathbf{x}_t)$.

7. The values and meaning of $\omega$ in Figure 4.

In Figure 4, the ablated parameter $\omega$ refers to the weight of the LPIPS loss in Equation (20). We tested $\omega = \{0.0, 0.025, 0.05, 0.075, 0.1, 0.2, 0.3, 0.4, 0.5\}$, which correspond in the exact order to the points from the top-right to the bottom-left in Figure 4. This result suggests that smaller values of $\omega$ encourage LLE to achieve higher PSNR, while larger values lead to better LPIPS. This provides a direct and practical way to control the trade-off between PSNR and LPIPS, depending on application-specific preferences.

We thank the reviewer for the helpful comments. We will annotate the corresponding $\omega$ values in Figure 4 and add this explanation and discussion to the figure caption in the camera ready version. Sorry for the inconvenience as we are unable to include any updated figures during the rebuttal phase.

8. Limitations-"The method is dependent on the specific diffusion algorithm employed."

We agree that LLE requires training for each specific inverse algorithm. This is because different algorithms vary in their design principles, formulations, and posterior score estimations, which lead to different sampling trajectories. Consequently, the high-order approximation coefficients of these trajectories differ, necessitating separate training for each algorithm. Designing a more general approach, including aligning sampling trajectories across algorithms or developing unified high-order approximation parameters, is a promising direction for future work. We thank the reviewer for this comment.

We would like to express our gratitude once again for the reviewer's comments and suggestions for our work. Should there be any further questions, we welcome continued discussion.

[R.1] Fornberg, Bengt. "Generation of finite difference formulas on arbitrarily spaced grids." Mathematics of computation 51.184 (1988): 699-706.

[R.2] Süli, Endre, and David F. Mayers. "An introduction to numerical analysis". Cambridge university press, 2003.

Dear Reviewer enUs,

Thank you once again for your helpful and constructive review.

As the Discussion phase is approaching its end in less than three days, we would like to kindly remind you to take a look at our rebuttal. We have added substantial experiments and clarifications in response to your valuable comments and concerns.

If there are any points we may have not fully addressed, we would be grateful if you could let us know.

Looking forward to your feedback.

Best regards,

The Authors.

We sincerely thank the reviewer for acknowledging the effectiveness and generality of our method. Below, we provide responses to the concerns and questions raised. All additional experimental results are reported in terms of PSNR ↑ / SSIM ↑ / LPIPS ↓.

1. "LLE requires additional finetuning and cannot be used out of the box."

We appreciate the reviewer's comment. We'd like to highlight that the training process of LLE is highly efficient, and LLE demonstrates cross-domain generalization ability. First of all, as shown in Table 6 on page 19, the training phase of LLE takes only a few minutes, making it practical and feasible in real-world scenarios. Moreover, as presented in Table 4 on page 18, LLE exhibits a certain degree of cross-domain generalization. We further validate this in a more challenging ImageNet-to-FFHQ cross-domain setting, where LLE still achieves good performance (see the table below, where we test DDNM with 5 steps). This indicates that domain shift does not significantly undermine the effectiveness of LLE. Lastly, we acknowledge that LLE currently requires retraining for different inverse problems. Given the minimal training cost, we consider this an acceptable trade-off. Designing more efficient ways to enhance LLE's performance on cross-task generalization is a promising direction for future work.

2. "The SDE version of the derivation should be highlighted more in the main paper."

We appreciate the reviewer's valuable suggestion. As the reviewer pointed out, most diffusion-based inverse algorithms are formulated in the SDE framework, and the injected random noise often plays a significant role in improving solution quality. Thus, the SDE-based derivation provided in Appendix B.3 offers a more complete explanation of the motivation and formulation of LLE. We will highlight the SDE formulation in the main text and explicitly clarify its connection to the motivation behind LLE in the camera ready version.

3. "How large are the gaps to methods taking significantly more steps?"

Regarding the comparison between the few-step LLE-enhanced method and the original algorithm with significantly more steps, we kindly refer the reviewer to Table 8 in Appendix C.5 (page 20), where we compare the performance of DDNM with LLE within 15 steps and the default DDNM with 100 steps. The results show that in many cases, LLE-enhanced DDNM with only 15 steps achieves performance comparable to the 100-step baseline.

Regarding the comparison between the LLE-enhanced method and the original algorithm both under significantly more steps, we supplement an additional experiment comparing DDNM with and without LLE under a larger number of sampling steps (25, 50, and 100) on the FFHQ dataset, as presented in the table below. To avoid the need to store an excessive number of past results when the number of steps is large, we limit the lookback window to 10 steps. The results show that as the number of steps increases, the performance gap between methods with and without LLE becomes smaller. This is consistent with our expectation as LLE searches for an improved sampling trajectory within the linear subspace spanned by previous samples. When the original algorithm takes more steps, the discretization error is reduced, and its trajectory already closely approximates the optimal one. In such cases, LLE tends to follow the original trajectory more closely, resulting in smaller gaps.

Therefore, LLE brings more significant improvements in few-step regimes, which aligns with our primary objective of enhancing the performance of inverse algorithms under a limited number of steps. We will include these results and the accompanying discussions in the appendix as an additional ablation study. We appreciate the reviewer's thoughtful question.

4. "Evaluating the diversity of solutions/full posterior."

Regarding the diversity on the full test set, we have included results in Appendix H that report FID scores, which serve as a quantitative measure of diversity across the dataset. Regarding the diversity of individual posteriors, we would like to clarify that the image restoration tasks considered in this paper typically have highly concentrated posteriors. For instance, in the 4$\times$ super-resolution task on 256$\times$256 images, the solution is generally close to deterministic and does not deviate significantly from the ground truth. As a result, a common practice in such settings is to evaluate performance using reconstruction metrics such as PSNR, SSIM, and LPIPS. We will include visualizations of solutions corresponding to different initial noises in the camery ready version to better illustrate the diversity of solutions. Sorry for the inconvenience that we are unable to include any visualizations during the rebuttal phase due to submission constraints. We thank the reviewer for this insightful suggestion.

5. "Can the LLE method be seen as learning a type of non-stationary solver?"

Yes, LLE can be viewed as learning a broad class of non-stationary solvers as defined in the provided paper. Specifically, the non-stationary solver in the paper is formulated as a linear combination of all previous steps' estimates and velocity fields (or scores equivalently). Accordingly, LLE can be interpreted as a data-driven approach to learn the optimal coefficients of this linear combination. We will include this paper as a reference and discuss the connection between LLE and these non-stationary solvers in the camera ready version. We thank the reviewer for the valuable comment!

6. "How are the time steps chosen? Are they just uniformly spaced? Are there methods to choose them in a more optimal way?"

Yes, in this paper, the time steps are uniformly spaced, as stated in Appendix F.2 (line 716 on page 25). We will add this clarification in the main text to make this setup clearer.

Adaptive selection of time steps may indeed further improve the performance of inverse algorithms. For example, by leveraging LD3 [R.1], the time steps of inverse algorithms could be parameterized and optimized in a data-driven manner to find more optimal schedules. Combining such adaptive time step schedules with LLE is a promising direction for future work. We appreciate the reviewer's suggestion and will include this discussion in the camera ready version.

7. "What happens if we choose a large value for the number of steps for LLE?"

See our response to point 3 above.

We would like to express our gratitude once again for the reviewer's comments and suggestions for our work. Should there be any further questions, we welcome continued discussion.

[R.1] Tong, Vinh, et al. "Learning to Discretize Denoising Diffusion ODEs." ICLR, 2025.

Dear Reviewer hxDr,

Thank you once again for your helpful and constructive review.

As the Discussion phase is approaching its end in less than three days, we would like to kindly remind you to take a look at our rebuttal. We have added substantial experiments and clarifications in response to your valuable comments and concerns.

If there are any points we may have not fully addressed, we would be grateful if you could let us know.

Looking forward to your feedback.

Best regards,

The Authors.