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

TL;DR

The paper tackles the inefficiency of diffusion-based inverse problem solvers by introducing a canonical three-module form that unifies a broad class of methods and a lightweight Learnable Linear Extrapolation (LLE) to boost performance in the few-step regime. LLE learns per-timestep linear coefficients to combine current and past estimates, improving trajectory alignment with higher-step solvers while remaining computationally cheap. For linear problems, LLE further decouples coefficients into range-space and null-space components, leveraging observation structure to enhance accuracy. Extensive experiments across nine inverse algorithms, five tasks, and multiple datasets demonstrate robust, consistent improvements with minimal overhead, suggesting LLE as a practical, general accelerator for diffusion-based inverse solvers.

Abstract

Diffusion-based inverse algorithms have shown remarkable performance across various inverse problems, yet their reliance on numerous denoising steps incurs high computational costs. While recent developments of fast diffusion ODE solvers offer effective acceleration for diffusion sampling without observations, their application in inverse problems remains limited due to the heterogeneous formulations of inverse algorithms and their prevalent use of approximations and heuristics, which often introduce significant errors that undermine the reliability of analytical solvers. In this work, we begin with an analysis of ODE solvers for inverse problems that reveals a linear combination structure of approximations for the inverse trajectory. Building on this insight, we propose a canonical form that unifies a broad class of diffusion-based inverse algorithms and facilitates the design of more generalizable solvers. Inspired by the linear subspace search strategy, we propose Learnable Linear Extrapolation (LLE), a lightweight approach that universally enhances the performance of any diffusion-based inverse algorithm conforming to our canonical form. LLE optimizes the combination coefficients to refine current predictions using previous estimates, alleviating the sensitivity of analytical solvers for inverse algorithms. Extensive experiments demonstrate consistent improvements of the proposed LLE method across multiple algorithms and tasks, indicating its potential for more efficient solutions and boosted performance of diffusion-based inverse algorithms with limited steps. Codes for reproducing our experiments are available at https://github.com/weigerzan/LLE_inverse_problem.

Figures (15)

• Figure 1: The proposed canonical form of diffusion-based inverse algorithms and the workflow of our LLE method. Our canonical form decomposes an inverse algorithm into three key modules: Sampler, Corrector, and Noiser, providing a unified framework that encompasses a wide range of existing approaches. The LLE method learns a linear combination of the current corrected estimate and previous results to obtain a better estimation of the original image, therefore universally enhancing diffusion-based inverse algorithms' performance under few steps.
• Figure 2: PSNR and LPIPS metrics under various steps for four algorithms on the FFHQ linear inverse problems.
• Figure 3: Qualitative comparison of restoration results with and without LLE on the FFHQ dataset. The top two rows present the results of $\Pi$GDM on the noiseless compressed sensing task, while the bottom two rows present the results of DPS on the noiseless inpainting task.
• Figure 4: PSNR v.s. LPIPS for varying $\omega$ in LLE (from bottom-left to top-right: $\omega = 0.5$, $0.4$, $0.3$, $0.2$, $0.1$, $0.075$, $0.05$, $0.025$, $0.0$).
• Figure 5: Visualization of learned coefficients of DDNM on the noisy inpainting task.

Theorems & Definitions (2)

• Proposition 1
• Corollary 2