Memory-Efficient Fine-Tuning via Structured Neural Network Pruning

Thanks for the valuable review!

1. We had included sparse fine-tuning (Lines 94–107), FISH-Mask (Line 104), and NN pruning (Lines 122–134) in the initial version. In the revision, we expanded the related work and introduction on the techniques of efficiently training with sparse matrices and we removed the description of “ours is the first work……”.
   For RoSA, we tested its performance under various settings but did not observe a significant advantage over LoRA. A comparison with RoSA (using mixed precision training with bfloat16 for RoSA) is included in Table 10 in Appendix D.6 of the revised paper. While we include a comparison with RoSA in the paper, we believe that this would not alter our main conclusion.
   The key insight we aim to emphasize is that replacing unstructured sparse matrices with sparse ones using our method results in more memory-efficient fine-tuning than the popular LoRA, without requiring additional implementations for sparse matrices. We have revised the paper to underscore this point more explicitly.
   We tested RoSA with $r=32$ and $d=1.2$% in full precision, following the settings suggested in the RoSA paper, and compared it to LoRA with r=64 under our fine-tuning process. Across seven evaluation datasets, RoSA's accuracies were similar to LoRA's—slightly better on some tasks and slightly lower on others. However, RoSA’s memory usage was significantly higher than LoRA’s under these settings.
   In a separate experiment, we attempted to fine-tune all linear layers of Llama-3-8B using RoSA but encountered out-of-memory (OOM) errors, requiring over 80GB of memory. Notably, under these settings, LoRA’s trainable parameters exceeded RoSA’s by 1M, yet LoRA’s peak memory usage during training was only 64.22GB. In another test with attention-frozen fine-tuning, RoSA's peak memory usage was 76.09GB, compared to LoRA’s 58.41GB. While RoSA's trainable parameters exceeded LoRA's by approximately 11M, this difference would contribute less than 0.2GB to LoRA's memory usage, highlighting a memory inefficiency of RoSA.
   RoSA relies on custom C++ implementations for its sparse matrix operations, complicating its application and making memory tracing difficult. The latest version of ‘spops’, released by the RoSA authors, failed to install on our device. We had to rely on an older version of ‘spops’ to execute RoSA. This discrepancy in software versions might partly explain the inconsistent conclusions presented in the RoSA paper regarding memory usage.
   Thanks for bringing the idea of ‘a combination of this paper’s pruning approach with a low-rank component’. We will leave the study of compatibility with other LoRA variants to be a future work, such as RoSA[1], QLoRA[2], and VeRA[3].
2. We added a footnote to the revision: “The Taylor importance here refers to computing the exact value without relying on approximations of the importance score or the gradient matrix used for deriving the importance score.”
3. Incorporating feedback from reviewer H135, we have included the results on Caltech101 in the revised version. This dataset is also an out-of-distribution (OOD) benchmark referenced in this paper [4].
4. Parameter dependency was commonly used in structured pruning (with some citation at line 127), this was the reason we mentioned this paragraph. We cited some papers of structured pruning at section 4.3 again in revision.
   Maybe there is some part unclear, $W_{\cdot j}^a$ represents the $j$-th input feature of layer $a$, while $W_{i \cdot}^b$ represents the $i$-th output feature of layer $b$, where the two are connected. In structured pruning, these features, $W^a_{\cdot j}$ and $W^b_{i \cdot}$, are pruned collectively. Similarly, in our approach, this dependency becomes fine-tuning these features together. For a visual representation of this relationship, please refer to Figure 2 in Appendix B.
5. In the revised version, we have mentioned the datasets on which the models are pre-trained. Additionally, we have reworded the paper to emphasize that "Without requiring additional implementations for sparse matrix operations, our novel SPF framework achieves memory efficiency while maintaining accuracies comparable to the popular LoRA."

[1] Nikdan, Mahdi et al. “RoSA: Accurate Parameter-Efficient Fine-Tuning via Robust Adaptation.” ArXiv abs/2401.04679 (2024)
[2] Dettmers, Tim, et al. "Qlora: Efficient finetuning of quantized llms." Advances in Neural Information Processing Systems 36 (2024).
[3] Kopiczko, Dawid Jan, Tijmen Blankevoort, and Yuki Markus Asano. "Vera: Vector-based random matrix adaptation." arXiv preprint arXiv:2310.11454 (2023).
[4] Kornblith, Simon et al. “Do Better ImageNet Models Transfer Better?” 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) (2018)

Thanks for the thorough review! We discuss the questions/comments raised in the “Weakness” section.

1. Thanks for pointing this out, we added clarification and did a thorough review in revision.
2. We selected these datasets because they are the standard ones used to benchmark fine-tuning. We also ensured that the size of datasets is small enough to train within a day on a single A100. For vision tasks, the main benefit of our approach is using significantly less trainable parameters, i.e. less FLOPs, to achieve the comparable results to dense fine-tuning. (For small-scale models, the models themselves are quite small, and most of the memory footprint comes from the datasets.) These image models have been pre-trained on ImageNet-1k, and the “pre-train then fine-tune” paradigm is often applied when the computation resources are limited. Given these facts, we didn’t fine-tune the models on ImageNet-1k.
3. We will try to add this to the final version. We were not able to do this within the rebuttal period given the timeline due to the limited computation resources. However, we've incorporated 7 commonly used datasets in the evaluation of Llama. ‘BoolQ’, ‘HellaSwag’, ‘WinoGrande’, ‘ARC-e’, ‘ARC-c’, and ‘OpenbookQA’ are included in the standard 8 zero-shot datasets, while ‘rte’ is also a challenging dataset. The experimental results of existing works, e.g. DoRA [1], show that Llama performs well on ‘PIQA’ and ‘SIQA’ by fine-tuning with most PEFT methods. So we exclude the two dataset and select other challenging datasets due to the limitation of computation resources.
4. In revision, we compared the results of our approach with rank 128 to the result of LoRA with rank 64, where the number of trainable parameters were similar. We also added the memory footprint for Llama and provided the exact number of trainable parameters but did not emphasize memory efficiency in small-scale model scenarios. In addition, we added the ablation study for different rank settings in Appendix D.
5. Maybe there was some part unclear in the initial version, we aimed to clarify that sparse fine-tuning can achieve even lower memory usage than the popular LoRA under comparable settings of trainable parameters by replacing unstructured sparse matrices with structured sparse ones using our approach. This is accomplished without requiring any complex implementations of forward passes, backward passes, or backpropagation, which often rely on custom C++ implementations to enable efficient sparse tensor operations. Consequently, we selected LoRA, the most widely used method, as our baseline for comparison.
   In revision, we attempted to include results for DoRA, a method that appears to be more widely adopted and validated than other LoRA variants. However, we encountered significant computational challenges. We included the memory usage of DoRA and RoSA[2] (see reviewer 3hdW) in Appendix D.2. Briefly, DoRA’s memory usage and fine-tuning time were substantially higher than those of LoRA. For example, in our experiments with Llama-3, DoRA consistently required over 80GB of memory for training and exhibited training times approximately 50%-100% longer than LoRA. Given our focus on highlighting the computational efficiency of our approach, we have opted to limit our comparisons to LoRA.
6. We added Figure 3 in appendix D4. We re-wrote section 5.4 and added appendix D to describe more on the memory benefit of our approach in detail.

[1] Liu, Shih-Yang, et al. "Dora: Weight-decomposed low-rank adaptation." arXiv preprint arXiv:2402.09353 (2024).
[2] Mahdi Nikdan, Soroush Tabesh, Elvir Crn ˇcevi ´c, and Dan Alistarh. RoSA: Accurate parameter- efficient fine-tuning via robust adaptation. In Forty-first International Conference on Machine Learning, 2024. URL https://openreview.net/forum?id=FYvpxyS43U

We thank the reviewer for the detailed comments and clarifications.

1.1. We added the memory footprint in tables of Llama (please see the revision). For image models and debertav3, we do not emphasize the memory footprint because the models themselves are quite small, and most of the memory footprint comes from the datasets. The main benefit of our approach is using significantly less trainable parameters, i.e. less FLOPs, to achieve the comparable results to dense fine-tuning.
1.2. For Llama, we compared the results of our approach with those of LoRA’s with (approximately) the same number of trainable parameters. (The “approximate” is because the rank parameter is chosen as a power of 2 for convenience.) We will highlight the fact that “even then the trainable parameters of our approach is 3 times that of LoRA, the memory usage of our method is smaller than that of LoRA.” This is because the memory usage consists of (a) the “base” memory needed to hold the model, (b) the parameters to be trained, (c) auxiliary memory for storing gradients, etc., (d) memory corresponding to intermediate layer activations. Even with (a) and (b) being the same and using the same optimizer, we see our method outperforming LoRA by a significant amount from (c) and (d).
1.3. “Curve of performance vs. parameter count/memory consumption”: thanks for the suggestion – we add this to the revision. We add Appendix D for the detailed study of memory footprints and rank settings.
2.(a) While we made efforts to position all tables and figures at the top of the pages in the revised version, some had to be placed differently to optimize the use of space in the main text.
2.(b) Regarding the matrix $\boldsymbol{M}$, this is to make our notation consistent with other sparse fine-tuning and LoRA papers. In Eq.6, we replace the notation $F-\boldsymbol{c}_i$, see the revision.

Thanks for the comments and suggestions.

1. (Layout, abbreviations): We will make these changes; the spacing was an unfortunate typesetting issue.
2. (More extensive experiments): Following the suggestion, we have done experiments on Caltech-101 and included the results in the updated draft. Notice that while our experiments on image datasets are somewhat limited, we have done much more extensive tests on language datasets. Given that many applications focus on fine-tuning language models like LLaMA, we focused more on them.
3. Comparison with LoRA-Prune: note that this is not actually a paper on fine-tuning; instead, their goal is a better compression procedure for large models. That said, there are other methods such as RoSA (see also the response to Reviewer 3hdW), that we have recently compared to. (In short, our procedure maintains accuracy, while not requiring any sparse matrix libraries, and using a considerably smaller amount of memory.)
4. Novelty: Maybe there is some part unclear, we aim to improve the memory challenge of sparse fine-tuning (SFT), which is a parallel line of LoRA and its variants. The key challenge of many SFT is that the unstructured sparse matrix requires additional implementation for forward pass, backward pass, and backpropagation computation (e.g. [1], [2], [3], [4], [5]). This often involves optimizing tensor computations by selectively processing only non-zero elements, e.g. torch.sparse [3], compressed sparse column/row (CSC/CSR) [4], semi-structured formats [5] The tradeoff in these approaches lies in the fact that they all increase time complexity to achieve reduced memory complexity. Therefore, some approaches also leverage C++ for acceleration, as seen in works like [1], [2]. (see their implementation: [6]) This complicates the practical application of these methods and increases the difficulty of further advancing this field. Our method obtains a memory- and parameter- efficiency without any additional implementations of sparse operations in forward pass, backward pass and back-propagation. We added paragraphs in section1 (line 60-68, 87-89) and section 2 (line 123-137) to more clarify the challenge of sparse fine-tuning. We also added some words at the research questions (line 69-72).
   Additionally, LoRA and its variants leverage matrix decomposition techniques, which are not directly applicable to vector parameters such as LayerNorm and BatchNorm. In contrast, our approach identifies a subset of parameters for fine-tuning by focusing on the most important neurons or channels, making it compatible with vector parameters. We believe this flexibility is another key strength of our work.
   By comparing our method with LoRA, we aim to demonstrate that structured sparse fine-tuning is a straightforward and memory-efficient solution. It relies solely on standard tensor operations available in regular libraries, all while maintaining competitive performance.

[1] Mahdi Nikdan, Tommaso Pegolotti, Eugenia Iofinova, Eldar Kurtic, and Dan Alistarh. SparseProp: Efficient sparse backpropagation for faster training of neural networks at the edge. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 26215–26227. PMLR, 23–29 Jul 2023. URL https://proceedings.mlr.press/v202/nikdan23a.html
[2] Mahdi Nikdan, Soroush Tabesh, Elvir Crn ˇcevi ´c, and Dan Alistarh. RoSA: Accurate parameter- efficient fine-tuning via robust adaptation. In Forty-first International Conference on Machine Learning, 2024. URL https://openreview.net/forum?id=FYvpxyS43U
[3] PyTorch is actively developing tools for sparse matrices, with the beta version available at https://pytorch.org/docs/stable/sparse.html
[4] Mohammad Hasanzadeh Mofrad, Rami Melhem, Yousuf Ahmad, and Mohammad Hammoud. Multi-threaded layer-wise training of sparse deep neural networks using compressed sparse column. In 2019 IEEE High Performance Extreme Computing Conference (HPEC), pp. 1–6. IEEE, 2019.
[5] Connor Holmes, Minjia Zhang, Yuxiong He, and Bo Wu. Nxmtransformer: semi-structured sparsifi-cation for natural language understanding via admm. Advances in neural information processing systems, 34:1818–1830, 2021.
[6] https://github.com/IST-DASLab/spops