Source-linked AI summary
Deep Equilibrium Architectures for Inverse Problems in Imaging
Davis Gilton, Gregory Ongie, Rebecca Willett
TL;DR
Ill-posed imaging inverse problems and fixed-depth unrolled solvers limit reconstruction flexibility and make deeper training difficult. The paper uses deep equilibrium models with forward-model-informed fixed-point iterations, reporting consistently better reconstructions and test-time accuracy–computation trade-offs. The approach also supports fixed-point acceleration, while some convergence guarantees require a trivial nullspace assumption on A.
Problem
Ill-posed imaging inverse problems require estimating images from noisy measurements, while deep unrolled solvers are typically trained for few iterations and are difficult to extend reliably.
Method
The paper constructs deep equilibrium inverse solvers whose fixed-point iterations incorporate the known forward model and classical optimization-based reconstruction structure.
Results
The proposed methods consistently outperform deep unrolled approaches across gradient, proximal, and ADMM base algorithms, while matching their performance with no more computation and improving with further computation.
Takeaways & Limitations
Users can select inference computation to navigate reconstruction accuracy–computation trade-offs, and fixed-point accelerators can speed the resulting inference computations.
Takeaways & Limitations
Convergence guarantees for DE-Prox and DE-ADMM assume λmin(AᵀA) > 0, excluding many compressed-sensing and undersampled-MRI settings.
Abstract
from arXiv · showhide
Recent efforts on solving inverse problems in imaging via deep neural networks use architectures inspired by a fixed number of iterations of an optimization method. The number of iterations is typically quite small due to difficulties in training networks corresponding to more iterations; the resulting solvers cannot be run for more iterations at test time without incurring significant errors. This paper describes an alternative approach corresponding to an infinite number of iterations, yielding a consistent improvement in reconstruction accuracy above state-of-the-art alternatives and where the computational budget can be selected at test time to optimize context-dependent trade-offs between accuracy and computation. The proposed approach leverages ideas from Deep Equilibrium Models, where the fixed-point iteration is constructed to incorporate a known forward model and insights from classical optimization-based reconstruction methods.
1 Introduction
Imaging inverse problems require estimating images from noisy, potentially ill-posed measurements. Deep unrolling addresses this with learned iterative solvers, while the paper proposes deep equilibrium networks for flexible, higher-performing reconstruction.
- Ill-posed inverse problems can make least-squares estimates nonunique or highly sensitive to noise.
- Deep unrolling learns effective regularizers by augmenting iterative inverse-problem solvers with neural networks.These architectures typically mimic only a small number of algorithmic iterations.
- 5-10 iterations are typical because stability, memory, and numerical issues complicate training deeper unrolled networks.
- Deep equilibrium models train inverse solvers to produce good reconstructions at convergence rather than at an arbitrary iteration count.The resulting computational budget can be selected at test time to trade accuracy against computation.
- The paper reports consistent performance improvements above state-of-the-art alternatives while enabling flexible inference computation.
2 Relationship to Prior Work
Prior work includes fixed-depth unrolled optimization networks, implicitly defined equilibrium networks, and denoiser-based PnP/RED methods. The paper positions its approach as an equilibrium formulation that avoids fixed-iteration deployment constraints while retaining the forward model in end-to-end reconstruction.
- Deep Unrolling: Deep unrolling uses fixed numbers of architecturally identical optimization-inspired blocks and is widely applied to MRI and other imaging inverse problems.
- Deep Equilibrium Models: The paper incorporates learned networks into classical gradient-based reconstruction updates, training the resulting estimate end-to-end with the forward model.
- Deep Unrolling: Unrolled networks usually use 5-10 iterations because GPU backpropagation memory grows linearly with iteration count.
- Deep Unrolling: Running a learned unrolled regularizer beyond its training depth performs poorly, so choosing the inference budget during training is essential.
- Deep Equilibrium Models: Deep equilibrium networks use weight-tied transformations whose infinite-depth limit, when it exists, is a fixed point.Implicit differentiation enables constant-memory training for these networks.
- PnP and RED: PnP replaces an explicit optimization step with a plugged-in denoiser, whereas RED uses an explicit regularizer related to a denoising function.
- PnP and RED: Unlike unrolling, recent PnP and RED denoisers are trained independently of the inverse problem and its forward operator.This decoupling makes the reconstruction system flexible across forward models without retraining the learned component.
3 Proposed Approach
The proposed approach designs forward-model-aware iteration maps whose fixed points estimate the image, effectively extending deep unrolling to infinite depth. It instantiates this framework with gradient, proximal-gradient, and ADMM-inspired choices, including a learned proximal replacement and fixed-point formulation for DE-ADMM.
- Framework: The method designs fθ(·; y) so its fixed point x(∞) estimates the image x⋆ from measurements y.The iteration map explicitly accounts for the forward model A and the inverse problem.
- Optimization-inspired maps: The proposed choices are connected to classical optimization schemes, including proximal gradient descent and ADMM for regularized least-squares reconstruction.ADMM is introduced through an equivalent constrained formulation with z = x and an augmented Lagrangian parameter α > 0.
- Framework: Unlike deep unrolling, the proposed architecture uses an effectively infinite number of iterations, while allowing other iteration-map choices beyond the three presented.The paper frames this as a previously unreachable paradigm for deep unrolling architectures in imaging.
- DE-Prox: For proximal-gradient reconstruction, the method replaces proxηr with a trainable network Rθ to form the iteration map.This mirrors deep unrolling while retaining the proximal-gradient update structure.
- DE-ADMM: For ADMM, replacing proxαr with Rθ yields updates depending only on previous x and u, enabling a fixed-point map on q = (x, u).The estimated image is the x component of the fixed point q(∞) = (x(∞), u(∞)).
4 Calculating forward passes and gradient updates
Deep equilibrium models compute fixed points for inference and training, while acceleration can improve fixed-point computation. Their implicit differentiation avoids storing arbitrarily many iterations by solving an inverse Jacobian-vector product.
- Fixed-point computation: Both training and inference require calculating a fixed point of the iteration map fθ(·; y).
- Fixed-point computation: Fixed-point accelerators can accelerate inference because unrolled deep networks can be viewed as fixed-point iterations.The paper compares Anderson acceleration, standard fixed-point iterations, and Broyden’s method.
- Fixed-point computation: Anderson acceleration uses past iterates to identify promising update directions through a least-squares optimization over residuals.The least-squares problem adds negligible computational overhead when the history size is small, such as m = 5.
- Practical considerations: Accelerating fixed-point iterations with auxiliary variables such as ADMM is non-trivial, although acceleration is not generally required for learning inverse-problem solvers.Standard fixed-point iterations may remain attractive for their implementation simplicity in other algorithms or settings.
- Training: Training uses stochastic gradient descent to find network parameters θ that locally minimize a cost over reconstructed training images.The reconstruction for each measurement is defined as the fixed point of fθ(·; yi).
- Gradient calculation: Once the equilibrium point is known, the gradient follows from the chain rule and the loss gradient with respect to the equilibrium state.For mean-squared error, the loss gradient is the residual x(∞) − x⋆.
- Gradient calculation: Implicit differentiation converts backpropagation through arbitrarily many fixed-point iterations into computing an inverse Jacobian-vector product.The auxiliary vector β(∞) is itself computed as a fixed point, using the same machinery as x(∞), with convergence related to a Neumann series.
5 Convergence Theory
The convergence analysis establishes conditions under which deep equilibrium iteration maps are contractive, ensuring convergence to fixed points and supporting gradient calculations. The guarantees depend on the regularization network, step sizes, and properties of the forward operator.
- General conditions: A contractive iteration map guarantees convergence of iterates to a unique fixed point.The analysis imposes conditions on the regularization network Rθ to ensure contractivity.
- General conditions: The regularization assumption is that Rθ − I is ϵ-Lipschitz.With a residual implementation Rθ = I + Nθ, this is equivalent to assuming Nθ is ϵ-Lipschitz.
- DE-Prox: DE-Prox is contractive under a step-size condition involving ϵ, L, and µ, and such a step size exists if ϵ < 2µ/(L − µ).The theorem assumes µ = λmin(A⊤A) > 0.
- DE-ADMM: DE-ADMM is contractive under a corresponding condition on its ADMM step-size parameter α involving ϵ and µ.
- Scope and limitations: The DE-Prox and DE-ADMM guarantees assume A has a trivial nullspace, excluding compressed sensing and undersampled MRI, although convergence is observed there in practice.The assumption is satisfied for some denoising and deblurring problems but violated by many other inverse problems.
- DE-Grad: The gradient-related truncated Neumann series converges when the Jacobian of fθ has spectral norm strictly below 1, equivalently when the iteration map is contractive.
6 Experimental Results
Experiments compare Deep Equilibrium, Deep Unrolled, Plug-and-Play, RED, TV, and related methods across deblurring, compressed sensing, and accelerated MRI. Deep Equilibrium methods generally outperform Deep Unrolled methods while supporting computation–accuracy trade-offs at inference.
- Comparison Methods and Inverse Problems: Experiments compare reconstruction methods across deblurring, compressed sensing, and accelerated MRI, including TV, PnP, RED, Deep Unrolled, and Deep Equilibrium approaches.The comparisons include multiple base algorithms and use datasets including CelebA and fastMRI.
- Main Results: Deep Equilibrium approaches consistently outperform Deep Unrolled approaches across gradient, proximal-gradient, and ADMM base algorithms.The experiments do not identify an obvious winner among iterative reconstruction architectures for fθ, suggesting the best choice may depend on the problem or setting.
- Effect of Acceleration: Deep Equilibrium generally achieves higher PSNR than Deep Unrolled solvers across a broad range of iteration counts.Deep Unrolled ProxGrad was trained for 10 iterations, whereas Deep Equilibrium performance remains effective beyond a single trained iteration count.
- Effect of Acceleration: Anderson acceleration reduces net inference time by reaching fixed-point convergence faster than standard fixed-point iteration.Although accelerated methods take more time per iterate, Anderson acceleration was selected for the remaining experiments because Broyden’s method produced slightly lower PSNR.
- Practical Considerations: Deep Equilibrium methods require only gradient-computation memory during training, enabling training at convergence rather than storing all unrolled iterations.Deep Unrolled training memory scales linearly with iteration count, while inference for all compared methods requires at most a constant number of stored iterations.
- Effect of Pre-Training: Pretraining improves DE-Prox reconstruction quality compared with random initialization and simplifies selecting algorithm-specific hyperparameters.The reported example concerns the initial internal step size for DE-Prox.
7 Conclusions
The paper reports quantitative benefits from implicitly defined infinite-depth networks for linear inverse problems in imaging. Extending the framework to nonlinear inverse problems remains an open area because the relevant conditions are still under investigation.
- Conclusions: Implicitly defined infinite-depth networks provide non-trivial quantitative benefits for solving linear inverse problems in imaging.These empirical benefits complement convergence guarantees unavailable to widely used Deep Unrolled methods.
- Limitations and Future Work: The paper’s exposition focuses on linear inverse problems, while the conditions for applying the proposed Deep Equilibrium methods to nonlinear iterative approaches remain under investigation.Whether these models lead to additional advances in image reconstruction is also left as an open question for future work.
8 Appendix
Appendix visualizations show intermediate Deep Equilibrium iterates converging to reasonable, high-quality reconstructions over many iterations. The appendix also documents the reconstruction inputs, MRI sampling, compressed-sensing matrices, and experimental settings.
- Further Qualitative Results: Appendix figures contain ground-truth images, measurements, and DU-Prox and DE-Prox reconstructions for deblurring, compressed sensing, and MRI.The measurements are projected back to image space for MRI and compressed sensing.
- Further Qualitative Results: Deep Equilibrium models converge quickly to reasonable reconstructions and maintain high-quality reconstructions after more than 100 iterations.Intermediate fixed-point iterates are visualized to demonstrate convergence behavior.
- Further Qualitative Results: DE-Prox visualizations show reconstructions converging to good results across forward problems, including 90 compressed-sensing iterations and 31 MRI iterations.The MRI example terminates at 31 iterations.
- Experimental Setup: For deblurring, the algorithm input is (A^T A + λI)^-1 A^T y with λ equal to σ; MRI and compressed sensing instead use A^T y.MRI masks use Cartesian sampling with the center 4% of frequencies fully sampled for both 4× and 8× acceleration.
- Experimental Setup: Compressed-sensing matrices use entries drawn from a Gaussian distribution with variance 1/m, and the same matrix is used for all learned methods.The matrices are scaled with A ∈ R^m×n.
- Experimental Setup: All Deep Unrolled methods were trained for 10 iterations, while optimization parameters for RED, PnP, and Deep Equilibrium methods were selected by logarithmic grid search.The grid spans 10^-4 to 10^1 with 20 elements in each dimension.