Stochastic Fine-Tuning of Language Models Using Masked Gradients (2024)

Abstract: Large Language Models (LLMs) have emerged as the dominant paradigm in Natural Language Processing owing to their remarkable performance across various target tasks. However, naively fine-tuning them for specific downstream tasks often requires updating a vast number of parameters, resulting in high computational costs and overfitting when training data is limited. In this paper, we propose a novel approach, called *Stochastic Tuning*, that addresses these challenges by selectively updating a small subset of parameters in each step of the tuning process. Our approach is characterized by its customization of updates based on task-specific partial gradients with respect to stochastic sub-networks. The advantage of Stochastic Tuning over existing solutions lies in its ability to consider both parameter weights as well as forward values which guarantees a context-sensitive fine-tuning. Our experiments demonstrate that Stochastic Tuning outperforms existing lightweight fine-tuning methods, improving average performance by over two points on RoBERTa across several tasks in the GLUE benchmark while updating merely **0.08**% of the model’s parameters. The code for our implementation can be found at https://github.com/m-Tajari/StocTuning_LLMs.

## introduction and related work full fine-tuning of a pre-trained language model (plm), a widely adopted approach in modern nlp, can be computationally expensive due to the need for updating all parameters of the model. by limiting the number of updatable parameters during the fine-tuning process, lightweight mechanisms reduce the computational cost by a large factor while retaining the performance. a class of lightweight methods augments models with small trainable modules. adapters @xcite are a prominent technique in this category. they inject a few modules into the transformer blocks, enabling comparable performance to the full fine-tuning scenario. instead of introducing new modules, prefix-tuning (li and liang, 2021) optimizes virtual tokens, called the prefix, which are prepended to the normal activation vectors of transformers. this method excels in low-data settings. in contrast to both these techniques, low-rank adaptation @xcite avoids adding new parameters. instead, it freezes the pre-trained ones and optimizes low-rank weights inserted into existing layers, achieving better results across various tasks. however, a drawback of all these methods is that they still require extra parameters to be injected into already large models. according to the lottery ticket hypothesis @xcite , large transformer-based models consist of sparse sub-networks, fine-tuning of which results in competitive performance to the full fine-tuning. various studies have tried to find optimal sub-networks which effectively transfer knowledge from a pre-trained model to different downstream tasks. notably, bias-term fine-tuning @xcite freezes all the transformer-encoder parameters but the biases. akbar-tajari et al. (2022) take a step further to generalize this approach by fine-tuning similar transformer modules across layers, including layernorms, the expressive power of whose variants has been theoretically studied @xcite . the results suggest that each transformer module can act as a winning ticket due to its ability in effectively transferring knowledge. however, while these localized fine-tuning techniques achieve acceptable performance across a variety of tasks, they rely too much on tuning predefined modules and neglect the importance of adapting parameters based on the specific requirements of the task at hand. this limitation prevents more finegrained, task-specific parameter updates, which are crucial for optimal performance on diverse downstream tasks. another branch of research takes a non-localized approach. these methods form sub-networks where any parameter of the model can be updated, without considering their role in different transformer modules. influenced by dropout @xcite , mixout @xcite ) regularizes fine-tuning process by randomly freezing subnetworks during sequential iterations, ignoring the significance of parameters. in a more efficient way, dynamic parameter selection @xcite adaptively selects promising sub-networks composed of important parameters. despite being dynamically selected, these sub-networks usually contain a substantial portion of model's parameters. in addition, dps only utilizes gradients of backpropagation without directly considering the value of parameters, which have successfully been used in similar pruning strategies @xcite . to mitigate these issues, we propose stochastic tuning, a method that estimates the importance of individual parameters by using gradients of back-propagation and the value of parameters together with forwarded values. based on the importance scores, our approach stochastically forms a task-specific binary mask of a predefined size. the optimization process is then constrained to sub-networks of parameters with highest importance scores. stochastic tuning utilizes both taskspecific data and the encoded knowledge of pretrained models to select sub-networks. this approach enables plms to seamlessly adapt to diverse tasks while retaining their inherent generalization capabilities. moreover, the randomness in our method acts as a regularizer; therefore, it can prevent models form overfitting, bringing about better generalization. stochastic tuning provides the following three key advantages compared to existing solutions: (1) it utilizes the encoded knowledge of plms by taking into account the value of parameters in the masking process, which ultimately reduces the time required for the fine-tuning process; (2) it tailors its selection of update sub-networks to the specific downstream task at hand; and (3) the stochasticity in the selection of the binary mask ensures that less important parameters also have the opportunity to participate in the fine-tuning process, bringing about regularization effect. our experiments on the glue benchmark @xcite demonstrate that stochastic tuning yields consistent improvements of about two points on average over previous state-of-the-art fine-tuning methods. for @xmath0, 1, • • • , t -1} and b t ⊆ d. note that equation (1) makes use of learning rate η ∈ r + and in-batch training loss: l(θ t , b t ). fine-tuning sub-networks corresponds to imposing a restriction on the number of parameters that are modified in each step of sgd. under these circumstances, the optimization problem can be formulated as: ## experiments we assess the performance of stochastic tuning on roberta base and roberta large @xcite . 5 following previous work @xcite @xcite , we undertake an extensive set of experiments on six datasets from the glue benchmark. to provide a thorough analysis, we report the average and standard deviation of the results obtained from seven models trained with distinct random seeds. we compare a range of fine-tuning methods for language models, including but not limited to dps mix @xcite , which dynamically selects task-relevant sub-networks; multi-head (akbar-tajari et al., 2022), which fine-tunes only attention module parameters; and bitfit (ben zaken et al., 2022), which updates only the bias terms. detailed information on datasets and prior methods is provided in appendices. ## conclusions we introduce stochastic tuning, a novel approach for highly efficient fine-tuning of pre-trained language models (plms). by utilizing masked gradients to update stochastic sub-networks, our method outperforms previous state-of-the-art fine-tuning techniques in terms of overall performance across multiple glue tasks. notably, we achieve this improvement by updating parameters solely within a small sub-network during each iteration, resulting in substantial computational cost reduction. the inherent randomness in our method acts as a form of regularization, effectively mitigating overfitting and consistently promoting better generalization. our results demonstrate the significant performance gains achieved with stochastic tuning on two distinct plms. we anticipate that stochastic tuning of plms holds promise for a wider range of application scenarios like multimodal model finetuning @xcite , leaving room for further exploration for future research endeavors. in line with the methodological choices of previous work, our research did not include any generative models. however, it is worth noting that future studies could broaden the scope of comparative analysis by incorporating fine-tuning approaches using generative models such as gpt @xcite , bart @xcite , and t5 @xcite . such an expansion would offer a more comprehensive understanding of the capabilities and trade-offs of different model architectures in natural language processing tasks.