TokenDrop + BucketSampler: Towards Efficient Padding-free Fine-tuning of Language Models
Authors: Amrit Nagarajan
Anand Raghunathan
Venue: EMNLP 2023 Findings (2023)
Keywords: Efficient training of LLMs
Efficient fine-tuning of LLMs
Padding-free variable sequence length batching
Token pruning
Regularizer for LLM training
TokenDrop
BucketSampler

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.

.

TokenDrop + BucketSampler: Towards Efficient Padding-free Fine-tuning of Language Models
Amrit Nagarajan School of ECE, Purdue University a edu
Anand Raghunathan School of ECE, Purdue University raghunathan@purduee edu
Abstract
(NLP), with recent models showing remarkable per- formance on a wide range of tasks (Bubeck et al., 2023). LMs are typically charactcrized by large model sizes, and are pre-trained on very large text corpora. Pre-trained LMs are subsequently fine- tuned to solve 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 pre- trained LMs such as BERT (Devlin et al., 2019) and OPT (Zhang et al., 2022) suggests that these LMs have been fine-tuned millions of times by dif- ferent users on a diverse range of downstream tasks. Further, LMs are notoriously sensitive to the ini- tialization of the task-specific final layer and the order in which training data is presented during fine-tuning (Dodge al., 2020). As a result, multi- ple runs fine- fine-tuning with different hyperparam- eters are often required to achieve high accuracy on given downstream task. Finally, due to large model sizes of LMs (that continue to grow), even 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 lot of attention, likely that compara- ble not more compute time and energy have been spent on fine-tuning compared to pre-training since the inception of LMs.
The great success of Language Models (LMs) for various Natural Language Processing (NLP) tasks is accompanied by computational chal- lenges during both pre-training and iin-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 downstream tasks. During fine- tuning, the presence of variable-length input sequences necessitates the use of padding to- kens when batching sequences. These padding tokens lead to ineffectual computations, ad- versely impacting the efficiency of fine-tuning. We also observe that LMs memorize the lim- ited task-specific training data despite the use of known regularization methods. Based these insights, we present Trop + BucketSam- pler, a framework that simultaneously improves efficiency and accuracy of LM fine-tuning, BucketSampler generates batches of samples with lower variance sequence lengths re- duce the number of padding tokens, but does so without the accompanying accuracy drop seen in previous approaches. TokenDrop is new regularizer that prunes random subset of in- significant tokens from each input sequence in every epoch prevent overfitting. TokenDrop drops more tokens from the longer sequences each batch further reduce variance in input lengths and the need for padding. TokenDrop + BucketSampler accelerates fine-tuning on di- verse downstream tasks by up 10.61 while also producing models that are up to 1.17% more accurate compared to conventional fine- tuning. Code is available https:: //gi thub. com/ Conm/amrnaa/ToeeDrrog BecoopppBuukttttee Sampler.
While several prior efforts have focused on re- ducing the costs of pre-training, relatively little attention has been paid to the computational chal- lenges of fine-tuning. We observe that LM fine- tuning presents a unique set of challenges. First, fine-tuning is performed on variable-length text quences, with significant spread in lengths. When batches of variable-length sequences are generated for fine-tuning, shorter sequences in batch are "padded" to the length of the longest sequence in the batch by adding padding tokens. However,
1 Introduction
Language Models (L.Ms) derived from attention- based Transformer networks have significantly ad- vanced the field of Natural Language Processing

---

computations on padding tokens are useless, and adversely affect throughput during fine-tuning. Sec- ond, fine-tuning is performed in supervised man- ner and requires expensive human annotations. As
95
Accuracy
85
a result, fine-tuning datasets are several orders of magnitude smaller than pre-training datasets. In addition, when overparameterized LMs are trained small task-specific datasets, overfitting leads to sub-optimal generalization performance (Fig. 1), even with the use of popular regularizers such as Dropout (Srivastava al., 2014). Consequently, fine-tuning performed over a very small number epochs (typically 55). Finally, we observe that fine-tuned LMs are adversely impacted by even minor grammatical errors in inputs during infer- ence 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, espe- cially when LM inputs are provided by users with different levels of language proficiency.
75
SST-2
MNU
QQP
Grammatically inputs
Inputs with minor grammatical
Figure Impact of minor grammatical errors on a Koberta-Base model fine-tuned using only grammati- cally correct examples. We generate inputs with minor grammatical errors by pruming articles ("a", 'an', 'the') and punctuation marks (comma, fullstop, apostrophe, etc.) from samples in the development set.
small datasets for very few epochs, find that these prior batching methods lead to signif- icant drop in accuracy. BucketSampler includes key algorithmic enhancements that tune the opti- mal 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 random sub- 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 im- proves the resilience of LMs to grammatical errors during inference by exposing the model to gram- matically 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 fine- tuning, such as freezing layers (Lee et al., 2019; Zaken et al., 2022) and the use of adapters (Houlsby et al., 2019), to achieve further efficiency gains. We summarize main contributions follows:
1
2
9
10
Training epoch
Training Loss
Loss
Figure I: Training curve obtained from fine-tuning Roberta-base RTE, language understanding task with 2.5K training samples, with dropout rate 0.1. We report loss averaged across 10 random seeds.
To address the aforementioned challenges, we present TokenDrop + BucketSampler, the synergis- tie combination of a 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 batch- ing strategies have been incorporated into popular NLP libraries such HuggingFacc Transformers (Wolf et al., 2019) and Lingvo (Shen al., 2019). However, these methods are optimized for training LMs from scratch. When applied to fine-tuning
We propose TokenDrop BucketSampler, a framework for accurate and efficient LM fine- tuning.
BucketSampler is a length-aware batching method that generates batches of sequences with similar lengths to reduce padding. Bucket- Sampler incorporates optimizations that enable fine-tuning convergence while maintaining high

---

throughput.
In addition, most fine-tuning datasets are only com- posed of sentences and phrases that are grammat- ically correct. As a result, minor errors in user- provided inputs during inference (such as missing punctuation marks) can degrade the quality of out- puts produced by the LM (Fig. 2).
Troe is a regularizer that drops a random subset of insignificant tokens from each sequence in every epoch to prevent LMs from memorizing fine-tuning data, while also further minimizing the need for padding.
We demonstrate that TooenDrop can be synergis- tically combined with BucketSampler to simulta- neously improve both accuracy (up to 1.2%) and efficiency (up to 10.61 x) of LM fine-tuning
To address these challenges, we propose Tok- enDrop, regularizer that drops random subset of words from each sequence in every training epoch. Unlike Dropout, TokenDrop introduces data diver- sity between the different training epochs, thereby making the models unlikely to see the same se- quence twice over the course of fine-tuning, and hence, less likely to overfit. We find that the choice of words selected for dropping has a significant im- pact on training convergence, since the semantics of the input text sequence can change if important tokens are dropped (see Appendix C). For instance, if the word "not" in the sequence "the movie not good" is dropped, the sentiment of the sen- tence changes from negative to positive. order to overcome this challenge, TokenDrop only drops stopwords from sequences. Stopwords are words in any language that do not contribute to the mean- ing of a sentence, but are added to make sentences grammatically correct. For instance, the words "the" and "was" in the aforementioned sequence stopwords. In effect, TokenDrop provides stronger regularization effect compared to Dropout (see Appendix D), and also improves the efficiency of fine-tuning by reducing the sequence length with- out affecting the meaning of the sequence. To- kenDrop also improves the resilience of LMs to grammatical errors during inference by exposing the model to grammatically incorrect sequences during fine-tuming In particular, TokenDrop gener- ates grammatically incorrect sequences by pruning stopwords from grammatically correct sequences.
2 Method
TokenDrop + BucketSampler is a framework for ac- curate and efficient fine-tuning of pre-trained LMs. The first component, TokenDrop. is a regularizer that randomly drops a subset of insignificant to- kens from each sequence in every epoch. The sec- ond component, BucketSampler, is length-aware batching method that generates batches of sam- ples with lower spread in sequence lengths. We demonstrate how Troe and BucketSampler can be synergistically combined (hence the name Drop + BucketSampler) to simultaneously re- duce overfitting and climinate the need for padding tokens, resulting in both accuracy and speed im- provements (Fig. 3).
Handomsanppier
Heavy Padding
Padding tokens
Training Data
3xche.
Padding
Tokandrop
BucketSammlnn
Near-zero Padding
Dropped tokens (Stopwords)
Figure 3: Illustration of different batching methods.
The procedure for applying TokenDrop given dataset is described in Algorithm 1. Given a dataset and a list of stopwords, we identify and prune a random subset of stopwords in each sequence. The number of stopwords pruned in each sequence is determined by the parameter. When
2.1 TokenDrop
Fine-tuning datasets for LMs are often small due to the need for expensive human annotations. As result, we find that pre-trained LMs quickly mem- orize the training samples during fine-tuning (Fig. 1). Dropout (Srivastava al., 2014) 1)) which drops random subset of neurons in each training epoch, is currently the most widely used regularizer. How- ever, we find that dropout is ineffective at prevent- ing overfitting during fine-tuning of LMs (Fig. 1).
T Tokens_te Lakens_to_aae too_tto_dro global TokenDrop Rate (analogous Dropout_Ratl in Dropout) used, i.e., the same fraction of stopwords are dropped in each sequence, then Tokens_to to_drop(ddoppooo (cquuncc)))))) number lemmber_of__ee sto_soowoddd (sequuenc)) Token Drop Rate. Dropp_ote.. We note that it is not
*

---

necessary use the same TropnDrop Rote for all sequences, and a different fraction of stopwords can be dropped in each sequence, as described in Section 2.3. The complete list of stopwords used in our experiments is provided in Appendix C.
mentally determined by profiling on-chip memory usage for different sequence lengths.
While the BucketSampler algorithm described above provides some improvements in the effi- ciency of fine-tuning, it still leaves room for fur- ther improvement along two directions. (1) The constraint that batches can only be formed using samples from the same bucket leads to "residual" batches with small numbers of samples, adversely affecting hardware utilization. Therefore, we pro- pose Residual Batch Merging (RBM) to merge residual batches in different buckets into larger batches. (2) From an accuracy standpoint, we find that the convergence of fine-tuning is adversely im- pacted when using (a) very large batch sizes, and (b) single learning rate schedule with variable batch sizes, since fine-tuning is often performed on small datasets for a very small number of epochs. To overcome these challenges, we propose Batch- Cap to progressively increase the maximum batch size over the course of fine-tuning. We also propose Learning Rate Modulation (LRM) to dynamically scale the learning rate based on the batch size. We explain these optimizations in greater detail in the following subsections.
Algorithm 1: Trop
Input
Input Dataset (Datoset), List of stopwords (Stopwords)) Number of stopwords to drop in each sample (Tokens_Le) to _to__rop)
Output: Dataset after applying TokenDrop
Settereel_Datasee
Filtered Dutaset = Dataset
for somple in Filtered Dataset do
Stopwords = random Sluffle(Stopwordd
while T Tokens__ to dirop(sample)) 0 do for word in Stopwords do if word in somple then
sample.pop(wore Tokens_lo satesss_odddoosseeeeee =
return Filtered Datasct
2.2 BucketSampler
BucketSampler combines a sequeece--enntth-awan aware, variable-batch-ae batching strategy with algorith- mic optimizations to enable faster fine-tuning. The method for generating batches with BucketSampler is illustrated in Algorithm/ 2. All samples in the dataset are divided into buckets such that sequences that fall into the same bucket have similar se- quence lengths. Each bucket is defined by a triplet (minnseq_len, max_seq_len, HW (ap)). Here, min_seq_les and mox__eq_____ denote the mini- mum and maximum sequence lengths of sequences in the bucket, respectively, i.c., all sequences whose lengths lie between min_seq_len (inclusive) and max_ mar_seq_len (exclusive) of bucket fall into that bucket (lines 20-23 in Alg. 2). Then, batches are generated by only combining sequences from the same bucket (lines 24-32 in Alg. 2). In effect, BucketSampler reduces the spread of sequence lengths in a batch, thereby reducing the need for padding tokens. BucketSampler also provides sup- port for variable batch sizes. In particular, since the peak memory requirements for processing a batch scale quadratically with sequence length due to the quadratic complexity of self-attention, we propose using large batch sizes when generating batches from buckets with small max_seq_lee and vice- versa. The HIW Cap parameter associated with each bucket encodes the maximum batch size that can be used for generating batches from that bucket on a given hardware platform. CCp is experi-
Residual Batch Merging (RBM): Peak hardware utilization is achieved when the number of sam- ples in each bucket is exact multiple of the bucket's HW Cap. However, in reality, it is highly likely that the aforementioned condition is not sat- isfied, resulting in one batch in most buckets with batch_sice < WWCap. We term these batches as "residual" batches, and the hardware under- utilized when processing residual batches. In order to reduce the impact of residual batches, we pro- pose merging residual batches from different buck- ets into larger batches. We present the Residual Batch Merging (RBM) algorithm to maximize hard- ware utilization, while also minimizing the number of additional padding tokens introduced as a result of merging sequences from different buckets (lines 5-18 in Algorithm 2). New batches are created by appending samples one-by-one from residual batches in each bucket, with buckets processed in increasing order of maz_seq_ (lines 7, 1 11-12). When the number of samples in the new batch be- comes an exact multiple of the Mucket_batch.. of the bucket corresponding the longest sequence in this new batch, the new batch is added to the list of batches used for fine-tuning (lines 13-15). Lines 8- 10 account for corner cases where the num-

---

Algorithm 2: BucketSampler
Input :Input Dataset (Dataset), Training epoch (epoch)
:utches generated from Datasel (batches)
Function (aWCa. epoch);
EpochCup= soaling__aatattnnnn
x base batch_sizze
bucket bucket_hatch_ min(HWCap, EpochCap)
retur bucket_ batch_s_iteeeeeeeee__
Function Beggal__eeeccMMMggggggniiiie
residual batches);
new batch =
for (residual batch, bucket batch (atch_size)) in residual batches do
if cardinality(new batch) >= bucket batch size then
aodchee.append(one s_athhhhh
new_ botchatch= [1]
11
sequence residual batch
12
new hwwp_batc.aapp apauchhappeniiiee (equenc))))
13
(cardinality new batch) > 0) and (cardinalitynew batch) bucket batch size 0) then batches batches..apendkeeee append(new batch)
new batch = []
16
new batch is not empry then patches.append(nes. batch)
17
18
retrr batches
buckets: List of buckets, where bucket defined the triplet (min seq len, seq Jen, HWCap) for sample in Datasel do for bucket in buckets do bucket.min, seg seq_len seq_len(sample) bucket.max-_eq _eegeageeeeeeeen then buuckt..apeenddsmmml.
batches= =
residual Battual_btchhs==]1]]
26
for backet in buckets backet Bat.uuuuiieee-BBaaacuuuuuu..... BatchCapluucee..WWW..... epoch)
27
28
barket_batcc_sieee size
29
index =
while index < ffllorddivisions diviision.curddinddiyyyuuuu (urdinddyyybbuuekel bucket butch buzch_size) bbucke_ batch_size) natehes..pennd(uuceelinne. index = index + bucket batch_size
30
31
index bucket batch size))
32
33
residual__aatches.pppcee haiichhs..ppeedddhcceeffiiieeeee :], bucket_hatell sastt__athh__spppeeeee size))
batches
Residual Batch Merging((utchees.reiin residual (ssidual_batchhs)
batches
ber of samples in the new batch is less than the bucket_batch__tt batbb_atchhic corresponding to the longest sequence in the new batch, but larger than the bucket_ sasee_bbach__sss_eeeeeeee corresponding to the next se- quence to added to the new batch. In this case, the new batch is considered to be "full", and is added to the list of batches used for fine-tuning. In effect, RBM reduces the number of possible residual buckets from n., where n is the number of buckets, to (only the final batch can be resid- ual batch, see lines 16-17). Since buckets are in- spected increasing order maar_se-_een, RBM also ensures that only residual samples from adja- cent buckets are combined to form larger batches, thereby minimizing the number of padding tokens used in merged batches.
small datasets for small numbers of epochs, find that fine-tuning does not converge when very large batch sizes are used due the sparsity of weight updates. For instance, RTE (one of the GLUE (Wang et al.) datasets) contains 2.5K training samples, and fine-tuning on RTE is typically performed over epochs. When a batch size of 1024 is used, only 3 weight updates are performed in cach epoch, for total of 9 weight updates during fine-tuning. As result, fine-tuning does not converge, and models exhibit high training loss and low test accuracy the end of fine-tuning. With BucketSampler, the HW Cap for buckets with small mar_ laas__se_lee_ can be large (for instance, HI Cap = 2525 for mar mar_seq see_len len = 5 when fine-tuning Goberta-Base on a NVIDIA RTX 2080Ti GPU). To ensure fine-tuning convergence.
BatchCap: Since fine-tuning is performed on

---

BucketSampler uses smaller batch sizes in the first epoch of fine-tuning, and progressively increases the cap on the maximum allowable batch size in the later epochs of fine-tuning (Smith et al., 2018). We introduce BatchCap to find the maximum allowable batch size in each epoch (lines 1-4 in Alg. 2). In every epoch, BatchCap limits the maximum batch size of each bucket as bucket bucket_batch_aneeeteetee batchaaiize min( W Cap(bucket), SappochCap(epochh) where EpochCap(epnch))
GPU clusters. In particular, when batch size is changed k * batch_size the learning rate is changed to learning
* sqr((k) to find the optimal learning rate for a given batch size when the optimal learning rate is known for a different batch size. In this work, we use this formulation for modulating learning rate a batch-by-bateh basis, i.e., changing learning rate for each batch based the size of the batch.
epoch)
2.3 TokenDrop + BucketSampler
ecaling_facttorpodt base baasebbbacll wwwcbbatcthhhee size. Here, and scaling_factor are hy- perparameters that control the batch size used in the first epoch of fine-tuning, and the growth rate of the maximum batch size across epochs, respectively. We note that BatchCap leads to hardware under-utilization in the early epochs of fine-tuning. However, BatchCap is necessary to achieve convergence when fine-tuning with BucketSampler, and our exponential scaling rule ensures high utilization in the majority of epochs.
When batches are generated using BucketSampler, the sequence lengths of all samples in batch lie between the mön_see sen_een and max_ seg_ee__ee of the bucket the batch was drawn from (ex- cept in merged residual batches). TokenDrop can be synergistically combined with BucketSam- pler to further equalize the lengths of all se- quences in a batch by pruming stopwords from longer sequences in the batch, thereby eliminat- ing the need for padding. To achieve this, we pro- pose defining T TokenDrop_Iie Rate a sequence- by-sequence basis, rather than having global Token Drop_ rate for all sequences. In par- ticular, the number of stopwords to drop in a se- quence is computed as faakens_to_dreee drop(samppl sample) cardinality! sample) min_seq sinnssse_eend_ (batch). Consequently, all sequences in a batch are pruned to min_ siii__se_lee(bbach by dropping a random subset of stopwords (Algorithm 3), thereby elim- inating padding tokens (except in merged resid- ual batches, where tokens_t to _ddorpp sample) may be larger than num num_ stopwords (sample)) due larger variance in sequence lengths across merged buckets).
Learning Rate Modulation (LRM): The use of BucketSampler leads to large variance batch sizes during fine-tuning. For instance, when fine- tuning Roberta-Base on a NVIDIA RTX 2080Ti GPU, HWCap 2525 for max_seq se___eee_lee_______ee_______see_ = 5 and 128. Since the choice of learning rate is highly sensitive to the batch size used during training (Krizhevsky, 2014; Smith et al., 2018), we find that fine-tuning using BucketSampler fails to converge with a single learning rate schedule, even when a grid search is performed to find the best learning rate. We propose Learning Rate Modulation (LRM) to over- come the limitations of using single learning rate schedule when training with variable batch sizes, LRM dynamically modulates the base learning rate based on the batch size of each batch. LRM scales the base lewrnenerrrraaaeeeee for each batch as dearning_rate(utte)) = leae_elanninngng rate
HWCap 64 for moo__seqleest
A Algorithm 3: TokenDrop + BucketSampler
Input
:Input Dataset (Dataset), Training epoch (epoch). List of stopwords (Stopwords)
Output:Batches generated from Dataset (batches)
batches BucketSamplert Datasct, cpoch) for batch in batches do for sample batch do Tokens sample))
spri@batch_sizec(batthh laabbbbbasch____aaaaaaaaaaa___hh (rutc__size)
cardinalityssample) min_ mni__se__enttaee len_hen(h)))h)
Here, base @use_learning_ree rrtt is the optimal learning rate schedule when training with fixed batch size, where all batches (except the last batch) have batch_size
s batches TokenDroop(utthhss BoeenDDooouuteeessttppoodds
Tokens to drop)
return batches
mass_batch_sse size. The formula for computing learning_rond (raniing_ate((aacch) is derived from the square-root scaling law relating learning rate and batch size (Krizhevsky, 2014), which is a popular trick for parallelizing training of deep neural networks using large batch sizes on
3 Experiments and Results
We implement TokenDrop + BucketSampler in Py- Torch using Huggingface Transformers (Wolf et al.,

---

2019). We perform experiments on NVIDIA RTX 2080 Ti GPU with 11 GB memory, and report re- sults averaged across 10 runs with different random seeds. We perform epochs of fine-tuning for our method and all baselines. The details of all hyper- parameters used in our experiments are described in Appendix A. We note that TokenDrop is not used when fine-tuning on CoLA, since the task involves identifying if given sentence is linguis- tically acceptable or not, and TokenDrop makes linguistically acceptable sequenees unacceptable.
note that TooenDrop is only applied to input se- quences, and bucketing performed based on input sequence lengths for generation tasks. As result, padding is still necessary for the target sequences. Resilience minor grammatical errors in in- puts: We find that training with TokenDrop signifi- cantly enhances the resilience of fine-tuned mod- els to minor grammatical errors in inputs. For in- stance, when articles ("a', "an", 'the") and punctu- ation marks are removed from the test sequences, the average accuracy on GLUE (except CoLA) and SQUAD drops by 5.2% Sobeeta-BBsse) and the average ROUGE-1 score drops by 3.1 points (T- small) for the baseline models. On the other hand, models fine-tuned with TokenDrop incur only 0.3% and 0.06 points drop in average accu- racy and ROUGE-1 scores, respectively, thereby demonstrating significantly higher resilience to mi- nor grammatical errors. The enhanced resilience to grammatical errors in models fine-tuned with TokenDrop can also be observed in Fig. 4, where there is negligible loss in accuracy even when 40% of all stopwords in each sequence are randomly chosen and deleted during inference.
3.1 TokenDrop BucketSampler improves the accuracy and efficiency of fine-tuning
Classification tasks: We present results of fine- tuning the popular Roberta (Liu et al., 2019) and Electra (Clark et al., 2020) models the GLUE (Wang et al.) and SQUADv1.1 (Rajpurkar et al.) datasets in Table 1. We find that fine-tuning with TokenDrop + BucketSampler consistently produces more accurate models compared to conventional fine-tuning with random batches (RandomSam- pler). In addition, TokenDrop- + BucketSampler also reduces the wall-clock fine-tuning time by up to 10.61 x compared to conventional fine-tuning. with an average speedup of 5.9x across the 10 GLUE and SQUAD tasks. With RandomSampler, 38. 38.9% of all tokens used for training are padding tokens, which reduces to just 0.2% with Tooe BucketSampler (not exactly 0%, since padding is needed in merged residual batches). We find that the speedup from using TokenDrop + BucketSam- pler on given task is dependent on two factors: (I) the sequence length histogram, and (2) the size of the dataset. We provide detailed analysis of the relationship between the statistics of the fine- tuning dataset and the speedup achieved from using TokenDrop + BucketSumpler in Appendix E. We also provide supplementary results on fine-tuning Robgrta-Large in Appendix B to demonstrate the benefits of using TokenDrop + BucketSampler for fine-tuning larger models.
3.2 TokenDrop + BucketSampler enables accurate and efficient inference
While the primary objective of TokenDrop + Buck- etSampler is to improve fine-tuning efficiency, we describe how they can also be used to improve the efficiency of both real-time and server-mode inference in the following subsections.
Real-time inference (batch size = 1): Real-time inference workloads have strict latency require- ments and bursty input rates, and hence, inputs are typically processed as soon as they arrive with batch size of 1. To reduce the latency of real- time inference, we propose filtering stopwords in the input text sequence by applying TokenDrop during inference also. Inference-tin time TokenDrop offers a promising approach for accelerating real- time inference, enabling speedups of 2.2X when all stopwords are pruned (T Drop Rate- = in Fig. 4). However, we find that models fine- tuned without TokenDrop suffer from large accu- racy drop when TokenDrop is applied at inference time. On the other hand, models trained with To- kenDrop exhibit significantly higher resilience inference-time TokenDrop, enabling 1.5x reduc- tion in inference latency with no loss in accuracy, and 2.2x speedup with < 1% loss in accuracy (Fig.
Generation tasks: We present results of fine- tuming the Tmll seq2seq model (Raffel et al., 2020) on text summarization using the XSum (Narayan et al., 2018) and CNN//ailyMail (Nal- lapati et al., 2016) datasets in Table 2. We find that fine-tuning with TokenDrop BucketSampler improves the ROUGE-1 score by up to 0.3 points, while also reducing the wall-clock fine-tuning time by up to 8.62x compared to RandomSampler. We

---

Table 1: Results of fine-tuning Roberta-base and Electra-base on the GLUE and SQUAD v1.1 development sets. We report Fl score for SQUAD, Matthews correlation for CoLA, Pearson Correlation for STS-B and accuracy for all other tasks. We report only "matched" accuracy for MNLI. Subscripts indicate standard deviation.
Model
Batching Method
SQUAD
MRPC
STS-B
SST-2
CoLA
QQP
QNLI
RTE
MNLI
WNLI
Avg
Roberta
RandonSampler
90.460,2
88.0910
89.85.3
94.150.
59.35as
91.36..1
92.58a.1
69.3118
$7.56..
53.51ga
81.62
Base
TokenDrop+ BacketSampler
90.5%
88663s
90.06.,4
94,3302
59,49I1
91.48
92.64.11
70.4gs
87.660.1
54.572a
$1.98
Electra-
RandomSampler
90.74o.5
88 28as
89.74a.4
95.01ea
65.48
91.9%a
92.411
77.070
88.330.6
54.2220
83.32
Base
Token Bucepppetttappppeeee
90.8%..
88.51a3
89,79.
96.12a3
66.03
92.1ea
93,09a
78.24
88.74,3
54.38
83.7
Average mee-ttnnng Speodup
9.87X
2,75X
7.73X
10.6X
9.39X
5.3XX
3,18X
2008X
4.09X
3.8X
5.9X
Table 2: Results of fine-tuning the T5-small seq2seq model on text summarization. We report the ROUGE- score. Subscripts indicate standard deviation.
in models fine-tuned with Troe to achieve 4.9x speedup with no loss in accuracy (Table 3).
3.3 Ablation: Breakdown of benefits from the different BucketSampler optimizations
Batching Method
XSum
CNNiDailyMail
RamdomSampler
32.58
24.5
24.81ss
We study the impact of the different BucketSam-- pler optimizations on accuracy and efficiency of fine-tuning in Table 4 and Figure 5. When no opti- mizations are used, BucketSampler achieves 5.3 x
Speedup
x.62x
7.77x
4). In addition, TooenDrop can be combined with progressive token pruming methods that prune the least-important tokens in each layer based on atten- tion scores (Wang et al., 2021; Goyal al., 2020), to achieve further gains in efficiency.
speedup, but incurs substantial accuracy drop due to insufficient training convergence. When Resid- ual Batch Merging (RBM) used, "stray" batches from different buckets are combined to form larger batches. In addition to improving efficiency by reducing hardware under-utilization, RBM also im- proves accuracy by reducing variability in batch sizes, thereby enabling better convergence with a single learning rate schedule. BatchCap further improves accuracy by using small batch sizes in early epochs, thereby ensuring sufficient numbers of weight updates to achieve training convergence. While BatchCap is necessary for achieving conver- gence, it leads to small drop in training efficiency due to hardware under-utilization in early epochs. Finally, the use of LRM to dynamically adjust the learning rate for each batch ensures that fine-tuning with BucketSampler results in near-identical train- ing curves (Fig. 5) and hence, accuracy (Table 4) compared to fine-tuning with RandomSampler.
2.5
84
2
Accurayy
81
1.5
0,2
0.4
0.8
TokenDrop Rate (Inference)
No TokenDrop
With TokenDrop
Speedup
Figure 4: Accuracy drop and speedups from using To- kenDrop during inference with Roberta-band (batch size 1). We report the average score on the GLUE tasks and SQUADv1.1. For 0 < TokenDrop Rate < I, we randomly prune touu_ssoowworde ssopwordsssecquenee)) x T TokenDrop_Rate stopwords in each sequence.
Related Work
Server-mode inference (batch size 1): In the server-mode inference setting. inputs arrive simul- taneously from several sources. Consequently, in- puts are processed in batches, with the goal of max- imizing throughput. In the server-mode inference setting, we propose utilizing TokenDrop + Bucket- Sampler to batch the inference queries. Here, we set BatchCap(bucker
Prior works have developed batching strategies for variable-length inputs to improve the efficiency of LM training. RandomSampler is the most com- monly 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 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 Huggingface Transformers
HW (Capbbuceet) for all buckets maximize hardware utilization, and hence, throughput. We find that using BucketSam- pler leads 4.5x speedup over random batching. addition, TokenDrop can also be synergistically combined with BucketSampler at inference time

---

Table 3: Results of using TokenDrop + BucketSampler during batched inference with Roberta-base. We report the average score across the GLUE tasks and SQUADvI.I. We assume that all samples in the test dataset arrive simultancously, and speedup is computed by comparing the wall-clock time taken to infer on all test samples.
Batching Method
Accuracy (Trained
Accuracy (Trained
Fraction pudding tokens
Speedup
(Inference)
without TokenDrop)
with TropnDrop)
RamdomSampler
81.62
82.02
40.1%
1X
BucketSampler
81.64
81,98
9.3%
4.5X
80,21
82.01
0.3%
4.9X
Table Impact of the different BucketSampler opti- mizations when fine-tuning Roberta-base. We report the average score across GLUE and SQUAD.
2.5
Randondamalar
BucketSompler
2
(RBM)
BucketBermler BatchCap)
Bunketioopplr (SBM BusiiCap LSM
1.5
Optimizations
Accuracy
Speedup
None
76.9
5.3x
1
RBM
5.8x
0.5
RBM
ButchCup
81.29
5.2x
RBM + BatchCap LRM
81.64
5.2x
0
Epoch
(Wolf et al., 2019)) sorts sequences in order of in- creasing sequence length, and generates batches of adjacent sequences in the sorted list. While LengthGroopcdSampler reduces padding, it uses fixed batch sizes, resulting in hardware under- utilization. Packing (Krell et al., 2022) generates batches by concatenating different inputs along the sequence length dimension. However, we find that unlike when training from scratch, the overheads of packing and the additional computations in atten- tion layers (computing irrelevant scores, followed by masking to prevent cross-contamination be- tween different sequences) not amortized over the small number of fine-tuning epochs. Finally, Tensorflow's tf.data.bucke__b be_ee__ue_ssee_e____eee_e divides the training dataset into buckets based on sequence length, and generates batches only from sequences in the same bucket, sim- ilar to BucketSampler. While packing and tf. bf.cketa.bcckt_bby ba_ii_y________ www.:ucceeebbyyssueeeeeeeee lenuthce_length support vari- able batch sizes to maximize hardware utilization, they incur accuracy drop during fine-tuning. In summary, while prior works have demonstrated ef- ficiency gains when training LMs from scratch, the unique characteristics of fine-tuning (small datasets, very few epochs) make these methods in- effective. As a result, TokenDrop + BucketSampler significantly outperforms prior methods in terms of both accuracy and fine-tuning efficiency (Table 5).
Figure 5: Average training loss across GLUE and SQUAD when fine-tuning Roberta-base.
Table 5:
Accuracy and efficiency of fine-tuning Roberta-base with different batching strategies. Re- sults are averaged across the GLUE tasks and SQUAD.
Batching Strategy
Accuracy
Speedup
RandomSanpler LennghhCroupeddamplerr
81.62
IX
81.65
3.8X
Packing
80.69
4,2X
tf. tf.bucket by sequequqecce_ccc length
76.9
5.3X
BucketSumpler
81.67
5.1X
TokkenDrop--BucctttSnppeer
81.98
5.9X
overfitting, while BucketSampler creates batches of sequences with similar sequence lengths to reduce the number of padding tokens. We also introduced algorithmic optimizations Residual Batch Merg- ing (RBM), BatchCap and Learning Rate Modula- tion (LRM) to maximize throughput and enable convergence when fine-tuning with BucketSampler. We demonstrated that TokenDrop can be syner- gistically combined with BucketSampler drop more tokens from longer sequences in each batch, thereby further reducing the need for padding. effect, fine-tuning with TokenDrop + BucketSam- pler produced more accurate models in a shorter time compared to conventional fine-tuning.
Acknowledgement
5 Conclusion
This work was supported in part by the Center for the Co-Design of Cognitive Systems (CoCoSys), a JUMP2.0 center sponsored by the Semiconductor Research Corporation (SRC) and DARPA, and in part by SRC under the AIHW program.
In this work, we presented TokenDrop Bucket- Sampler for accurate and efficient fine-tuning of LMs. TokenDrop prunes a random subset of stop- words in each sequence in every epoch to reduce

---

7 Limitations
Weight initializations, data orders, and early stop- ping. CORR, abs/2002.06305.
(I) TokenDrop is not universally applicable to all fine-tuning tasks. In particular, Too cannot be used when dropping stopwords can potentially change the labels associated with sequences, For instance, when fine-tuning LM to identify if given sequence is grammatically correct or not (as in CoLA), dropping stopwords from sequences will make all sequences grammatically incorrect. (2) BucketSampler is not useful for self-supervised training with the Masked Language Modelling (MLM) objective. When LMs are trained with MLM, the text corpus is divided into fixed-size blocks and fed to the model, thereby resulting in constant input sizes. (3) The use of Batch- Cap in BucketSampler introduces hardware under- utilization in the first epoch of fine-tuning. We find that BatchCap is necessary to achieve convergence when fine-tuning with TokenDrop + BucketSam- pler with the same set of hyperparameters as Ran- domSampler. However, it is possible that the use of advanced optimizers and/or learning rate schedules enable convergence without the need for Batch- Cap, thereby further increasing efficiency. We plan explore this as part of future work.
Saurabh Goyal, Anamitra Roy Choudhury, Saurabh Raje, Venkatesan T. Chakaravarthy, Yogish Sabhar- wal, and Ashish Verma. 2020. Power-bert: Accel- erating BERT inference via progressive word-vector elimination. In Proceedings of the 37th International Conference on Machine Learning. 1CML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 3690-3699. PMLR.
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin de Laroussilhe, Andrea Ges- mundo, Mona Attariyan, and Sylvain Gelly. 2019. Parameter-efficient transfer learning for NL.P. In Pro- ceedings of the 36th International Conference on Ma- chine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 Proceedings of Machine Learning Research, pages 2790-2799. PMLR.
Mario Michael Krell, Matej Kosec, Sergio P. Perez, and Andrew Fitzgibbon. 2022. Efficient sequence pack- ing without cross-contamination: Accelerating large language models without impacting performance.
Alex Krizhevsky. 2014. One weird trick for par- allelizing convolutional neural networks. CoRR, abs/1404.5997.
Jaejun Lee, Raphael Tang, and Jimmy Lin. 2019. What would elsa do? freezing layers during transformer fine-tuning. CoRR, 0tt//9911.0399
References
Sébastien Bubeck, Varun Chandrasckaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott M. Lundberg, Harsha Nori, Hamid Palangi, Marco Túlio Ribeiro, and Yi Zhang. 2023. Sparks of artificial general intelligence: Early experiments with GPT-4. CoRR, abs/2303. 12712.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Man- dar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer. and Veselin Stoyanov. 2019. Roberta: A robustly optimized BERT pretraining approach. CoRR, abs//107.11666
Ramesh Nallapati. Bowen Zhou, Cícero Nogueira dos Santos, Çaglar Gülçehree and Bing Xiang. 2016. Abstractive text summarization using sequence-to to- sequence rmns and beyond. In Proceedings of the 20th SIGNLL Conference Computational Natural Language Learning, CoNLL 2016, Berlin, Germany, August 11 11-12, 2016, pages 280-290. ACL.
Kevin Clark, Minh- Thang Luong, Quoc V. Le, and Christopher D. Manning. 2020. ELECTRA: pre- training text encoders as discriminators rather than generators. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa. Ethiopia, April 26-30, 2020. OpenReview.net.
Jacob Devlin, Ming-Wa Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: pre-training of deep bidirectional transformers for language under- standing. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Tech- nologies, NAACL-FILT 2019, Minneupolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers), Association for Computational
Shashi Narayan, Shay B. Cohen, and Mirella Lapata. 2018. Don't give me the details, just the summary! topic-aware convolutional neural networks for treme summarization. In Proceedings of the 2018 Conference on Empirical Methods in Natural Lan- guage Processing, Brussels, Belgium, October 31 November 2018, pages 1797-1807. Association for Computational Linguistics.
pages Linguistics.
Colin Raffel, Noam Shazeer, Adam Roberts. Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. 2020. Exploring the limits of transfer learning with a unified text-to-text trans- former. J. Mach. Learn. Res., 11:140:1-1400
Jesse Dodge, Gabriel Ilharco, Roy Schwartz, Ali Farhadi, Hannaneh Hajishirzi, and Noah Smith. 2020. Fine-tuning pretrained language models:

---

Pranav Rajpurkar. Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100, 000+ questions for ma- chine comprehension of text. In EMNLP 2016.
when training with Tooe BucketSampler compared to training with RandomSampler in or- der to account for the regularization effect of To- kenDrop. In particular, we set the dropout rate to on small datasets (<10K training samples) to enable convergence within 3 epochs. For larger datasets with >> 10K training samples, we reduce the dropout rate to 0.025, since sufficient weight updates performed to enable convergence even when small amount of dropout is used, and the ad- ditional regularization from dropout leads to higher test accuracy.
Jonathan Shen, Patrick Nguyen, Yonghui Wu, Zhifeng Chen, Mía Xu Chen, Ye Jia, Anjuli Kannan, Tara N. Sainath, Yuan Cao, Chung- Cheng Chiu, Yanzhang He, Jan Chorowski, Smit Hinsu, Stella I Laurenzo, and James Qin ct.al. 2019. Lingvo: modular and scal- able framework for sequence-to-squeeee modeling. CoRR, 1b2/0999...295.
Samuel I.. Smith, Pieter-Jan Kindermans, Chris Ying, and Quoe V. Le. 2018. Don't decay the learning rate, increase the batch size. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenReview.net.
Choice of buckets: and mar_seq_les of a bucket determine the se- quence length range of all samples in the bucket. When TokenDrop is not used, we simply define buckets in length increments of I i.e., bucket 1 will have sequences of length 1, bucket 2 will have quences of length 2, and on. Therefore, the num- ber of buckets will be equal to the max_seq of the dataset, and the [min_seq ses__seeleenn____ mar_ saas__se_leen of bucket i will be (i,)). This will eliminate the need for padding in non-residual batches (since all sequences in a bucket will have the same se- quence length), and RBM will keep hardware utilization as high as possible while minimizing padding in residual batches. When TokenDrop is used, the sequence length spread (mar_seq_len mmnn_seq_ee peq_len + 1) of each bucket must be >1 to enable token dropping in non-residual batches also. However, if the sequence length spread of each bucket is too large, then it is highly likely that piatensslo_droppp (sspssspple will be larger than number_of _f_o___o__ttpooddss__r_ (sample) for most sam- ples. This has two drawbacks: (1) it necessitates the use of padding, and (2) all stopwords in many samples will be dropped in every epoch, thereby reducing data diversity between epochs and leading to weaker regularization. In our experiments, we find that defining buckets in increments of 5, where bucket 1 has (min_seq_len) max_ mar_seq_len) = [0,5), bucket 2 has (min_seq_len) max_seaee seq_len) [5,10) and so on, provides good trade-off be- tween minimizing padding and ensuring sufficient regularization for high test accuracy.
The parameters min_seq_les
Nitish Srivastava, Geoffrey E. Hinton, Alex Krizhevsky, Πya Sutskever, and Ruslan Salakhutdinov. 2014. Dropout: simple way to prevent neural networks from overfitting, J. Mach. Learn. Res., 15(1)11929- 1958.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman, GLUU:: A multi-task benchmark and analysis platform for natural language understanding. In ICLR 2019.
Hanrui Wang. Zhekai Zhang, and Song Han. 2021. Spat- ten: Efficient sparse attention architecture with cas- cade token and head pruning. In IEEE International Symposium on tecture, HPCA 2021, Scoul, South Korea, February 27- March 2021, pages 97-110. IEEE.
Computer Archi-
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pier- ric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, and Jamie Brew. 2019. Huggingface's 's transformers: State-of-the-an natural language processing. CoRR, Ab0./991..3771.
Elad Ben Zaken, Yoav Goldberg, and Shauli Ravfogel. 2022. Bitfit: Simple parameter-efficien mne-tuning for transformer-based masked language-models. In Proceedings of the 60th Anmual Meeting of the As- sociation for Computational Linguistics ( (Volume 2: Short Papers), ACL 2022, Dublin, Ireland, May 22- 27, 2022, pages 1-9. Association for Computational Linguistics.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona T. Diab, Xian Li. Xi Victoria Lin, Todor Mihaylov, Myle Ott. Sam Shleifer, Kurt Shus- ter, Daniel Simig, Punit Singh Koura, Anjali Srid- har, Tianlu Wang, and Luke Zettlemoyer. 2022. OPT: open pre-trained transformer language mod- els. CoRR, ahs/2205.01063
List of stopwords used for TokenDrop: [',,, 'me', "my', 'myself", 'we", 'our', 'ours', 'ourselves', "you', "you're", "you've" "oou've". "you'll" 11",, "you'd", 'your', 'yours', 'yourself", 'yourselves', 'he', 'him', 'his', "himself', 'she", "she's", 'her', 'hers", 'herself', 'it', "it's", 'its', 'itself", 'they', 'them', 'their', 'theirs',
A List of hyperparameters
The hyperparameters used in our experiments are described in Table 6. We reduce the dropout rate

---

Table 6: List of hyperparameters used in our experiments. We use the same set of hyperparameters for all GLUE and SQUAD tasks. We also the same set of hyperparameters when fine-tuning both Roberta-base and Electra-base. HWCap is experimentally determined on a NVIDIA RTX 2080Ti GPU with 1 IGB memory.
Hyperparameter
Value (RandomSamplen)
Value (TokenDrop BucketSampler)
Batch Size
Variable
Learning Rate
Best among [3, c
Variable
for IMRPC, STS-B, CoLA, RTE, WNLII.
Dropout Rate
0.1
0,025 ISST-2, QQP, QNL, MNLI, SQUAD!
Buckets (min (min__sg__les Jen.
Defined in increments 5, L... (0,5.2525), (5.10.1250). (10,15,800) (005,000, ... (22.......4).. (125,128... 122.61) (2, 64); base butch size = Batch Size used for RandomSampler (Learning Rate, Batch Size) used for RandomSanpler
N/A
max max_soq Jen, HWCsp)
BatchCap: (scaling factor, base hutch size)
N/A
LRM: (base_learmin leariing_attee
N/A
batch
ie. (best among (3, 4, 5) e 64)
'themselves', 'what', 'which', 'who", 'whom', tecturc. "this", 'that', "Ihat'll", 'these", 'those", 'am', 'is'. "are', "was", 'were', 'be", "been', "being'. 'have'. "has", 'had', 'having', 'do', 'does', 'did', 'doing' "a', "an', 'the', 'and", 'but', 'if", 'or'. 'because', "as', "until", 'while', 'of", "at', 'by'. 'for', 'with', "about", "aggaist'', 'between', 'into', 'through', 'dur- ing', "before', "after', "above', "below', 'to", 'from', "up'. 'down'. 'in', 'out', 'on", 'off", 'over', 'under", 'again', 'further', 'then'. 'once', 'here', 'there', "when', 'where', "why', 'how", 'all', 'any', 'both', "each', 'few', "more', 'most', 'other', 'some', "such', "only', 'own'. 'so', 'than', 'too', 'very'. "can', 'will', 'just', 'should', "should've" 'now"." "..', '(, ""). We remove all pronouns from the stopword list on tasks where they are required (such as WNLI, where the task is to select the ref- erent of a given pronoun from list of choices).
C Ablation: Necessity of dropping a random subset of stopwords with TokenDrop
We demonstrate the need for dropping a random subset of stopwords from sentences in Table 8. When all tokens considered for random drop- ping, fine-tuning does not converge in the small number of epochs typically used, since dropping important tokens changes the meaning of the input sequences. On the other hand, dropping only a ran- dom subset of stopwords does not alter the meaning of sentences, thereby enabling convergence while also preventing overfitting. We also find that drop- ping all stopwords in each sequence in every epoch leads to loss in accuracy. This is because there is no longer regularization effect, since the same sequences (with all stopwords dropped) are pre- sented to the model in every epoch, and the pruning of stopwords makes sequences shorter, and poten- tially easier for the model to memorize (Table 8).
B Results of fine-tuning Roberta-Large
We present results of fine-tuning Roberta-Large (Liu et al., 2019) on GLUE (Wang et al.) and SQUADv1.1 (Rajpurkar et al.) in Table 7. We find that fine-tuning with TooenDrop BucketSampler produces models that are up to 1.96% more accu- rate, while also reducing the wall-clock fine-tuning time by up to 10.12x compared fine-tuning with RandomSampler. The average accuracy gain from using TokenDrop + BucketSampler on Roberta- Large is more than the gain on Coberta-Base, since overfitting is a bigger problem on larger models. In particular, Roberta-Large has approximately 3 the number of parameters as Roberta-Baan Base, making it more susceptible to memorizing the limited train- ing samples during fine-tuning. On the other hand, the fine-tuning speedup from using TokenDrop + BucketSampler on Roberta-Large is similar to the speedup on Roberta- Base, since speedup depends on the dataset statistics rather than the model archi-
D Reducing overfitting with TokenDrop
We analyze the effectiveness of TokenDrop reducing overfitting when fine-tuning on three datasets of different sizes small (RTE, with 2.5K training samples), medium (SST-2, with 67K train- ing samples), and large (MNLI, with 393K training samples) in Fig. 6. We find that the training loss decreases more gradually when TokenDrop is used, since different data samples are presented in each epoch. As a result, the model does not simply memorize the training data within a few epochs, leading to better generalization performance.

---

Table Results of fine-tuning Roberta-Large on the GLUE and SQUAD v1.1 development sets. We report FI for SQUAD, Matthews correlation for CoLA, Pearson Correlation for STS-B and accuracy for all other tasks. We report only "matched" accuracy for MNLI. Subscripts indicate standard deviation.
Batching Method
SQUAD
MRPC
STS-B
SST-2
CoL.A
QQP
QNLI
RTE
MNLI
WNLI
Avg
RamiemSampler
91. 91.33pa
89.71ps
91.65a.7
96. 26.281.
65,491.4
92.38o
93.72p.1
87.1p.1
89.58
72.124.2
86.94
Trop TokeeDrop-BcckeeSamppe Backeer
91.68a2
90.56os
91.826.7
96,51a
65,54),s
92.74o.s
93.99
88.342.
89.88
74.08,,
81,98
Speedup
9.46X
2.91X
7.28X
10.12X
8,96X
5.6X
3.14X
2.02X
4.19X
3.92X
5.76X
RTE
SST-2
MNAI
15
86
0.5
3
10
30
TokenDrepl
Dropl
TokanCreal
Tolker Drlonrrapp)
TTckenrragl
Test (TokenDrop)
(TekeDrrap)
Trabhing (Taieeerrg
(rilee
Epoch
Figure 6: Training curves obtained from fine-tuning Roberta-base. We report loss averaged across 10 random seeds. Dropout rate = 0.1 is used for all datasets when fine-tuning without Tropn When fine-tuning with TokenDrop. Dropout Rate for RTE, 0.025 for SST-2 and MNLI), as listed Table 6.
Table 8: Evaluation of regularization strategies. We report the average score across GLUE and SQUAD.
of fine-tuning, the other steps of BucketSampler are very similar to those in RandomSampler (with the exception of RBM, which accounts for very small fraction of the runtime), and hence, the over- heads are negligible. We analyze the overheads of TokenDrop BucketSampler on three fine-tuning datasets of different sizes small (RTE, with 2.5K training samples), medium (SST-2, with 67K train- ing samples), and large (MNLI, with 393K training samples) - in Figure 7. We find that the most time- consuming parts of TooenDrop (identifying all stop- words in each sequenee) and BucketSampler (di- viding the dataset into buckets) are performed only once at the start of fine-tuning, On the other hand, the operations performed each epoch account for only a small fraction of the total runtime. Con- sequently, the overheads TokenDrop and Buck- etSampler are better amortized over the course of Line-tuning when training on larger datasets, lead- ing to higher speedups.
Regularization Method
Accuracy
Dropout (Dropout rate = 0.1)
81.62
Drop any random token (Drop rate 0.3)
78.44
TooenDrop: Drop only stopwords (TokenDrup rate 0.3)
Drop stopwords (TokenDrop rate
E Factors impacting speedup from using TokenDrop BucketSampler a given task
We find that the size of the fine-tuning dataset and the sequence length spread play key roles in de- termining the speedup achieved when using To- kenDrop + BucketSampler analyze these relationships in the following subsec- tions.
given task. We
E.1 Dataset size
Both TokenDrop and BucketSampler add small overheads at fine-tuning time. When TokenDrop is used, stopwords are first identified in each se- quence. Then, a random subset of stopwords are pruned from each sequence in every epoch. When BucketSampler is used, the training dataset is first divided into buckets. Then, batches randomly generated from each bucket in every epoch, and residual batches from different buckets are merged. Finally, the generated batches are shuffled to ran- domly order batches from different buckets. As a result, once the buckets are generated at the start
Further accelerating hyperparameter search during fine-tuning. When hyperparameter tun- ing is necessary, it is sufficient to split the dataset into buckets and perform stopword identification only once, during the first epoch of the first fine- tuning run. They can then be re-used for subse- quent runs with different hyperparameters. As a result, if we compare the wall-clock time taken perform tine-tuning with 10 different random seeds, TokenDrop + BucketSampler achieves an average

---

speedup of 6.8x over RandomSampler across the 9 GLUE tasks and SQUAD. We also note that all speedups reported in Section 3 are not computed this way. Instead, is assumed that both steps (stopword identification and bucket generation) are performed in every fine-tuning run, even when re- sults are averaged across multiple random seeds.
50% of all tokens are expected be padding to- kens with RandomSampler), In addition, SST-2 has Lavg = 14, and hence, the majority of batches can be processed with batch sizes >800 (HW Cap = 800 for the bucket with mar_seq_ = 15), lead- ing to large speedups over RandomSumpler (where the batch size of all batches is determined by the longest sequence in the dataset, leading to hard- ware under-utilization when processing batches with short sequences only).
Percentage of wall llll-ccck
100
fine-tuning time (%)
80
60
12
40
DN Speee Soupp
20
STS-0
00P
RTE
SST-2
MRPC
MNU
WNU
ONLI
@ TokenDrop
BucketSampler
Training
RTE
2
0.15
0.25
0.35
0.45
Figure 7: Overheads of TokenDrop + BucketSampler. Times are measured on a NVIDIA RTX 2080 Ti GPU with 11 GB memory.
Expected fraction of padding tokens
Figure 8:
Impact of sequence length spread of dataset on fine-tuning speedup achieved using To- kenDrop BucketSampler.
E.2 Sequence length spread of dataset
We quantify the sequence length spread a dataset using two parameters: Lavg and Lg. Lavg is the average sequence length of all sequences in the dataset, while LK is the maximum possible sequence length such that at least K% of all se- quences in the dataset are longer than Lg. When K 87.5% and assuming batch size of 64 (HW Cap 64 when max__eq_ NVIDIA RTX 2080 Ti GPU), the probability of each batch having at least one sequence with length 0.9999 when batch-
128 on
LK = (1 - (87 1/107/5/100// ing with RandomSampler. As a result, it is ex- pected that each batch will be padded to at least L87.5- Consequently, ((187.5 I=awg// / Lgg.s) s) is conservative estimate of the fraction of all tokens that are expected to be padding tokens when batching with RandomSam- pler. We observe that falding_fraclinn a direct correlation with speedup achieved when using TooenDrop + BucketSampler (Fig. 8). We observe some outliers when the datasets are very small, as in the case of RTE (2.5k training samples) and WNLI (634 training samples), since the over- heads of TokenDrop BucketSampler account for larger fraction of the wall-clock fine-tuning time (see Fig. 7). We achieve maximum speedup on SST-2, relatively large dataset (67K training sam- ples) with high madding_fraction_ee (nearly
has


Reviews:
[{"review_id": "JKE3atLJDa", "reviewer": "Reviewer_QVce", "summary": "", "questions": "1. Figure 1 is used to show that there is still overfitting during fine-tuning without Dropout. It is also claimed that TokenDrop is better in preventing overfitting. Are there any statistics/charts to support this?\n2. Line 186, it states that “Unlike Dropout, TokenDrop introduces data diversity between the different training epochs”. This sounds inaccurate to me, since Dropout also causes data diversity by random sampling in different training epochs.\n3. Algorithm 1, Line 8: is this line incomplete?\n4. Algorithm 3, Line 4: does this mean for non-residual batches, TokenDrop is not applied? If that’s true, what’s the ratio of tokens dropped during fine-tuning?", "limitations": "", "rating": 3, "confidence": 3, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the insightful review. We address the questions and comments below.\n\n**Re:  How does the speedup change with different fine-tuning set sizes, different length distributions?**\n\n**Dataset size**: Speedup is typically higher on larger datasets, since the overheads introduced by TokenDrop+BucketSampler are better amortized over the course of fine-tuning. For instance, the overheads of TokenDrop+BucketSampler account for approximately 20% of the wall-clock fine-tuning time when training on RTE, which has only 2.5K training samples. On the other hand, when fine-tuning on MNLI, which is a relatively large dataset with 393K training samples, the overheads of TokenDrop+BucketSampler account for only around 5% of the wall-clock fine-tuning time. This can be seen in Fig. 5 in Appendix C (page 11) in the paper. We also note that in spite of the overheads, TokenDrop+BucketSampler achieves substantial speedups over prior methods even on small datasets. \n\n**Length distribution**: Potential for speedup is maximum when a majority of sequences in the dataset are short, but there are a substantial number of long sequences as well. For example, in the SST-2 dataset, approximately 80% of all sequences have length <20, but approximately 10% of all sequences have length >35. As a result, when using RandomSampler with a batch size of 64, there is a >99% chance that each batch will have at least one sequence with length >35, leading to all samples in the batch being padded to >35 tokens. In contrast, when using TokenDrop+BucketSampler, the majority of batches can be processed with length <20, leading to large speedups over RandomSampler. This can be seen in Table 1, where speedup is maximum on SST-2. We will add a detailed analysis in the final paper about the correlation between the sequence length spread and size of a dataset, and the speedup achieved.\n\n**Re: Stronger baselines (e.g., LengthGroupedSampler, or tf.bucket_by_sequence_length without varying batch sizes).**\n\nWe compare TokenDrop+BucketSampler with LengthGroupedSampler and tf.bucket_by_sequence_length with variable batch sizes in Table 4 (page 8 in our paper) in terms of both accuracy and efficiency. In summary, we find that TokenDrop+BucketSampler achieves ~0.3% higher accuracy along with 1.6X speedup over LengthGroupedSampler, and 5% higher accuracy along with 1.1X speedup over tf.bucket_by_sequence_length with variable batch sizes (batch_size = HWCap for each bucket) averaged across the GLUE tasks and SQUAD with Roberta-Base. Following the reviewer’s suggestion, we performed a new experiment and found that TokenDrop+BucketSampler achieves ~0.3% higher accuracy along with 1.5X speedup over tf.bucket_by_sequence_length without variable batch sizes averaged across the GLUE tasks and SQUAD. We can also report speedups over these methods for each task (similar to Table 1) in the final paper. We chose RandomSampler as baseline for the primary results since it is by far the most widely used method in practice for fine-tuning LLMs.\n\n**Re: The ablation study does not clearly reveal the contribution from TokenDrop.**\n\nFollowing the reviewer’s suggestion, we have evaluated TokenDrop’s contribution. When using only BucketSampler, the average accuracy is 81.67% and speedup (over RandomSampler) is 5.1X on GLUE and SQUAD with Roberta-Base. When both TokenDrop and BucketSampler are used, the accuracy and speedup increase to 81.98% and 5.9X, respectively. Therefore, the use of TokenDrop clearly improves both the accuracy and efficiency of fine-tuning. We thank the reviewer for pointing out this missing entry, and we will add this to Table 4 in the final paper.\n\n**Re: Figure 1 is used to show that there is still overfitting during fine-tuning without Dropout. It is also claimed that TokenDrop is better in preventing overfitting. Are there any statistics/charts to support this?**\n\nAs we are unable to attach figures to the rebuttal, we use the table below to compare the training and validation losses with and without TokenDrop while fine-tuning Roberta-Base on RTE. The training loss decreases more gradually when TokenDrop is used, since different data samples are presented in each epoch. As a result, the model does not simply memorize the training data within a few epochs, leading to better generalization performance. We will add graphs similar to Figure 1 to demonstrate that TokenDrop does indeed reduce the impact of overfitting when fine-tuning on datasets of different sizes in the final paper using both Roberta-Base and Roberta-Large. \n\n|          | Without TokenDrop   |              | With TokenDrop |            |\n|--------| -------- | ------- | ------- | ------- |\n| Epoch| Train Loss | Validation Loss | Train Loss | Validation Loss |\n| 1 | 0.59 | 0.83| 0.57 | 0.79|\n| 2 | 0.32 | 0.68 | 0.42 | 0.57|\n|3| 0.16|0.59|0.28|0.5|\n|4|0.11|0.61|0.19|0.47|\n|5|0.09|0.69|0.17|0.51|\n|6|0.08|0.77|0.14|0.52|\n|7|0.06|0.98|0.1|0.54|\n|8|0.04|1.16|0.09|0.61|\n|9|0.04|1.36|0.08|0.74|\n|10|0.03|1.39|0.05|0.87|\n\n**Re: Line 186, it states that “Unlike Dropout, TokenDrop introduces data diversity between the different training epochs”. This sounds inaccurate to me, since Dropout also causes data diversity by random sampling in different training epochs.**\n\nDropout only randomly drops a random subset of neurons in each training iteration. The training set is not changed between epochs. On the other hand, TokenDrop changes the training set itself in each epoch, by dropping a random subset of stopwords from each sample. This is what we meant by data diversity between different epochs, since a different subset of stopwords is dropped from each training sample in every epoch, leading to different training samples being presented to the model in every epoch. We will clarify this in the paper.\n\n**Re: Algorithm 1, Line 8: is this line incomplete?**\n\nYes, it should be Tokens_to_drop(sample)--. We thank the reviewer for pointing out this mistake, and we will fix it in the paper.\n\n**Re: Algorithm 3, Line 4: does this mean for non-residual batches, TokenDrop is not applied? If that’s true, what’s the ratio of tokens dropped during fine-tuning?**\n\nTokenDrop is applied to non-residual batches also. Since batches are only formed from sequences within the same bucket, all sequences in a batch will have min_seq_len(bucket) < length < max_seq_len(bucket). Therefore, if a sequence in a non-residual batch has length greater than the minimum length of all sequences in that batch, then {length - min_seq_len(batch)} tokens are dropped from that sequence. In our experiments, we define buckets in length increments of 5, i.e, bucket 1 has (min_seq_len, max_seq_len) = [0,5), bucket 2 has (min_seq_len, max_seq_len) = [5,10) and so on. If, for example, a batch drawn from bucket 2 has 4 sequences with lengths (5, 6, 7 and 9), then the number of tokens dropped from these sequences are (0, 1, 2, 4), respectively. The ratio of tokens dropped varies across datasets, and is typically between 15-20% of all tokens. We will make this clearer in the final version of the paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. This paper studies the improvement of an important task. Fine-tuning LLM is the most widely adapted way to leveraging LLMs. Improving the fine-tuning efficiently can have great impact in both academia and industry.\n2. The proposed techniques can help both LLM fine-tuning and inference. Especially for the real-time inference scenario, the experiments show after the fine-tuning, dropping stop words during inference time can bring a speedup of 2.2X and almost without accuracy loss.", "reasons_to_reject": "1. It would be good to show some more detailed analyses. For example, how does the speedup change with different fine-tuning set sizes, different length distributions?\n2. The random sampler baseline used in empirical evaluation is weak. Maybe a simple sequence length bucketing baseline (e.g., `LengthGroupedSampler`, or `tf.bucket_by_sequence_length` without varying batch sizes) could be a better baseline?\n3. The ablation study does not clearly reveal the contribution from TokenDrop.", "questions_to_authors": "1. Figure 1 is used to show that there is still overfitting during fine-tuning without Dropout. It is also claimed that TokenDrop is better in preventing overfitting. Are there any statistics/charts to support this?\n2. Line 186, it states that “Unlike Dropout, TokenDrop introduces data diversity between the different training epochs”. This sounds inaccurate to me, since Dropout also causes data diversity by random sampling in different training epochs.\n3. Algorithm 1, Line 8: is this line incomplete?\n4. Algorithm 3, Line 4: does this mean for non-residual batches, TokenDrop is not applied? If that’s true, what’s the ratio of tokens dropped during fine-tuning?", "paper_topic_and_main_contributions": "This paper presents two techniques to improve the efficiency of LLM fine-tuning. The first one is to randomly drop out a subset of stop words in each sequence in every epoch to reduce the number of tokens to encode/decode and alleviate overfitting. The second one is based on  bucketing, which include several improvements: 1) batch size is determined based on hardware limit and bucket sequence length, 2) residual buckets are merged together, 3) adaptive learning rate based on batch size. Empirical evaluations show the combination of the two techniques can significantly speedup the fine-tuning without loss of accuracy, comparing to a random sampler baseline.", "excitement": 3, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "63qqO9ldzZ", "reviewer": "Reviewer_4LK5", "summary": "", "questions": "A) The title and the build-up to the experimental section is slightly misleading. The title says \"large language models\" but the term \"large\" is very subjective. The introduction mentions \"LLMs are characterized by large model sizes (exceeding 1 Trillion parameters), and are pre-trained on very large text corpora\", but experiments are conducted on base models (~110 million parameters). While it is definitely not an issue if the experiments are done on base models, it would benefit the reader if the wording aligns with the experiments, and removing the word \"large\" throughout would help achieve that. \n\nB) L377-381: How does this formulation of LRM come into place? Is it borrowed from previous literature or motivated due to other reasons?  Motivating its formulation might help add clarity to the reader.\n\nC) Since there are multiple components involved in the framework, laying out its scope in terms of tasks and models it would benefit would be really helpful for its usability in the real world.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful feedback. We address the main questions and comments below and provide new results on larger models.\n\n**Re: Applicability to generation tasks**\n\nWhile the reviewer is correct that generalization tasks do not allow ungrammatical outputs, TokenDrop is applied only on the input text sequences during fine-tuning, and never on the output sequences. As a result, TokenDrop can be used for generalization tasks, as long as applying TokenDrop on the input does not change the label (like in the case of CoLA, where grammatically correct sequences become incorrect when TokenDrop is used). We validate this by using TokenDrop+BucketSampler to fine-tune the T5-small seq2seq model on xsum, a text summarization task. We achieve an average ROGUE-1 score of 32.84 when fine-tuning with TokenDrop+BucketSampler, which is 0.26 points higher than the score obtained using RandomSampler. In addition, fine-tuning with TokenDrop+BucketSampler is 8.62X faster than RandomSampler. \n\nWe also find that fine-tuning with TokenDrop **significantly increases the resilience of LLMs to grammatical errors** in text sequences during inference. LLMs trained with RandomSampler are adversely impacted by even minor grammatical errors. For instance, if only articles (“a”, “an”, “the”) were dropped from the test inputs, the average ROUGE-1 score drops from 32.58 to 29.98. On the other hand, when LLMs are trained with TokenDrop, the drop in ROUGE-1 score is only 0.05 (32.84 -> 32.79) when articles are dropped, indicating that they are resilient to minor grammatical errors. Since conventional fine-tuning (RandomSampler) uses only grammatically perfect sequences to train the model, high accuracy is obtained only on grammatically perfect test sequences. This can also be seen in Fig 3 (page 7 in our paper), where dropping only a small subset of stopwords that do not change the meaning of sequences leads to a large drop in accuracy on models trained with RandomSampler. This is important because the assumption of seeing only grammatically correct inputs during inference does not always hold in real-world scenarios, especially since non-native speakers also use these tools. \n\n**Re: Results on larger models**\n\nFollowing the reviewer’s suggestion, we experiment with fine-tuning Roberta-Large on GLUE and SQUAD. When RandomSampler is used, we obtain an average score of 86.94 across the 10 tasks (averaged across 10 random seeds). With TokenDrop+BucketSampler, the average score increases to 87.51. Therefore, TokenDrop+BucketSampler leads to 0.57% accuracy gain on Roberta-Large. The accuracy gain on Roberta-Large is more than the gain on Roberta-Base, since overfitting is a bigger problem on larger models. Fine-tuning with TokenDrop+BucketSampler is also 6X faster than fine-tuning with RandomSampler. The speedup on Roberta-Large is very similar to the speedup on Roberta-Base, since speedups depend on the dataset statistics rather than the models (max batch size with RandomSampler and HWCap for the buckets are affected by similar amounts when moving to larger models, so speedups remain approximately constant across model sizes). We will add task-wise results on Roberta-Large to Table 1 in the final paper.\n\n**Re: Applicability to different languages, and the need to know stopwords apriori**\n\nThe reviewer is correct about needing to know the stopwords apriori for different languages.  Lists of stopwords in most major languages are publicly available (https://github.com/6/stopwords-json, https://advertools.readthedocs.io/en/master/advertools.stopwords.html etc.). However, we acknowledge that these lists may not be available for low-resource and understudied languages. We will add this to the limitations section in the final paper.\n\n**Cross-lingual generalizability** – Intuitively, we do not expect TokenDrop to affect cross-lingual generalizability since (a) it only drops words that do not contribute to the meaning of the sentence, and (b) only a random subset of stopwords are dropped in each iteration i.e., not all stopwords are dropped from every sequence, and a different set of stopwords are dropped from a given sequence in each epoch. However, we are unable to provide empirical evidence for this in the short rebuttal period. We will either empirically show that TokenDrop does not affect cross-lingual transfer in the final paper, or call it out in the limitations section.\n\n**Re: The title and the build-up to the experimental section is slightly misleading**\n\nWe used the term LLM to refer to a particular class of models – pre-trained foundation models that are subsequently fine-tuned for different NLP tasks – rather than models that are enormous. However, after re-reading the paper in light of the reviewer’s comments, we agree that the title and intro may be slightly misleading. We will remove references to “large” in the final version, and place more emphasis on the class of models we are targeting (foundation models or language models)\n\n**Re: How does this formulation of LRM come into place? Is it borrowed from previous literature or motivated due to other reasons? Motivating its formulation might help add clarity to the reader.**\n\nScaling learning rate by a factor of sqrt(k) when the batch size is scaled by a factor of k is a popular method for parallelizing training of deep neural networks on large GPU clusters (https://arxiv.org/abs/1404.5997). In particular, when batch size is changed to k * batch_size, the learning rate is changed to learning_rate * sqrt(k). Prior works have used this square root scaling for fixed batch sizes (i.e., batch size is determined at the start of training based on the number of GPUs available) to find the optimal learning rate for a given batch size when the optimal learning rate is known for a different batch size. In this work, we find that this formulation works well for variable batch sizes also (i.e., changing learning rate for each batch based on the size of the batch). We thank the reviewer for pointing out this missing motivation, and we will add this to the final paper along with the relevant references.\n\n**Re: Since there are multiple components involved in the framework, laying out its scope in terms of tasks and models it would benefit would be really helpful for its usability in the real world.**\n\nBucketSampler is useful when fine-tuning is performed on datasets with variable length sequences. TokenDrop is useful on tasks where dropping stopwords from input sequences does not change the output labels. We will make this clear in the final paper.\n\n**Re: Typos Grammar Style And Presentation Improvements**\n\nWe thank the reviewer for the suggestions. We will incorporate them in the final paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "A) Writing: The paper is easy-to-understand and the framework is detailed. The motivation is clear and the shortcomings of past works are well laid-out.\n\nB) Code and Soundness: The code is released and results reported are averaged across 10 seeds. The framework is compared to several past methods and lay out the advantages over them.\n\nC) Ablation: Each piece of the framework and its contribution is studied in Table 3.\n\nD) The efficiency speedups are impressive, all while maintaining accuracy.", "reasons_to_reject": "A) Scope and Generalizability: The major weakness in the framework's design is its limited scope to multiple tasks and languages. With regards to tasks, it is not applicable to a few classification- and all generation-tasks. Amongst classification tasks, TokenDrop cannot be used for COLA (which judges grammaticality), and generation tasks (like machine translation or summarization) wouldn't allow for ungrammatical outputs either (although the paper does not mention generation tasks but clarification on the scope would be beneficial). In general, the method presents a mismatch between pre-training and finetuning, which can be problematic in a few settings like zero-shot cross-lingual transfer in multilingual models. If a multilingual model is fine-tuned in English with stopwords dropped it might affect cross-lingual generalization. One would also require knowing stopwords apriori if finetuning on different languages. \nWith regards to models, the framework is tested on base models (110M params), while being motivated as an efficient LLM fine-tuning framework. I do not see the experimental setup of testing only base models to be problematic, but its framing is misleading (elaborated in the questions section). However, I bring up this point here because the method involves optimization tricks like learning rate modulation and BatchCap, whose generalizability to larger models is uncertain. If instead of electra-base for example, the paper would include roberta-base and roberta-large (300-350M params?), it would alleviate some of these concerns.", "questions_to_authors": "A) The title and the build-up to the experimental section is slightly misleading. The title says \"large language models\" but the term \"large\" is very subjective. The introduction mentions \"LLMs are characterized by large model sizes (exceeding 1 Trillion parameters), and are pre-trained on very large text corpora\", but experiments are conducted on base models (~110 million parameters). While it is definitely not an issue if the experiments are done on base models, it would benefit the reader if the wording aligns with the experiments, and removing the word \"large\" throughout would help achieve that. \n\nB) L377-381: How does this formulation of LRM come into place? Is it borrowed from previous literature or motivated due to other reasons?  Motivating its formulation might help add clarity to the reader.\n\nC) Since there are multiple components involved in the framework, laying out its scope in terms of tasks and models it would benefit would be really helpful for its usability in the real world.", "paper_topic_and_main_contributions": "This paper proposes a framework to improve efficiency during fine-tuning of language models, without harming the accuracy. The efficiency speedup is made possible by BucketSampler, which creates batches of samples having similar sequence lengths. BucketSampler has three additional components, the first two of which are necessary to achieve convergence: a) modulating the learning rate based on the batch size; b) having a cap on the maximum batch size (which can grow large for batches that are composed of small sequences) depending on the training epoch to prevent sparse updates; c) a way to merge residual samples from multiple buckets. TokenDrop is a method that drops tokens (specifically, stopwords only) from sequences, which serves the dual purpose of acting as a regularizer and creating batches of equal length (once they have been bucket-ed). Their framework is tested on roberta-base and electra-base for all the GLUE tasks and SQuAD. They achieve an average speedup of 5.9x with slightly better (~0.3% gain) accuracies.", "excitement": 4, "reproducibility": 5, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Missing_References": "not extremely well-versed with the literature so not sure of this if any"}, "extra_text": {"Typos_Grammar_Style_and_Presentation_Improvements": "A) The introduction and Section 2 (L158-171) while clear, seems repetitive. For example, the method and its motivation is described multiple times (L010-L32, L076-L129, L141-L151, L158-L171). This can be cut down to two overall places of description instead of 4, one in the abstract and one the introduction. The saved space can be used to create a visualization for Section 2.2 which describes BucketSampler and the three components with a toy example?"}}, {"review_id": "pdZTYI95PZ", "reviewer": "Reviewer_4fLH", "summary": "", "questions": "1. At line 8, authors mention \"which is expected to be performed much more broadly albeit at lower computational cost.\", what do they mean by more broadly?\n2. In line 442, authors mention that \"the datasets with larger numbers of short sequences benefit more from BucketSampler\".  Why that is the case? My intuition is that when the distribution of the length is spread, with sequences of varying lengths, we would benefit more from this technique. Would it be possible to see where the higher percentage of shorter sequences has a correlation with speed up gains?\n3. In section 2.3, authors mention min_seq_len and max_seq_len. How are these selected? \n4. How did you choose the dynamic learning rate for a batch in line 379? What was the intuition, did you try different formulas?\n5. In line 316, O(n) has been used. Why do we need the complexity notation? Why not just _n_?\n6. In line 258 authors mention \"HWCap is experimentally determined by profiling on-chip memory usage for different sequence lengths.\". Does this mean that for every hardware or model, one has to do manual inspection. If so, is this a limitation on using the method out of the box?\n7. Authors claim (line 180) that \"we find that dropout is ineffective at preventing overfitting during fine-tuning (Fig. 1). The Figure depicts only RoBERTa base model on RTE task. Can we make such a broad claim based on one setting?", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed feedback. We address the main questions and comments below, including a new comparison and new results on larger models.\n\n**Re: While authors mention LengthGroupedSampler in the related work section, there is no comparison in the experiments.**\n\nWe compare TokenDrop+BucketSampler with LengthGroupedSampler in Table 4 (page 8 in our paper) in terms of both accuracy and efficiency. In summary, we find that TokenDrop+BucketSampler achieves ~0.3% higher accuracy and 1.6X speedup over LengthGroupedSampler averaged across the GLUE tasks and SQUAD with Roberta-Base. We can also report the speedup over LengthGroupedSampler for each task (similar to Table 1) in the final paper. We chose RandomSampler as baseline for the primary results since it is by far the most widely used method in practice for fine-tuning LLMs.\n\n**Re: The authors have only experimented with RoBERTa base. It will be good to include experiments with seq2seq models and also slightly larger models, for instance RoBERTa large, to see how different architectures or size affect the speed up or performance drop.**\n\nIn addition to Roberta-Base, we also provide results on Electra-Base in Table 1 (page 7 in our paper). \n\n**Larger models** – Following the reviewer’s suggestion, we performed experiments with Roberta-Large on GLUE and SQUAD. When RandomSampler is used, we obtain an average score of 86.94 across the 10 tasks (averaged across 10 random seeds). With TokenDrop+BucketSampler, the average score increases to 87.51. Therefore, TokenDrop+BucketSampler leads to 0.57% accuracy gain on Roberta-Large. The accuracy gain on Roberta-Large is more than the gain on Roberta-Base, since overfitting is a bigger problem on larger models. Fine-tuning with TokenDrop+BucketSampler is also 6X faster than fine-tuning with RandomSampler. The speedup on Roberta-Large is very similar to the speedup on Roberta-Base, since speedups depend on the dataset statistics rather than the models (max batch size with RandomSampler and HWCap for the buckets are affected by similar amounts when moving to larger models, so speedups remain approximately constant across model sizes). We will add task-wise results on Roberta-Large to Table 1 in the final paper.\n\n**seq2seq models** – We experiment by using TokenDrop+BucketSampler to fine-tune the T5-small seq2seq model on xsum, a text summarization task. We achieve an average ROGUE-1 score of 32.84 when fine-tuning with TokenDrop+BucketSampler, which is 0.26 points higher than the score obtained using RandomSampler. In addition, fine-tuning with TokenDrop+BucketSampler is 8.62X faster than RandomSampler. \n\nWe also find that fine-tuning with TokenDrop **significantly increases the resilience of LLMs to grammatical errors** in text sequences during inference. LLMs trained with RandomSampler are adversely impacted by even minor grammatical errors. For instance, if only articles (“a”, “an”, “the”) were dropped from the test sequences, the average ROUGE-1 score drops from 32.58 to 29.98. On the other hand, when LLMs are trained with TokenDrop, the drop in ROUGE-1 score is only 0.05 (32.84 -> 32.79) when articles are dropped, indicating that they are resilient to minor grammatical errors. Since conventional fine-tuning (RandomSampler) uses only grammatically perfect sequences to train the model, high accuracy is obtained only on grammatically perfect test sequences. This can also be seen in Fig 3 (page 7 in our paper), where dropping only a small subset of stopwords that do not change the meaning of sequences leads to a large drop in accuracy on models trained with RandomSampler. This is important because the assumption of seeing only grammatically correct sequences during inference does not always hold in real-world scenarios, especially since non-native speakers also use these tools. \n\n**Re: At line 8, authors mention \"which is expected to be performed much more broadly albeit at lower computational cost.\", what do they mean by more broadly?**\n\nCurrently, pre-training of LLMs can only be performed by those with access to large GPU/ TPU clusters. As a result, fine-tuning pre-trained LLMs is the primary way of leveraging LLMs for the vast majority of users. This is what we meant by “more broadly”. Pre-trained LLMs that are publicly available (such as BERT, OPT, etc.) are fine-tuned millions of times by large numbers of users on a diverse range of downstream tasks. We will make this clearer in the final paper.\n\n**Re: In line 442, authors mention that \"the datasets with larger numbers of short sequences benefit more from BucketSampler\". Why that is the case? It will be nice to see the relationship between the spread of length across a dataset and the speed up gain.**\n\nThat line was incomplete. We thank the reviewer for pointing this out, and we will fix it in the final paper. Potential for speedup is maximum when a majority of sequences in the dataset are short, but there are a substantial number of long sequences also. For example, in the SST-2 dataset, approximately 80% of all sequences have length <20, but approximately 10% of all sequences have length >35. As a result, when using RandomSampler with a batch size of 64, there is a >99% chance that each batch will have at least one sequence with length >35, leading to all samples in the batch being padded to >35 tokens. In contrast, when using TokenDrop+BucketSampler, the majority of batches can be processed with length <20, leading to large speedups over RandomSampler. This can be seen in Table 1, where speedup is maximum on SST-2. We will add a detailed analysis in the final paper about the relationship between the sequence length spread and size of a dataset, and the speedup achieved.\n\n**Re: In section 2.3, authors mention min_seq_len and max_seq_len. How are these selected?**\n\nmin_seq_len and max_seq_len of a bucket determine the sequence length range of all samples in the bucket. When TokenDrop is not used, we simply define buckets in length increments of 1 i.e., bucket 1 will have sequences of length 1, bucket 2 will have sequences of length 2, and so on. Therefore, the number of buckets will be the max_seq_len of the dataset, and the [min_seq_len, max_seq_len) of bucket i will be [i,i). This will eliminate the need for padding in non-residual batches (since all sequences in a bucket will have the same sequence length), and RBM will keep hardware utilization as high as possible while minimizing padding in residual batches. When TokenDrop is used, the sequence length spread (max_seq_len - min_seq_len + 1) of each bucket must be >1 to enable token dropping in non-residual batches also. However, if the sequence length spread of each bucket is too large, then it is highly likely that Tokens_to_drop(sample) will be larger than number_of_stopwords(sample) for most samples (Algorithm 3, line 4, page 6). This has two drawbacks: (1) it necessitates the use of padding, and (2) all stopwords in many samples will be dropped in every epoch, thereby reducing data diversity between epochs and leading to weaker regularization. In our experiments, we find that defining buckets in increments of 5, where bucket 1 has (min_seq_len, max_seq_len) = [0,5), bucket 2 has (min_seq_len, max_seq_len) = [5,10) and so on, provides good trade-off between minimizing padding and ensuring sufficient regularization for high validation accuracy. We thank the reviewer for pointing out this missing detail, and we will add an additional subsection in Appendix A describing how hyperparameters were chosen to complement the list of hyperparameters in Table 5.\n\n**Re: How did you choose the dynamic learning rate for a batch in line 379? What was the intuition, did you try different formulas?**\n\nScaling learning rate by a factor of sqrt(k) when the batch size is scaled by a factor of k is a popular method for parallelizing training of deep neural networks on large GPU clusters (https://arxiv.org/abs/1404.5997). In particular, when batch size is changed to k * batch_size, the learning rate is changed to learning_rate * sqrt(k). Prior works have used this square root scaling for fixed batch sizes (i.e., batch size is determined at the start of training based on the number of GPUs available) to find the optimal learning rate for a given batch size when the optimal learning rate is known for a different batch size. In this work, we find that this formulation works well for variable batch sizes also (i.e., changing learning rate for each batch based on the size of the batch). We thank the reviewer for pointing out this missing motivation, and we will add this to the final paper along with the relevant references.\n\n**Re: In line 316, O(n) has been used. Why do we need the complexity notation? Why not just n?**\n\nThe reviewer is correct. There is no need for the complexity notation. We will change it to just n in the final paper.\n\n**Re: In line 258 authors mention \"HWCap is experimentally determined by profiling on-chip memory usage for different sequence lengths.\". Does this mean that for every hardware or model, one has to do manual inspection. If so, is this a limitation on using the method out of the box?**\n\nWe automate this entire process through a simple script that tracks the peak memory usage of a single forward + backward pass for a given model and sequence length with a random base batch size. We then increase or decrease the batch size depending on the ratio of peak memory usage to hardware memory capacity, and this process is repeated till the peak memory usage is less than but close to the hardware memory capacity. These steps are then repeated for different sequence lengths under consideration. Since this process is highly efficient and takes less than a minute for all models in our experiments, we do not believe that this will be a limitation on using our methods out of the box.\n\n**Re: Authors claim (line 180) that \"we find that dropout is ineffective at preventing overfitting during fine-tuning (Fig. 1). The Figure depicts only RoBERTa base model on RTE task. Can we make such a broad claim based on one setting?**\n\nWe observe similar trends on other tasks also, with both Roberta-Base and Electra-Base. We also observe that the overfitting problem is even more severe on Roberta-Large. We will add more curves when fine-tuning on datasets of different sizes in the final paper. We chose RTE to demonstrate how severe the overfitting problem can be during fine-tuning even with base models, and to motivate the need for stronger regularization.\n\n**Re: Missing References, Typos Grammar Style And Presentation Improvements**\n\nWe thank the reviewer for the suggestions and for pointing out a source of relevant references. We will incorporate them in the final paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "The techniques authors have introduced and their experiments show that fine-tuning and inference of ReBERTa can be made faster by up to 10x on some datasets, specifically GLUE and SQUAD.\n\nWhile authors have experimented with RoBERTa, this method cab be applicable to other transformer models.\nThe authors have carried out an ablation study to show which techniques can lead to speed up or prevent performance drop.\n", "reasons_to_reject": "- While authors mention LengthGroupedSampler in the related work section, there is no comparison in the experiments.\n- I would like to have seen more analysis of the results. for instance, it will be nice to see the relationship between the spread of length across a dataset and the speed up gain. In Table 1, we see that speed up gain is different for different dataset but it will be more helpful to see why that is. \n- The authors have only experimented with RoBERTa base. It will be good to include experiments with seq2seq models and also slightly larger models, for instance RoBERTa large, to see how different architectures or size affect the speed up or performance drop.", "questions_to_authors": "1. At line 8, authors mention \"which is expected to be performed much more broadly albeit at lower computational cost.\", what do they mean by more broadly?\n2. In line 442, authors mention that \"the datasets with larger numbers of short sequences benefit more from BucketSampler\".  Why that is the case? My intuition is that when the distribution of the length is spread, with sequences of varying lengths, we would benefit more from this technique. Would it be possible to see where the higher percentage of shorter sequences has a correlation with speed up gains?\n3. In section 2.3, authors mention min_seq_len and max_seq_len. How are these selected? \n4. How did you choose the dynamic learning rate for a batch in line 379? What was the intuition, did you try different formulas?\n5. In line 316, O(n) has been used. Why do we need the complexity notation? Why not just _n_?\n6. In line 258 authors mention \"HWCap is experimentally determined by profiling on-chip memory usage for different sequence lengths.\". Does this mean that for every hardware or model, one has to do manual inspection. If so, is this a limitation on using the method out of the box?\n7. Authors claim (line 180) that \"we find that dropout is ineffective at preventing overfitting during fine-tuning (Fig. 1). The Figure depicts only RoBERTa base model on RTE task. Can we make such a broad claim based on one setting?", "paper_topic_and_main_contributions": "The paper introduces two techniques or tricks to utilities the hardware for training large language models, while avoiding degrading of accuracy. These two methods are:\n- BucketSampler: they create batches based on the size of the sequences. Longer sequences are grouped together and shorter ones are grouped together.\n- Token Drop: as a form of regulariser and also to reduce padding, the authors introduce TokenDrop. By dropping tokens, at a higher rate for longer sequences, they reduce the need for padding and therefore, making the training faster and more efficient.\n- They also use a few more tricks to avoid large batches (batch cap) or small batches (residual batch merging)\n\nTheir results on GLUE and SQUAD dataset, using Roberta model only, shows that using all the techniques introduce, there could be a speed up gain of up to x10, while avoiding performance drop. The techniques are more effective when a dataset has sequences of a varied length, especially higher number of short sequences.", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Missing_References": "- Bucketing was a technique that was (perhaps introduced) and used widely when RNNs and LSTMs became very popular around 2015/6. There are no references to those earlier works. "}, "extra_text": {"Typos_Grammar_Style_and_Presentation_Improvements": "in the second page, when authors are listing their contributions, there are outlined in 4 points. However, those 4 points are repetition of the TokenDrop and BucketSampler. I suggest authors make it more concise and reduce the contribution two two points 1. TokenDrop 2. BucketSampler. Authors can add another bullet point for the extra techniques they introduced such as dynamic learning rate or residual batch merging."}}]

Author rebuttal:

Meta-review: This paper presents two methods to improve the efficiency of fine-tuning, one relates to randomly dropping some of the input tokens, and the second is based on bucketing.  The reviewers appreciated the quality of the results, and the ablations. Some of the reviewers initially had some concerns around the analysis and the baselines, as well as the scope of the experiments. The rebuttal helped alleviate many of these concerns.
Decision: Accept-Findings