Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

A primal-dual algorithm for image reconstruction with input-convex neural network regularizers

Matthias J. Ehrhardt, Subhadip Mukherjee, Hok Shing Wong

TL;DR

The paper addresses variational image reconstruction with ICNN-based regularizers and introduces a convex reformulation via epigraphical projections that decouples the nested network structure. This reformulation is proven equivalent to the original problem and solved efficiently with a tailored primal-dual algorithm featuring a specialized step-size scheme and parallelizable updates. Empirical results show faster convergence and greater stability than subgradient methods, with competitive performance in the smooth setting and clear benefits when used as a lower-level solver in bilevel learning to train the regularizer. The framework demonstrates strong reconstruction performance across denoising, CT, and inpainting tasks and is amenable to deeper networks and parallel computation, highlighting practical impact for data-driven regularization in inverse problems.

Abstract

We address the optimization problem in a data-driven variational reconstruction framework, where the regularizer is parameterized by an input-convex neural network (ICNN). While gradient-based methods are commonly used to solve such problems, they struggle to effectively handle non-smooth problems which often leads to slow convergence. Moreover, the nested structure of the neural network complicates the application of standard non-smooth optimization techniques, such as proximal algorithms. To overcome these challenges, we reformulate the problem and eliminate the network's nested structure. By relating this reformulation to epigraphical projections of the activation functions, we transform the problem into a convex optimization problem that can be efficiently solved using a primal-dual algorithm. We also prove that this reformulation is equivalent to the original variational problem. Through experiments on several imaging tasks, we show that the proposed approach not only outperforms subgradient methods and even accelerated methods in the smooth setting, but also facilitates the training of the regularizer itself.

A primal-dual algorithm for image reconstruction with input-convex neural network regularizers

TL;DR

The paper addresses variational image reconstruction with ICNN-based regularizers and introduces a convex reformulation via epigraphical projections that decouples the nested network structure. This reformulation is proven equivalent to the original problem and solved efficiently with a tailored primal-dual algorithm featuring a specialized step-size scheme and parallelizable updates. Empirical results show faster convergence and greater stability than subgradient methods, with competitive performance in the smooth setting and clear benefits when used as a lower-level solver in bilevel learning to train the regularizer. The framework demonstrates strong reconstruction performance across denoising, CT, and inpainting tasks and is amenable to deeper networks and parallel computation, highlighting practical impact for data-driven regularization in inverse problems.

Abstract

We address the optimization problem in a data-driven variational reconstruction framework, where the regularizer is parameterized by an input-convex neural network (ICNN). While gradient-based methods are commonly used to solve such problems, they struggle to effectively handle non-smooth problems which often leads to slow convergence. Moreover, the nested structure of the neural network complicates the application of standard non-smooth optimization techniques, such as proximal algorithms. To overcome these challenges, we reformulate the problem and eliminate the network's nested structure. By relating this reformulation to epigraphical projections of the activation functions, we transform the problem into a convex optimization problem that can be efficiently solved using a primal-dual algorithm. We also prove that this reformulation is equivalent to the original variational problem. Through experiments on several imaging tasks, we show that the proposed approach not only outperforms subgradient methods and even accelerated methods in the smooth setting, but also facilitates the training of the regularizer itself.

Paper Structure

This paper contains 18 sections, 4 theorems, 40 equations, 11 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Overview of learned regularizers
  3. Existing methods for solving variational problems
  4. Contributions
  5. Background
  6. Architecture of ICNNs
  7. Training Methods
  8. Adversarial regularization framework
  9. Bilevel Learning
  10. Constrained Convex Reformulation
  11. Primal-Dual framework
  12. Experiments
  13. Salt and pepper denoising
  14. CT with Poisson noise
  15. Deeper network
  16. ...and 3 more sections

Key Result

Proposition 3.1

Under Assumption 1, $R_{\boldsymbol{\theta}}$ defined by (gen-eq) is convex with respect to $\mathbf{x}$.

Figures (11)

  • Figure 1: Example illustrating the non-convexity of (\ref{['var-g']}).
  • Figure 1: Denoising: Ablation study of proposed method for step-size hyperparameters. The markers on the left corresponds to those depicted in the energy versus iterations plots on the right.
  • Figure 2: Denoising: Comparison to subgradient methods. Note that SM-C either converges only to a suboptimal solution or is very slowly. SM-D converges, but remains slightly slower than the proposed method.
  • Figure 3: Denoising: Visual comparison of reconstructions, with PSNR values shown in the top right corner. The proposed method achieves a visually satisfactory reconstruction within 15 iterations, while that of SM-C remains noisy.
  • Figure 4: CT: Comparison to subgradient methods. Note that SM-C is faster than SM-D in this case, and at times even faster than the proposed method initially.
  • ...and 6 more figures

Theorems & Definitions (8)

  • Proposition 3.1
  • Proof 1
  • Proposition 3.2
  • Proof 2
  • Theorem 3.3
  • Proof 3
  • Corollary 3.4
  • Proof 4