Table of Contents
Fetching ...

Deep Equilibrium models for Poisson Imaging Inverse problems via Mirror Descent

Christian Daniele, Silvia Villa, Samuel Vaiter, Luca Calatroni

TL;DR

This work introduces DEQ-MD, a Deep Equilibrium Model for Poisson imaging that uses Mirror Descent in Burg’s entropy geometry to minimize the KL data term. It provides a convergence analysis under a Kurdyka–Łojasiewicz framework for subanalytic functions with non-closed domains, ensuring fixed-point convergence of the learned reconstruction even without convexity or standard Lipschitz-smoothness. The method yields competitive results against model-based and Bregman-PnP approaches while using lighter architectures and offering parameter-free inference at test time, with robust generalization across noise levels and even to super-resolution. The practical impact lies in stable, efficient Poisson inverse problem solvers that better exploit non-Euclidean geometry and learned regularizers in imaging applications such as microscopy and PET.

Abstract

Deep Equilibrium Models (DEQs) are implicit neural networks with fixed points, which have recently gained attention for learning image regularization functionals, particularly in settings involving Gaussian fidelities, where assumptions on the forward operator ensure contractiveness of standard (proximal) Gradient Descent operators. In this work, we extend the application of DEQs to Poisson inverse problems, where the data fidelity term is more appropriately modeled by the Kullback--Leibler divergence. To this end, we introduce a novel DEQ formulation based on Mirror Descent defined in terms of a tailored non-Euclidean geometry that naturally adapts with the structure of the data term. This enables the learning of neural regularizers within a principled training framework. We derive sufficient conditions and establish refined convergence results based on the Kurdyka--Lojasiewicz framework for subanalytic functions with non-closed domains to guarantee the convergence of the learned reconstruction scheme and propose computational strategies that enable both efficient training and parameter-free inference. Numerical experiments show that our method outperforms traditional model-based approaches and it is comparable to the performance of Bregman Plug-and-Play methods, while mitigating their typical drawbacks, such as time-consuming tuning of hyper-parameters. The code is publicly available at https://github.com/christiandaniele/DEQ-MD.

Deep Equilibrium models for Poisson Imaging Inverse problems via Mirror Descent

TL;DR

This work introduces DEQ-MD, a Deep Equilibrium Model for Poisson imaging that uses Mirror Descent in Burg’s entropy geometry to minimize the KL data term. It provides a convergence analysis under a Kurdyka–Łojasiewicz framework for subanalytic functions with non-closed domains, ensuring fixed-point convergence of the learned reconstruction even without convexity or standard Lipschitz-smoothness. The method yields competitive results against model-based and Bregman-PnP approaches while using lighter architectures and offering parameter-free inference at test time, with robust generalization across noise levels and even to super-resolution. The practical impact lies in stable, efficient Poisson inverse problem solvers that better exploit non-Euclidean geometry and learned regularizers in imaging applications such as microscopy and PET.

Abstract

Deep Equilibrium Models (DEQs) are implicit neural networks with fixed points, which have recently gained attention for learning image regularization functionals, particularly in settings involving Gaussian fidelities, where assumptions on the forward operator ensure contractiveness of standard (proximal) Gradient Descent operators. In this work, we extend the application of DEQs to Poisson inverse problems, where the data fidelity term is more appropriately modeled by the Kullback--Leibler divergence. To this end, we introduce a novel DEQ formulation based on Mirror Descent defined in terms of a tailored non-Euclidean geometry that naturally adapts with the structure of the data term. This enables the learning of neural regularizers within a principled training framework. We derive sufficient conditions and establish refined convergence results based on the Kurdyka--Lojasiewicz framework for subanalytic functions with non-closed domains to guarantee the convergence of the learned reconstruction scheme and propose computational strategies that enable both efficient training and parameter-free inference. Numerical experiments show that our method outperforms traditional model-based approaches and it is comparable to the performance of Bregman Plug-and-Play methods, while mitigating their typical drawbacks, such as time-consuming tuning of hyper-parameters. The code is publicly available at https://github.com/christiandaniele/DEQ-MD.

Paper Structure

This paper contains 46 sections, 15 theorems, 79 equations, 12 figures, 6 tables, 1 algorithm.

Key Result

Lemma 4.2

Let $f: x \mapsto \text{KL}(y,Ax)$ and $h$ be the Burg's entropy, defined in burg's entropy. Then, for any $L$ satisfying $L \geq \|y\|_1 = \sum_{i=1}^{m} |y_i|,$ the function $Lh - f$ is convex on $\mathbb{R}_{>0}^n=\operatorname{int}(\operatorname{dom}(h))$, i.e eq nolip holds.

Figures (12)

  • Figure 1: The 3 blur kernels used in the experiments. (a) is a real-world camera shake kernel, see Levin. (b) is a $11 \times 11$ Gaussian kernel with standard deviations $\sigma$ =1.2. (c) is a $9 \times 9$ uniform kernel.
  • Figure 2: Illustration of the two different networks employed for the parameterization of the regularization function.
  • Figure 3: DEQ-MD forward pass at test time: functional values decrease and PSNR improvement along iterations.
  • Figure 4: DEQ-MD forward pass at test time: step-size variations due to backtracking and relative error behavior. The estimation of the theoretical step-size is computed as $\overline \tau =\frac{1}{L}$, where $L=\|y\|_1$ is the constant satisfying \ref{['eq nolip']} for the pair $(\text{KL}(y,A\cdot),h)$. This value provides an upper bound to the theoretical step-size $\tau<\frac{1}{L+L'}<\frac{1}{L}= \overline{\tau}$, where $L'$ is the constant satisfying \ref{['eq nolip']} for the pair $(R_{\theta},h)$.
  • Figure 5: Number of iterations performed in the forward pass to reach the stopping criterion threshold for four different initializations, averaged on the test set, varying with epochs.
  • ...and 7 more figures

Theorems & Definitions (44)

  • Definition 2.1: Well-posedness of a DEQ
  • Definition 4.1: NoLip, $L$-SMAD, Relative smoothness nolip,Bolte2018,RelativeSmoothness
  • Lemma 4.2: nolip Lemma 7
  • Proposition 4.3: KŁ with non-closed domain
  • proof
  • Corollary 4.4
  • proof
  • Proposition 4.5: Convergence of DEQ-MD with learnable regularization
  • proof
  • Remark 4.6: Analyticity is not restrictive
  • ...and 34 more