Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Equivariant Deep Equilibrium Models for Imaging Inverse Problems

Alexander Mehta, Ruangrawee Kitichotkul, Vivek K Goyal, Julián Tachella

TL;DR

This work tackles the challenge of training deep equilibrium models (DEQs) for imaging inverse problems under equivariant imaging (EI) without ground-truth data. It introduces a modular implicit differentiation (ID) approach that simplifies backpropagation through fixed-point computations, enabling effective EI losses to be used with DEQs. Empirically, DEQs trained with ID outperform Jacobian-free backpropagation and deliver reconstruction quality close to supervised training, while EI-trained DEQs exhibit equivariance and implicitly learn a symmetry-invariant proximal map. The results demonstrate the practicality of self-supervised, end-to-end DEQ training for complex imaging tasks like sparse-view CT and accelerated MRI, with strong performance gains over baselines such as deep unrolling and PnP/RED. Overall, the work highlights the potential of EI-powered DEQs to enable high-quality reconstructions in settings where ground-truth data are scarce or unavailable, aided by principled modular differentiation and symmetry considerations.

Abstract

Equivariant imaging (EI) enables training signal reconstruction models without requiring ground truth data by leveraging signal symmetries. Deep equilibrium models (DEQs) are a powerful class of neural networks where the output is a fixed point of a learned operator. However, training DEQs with complex EI losses requires implicit differentiation through fixed-point computations, whose implementation can be challenging. We show that backpropagation can be implemented modularly, simplifying training. Experiments demonstrate that DEQs trained with implicit differentiation outperform those trained with Jacobian-free backpropagation and other baseline methods. Additionally, we find evidence that EI-trained DEQs approximate the proximal map of an invariant prior.

Equivariant Deep Equilibrium Models for Imaging Inverse Problems

TL;DR

This work tackles the challenge of training deep equilibrium models (DEQs) for imaging inverse problems under equivariant imaging (EI) without ground-truth data. It introduces a modular implicit differentiation (ID) approach that simplifies backpropagation through fixed-point computations, enabling effective EI losses to be used with DEQs. Empirically, DEQs trained with ID outperform Jacobian-free backpropagation and deliver reconstruction quality close to supervised training, while EI-trained DEQs exhibit equivariance and implicitly learn a symmetry-invariant proximal map. The results demonstrate the practicality of self-supervised, end-to-end DEQ training for complex imaging tasks like sparse-view CT and accelerated MRI, with strong performance gains over baselines such as deep unrolling and PnP/RED. Overall, the work highlights the potential of EI-powered DEQs to enable high-quality reconstructions in settings where ground-truth data are scarce or unavailable, aided by principled modular differentiation and symmetry considerations.

Abstract

Equivariant imaging (EI) enables training signal reconstruction models without requiring ground truth data by leveraging signal symmetries. Deep equilibrium models (DEQs) are a powerful class of neural networks where the output is a fixed point of a learned operator. However, training DEQs with complex EI losses requires implicit differentiation through fixed-point computations, whose implementation can be challenging. We show that backpropagation can be implemented modularly, simplifying training. Experiments demonstrate that DEQs trained with implicit differentiation outperform those trained with Jacobian-free backpropagation and other baseline methods. Additionally, we find evidence that EI-trained DEQs approximate the proximal map of an invariant prior.

Paper Structure

This paper contains 24 sections, 1 theorem, 10 equations, 2 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Equivariant Imaging
  4. Deep Equilibrium Models
  5. Backpropagation for DEQ
  6. Difficulty in Training DEQ with EI
  7. Modular Implicit Differentiation
  8. Experiments
  9. Tasks
  10. Sparse-view CT
  11. Accelerated MRI
  12. Reconstruction Methods
  13. Deep Equilibrium Models
  14. Deep Unrolling
  15. U-Net Refinement
  16. ...and 9 more sections

Key Result

Proposition 1

The vector-Jacobian product between a row vector $g \in \mathbb{R}^{1 \times n}$ and a Jacobian ${\partial \widehat{x}} / {\partial \theta} \in \mathbb{R}^{n \times p}$ of a fixed point $\widehat{x} = \operatorname{F}_{\theta}(\widehat{x})$ with respect to the parameters $\theta \in \mathbb{R}^n$ is

Figures (2)

  • Figure 1: Reconstructions from Deep Equilibrium Models (DEQs) and U-Nets trained with full supervision (Sup.) and without ground truth via equivariant imaging (EI). Top -- Sparse View CT; Bottom -- 8x Accelerated MRI with Noise. PSNR is reported in image corners.
  • Figure 2: PyTorch implementation of DEQ forward pass and implicit differentiation, adapted from giltonDeepEquilibriumArchitectures2021. The fixed-point is computed in lines 3–4; the remaining lines implement ID. Lines 7-12 define backward_hook for backpropagating through the fixed point xhat given an upstream gradient grad. Line 6 performs an evaluation of $\mathop{\mathrm{F}}\nolimits_\theta$ so that $\partial \operatorname{F}_{\theta}(\widehat{x}) / \partial \theta$ is captured in the autograd graph.

Theorems & Definitions (1)

  • Proposition 1