PaperlayerView exact revision on arXiv

Source-linked AI summary

Learned Primal-dual Reconstruction

Jonas Adler, Ozan Öktem

arXiv:1707.06474v3math.OCcs.CVcs.NEmath.FA

TL;DR

Tomographic reconstruction needs methods that combine forward-model knowledge with data-driven learning, especially for clinical-sized inverse problems. The Learned Primal-Dual algorithm provides this framework and improves reconstruction quality over classical and learned alternatives for analytical and human phantoms.

  • Problem

    Data-driven reconstruction had not yet demonstrated an entirely data-driven approach applicable to clinical-sized tomographic problems, motivating methods that incorporate forward-model knowledge.

  • Method

    The Learned Primal-Dual algorithm unrolls a primal-dual optimization scheme and replaces proximal operators with learned operators in a framework that incorporates the forward model.

  • Results

    For analytical phantoms, the method improves over classical and post-processing algorithms by at least 6 dB while improving SSIM; for human phantoms, it improves over TV by 6.6 dB and learned post-processing by 2.2 dB.

  • Takeaways & Limitations

    The method works directly from tomographic data without an initial reconstruction and has reconstruction time comparable to filtered back-projection and learned post-processing for 512 × 512 human phantoms.

Abstract

from arXiv · show

We propose the Learned Primal-Dual algorithm for tomographic reconstruction. The algorithm accounts for a (possibly non-linear) forward operator in a deep neural network by unrolling a proximal primal-dual optimization method, but where the proximal operators have been replaced with convolutional neural networks. The algorithm is trained end-to-end, working directly from raw measured data and it does not depend on any initial reconstruction such as FBP. We compare performance of the proposed method on low dose CT reconstruction against FBP, TV, and deep learning based post-processing of FBP. For the Shepp-Logan phantom we obtain >6dB PSNR improvement against all compared methods. For human phantoms the corresponding improvement is 6.6dB over TV and 2.2dB over learned post-processing along with a substantial improvement in the SSIM. Finally, our algorithm involves only ten forward-back-projection computations, making the method feasible for time critical clinical applications.

I. INTRODUCTION

Tomographic reconstruction combines forward-model knowledge and prior information to recover images from indirect, noisy observations. The paper motivates a framework that integrates these model-driven elements with deep learning for inverse problems.

  • Inverse problems reconstruct system parameters from indirect observations, including interior images in CT and MRI.
  • Model-driven reconstruction relies on a forward model, data statistics, and prior information about the image.
  • Deep learning has largely succeeded where forward-model knowledge is unnecessary or of limited importance, while clinical-scale tomographic reconstruction remained undemonstrated.
  • The Learned Primal-Dual method addresses whether forward-model knowledge can be incorporated into neural-network reconstruction for ill-posed inverse problems.
  • Existing tomographic learning directions include post-processing approximate reconstructions, learning regularizers, and learned iterative schemes.
  • Learned iterative schemes update reconstructions using the current iterate and outputs from the forward operator and its adjoint.

II. CONTRIBUTION AND OVERVIEW OF PAPER

The paper proposes a learned iterative reconstruction framework that combines deep learning with model-based reconstruction. It learns the complete mapping from measured data to reconstruction rather than only post-processing or learning a prior.

  • The Learned Primal-Dual framework uses CNNs in reconstruction and data spaces connected by the forward operator and its adjoint.
  • The networks are trained to minimize reconstruction mean squared error.
  • The method learns the whole reconstruction operator from data, rather than only a post-processing transformation or an isolated prior.
  • The authors release code and learned parameters to support reproduction and application to other inverse problems.

III. THE LEARNED PRIMAL-DUAL ALGORITHM

The algorithm is motivated by primal-dual optimization for large-scale, non-smooth inverse problems. It alternates primal and dual updates while using the forward operator and its adjoint to connect data and reconstruction spaces.

  • Gradient methods are unsuitable for many non-differentiable imaging objectives, while smoothing introduces extra parameters and non-exact solutions.
  • Proximal methods replace gradient steps to handle non-smooth objective functionals directly.
  • Proximal primal-dual schemes introduce an auxiliary dual variable and alternate updates of primal and dual variables.
  • PDHG accommodates objectives with a possibly non-linear operator K, dual functional F, and primal functional G.
  • For the inverse-problem objective, the framework identifies K with the forward operator T, F with the data term, and G with the regularizer.
  • In CT, the forward operator is a ray transform, and the adjoint of the derivative corresponds to back-projection.

B. Learned PDHG

The Learned PDHG variant replaces primal and dual proximal operators with learned operators and runs for a fixed number of iterations. Its reported performance was comparable to traditional methods but below state-of-the-art deep-learning reconstruction methods.

  • Learned PDHG replaces proximal operators with parametrized operators whose parameters are learned from training data.
  • A fixed iteration count establishes the stopping criterion and fixes the computation budget before training for time-critical applications.
  • Algorithm 2 uses learned primal and dual proximal operators within an I-iteration PDHG scheme.
  • The learned algorithm infers proximal parameters, step lengths, and an overrelaxation parameter from training data.
  • The implemented method performed comparably to traditional methods but did not surpass state-of-the-art deep-learning image reconstruction.

C. Learned Primal-Dual

The Learned Primal-Dual algorithm extends learned PDHG by allowing learned memory, update combinations, evaluation points, and iteration-specific proximal operators. It can use zero initialization without relying on an earlier reconstruction, while pseudo-inverse initialization did not improve final results.

  • C. Learned Primal-Dual: The Learned Primal-Dual algorithm extends learned PDHG with memory in the primal and dual spaces between iterations.The primal and dual variables are expanded to retain multiple components across iterations.
  • C. Learned Primal-Dual: The network learns how to combine previous updates with forward-operator evaluations instead of enforcing fixed update forms.This replaces an explicitly prescribed dual update with a learned combination.
  • C. Learned Primal-Dual: The network learns the forward-operator evaluation point rather than using hard-coded over-relaxation.This allows the evaluation location to vary as part of the learned algorithm.
  • C. Learned Primal-Dual: Iteration-specific learned proximal operators increase the parameter space and notably improve reconstruction quality.Together with the other modifications, this defines the Learned Primal-Dual algorithm in algorithm 3.
  • C. Learned Primal-Dual: Zero initialization is the simplest starting point, and pseudo-inverse initialization marginally reduced training time without improving final results.Because pseudo-inverse initialization adds dependence on an earlier reconstruction, the reported values use zero initialization.

D. Connection to variational regularization

The Learned Primal-Dual framework contains several optimization methods as special cases and is evaluated on simplified and clinically realistic low-dose CT problems. The experiments use ellipse and human phantoms with linear and nonlinear forward-model settings.

  • D. Connection to variational regularization: Nprimal = 2 and Ndual = 1 reduce the Learned Primal-Dual algorithm to classical PDHG under specified parameter choices.This establishes a direct connection between the learned scheme and a classical primal-dual method.
  • D. Connection to variational regularization: Sufficient training data can approximate the explicit proximal equalities arbitrarily well through the universal approximation property of neural networks.The learned proximal operators need not explicitly access the true proximal operators.
  • D. Connection to variational regularization: Gradient descent and limited-memory BFGS are special cases obtainable through appropriate choices of learned proximal operators.The paper presents gradient descent as a concrete sub-case and identifies limited-memory BFGS as another.
  • D. Connection to variational regularization: With appropriate parameters, the proposed algorithm should perform at least as well as current variational regularization schemes at the same fixed number of iterates.This is stated as a comparison under matching stopping criteria.
  • D. Connection to variational regularization: The evaluation covers low-dose CT with ellipse phantoms and human abdominal CT phantoms using analytical, linearized, and nonlinear forward models.Ellipse data use a sparse 30-view parallel-beam geometry with Gaussian noise, while human data use fan-beam geometry and Poisson noise.
  • D. Connection to variational regularization: The nonlinear pre-log and linear post-log forward models are both implemented and compared, with modified Shepp-Logan and held-out patient data used for validation.The post-log transformation converts the model to the ray transform.

B. Implementation

The implementation combines ODL operator layers with TensorFlow neural networks and uses residual convolutional learned proximal operators. Ten unrolled iterations yield 60 convolutional layers while retaining approximately 2.4 · 10^5 parameters.

  • B. Implementation: Operator components were implemented in ODL and converted into TensorFlow layers, while neural-network layers and training were implemented in TensorFlow.The ray transform and its adjoint used the GPU-accelerated astra_gpu backend.
  • B. Implementation: The proximal networks use PReLU nonlinearities, 3 × 3 convolutions, five persistent primal and dual channels, and separate channel widths for primal and dual operators.The dual operators receive an additional input because the measured data g is supplied to them.
  • B. Implementation: CNNs were chosen because forward-operator and prior properties in CT and MRI are often approximately translation invariant.The resulting reconstruction operator is therefore expected to be approximately translation invariant.
  • B. Implementation: The learned proximal operators use residual networks formed by an identity operator and convolutional affine layers.The identity structure reflects that proximal operators are typically close to identity and can make networks easier to train.
  • B. Implementation: I = 10 unrolled iterations require ten evaluations each of the forward operator and the adjoint derivative, producing 60 convolutional layers and approximately 2.4 · 10^5 parameters.Each iterate contains two three-layer networks, one for each proximal operator.
  • B. Implementation: Training uses Xavier initialization, ADAM with cosine annealing, gradient-norm clipping, and batch sizes of 5 for ellipse data and 1 for human phantoms.No parameter regularization, dropout, or batch normalization was used.

2) Incorporating the forward operator in neural networks:

Training requires gradients through all unrolled network components, including the learned proximal operators and the forward and backward operators. Automatic differentiation supplies these gradients by combining TensorFlow and ODL derivative-adjoint computations.

  • 2) Incorporating the forward operator in neural networks:: The loss gradient depends on every neural-network component, the forward operator T, and the adjoint derivative [∂T(f)]* across all I iterations.This propagation through the unrolled reconstruction creates the main differentiation challenge.
  • 2) Incorporating the forward operator in neural networks:: TensorFlow automatic differentiation applies backpropagation, using TensorFlow for proximal derivatives and ODL for operator derivatives and their adjoints.The approach relies on the chain rule through the complete learned reconstruction operator.

C. Comparison

The study compares Learned Primal-Dual reconstruction with classical, learned, and simplified alternatives using runtime, PSNR, and SSIM. The comparison also examines how retaining or discarding the forward operator changes the reconstruction setup.

  • The evaluation includes standard FBP, isotropic TV reconstruction, partially Learned Gradient, U-Net post-processing, learned PDHG, and simplified Learned Primal-Dual variants.The simplified variants include a Learned Primal method and FBP + residual denoising.
  • FBP uses a Hann filter, while TV uses 1000 iterations of classical PDHG, with reconstruction settings selected to maximize PSNR.
  • The network receives data at the dual iterates and replaces classical proximal updates with CNN-based primal and dual blocks.Initial guesses enter from the left; in classical PDHG, the corresponding updates use proxτG and proxσF∗.
  • The human-phantom comparison considers both non-linear and linearized forward operators but compares only FBP, TV, and U-Net denoising because training times are noticeably longer.
  • All learned algorithms use the same training scheme and are evaluated by runtime, PSNR, and structural similarity index (SSIM).

V. RESULTS

The Learned Primal-Dual algorithm outperforms classical and learned reconstruction methods on ellipse and human phantom low-dose CT data, while remaining computationally competitive. Its reconstructions improve quantitative metrics and fine-detail recovery, though visual over-smoothing remains a limitation.

  • Ellipse data: 1.3 dB separates the Learned Primal-Dual algorithm from the closest Learned Primal method on the ellipse data.The Learned Primal-Dual reconstruction also avoids a halo artifact near the outer bone seen in Learned PDHG and Learned Primal results.
  • Runtime: At least 2 orders of magnitude faster than TV regularization are all learned methods, while forward-operator methods run ≈6 times slower than FBP-plus-denoising methods.For full-size data, Learned Primal-Dual is more competitive with FBP and U-Net denoising than on ellipse data because larger data increase FBP runtime.
  • Human phantom data: 10.5 dB over FBP, 6.6 dB over TV, and 2.2 dB over the U-Net denoiser are the Learned Primal-Dual improvements for human phantoms.The algorithm also produces a large SSIM improvement over TV and U-Net denoising.
  • Visual assessment: The Learned Primal-Dual reconstruction recovers finer detail and suppresses streak and noise-created artifacts, but its texture appears slightly over-smoothed.The over-smoothing is especially apparent in zoomed regions despite improved fine-detail reconstruction.
  • Overall results: The Learned Primal-Dual algorithm outperforms classical reconstruction methods in both PSNR and SSIM and improves upon learned post-processing for ellipse and human phantom data.For 512 × 512 human phantoms, reconstruction time is comparable with filtered back-projection and learned post-processing.
  • Discussion: Incorporating the forward operator is especially advantageous for very noisy, under-sampled ellipse data, whereas the improvement is smaller for less noisy human phantom data.The authors conjecture that direct-from-data learned schemes gain more as data quality decreases.

VII. CONCLUSIONS

The Learned Primal-Dual algorithm reconstructs directly from tomographic data without an initial reconstruction, combining learned operators with a primal-dual scheme. It achieves state-of-the-art CT results on analytical and human phantoms, while motivating further research and applications to other imaging modalities.

  • VII. CONCLUSIONS: The algorithm replaces proximal operators in a PDHG-inspired scheme with learned operators and works directly from tomographic data.It does not depend on an initial reconstruction.
  • VII. CONCLUSIONS: At least 6 dB improvement over FBP, TV, and post-processing methods was achieved for analytical phantoms, with improved SSIM.Human-phantom improvements were 6.6 dB over TV and 2.2 dB over learned post-processing.
  • VII. CONCLUSIONS: The authors hope the method will inspire further research in Learned Primal-Dual schemes and applications to other imaging modalities.This is presented as a direction for future work.
Loading 1707.06474v3…