S parse G rad: A Selective Method for Efficient Fine-tuning of MLP Layers (2024)

Abstract: The performance of Transformer models has been enhanced by increasing the number of parameters and the length of the processed text. Consequently, fine-tuning the entire model becomes a memory-intensive process. High-performance methods for parameter-efficient fine-tuning (PEFT) typically work with Attention blocks and often overlook MLP blocks, which contain about half of the model parameters. We propose a new selective PEFT method, namely SparseGrad, that performs well on MLP blocks. We transfer layer gradients to a space where only about 1% of the layer’s elements remain significant. By converting gradients into a sparse structure, we reduce the number of updated parameters. We apply SparseGrad to fine-tune BERT and RoBERTa for the NLU task and LLaMa-2 for the Question-Answering task. In these experiments, with identical memory requirements, our method outperforms LoRA and MeProp, robust popular state-of-the-art PEFT approaches.

## introduction due to the tendency to increase the size of transformer models with each new generation, we need efficient ways to fine-tune such models on downstream task data. the usual practice is fine-tuning a large pre-trained foundational model on a downstream task. the major problem that prevents efficient fine-tuning is a steady increase in the memory footprint. one of the best strategies is high-performance methods for parameter-efficient fine-tuning (peft). typically, such methods as lora @xcite focus on attention blocks and do not consider dense mlp blocks. since mlp blocks can take a significant fraction of the model parameters (see ## related work in the last few years, many approaches to peft have appeared. @xcite distinguishes three types of methods: additive, reparametrization-based, and selective. in additive peft, small neural networks called adapters are added to the main model to steer the outputs of its modules @xcite . adapters are trainable, therefore, the main model remains unchanged. @xcite adapt this approach to nlp. in reparametrization-based approaches low-rank representations of trainable parameters are used. for example, lora @xcite parameterizes the weight update by a trainable low-rank matrix decomposition. in the original paper, lora is applied to self-attention modules, but not to mlp ones. in the selective methods, parts of the model or sets of the parameters are chosen for fine-tuning using some heuristics. such methods include, for example, bit fit @xcite or meprop @xcite , where only top-k parameters are updated during backpropagation. the approach proposed in this paper is related to selective methods. ## method our aim is to reduce the amount of trainable parameters at the fine-tuning stage. taking into account that fine-tuning data is restricted to a limited scope, we assume there is a basis where the weight gradient matrix is very close to being sparse. to identify this basis, we applied a decomposition technique to the stacked weight gradient matrices. as a result, we introduce a new pytorch layer class, sparsegradlinear, which transitions weights to this sparse gradient space, accumulates gradients in sparse form, and enables the reverse transition back to the original space. ## time and memory consumption per training iteration we measure the peak memory allocated during training using the cuda memory allocator statistics. in terms of training time, lora demonstrates the fastest training, followed by sparsegrad, and then standard fine-tuning. ## experiments we conducted experiments on three transformerbased encoder models, bert and roberta base and large, on the glue @xcite benchmark, and the llama-2 decoder model on the openassistant conversations corpus @xcite . we compared the fine-tuning of the full model (regular ft scheme) with three peft methods, namely lora, meprop and sparsegrad, applyed to mlp blocks. to harness lora, we use an official repository code. for the meprop method, we kept the largest elements in the ∂l ∂w matrix. the proposed sparsegrad involves replacing layers in mlp blocks with its sparsegrad-linear equivalents. ## conclusion we propose a new selective peft method called sparsegrad, which identifies a space where the gradients exhibit a sparse structure and updates only its significant part. sparsegrad is validated through experiments conducted on the bert, roberta and llama-2 model models, demonstrating its superiority over the additive lora and selective meprop methods. leveraging the sparsity property significantly accelerated the calculations in sparsegrad. our method runs faster than standard fine-tuning but slower than lora, while yielding better performance than lora; the same trend applies to memory usage. in summary, our method serves as an alternative to lora in situations where the performance of the final model takes precedence over the execution time. the source code as well as links to pretrained models are available at repository. foot_1