Paper: A Structured Pruning Algorithm for Model-based Deep Learning (ICLR-2024)

Reviewer comment (Writing):
The explanation of the methods could be more clear. As the reviewer is not familiar with MBDL methods discussed in the paper or DepGraph, they cannot really follow how this pruning works. For example, "$f_{\theta,i}$ has dependency on $f_{\theta,j}$" is very briefly explained but seems to be a key ingredient of the proposed pruning approach.

Author response:
Prompted by your comment, the revised manuscript did a better job of explaining the pruning method. To be specific, we (1) detailed the application of SPADE to a simple network, (2) added a discussion of the inter-layer and intra-layer dependencies, and (3) provided a clear, step-by-step depiction of the pruning process. Please refer to the method section in the revised manuscript for more details.

Reviewer comment (Reproducibility):
It is unclear how the layers are actually removed. Probably it is not just taking a layer out of the network, as this would need additional assumptions about compatibility of the input and output dimensions of the layer, and it is uncertain whether the network would give any meaningful output after such removal.

Author response:
Prompted by your comment, the revised manuscript did a better job of explaining the pruning method. To be specific, we (1) detailed the application of SPADE to a simple network, (2) added a discussion of the inter-layer and intra-layer dependencies, and (3) provided a clear, step-by-step depiction of the pruning process. Please refer to the method section in the revised manuscript for more details.

Reviewer comment (Presentation):
The presentation could be improved. For example, the authors could take one MBDL method and describe in detail how their approach is applied to it. Table 2 and Figure 4 could be removed or moved to supplementary material if more space is needed.

Author response:
Prompted by your comment, the revised manuscript did a better job of explaining the pruning method. To be specific, we (1) detailed the application of SPADE to a simple network, (2) added a discussion of the inter-layer and intra-layer dependencies, and (3) provided a clear, step-by-step depiction of the pruning process. Please refer to the method section in the revised manuscript for more details.

Reviewer comment (Novelty):
It could be clarified what is the authors' contribution or novelty with respect to DepGraph.

Author response:
While pruning is a general idea, this topic has never been explored in the context of model-based deep learning (MBDL). MBDL architectures are inherently memory intensive, where pruning can have a big impact. Our work is the first to investigate the potential of pruning in MBDL  Prompted by feedback, we revised our manuscript to better clarify our contribution. We revised the title of our paper to _“Efficient Model-based Deep Learning via Network Pruning”_ to better highlight our contribution.

Reviewer comment (Evaluation):
The speedup seems quite low compared to the pruning ratio. For example, 10% pruning only seems to give 0–3% speedup, and E2EVar (Table 1) has only 24% speedup with a 65% pruning ratio. The reviewer asks whether pruning removes lower-complexity layers, whether times include possible compilation times (XLA, jit, etc.), or how this can be explained.

Author response:
The speed-up is not proportional to the pruning ratio because _MBDL networks consist of non-prunable data-consistency layers_ that enforce consistency between intermediate results and the raw measurements. These data-consistency layers correspond to the data-fidelity term $g(x)$ in the objective function (2). Note also recent studies have proposed to reduce the computational complexity of these data-consistency layers, such as “SGD-Net: Efficient Model-Based Deep Learning With Theoretical Guarantees”. SPADE is fully compatible with these methods.

Reviewer comment (Novelty):
While I appreciate the authors effort in improving the computation efficiency of deep learning methods for inverse problems, I wonder if ICLR is the correct venue for this paper. I personally think that the contributions maybe more of relevance to a compressed sensing, and super resolution community, as the pure technical contribution in this paper maybe limited. At the heart, the paper applies existing pruning method, and existing fine-tuning/training methods.

Author response:
Our main goal in this paper is to improve the efficiency of MBDL networks via network pruning. This directly relates to the computational efficiency and scalability issues of DL research, aligning with the interests of the ICLR community.

Reviewer comment (Novelty):
While the paper specialize in Model Based Deep Learning, and the introduction focuses on methods like RED, I fail to see how the proposed method SPADE is relevant to such methods, or how any specialized knowledge from this domain is used. It occurs to me that this can be applied to any CNN, not just the ones used within MBDL framework. Can authors confirm?

Author response:
While pruning is a general idea, this topic has never been explored in the context of model-based deep learning (MBDL). MBDL architectures are inherently memory intensive, where pruning can have a big impact. Our work is the first to investigate the potential of pruning in MBDL  Prompted by feedback, we revised our manuscript to better clarify our contribution. We revised the title of our paper to _“Efficient Model-based Deep Learning via Network Pruning”_ to better highlight our contribution.

Reviewer comment (Experiments):
Why was compressed sensing and super resolution selected as the tasks of choice? Showing more inverse problems, and including more networks can help establishing a stronger case for the paper.

Author response:
SPADE is applicable to any inverse problems and any MBDL. We considered CS-MRI and image super-resolution because they are two widely studied inverse problems. We agree that it is beneficial to validate SPADE on more inverse problems and MBDL networks. We can do it in our future work.

Reviewer comment (Novelty):
It is questionable that the proposed method is indeed for MBDL only. It could be seen as a mere combination of DepGraph and group sparsity based pruning.

Author response:
While pruning is a general idea, this topic has never been explored in the context of model-based deep learning (MBDL). MBDL architectures are inherently memory intensive, where pruning can have a big impact. Our work is the first to investigate the potential of pruning in MBDL  Prompted by feedback, we revised our manuscript to better clarify our contribution. We revised the title of our paper to _“Efficient Model-based Deep Learning via Network Pruning”_ to better highlight our contribution.

Reviewer comment (Experiments):
Experiments may not be as comprehensive as it should be: it is unclear if the sampling pattern in CS-MRI is fixed for all / super resolution was not evaluated on diverse test datasets, which were frequently used for super resolution literature.

Author response:
We used a uniform sampling pattern in CS-MRI. We conducted experiments on image super-resolution on a smaller testing set, simulating situations where only pre-trained models are accessible online and specific applications have limited data. The revised manuscript has been updated to clarify them.

Reviewer comment (Novelty):
Q2. It is unclear why the proposed three fine-tuning methods are novel: they are simply three popular losses that were used for full networks and there is no special treatment for pruned networks. Please clarify.

Author response:
We would like to highlight that our paper proposes _a novel application of structured pruning to MBDL networks_, an application _not explored in any previous studies_. While similar loss functions might have been proposed in other papers, those fine-tuning methods have never been investigated in the context of efficient MBDL networks via network pruning.

| Fine-tuning methods          | Need ground-truth? | Need trained original model?|
|:--------------------:|:------------:|:------------:|
| Supervised   | &#9745;          | &#9744;          |
| School    | &#9744;          | &#9745;          |
| Self-supervised        | &#9744;          | &#9744;          |

Reviewer comment (Experiments):
Q3. It may be true that "its potential has remained unexplored in the realm of imaging inverse problems," but I am not sure if exploring this potential is important considering that there are a number of drawbacks. For example, it is unclear if this work simulated diverse sampling patterns for CS-MRI, which is much more practical than using a single pattern. If different sampling patterns require a new pruning, then it may not be as convenient as using a full network. Moreover, it is unclear if 51-81% speed up is beneficial by sacrificing the performance by 0.77dB in CS-MRI, which may be critical for missing clinical information. More convincing arguments for the necessity on network pruning for inverse problems along with practical experiments should follow. For super resolution, 0.68dB is a huge gap and most prior works evaluated their methods in more diverse datasets.

Author response:
We would like to highlight that it is a trade-off in pruning between performance and network complexity. _One can tune this trade-off based on the specific application_. For example, in applications that are sensitive to performance drop, one can set the pruning ratio in CS-MRI to 20% such that the performance drop is mostly negligible. We will improve this trade-off in our future work. It is also worth mentioning that, in Figures 5 and 6, _the visual difference is unnoticeable_ between the images recovered by the unpruned model and SPADE for CS-MRI and super-resolution, respectively.

Reviewer comment (Novelty):
The contributions of the paper are mostly comprised of a combination of existing techniques such as the pruning algorithm and the fine-tuning techniques.

Author response:
The topic of pruning has never been explored in the context of model-based deep learning (MBDL). MBDL architectures are inherently memory intensive, where pruning can have an impact. Our work is really the first to investigate the potential of pruning in MBDL  Prompted by your remark, we revised our manuscript to better clarify our contribution. We revised the title of our paper to _“Efficient Model-based Deep Learning via Network Pruning”_ to better highlight our contribution.

Reviewer comment (Experiments):
Is it possible to combine fine-tuning losses, rather than view them as independent techniques, and could that help preserve performance?

Author response:
Thanks for your suggestions. We indeed have _experimented with combined fine-tuning losses before but did not find any improvement_. Our hypothesis is that the hierarchy of supervision quality—supervised, school, and self-supervised fine-tuning—may lead to one loss function dominating the supervision when combined. As an example, we formulated below a combined loss function with the addition of school loss and self-supervised loss.

$l_{combined}(\hat{\theta}) = l_{sc}(\hat{\theta}) + l_{ss}(\hat{\theta})$

We conducted experiments on the VarNet model with a 65% pruning ratio. The fine-tuning results 
below show that using combined loss does not have an advantage over solely employing school fine-tuning losses.

**Quantitative evaluation of three different losses: combined, school, and self-supervised in PSNR**
| Network | Pruning Ratio | School + Self-supervised | School | Self-supervised |
|:---------:|:---------------:|:--------------------------:|:--------:|:-----------------:|
| VarNet  | 0%            | 39.25                    | -      | -               |
| VarNet  | 65%           | **36.88**                   | 37.36  | 34.17           |

**Quantitative evaluation of three different losses: combined, school, and self-supervised in SSIM(%)**
| Network | Pruning Ratio | School + Self-supervised | School | Self-supervised |
|:---------:|:---------------:|:--------------------------:|:--------:|:-----------------:|
| VarNet  | 0%            | 97.7                     | -      | -               |
| VarNet  | 65%           | **96.6**                     | 96.7   | 95.3            |

Reviewer comment (Evaluation):
How does the memory complexity change at test-time? Memory complexity is also a quite important consideration for which discussion has not been included.

Author response:
Prompted by your comment, we evaluated the test-time memory complexity across different pruning ratios for DEQ, VarNet, and E2EVar networks. The results as shown below demonstrate a clear reduction in GPU memory usage as the pruning ratio increases.

| Network | Pruning Ratio | GPU Memory (MB) |
|:---------:|:---------------:|:-----------------:|
| DEQ     | 0%            | 3.82 (100.0%)   |
|         | 10%           | 3.37 (88.22%)   |
|         | 20%           | 3.04 (79.58%)   |
|         | 35%           | 2.44 (63.87%)   |
|         | 65%           | 1.37 (35.86%)   |
| VarNet  | 0%            | 76.37 (100.0%)  |
|         | 10%           | 68.06 (89.12%)  |
|         | 20%           | 61.16 (80.08%)  |
|         | 35%           | 47.96 (62.80%)  |
|         | 65%           | 26.77 (35.05%)  |
| E2EVar  | 0%            | 77.59 (100.0%)  |
|         | 10%           | 70.27 (90.57%)  |
|         | 20%           | 63.77 (82.19%)  |
|         | 35%           | 49.72 (64.08%)  |
|         | 65%           | 28.83 (37.16%)  |

Reviewer comment (Writing):
Typographical errors should be fixed via proofreading.

Author response:
The revised manuscript has done a better job of fixing the typographical errors.