Efficient Unrolled Networks for Large-Scale 3D Inverse Problems

TL;DR

This paper tackles the challenge of scaling unrolled networks to large-scale 3D inverse problems by introducing domain partitioning and a fast normal-operator approximation. Domain partitioning allows training on small patches while preserving full-volume inference, and the normal-operator factorization $\mathbf{A}^\top\mathbf{A} \approx \mathbf{H}(\mathbf{m},\boldsymbol{\lambda}) = \mathrm{diag}(\mathbf{m}) \mathbf{F}^{-1} \mathrm{diag}(\boldsymbol{\lambda}) \mathbf{F}$ enables FFT-based, memory-efficient data-consistency updates. The combination yields state-of-the-art performance on large-scale CBCT and MC-MRI with substantial memory and compute savings, including handling $501^3$ volumes on a single GPU. The work provides a practical path toward deploying high-quality, data-consistent reconstructions in resource-constrained environments and outlines clear avenues for extending to non-Cartesian and Poisson-noise settings.

Abstract

Deep learning-based methods have revolutionized the field of imaging inverse problems, yielding state-of-the-art performance across various imaging domains. The best performing networks incorporate the imaging operator within the network architecture, typically in the form of deep unrolling. However, in large-scale problems, such as 3D imaging, most existing methods fail to incorporate the operator in the architecture due to the prohibitive amount of memory required by global forward operators, which hinder typical patching strategies. In this work, we present a domain partitioning strategy and normal operator approximations that enable the training of end-to-end reconstruction models incorporating forward operators of arbitrarily large problems into their architecture. The proposed method achieves state-of-the-art performance on 3D X-ray cone-beam tomography and 3D multi-coil accelerated MRI, while requiring only a single GPU for both training and inference.

Figures (10)

• Figure 1: Peak video memory complexity (dashed lines) and global execution times (dotted lines) of isolated components used in unrolling. We show the cost of evaluating and back-propagating through a standard 3D data consistency step (using gradient descent) and a standard 3D network step (using a 3D DRUNet zhangPlugPlayImageRestoration2021). We see here that the bottleneck lies in the network step, which grows rapidly with the volume size, while the data-consistency step remains manageable even at high resolutions.
• Figure 2: Domain partitioning strategy: in case of a forward operator ${\bm{A}}$ that mixes the signal ${\bm{x}}$ in a non-trivial manner, we can still decompose the full domain $\mathbb{R}^{n}$ into two orthogonal subspaces $\mathbb{R}^{p}$ and $\mathbb{R}^{q}$. Then by linearity, we solve for the unknown smaller patch ${\bm{x}}_{\text{patch}} \in \mathbb{R}^{p}$ (red) given the context ${\bm{x}}_{\text{context}} \in \mathbb{R}^{q}$ (blue).
• Figure 3: Illustrations of sparse view reconstructions with [30/1200] projections on the Walnut-CBCT dersarkissianConebeamXrayComputed2019 dataset using the methods compared in Tab. \ref{['tab:walnut:quantitative']}. First row axial slices, second row vertical slices from the same sample. PSNR is computed per slice.
• Figure 4: Illustrations of MC-MRI reconstructions with acceleration rate of $5$ on the Calgary-Campinas dataset souza_open_2018 for the methods compared in \ref{['tab:calgary:quantitative']}. First row: axial slice, second row: coronal slice from the same sample. PSNR is computed per slice.
• Figure 5: Illustrations of the Walnut-CBCT dataset. Examples of two radiographies.

Theorems & Definitions (4)

• Remark 2.1
• Remark 4.1
• Remark 4.2
• Remark 4.3