T oken D rop + B ucket S ampler: Towards Efficient Padding-free Fine-tuning of Language Models (2023)

Abstract: The great success of Language Models (LMs) for various Natural Language Processing (NLP) tasks is accompanied by computational challenges during both pre-training and fine-tuning. Pre-training has attracted significant attention due to its huge computational footprint. We focus on the fine-tuning of pre-trained LMs, which is expected to be performed much more frequently as the pre-trained models are adapted to downstream tasks. During fine-tuning, the presence of variable-length input sequences necessitates the use of padding tokens when batching sequences. These padding tokens lead to ineffectual computations, adversely impacting the efficiency of fine-tuning. We also observe that LMs memorize the limited task-specific training data despite the use of known regularization methods. Based on these insights, we present TokenDrop + BucketSampler, a framework that simultaneously improves efficiency and accuracy of LM fine-tuning. BucketSampler generates batches of samples with lower variance in sequence lengths to reduce the number of padding tokens, but does so without the accompanying accuracy drop seen in previous approaches. TokenDrop is a new regularizer that prunes a random subset of insignificant tokens from each input sequence in every epoch to prevent overfitting. TokenDrop drops more tokens from the longer sequences in each batch to further reduce variance in input lengths and the need for padding. TokenDrop + BucketSampler accelerates fine-tuning on diverse downstream tasks by up to 10.61X, while also producing models that are up to 1.17% more accurate compared to conventional fine-tuning. Code is available at https://github.com/amrnag/TokenDrop-BucketSampler. .

## introduction language models (lms) derived from attentionbased transformer networks have significantly advanced the field of natural language processing (nlp), with recent models showing remarkable performance on a wide range of tasks @xcite . lms are typically characterized by large model sizes, and are pre-trained on very large text corpora. pre-trained lms are subsequently finetuned to solve a range of downstream tasks. while pre-training is the most expensive step by far in lm creation, and hence has attracted the most attention, we observe that the relatively high frequency of fine-tuning makes it an important challenge in its own right. data available from public-domain pretrained lms such as bert @xcite and opt @xcite suggests that these lms have been fine-tuned millions of times by different users on a diverse range of downstream tasks. further, lms are notoriously sensitive to the initialization of the task-specific final layer and the order in which training data is presented during fine-tuning @xcite . as a result, multiple runs of fine-tuning with different hyperparameters are often required to achieve high accuracy on a given downstream task. finally, due to large model sizes of lms (that continue to grow), even a single fine-tuning run is compute-intensive and can take several gpu-days. in summary, while the enormous computational costs of pre-training have attracted a lot of attention, it is likely that comparable if not more compute time and energy have been spent on fine-tuning compared to pre-training since the inception of lms. while several prior efforts have focused on reducing the costs of pre-training, relatively little attention has been paid to the computational challenges of fine-tuning. we observe that lm finetuning presents a unique set of challenges. first, fine-tuning is performed on variable-length text sequences, with significant spread in lengths. when batches of variable-length sequences are generated for fine-tuning, shorter sequences in a batch are "padded" to the length of the longest sequence in the batch by adding padding tokens. however, computations on padding tokens are useless, and adversely affect throughput during fine-tuning. second, fine-tuning is performed in a supervised manner and requires expensive human annotations. as a result, fine-tuning datasets are several orders of magnitude smaller than pre-training datasets. in addition, when overparameterized lms are trained on small task-specific datasets, overfitting leads to sub-optimal generalization performance (fig. 1 ), even with the use of popular regularizers such as dropout @xcite . consequently, fine-tuning is performed over a very small number of epochs (@xmath0. finally, we observe that fine-tuned lms are adversely impacted by even minor grammatical errors in inputs during inference when only grammatically perfect sequences are used for fine-tuning (fig. 2 ). this sensitivity is problematic since the assumption of seeing only grammatically correct sequences during inference does not always hold in real-world scenarios, especially when lm inputs are provided by users with different levels of language proficiency. to address the aforementioned challenges, we present tokendrop + bucketsampler, the synergistic combination of a new regularizer and batching method for accurate and efficient lm fine-tuning. our first contribution, bucketsampler, generates batches of samples with reduced spread in sequence lengths, thereby reducing computational overheads due to padding. in addition, the batch size is varied based on the lengths of the samples in the batch to maximize hardware utilization. improved batching strategies have been incorporated into popular nlp libraries such as huggingface transformers @xcite and lingvo @xcite . however, these methods are optimized for training lms from scratch. when applied to fine-tuning on small datasets for very few epochs, we find that these prior batching methods lead to significant drop in accuracy. bucketsampler includes key algorithmic enhancements that tune the optimal batch size and learning rate schedules, thereby enabling fine-tuning to achieve high accuracy while also maintaining high hardware utilization. our second contribution is tokendrop, a novel regularizer that identifies and drops a random subset of insignificant tokens from each sequence in every epoch. more tokens are dropped from the longer sequences in each batch, further reducing the need for padding (fig. 3 ). further, tokendrop reduces overfitting by ensuring that the model does not see the same sequence repeatedly over the course of fine-tuning. as a side effect, it also improves the resilience of lms to grammatical errors during inference by exposing the model to grammatically incorrect inputs during fine-tuning. since tokendrop + bucketsampler improves fine-tuning efficiency by eliminating ineffectual computations, it can be combined with previously proposed approaches for parameter-efficient finetuning, such as freezing layers @xcite and the use of adapters @xcite , to achieve further efficiency gains. we summarize our main contributions as follows: ## method tokendrop + bucketsampler is a framework for accurate and efficient fine-tuning of pre-trained lms. the first component, tokendrop, is a regularizer that randomly drops a subset of insignificant tokens from each sequence in every epoch. the second component, bucketsampler, is a length-aware batching method that generates batches of samples with lower spread in sequence lengths. we demonstrate how tokendrop and bucketsampler can be synergistically combined (hence the name tokendrop + bucketsampler) to simultaneously reduce overfitting and eliminate the need for padding tokens, resulting in both accuracy and speed improvements (fig. 3 ). ## related work prior works have developed batching strategies for variable-length inputs to improve the efficiency of lm training. randomsampler is the most commonly used batching technique, and is the default method in most nlp libraries. randomsampler randomly selects samples from the training dataset to generate batches, and pads all sequences in a batch to the maximum length of sequences in the batch. consequently, training with random-sampler requires substantial padding, and hence, substantial wasted computations. lengthgrouped-sampler (introduced in huggingface transformers however, we find that unlike when training from scratch, the overheads of packing and the additional computations in attention layers (computing irrelevant scores, followed by masking to prevent cross-contamination between different sequences) are not amortized over the small number of fine-tuning epochs. finally, tensorflow's tf.data.bucket_by_sequence_length divides the training dataset into buckets based on sequence length, and generates batches only from sequences in the same bucket, similar to bucketsampler. while packing and tf.data.bucket_by_sequence_length support variable batch sizes to maximize hardware utilization, they incur accuracy drop during fine-tuning. in summary, while prior works have demonstrated efficiency gains when training lms from scratch, the unique characteristics of fine-tuning (small datasets, very few epochs) make these methods ineffective. as a result, tokendrop + bucketsampler significantly outperforms prior methods in terms of both accuracy and fine-tuning efficiency ( ## conclusion in this work, we presented tokendrop + bucket-sampler for accurate and efficient fine-tuning of lms. tokendrop prunes a random subset of stopwords in each sequence in every epoch to reduce