Source-linked AI summary
MoDL: Model Based Deep Learning Architecture for Inverse Problems
Hemant Kumar Aggarwal, Merry P. Mani, Mathews Jacob
TL;DR
Inverse reconstruction must combine learned image information with measurement consistency across varied and often complex forward models. MoDL unrolls a model-based variational algorithm with CNN priors, shared weights, end-to-end training, and numerical data-consistency blocks; the reported experiments show improved performance and reduced sensitivity to acquisition settings, while the preliminary implementation remains limited by GPU memory and scope.
Problem
Existing reconstruction methods face challenges using learned priors across complex forward models, acquisition settings, and data-constrained imaging problems.
Method
MoDL unrolls alternating variational updates that combine a learned CNN prior with forward-model data consistency, shared weights, end-to-end training, and conjugate-gradient blocks.
Results
The experiments report improved performance over proximal-gradient and pre-trained-denoiser approaches, with reduced sensitivity to undersampling and noise acquisition settings.
Takeaways & Limitations
Weight sharing decouples iteration count from network complexity, supporting more iterations with fewer trainable parameters and reduced overfitting risk in data-scarce settings.
Takeaways & Limitations
The preliminary implementation was restricted to ten iterations and a five-layer CNN by GPU memory limitations, while proposed enhancements were outside the work's scope.
Abstract
from arXiv · showhide
We introduce a model-based image reconstruction framework with a convolution neural network (CNN) based regularization prior. The proposed formulation provides a systematic approach for deriving deep architectures for inverse problems with the arbitrary structure. Since the forward model is explicitly accounted for, a smaller network with fewer parameters is sufficient to capture the image information compared to black-box deep learning approaches, thus reducing the demand for training data and training time. Since we rely on end-to-end training, the CNN weights are customized to the forward model, thus offering improved performance over approaches that rely on pre-trained denoisers. The main difference of the framework from existing end-to-end training strategies is the sharing of the network weights across iterations and channels. Our experiments show that the decoupling of the number of iterations from the network complexity offered by this approach provides benefits including lower demand for training data, reduced risk of overfitting, and implementations with significantly reduced memory footprint. We propose to enforce data-consistency by using numerical optimization blocks such as conjugate gradients algorithm within the network; this approach offers faster convergence per iteration, compared to methods that rely on proximal gradients steps to enforce data consistency. Our experiments show that the faster convergence translates to improved performance, especially when the available GPU memory restricts the number of iterations.
I. INTRODUCTION
Model-based recovery combines an explicit forward model with regularization to address ill-posed reconstruction, while deep methods learn image priors. MoDL unrolls alternating CNN and data-consistency updates, sharing weights and using optimization blocks to support complex inverse problems.
- Motivation: Ill-posed recovery from sparse or noisy measurements is formulated with a forward model and regularization priors that favor natural images.Examples include total variation, patch-based methods, low-rank penalties, and learned priors.
- Proposed framework: MoDL provides a systematic framework that combines a variational data-consistency term with a learned CNN prior for arbitrary inverse problems.An alternating recursive algorithm is unrolled into interleaved CNN and data-consistency blocks.
- Proposed framework: For complex forward models such as multichannel MRI, MoDL enforces data consistency by solving a quadratic subproblem with conjugate-gradient optimization.Simpler problems can admit analytical solutions, whereas CG handles more complex operators within the network.
- Comparison with prior methods: Direct inversion networks require large receptive fields and extensive training data, while model-based schemes can use smaller networks because the forward model supplies measurement structure.Direct inversion models may also need separate trained networks for differing acquisition parameters.
- Contributions: CG blocks reduce data-consistency error more rapidly per iteration than proximal-gradient approaches without increasing training memory demand.The paper reports improved performance for this strategy, particularly under limited GPU memory.
- Contributions: MoDL combines numerical optimization blocks, a variational model-based formulation, shared weights across iterations, and end-to-end training.The paper identifies these components as its main novelties relative to related deep-learning schemes.
B. Deep learned image reconstruction: the state-of-the-art
Deep image-reconstruction methods use CNNs to learn inverse mappings or regularizers, often embedding the forward model through unrolled iterative updates. Their applicability is constrained by receptive-field, training-data, and operator-structure requirements.
- Direct CNN inversion: Many deep-learning methods recover images with a learned CNN applied after transforming measurements into the image domain.The CNN is trained to approximate inversion of the normal operator A^H A for signals in the image set.
- Operator structure: For Fourier sampling, blurring, and projection imaging, translation-invariant normal operators make convolutional CNN structures suitable for reconstruction.The receptive field must match the point-spread-function support.
- Limitations: Fourier and projection imaging may require CNN receptive fields spanning the image, motivating large networks such as UNET with many layers.Large parameter counts increase the need for extensive training data, while parallel MRI may violate translation invariance.
- Unrolled methods: Unrolled iterative algorithms alternate CNN-based regularization with forward-model updates, allowing the network to focus on image redundancies rather than full-image inversion.The cited approaches are related to the proposed scheme but differ in training and architecture.
III. PROPOSED METHOD
MoDL formulates reconstruction with a learned CNN prior that removes noise and alias artifacts while enforcing data consistency. The resulting architecture alternates denoising and data-consistency operations, with trainable regularization.
- The regularization formulation combines the CNN prior with a data-consistency term and uses a trainable regularization parameter λ.Substituting the residual denoiser expression yields the proposed penalized reconstruction formulation.
- MoDL formulates image reconstruction using a learned CNN estimator of noise and alias patterns as a regularization prior.The prior uses Nw(x), with Dw(x) representing the denoised image after artifact and noise removal.
- The proposed network alternates a CNN denoising block with a data-consistency layer in an unrolled architecture.The denoising blocks use residual learning, and the architecture is organized into layers and recursive blocks.
- Because the forward model is explicitly used, the network need not have a full-image receptive field and can use fewer parameters than direct inversion networks.End-to-end training distinguishes this approach from methods using pre-trained denoisers.
A. Unrolling the recursive network
MoDL derives an unrolled network by approximating the recursive optimization algorithm locally and sharing its denoising and regularization parameters across iterations. This sharing reduces parameter count while supporting robust training and reconstruction quality.
- A. Unrolling the recursive network: A local Taylor approximation of the nonlinear denoising map leads to an alternating algorithm that approximates the reconstruction objective.The approximation is valid in the vicinity of the current iterate and neglects a second-order term for small perturbations.
- A. Unrolling the recursive network: The data-consistency subproblem is solved using the normal equations.
- A. Unrolling the recursive network: The fixed-iteration algorithm becomes an unrolled deep network whose denoising weights and regularization parameter are shared across iterations.The network is initialized with z0 = 0, and the same denoising operator Dw is used at every iteration.
- A. Unrolling the recursive network: Sharing weights reduces the number of network parameters by a factor equal to the number of iterations.The authors report improved training robustness and reconstruction quality from this design.
- A. Unrolling the recursive network: The penalized formulation and alternating minimization algorithm were selected for simplicity rather than because convergence guarantees are central to the practical network.A constrained ADMM alternative can provide rigorous convergence guarantees, but the network is not practically iterated to convergence.
B. Data consistency layer
MoDL enforces data consistency through numerical optimization blocks tailored to the forward model. For complex operators such as multichannel MRI, the method uses conjugate-gradient optimization within the unrolled network.
- B. Data consistency layer: For single-channel undersampled MRI with A = SF, the data-consistency solution can be computed analytically.S selects sampled Fourier rows, allowing the Fourier transform at sampled locations to be evaluated directly.
- B. Data consistency layer: Complex forward operators such as multichannel MRI are not analytically invertible.
- B. Data consistency layer: MoDL solves the data-consistency subproblem with conjugate gradients, placing numerical optimization layers inside the unrolled deep network.This differs from proximal-gradient approaches that alternate CNN blocks with steepest-descent steps.
C. End-to-end training of the deep network
MoDL trains its shared CNN parameters end to end for the reconstruction task while backpropagating through numerical optimization blocks. The shared network learns to handle differing noise and alias statistics across iterations without prescribing iteration-specific noise levels.
- C. End-to-end training of the deep network: Pre-trained denoisers cannot fully capture undersampling artifacts because those artifacts differ from Gaussian noise statistics.MoDL instead trains the reconstruction network end to end for the acquisition setting.
- C. End-to-end training of the deep network: Training minimizes mean square error between the final reconstruction xK and the target image t.The target image is indexed as t(i).
- C. End-to-end training of the deep network: The CNN weight parameter w is shared across iterations and its gradient is computed through the unrolled network using the chain rule.The Jacobian Jw(z) contains derivatives of CNN outputs with respect to the shared weights.
- C. End-to-end training of the deep network: The shared CNN learns to denoise noise and alias patterns with different statistics across iterations, eliminating a prescribed noise-variance or regularization-selection recipe.The same Dw and λ are used at every iteration.
- C. End-to-end training of the deep network: Backpropagation includes numerical optimization blocks, and gradients through conjugate-gradient blocks can be evaluated recursively.The CG block has no trainable parameters, and the network parameters and λ are updated with Adam.
- C. End-to-end training of the deep network: Training data with different undersampling patterns are used to address changes in the acquisition scheme.The unrolled architecture is implemented with shared variables so parameters at all iterations are initialized and updated together.
D. Implementation details
The CNN implementation uses an N-layer architecture with convolution, batch normalization, and ReLU blocks, while recursive weight sharing supports efficient initialization. A two-step initialization strategy trains a one-iteration model before initializing the full unrolled network.
- CNN architecture: The Nw block uses N layers with 64 filters per layer, combining convolution, batch normalization, and ReLU activations.The final layer omits ReLU to preserve negative learned noise patterns.
- CNN architecture: The learned noise is added to the block input following a residual learning strategy.
- Training initialization: A one-iteration model is trained first, and its learned weights initialize the recursively shared unrolled network.This initialization is reported as faster and more reliable than random initialization of the full network.
IV. EXPERIMENTS AND RESULTS
The experiments compare MoDL with variants differing in optimization strategy, training approach, and weight sharing, using common training data where applicable. MoDL is characterized as a CG-ET-WS strategy.
- Experimental design: The validation compares variants that differ in optimization strategy, training approach, and network architecture.The compared methods were trained with the same dataset as MoDL and the same total number of slices.
- Compared variants: MoDL uses conjugate gradients for data consistency, end-to-end training, and shared CNN weights across iterations.This configuration is categorized as the CG-ET-WS strategy.
- Compared variants: The study compares CG-ET-WS against SD-ET-WS, CG-PD-NS, and CG-ET-NS variants.These variants respectively change the data-consistency solver, use iteration-specific pretrained denoisers, or remove CNN weight sharing.
- Datasets: The experiments include a training-data-size analysis using a CatDog dataset because medical-imaging training data are scarce.The CatDog data were not used to pre-train the network in the remaining MRI experiments.
B. Data for training and testing the algorithm
MRI experiments use multichannel 3D brain data with separate training and testing subjects, known coil sensitivities, and distinct sampling masks. Reconstruction quality improves with more iterations, while conjugate-gradient blocks reduce data-consistency cost faster than steepest descent.
- MRI data: The MRI dataset contains fully sampled 3D T2 CUBE brain images acquired with a 12-channel head coil from five volunteers.Four subjects supplied training data and the fifth supplied testing data.
- MRI data: Training uses 360 slices across 12 coils, while testing uses 164 slices, with coil sensitivity maps assumed known.
- Sampling: Testing uses sampling masks different from those used during training, while each training slice shares one mask across all 12 coils.
- Optimization strategy: Conjugate-gradient blocks reduce data-consistency cost faster than steepest-descent blocks, improving performance at the same iteration count.The CG block has no trainable parameters, so memory demand depends on iterations rather than CG steps per iteration.
- Iteration count: Average testing PSNR improves as model iterations increase and saturates around 8-10 iterations at 10x acceleration.The authors use this range for subsequent experiments.
D. Effect of dataset size
Weight sharing reduces the parameter burden and makes performance less sensitive to training-set size. The trained network also shows relative insensitivity to acquisition settings and progressively reduces aliasing across reconstruction iterations.
- Effect of dataset size: The WS architecture is compared with NS as training images increase from 50 to 2100.The comparison evaluates how weight sharing changes the effect of dataset size.
- Effect of dataset size: A 3-layer NS model has 199K parameters, approximately matching the 188K parameters of a 7-layer WS model.
- Effect of dataset size: A 21-fold increase from 100 to 2100 training samples yields only 1 dB PSNR improvement for WS, whereas NS requires substantially more data for comparable performance.The same-parameter NS strategy performs worse than WS at all tested training-set sizes.
- Acquisition settings: A model trained at 10x acceleration reconstructs images at 10x, 12x, and 14x acceleration and in super-resolution settings.The figure reports PSNR values in dB for these acquisition settings.
- Iterative reconstruction: Across iterations, the CNN estimates changing alias/noise components, and the reconstruction progressively reduces noise and aliasing.The denoised output is combined with data consistency to produce the next iterate, repeated for ten iterations.
F. Comparison with other deep learning frameworks
MoDL is compared with deep-learning, compressed-sensing, and parallel-imaging methods across acceleration factors and evaluation settings. The experiments report consistently better reconstruction performance, while attributing the gains to conjugate-gradient data consistency, end-to-end training, and shared weights.
- Quantitative comparison: MoDL’s average PSNR was compared with three deep-learning frameworks and CSTV across acceleration factors from 2x to 20x.The comparison used 164 testing slices; deep-learning models were trained at 10x acceleration with Gaussian noise.
- Intermediate reconstructions: The same network removed substantially different alias patterns in the 16-fold acceleration and super-resolution intermediate reconstructions.The figure compares iterations 2 and 4 for 16-fold acceleration with iterations 1 and 5 for super-resolution; the network was trained at 10x.
- Quantitative comparison: MoDL performed better than the other techniques for all evaluated acceleration factors.The paper attributes this consistent improvement to CG-based data consistency, end-to-end training, and weight sharing.
- Quantitative comparison: CG-ET-NS performed poorest because its non-sharing architecture has 10 times the capacity of the corresponding weight-sharing strategy and requires more training data.The comparison also reports MoDL outperforming CSTV, whose regularizer is fixed total variation rather than adaptively learned from data.
- Cross-validation: Five-fold cross-validation evaluated average PSNR at 6x and 10x acceleration by training on four subjects and testing on the fifth.The reported average results suggest that MoDL is robust enough for practical use.
V. DISCUSSION
The discussion presents MoDL as a model-based deep-learning framework whose shared weights, embedded conjugate-gradient optimization, and learned prior support flexible reconstruction with limited model complexity. It also identifies practical and theoretical limitations, including memory-constrained experiments, possible sampling bias, untested alternatives, and absent convergence analysis.
- Limitations: Experiments were restricted to ten iterations and a five-layer CNN because of GPU memory limitations.The authors state that more training data, augmentation, regularization priors, dropout, or deeper CNNs could improve performance, but leave these enhancements for future work.
- Limitations: Variable-density sampling may shift reliance toward lower k-space samples, causing a slight loss in high-frequency details.Fixing the sampling pattern might improve results but would increase sensitivity to acquisition settings; ADMM or momentum might also converge faster.
- Limitations: The network’s convergence was not theoretically analyzed, and the study used a finite fixed number of iterations.The authors defer evaluation of more iterations and detailed convergence analysis to future work.
- Contributions: MoDL combines data-driven learning with a physics-based model to design deep architectures for inverse problems with arbitrary structure.Embedding conjugate gradients within a network layer supports complex forward models such as multichannel MRI and facilitates additional image priors.
- Contributions: Sharing weights across iterations decouples convergence from network complexity and permits more iterations without increasing the degrees of freedom.The paper links this design to reduced overfitting risk and improved results with relatively few trainable parameters.