Make Pre-trained Model Reversible: From Parameter to Memory Efficient Fine-Tuning

Summary

>**Rebuttal:** The provided details satisfy my concerns. I think this paper should be accepted after applying the agreed changes. >**TL;DR:** The proposed MEFT reduces training activation memory, which is an important topic for the ML community. MEFT outperforms existing PEFT techniques while consuming up to 84% less memory compared to full fine-tuning. However, MEFT comes with significant computation overhead, which is not explored. Furthermore, some key comparisons, such as gradient checkpointing, are missing. Addressing my concerns and questions would improve my score. This paper proposes the Memory Efficient Fine-Tuning (MEFT) technique that achieves similar accuracy to that of traditional PEFT techniques while consuming up to 84% less memory compared to full fine-tuning. The paper discusses previous PEFT techniques and provides insights on key attributes for them to work. These insights include empirical results and detailed explanations. One of these key insights is the weights initialization. The paper shows the importance of initializing the PEFT weights such that they do not have any effect on the underlying LLM before starting the training process. The paper proposes to combine reversible models with PEFT existing techniques to create MEFT algorithms that gain from both worlds. The paper provides three different variants of MEFT, which are explained as well as illustrated. The three variants of MEFT are empirically validated on four different architecture backbones BERT, RoBERTa, BART and OPT. The MEFT variants are compared to several popular PEFT techniques as well as some memory efficient techniques. These comparison are conducted on the GLUE dataset.

Strengths

* **S.1.** The proposed MEFT algorithm tackles an important problem in existing PEFT techniques, which makes LLM PEFT training more accessible to researchers with low resources by reducing the activation memory up to 84% of full fine-tuning. * **S.2.** The paper provides interesting insights that are explained and backed by results and detailed illustrations. * **S.3.** The proposed MEFT variants outperform existing PEFT techniques on GLUE dataset with several architecture backbones. * **S.4.** Reproduction code is provided as part of the submission.

Weaknesses

* **W.1.** The training time of MEFT is not well explored and is only noted as "double the training time". Hence, the tradeoff between memory and computation efficiency is not well understood. * **W.2.** Popular existing memory efficient training techniques such as tensor rematerialization (gradient checkpointing) [2][3] and ZeRO [1] are not compared to. These techniques sacrifice training speed for reduction in memory consumption, similar to the goal of MEFT. Furthermore, these techniques are compatible with existing PEFT approaches. Finally, these techniques are widely available in deep learning frameworks such as PyTorch, DeepSpeed, and Hugging Face Accelerate which makes comparison relatively simple. * **W.3.** The font sizes in Figures 1,2,3,5, and 6 are too small which makes it very difficult to read. Furthermore, some figures are not self explanatory, such as Figure 5 which does not provide any explanation for the difference between the solid and dotted lines. These problems burden the reading flow of the paper. **Typos.** * Line #86: "for obtaining" → "to obtain" * Line #86: "though we" → "even though we" * Line #95: "in Figure 2a" → "in Figure 2a," * Line #97: "though with" → "even with" * Line #106: "from an modified" → "from a modified" * Line #114: "the added and" → "the added" [1] Ren, J., Rajbhandari, S., Aminabadi, R.Y., Ruwase, O., Yang, S., Zhang, M., Li, D. and He, Y., 2021, July. ZeRO-Offload: Democratizing Billion-Scale Model Training. In USENIX Annual Technical Conference (pp. 551-564). [2] Jain, P., Jain, A., Nrusimha, A., Gholami, A., Abbeel, P., Gonzalez, J., Keutzer, K. and Stoica, I., 2020. Checkmate: Breaking the memory wall with optimal tensor rematerialization. Proceedings of Machine Learning and Systems, 2, pp.497-511. [3] Beaumont, O., Eyraud-Dubois, L. and Shilova, A., 2021. Efficient combination of rematerialization and offloading for training dnns. Advances in Neural Information Processing Systems, 34, pp.23844-23857.

Questions

* **Q.1.** The paper proposes MEFT which sacrifices computation speed at the cost of lower peak memory. There are several existing common approaches (such as gradient checkpointing and DeepSpeed) for general memory efficient training which are compatible with PEFT techniques. Why are these comparisons not explored or detailed? * **Q.2.** In the LoRA modified initialization method, both $W_{up}$ and $W_{down}$ are initialized with a single constant. This causes unwanted symmetries, which makes it more difficult to learn. An alternative approach would be to initialize $W_{down}=1$ (as tested in this paper, but instead of initializing $W_{up}$ to a constant, initialize it with a random gaussian distribution where $c$ can act as the distribution mean. Does this make sense? * **Q.3.** In Line #175 it is noted that the adapter weights' $W_{down}$ and $W_{up}$ are initialized from a random gaussian distribution with zero mean. It is further noted, that the outputs of the adapter would be close to zero. However, in practice this does not hold true, and the actual outputs can significantly deviate from the desired zero output. What am I missing here? * **Q.4.** In Figure 6 it seems that LoRA only reduces slightly the peak memory in comparison to the full fine-tuning. However, the optimizer states, which are generally 3 times the size of the model, are significantly smaller in LoRA. Does this result hold for other LLMs and how does the batch size affect it as well? * **Q.5.** The MEFT technique does not seem to be only compatible with LLMs and can be compatible with other domains, as sted in the conclusion section. However, the experimental results are solely on LLMs. What is the reason for that? * **Q.6.** In Table 2 the memory usage of Y-Tuning is missing, why is that?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The limitations are discussed in paragraph line #315. These limitations include limited compatibility with quantization techniques due to training instabilities and significantly longer training time due to activation re-computation.

Summary

This paper points out that the existing parameter-efficient fine-tuning (PEFT) methods are not memory-efficient, since they need to cache most of the intermediate activations for the gradient calculation. To resolve this problem, the authors introduce to use the reversible model to reduce the activation memory. Specifically, this paper have these contributions: 1) it is important to preserve the weight of PLM at the initialization; 2) insert adapters and make it reversible without additional pre-training. Extensive results also demonstrate the effectiveness of the proposed method.

Strengths

1. This paper points out the importance of memory-efficient fine-tuning in current pre-trained models. And the authors introduce the idea of RevNet into the pre-trained language model, and design three reversible variants to PLM.

Weaknesses

1. Some conclusions seem naive. For example, preserving the starting point from the pre-trained model at the beginning of training is a common experience. Previous works like LoRA also point out the importance of this point, and this point is also a common method to avoid catastrophic forgetting. 2. Memory-efficient training is important, but the authors do not mention quantization in this paper. Quantization has been considered as the one of most effective methods in LLM to save memory. While for RevNet or the proposed architecture, these reversible-based neural networks actually are a trade-off between computation and memory (it will bring more computation to save memory). So, I do not think the proposed method is advantageous to Quantization and directly comparing the proposed methods with PEFT may also not be appropriate. 3. Essentially, the proposed method adopts an adapter structure and then uses the reversible structure into the transformer and presents three patterns (which formulate different $\cal{F}$ and $\cal{G}$) and trains each network from the start point. The idea is to combine existing work (RevNet) with current PEFT methods, and then combine the initialization that preserving the starting point from the pretrained model. So, I think the contributions are limited. 4. In experimental results, authors report MEFT with vanilla gradient. I think such reports are confusion, since the main contribution in this paper is the reversible gradient with MEFT structure.

Questions

1. This paper presents three variants of MEFT, but is there any conclusion that which one should be chosen under different scenarios?

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

The authors have presented limitations of their proposed works, and currently, these limitations are still a problem for the proposed methods. Besides, I also give some comments about the limitations of this paper (please see my comments in **Weaknesses**). I do not see any potential negative societal impact.

Summary

Authors propose a memory-efficient (MEFT) variant of parameter-efficient fine-tuning (PEFT), which utilizes reversible networks to save memory stored for activations during model forward propagation. The authors showed insights into the reasons for the effectiveness of PEFT. They evaluated several variants of MEFT with GLUE and QA datasets, indicating that the proposed method does not produce performance gaps while reducing memory usage.

Strengths

- The proposed method shows competitive results with benchmarks - Experiments with the initialization of LoRa are insightful - Generally, the proposed methods are useful for practitioners

Weaknesses

- The presentation could be improved. E.g., I struggled with Section 2.3 since $h^1, h^2 $ (L136) was not defined above. - It would be fascinating to compare the proposed method with gradient checkpointing used for training with LoRA/etc. Naively, the usage of invertible layers should be faster. However, seeing such an evaluation in the paper would be highly beneficial.

Questions

Please refer to the weaknesses section.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

–

Response to Rebuttal

Thank you for the detailed answers and results. The provided results and details satisfy my concerns. I will update my review accordingly.

Thank you for the increased score

We are really encouraged by the increased score. And thank you for the advice on incorporating the comparison to gradient checkpointing, which makes our paper more solid.