Efficient Model-Based Deep Learning via Network Pruning and Fine-Tuning

TL;DR

This work targets the high test-time cost of model-based deep learning (MBDL) for imaging inverse problems by introducing a structured pruning pipeline guided by the layer-group dependencies and the group $\ell_1$-norm, followed by three fine-tuning strategies (supervised, school, self-supervised) to recover performance. The method prunes CNN priors within DEQ and deep unfolding (DU) frameworks, achieving up to 50% and 32% speedups with minimal degradation. Extensive experiments on compressed sensing MRI and image super-resolution demonstrate substantial speedups with only small losses in reconstruction quality, and reveal that school and self-supervised fine-tuning can closely approach or match supervised performance under various data availability conditions. This approach enables scalable deployment of MBDL for large-scale imaging tasks and opens avenues for exploring different pruning rankings and self-supervised loss formulations.

Abstract

Model-based deep learning (MBDL) is a powerful methodology for designing deep models to solve imaging inverse problems. MBDL networks can be seen as iterative algorithms that estimate the desired image using a physical measurement model and a learned image prior specified using a convolutional neural net (CNNs). The iterative nature of MBDL networks increases the test-time computational complexity, which limits their applicability in certain large-scale applications. Here we make two contributions to address this issue: First, we show how structured pruning can be adopted to reduce the number of parameters in MBDL networks. Second, we present three methods to fine-tune the pruned MBDL networks to mitigate potential performance loss. Each fine-tuning strategy has a unique benefit that depends on the presence of a pre-trained model and a high-quality ground truth. We show that our pruning and fine-tuning approach can accelerate image reconstruction using popular deep equilibrium learning (DEQ) and deep unfolding (DU) methods by 50% and 32%, respectively, with nearly no performance loss. This work thus offers a step forward for solving inverse problems by showing the potential of pruning to improve the scalability of MBDL. Code is available at https://github.com/wustl-cig/MBDL_Pruning .

Figures (5)

• Figure 1: An illustration of our proposed pipeline. Our proposed pipeline consists of two components (see Section \ref{['sec-method']}): (a) a structured pruning algorithm that physically removes CNN filters based on the group $\ell_1$-norm, and (b) fine-tuning algorithms to minimize the performance loss between the pre-trained model and the pruned model. Each fine-tuning strategy has unique applicability depending on the presence of the pre-trained models and high-quality ground truth.
• Figure 2: Illustration of testing time evolution between MBDL models and the pruned variants using our proposed pipeline. Left: evolution of the distance between two consecutive images; Middle and Right: evolution of the testing PSNR values compared to the ground truth. We used the supervised and school strategies to fine-tune DEQ and USRNet, respectively. Note how models pruned by our proposed pipeline can significantly reduce the testing time while maintaining competitive performance.
• Figure 3: Degradation PSNR percentage of pruned MBDL models compared to the unpruned model in different pruning ratios and different fine-tuning strategies. These results correspond to experiments of CS-MRI at the sampling rate of $16.6\ \%$. Note how supervised fine-tuning method can reduce $65\%$ parameters while maintaining less than $4\%$ PSNR degradation.
• Figure 4: Degradation PSNR percentage of pruned MBDL models compared to the unpruned model in different pruning ratios and different fine-tuning strategies. These results correspond to experiments of CS-MRI at the sampling rate of $12.5\ \%$. Note how supervised fine-tuning method can reduce $65\%$ parameters while maintaining less than $4\%$ PSNR degradation.
• Figure 5: Visual results of USRNet and its pruned variants with the school fine-tuning strategy at scale of $\times$3. Note how USRNet has been pruned more than $65\%$ parameters but kept similar visual quality compared to the unpruned model.