Effective Pretraining Objectives for Transformer-based Autoencoders (2022)

Abstract: In this paper, we study trade-offs between efficiency, cost and accuracy when pre-training Transformer encoders with different pre-training objectives. For this purpose, we analyze features of common objectives and combine them to create new effective pre-training approaches. Specifically, we designed light token generators based on a straightforward statistical approach, which can replace ELECTRA computationally heavy generators, thus highly reducing cost. Our experiments also show that (i) there are more efficient alternatives to BERT’s MLM, and (ii) it is possible to efficiently pre-train Transformer-based models using lighter generators without a significant drop in performance.

## introduction transformer-based models @xcite require expensive hardware to be pre-trained @xcite . recently, many works focused on reducing pre-training cost @xcite @xcite . electra, for example, proposes to train bert as a discriminator rather than a generator @xcite . they replace the masked language modeling objective (mlm) @xcite with token detection (td). then, the discriminator detects if input tokens are original or fake created by a small generator network. on the one hand, the discriminator is much more efficient than mlm. on the other hand, the use of a generator requires the pre-training of a second transformer, increasing the pre-training cost. electra has been shown to be more accurate than bert. however, it is not clear if this superior performance is due to its innovative architecture or to the long and extensive training, which highly increases the computation cost for obtaining the final language model. ## related work many different objectives for self-supervised learning have been proposed in recent works, such as causal language modeling (clm) @xcite @xcite , masked language modeling (mlm) @xcite and token detection (td) @xcite , the latter used by electra, which is composed by a generator and a discriminator. while the generator is trained with mlm to find suitable candidates to replace the special mask tokens, the discriminator should recognize the replacements in the original text instead. after pretraining, the generator is removed, and the discriminator is used as the pre-trained language model. electra introduces many innovations: (i) the exploitation of the whole output of the discriminator to compute the loss function, thus having a stronger signal for the back-propagation; (ii) the usage of a generator network to find suitable replacements and (iii) the fact that the discriminator does not see spurious tokens such as the mask token. the latter is a main drawback of the original bert, as it creates input discrepancies between pre-training and fine-tuning, since the cls representation is dependent on all input tokens thanks to the self-attention mechanism. other research directions to reduce training time address the architecture instead of the learning objective. in albert @xcite , the authors tie the weights of every transformer layer to save gpu memory, thus enabling bigger batch sizes. however, since the expressive power of their models is reduced when layers are tied, they must train for much longer. @xcite and @xcite instead use distillation to reduce the model size, but the pre-training is still expensive because it requires a large teacher architecture. although pre-training is performed only once, it usually requires weeks and costly machines @xcite , so it is important to find alternative ways to pre-train transformers. @xcite provides an overview of many recent advancements in transformer efficiency. another successful mlm improvement regarding the objectives is spanbert @xcite , which proposes two new objectives: span-masking and span-boundary-objective (sbo). specifically, the span-masking objective is a refined version of mlm that masks contiguous spans of text instead of single tokens, while, with sbo, they predict the span content by considering only the output representation corresponding to the tokens on the boundaries. furthermore, in @xcite , the authors propose a technique to improve downstream performance by adapting the model to the final task while pre-training. similarly, in di liello et al. (2022) they do continuous pre-training with custom objectives to better adapt the model for answer sentence selection (as2). in t5 @xcite , the authors propose to use deshuffling @xcite to pre-train an autoregressive model. they shuffle random spans of text and ask the model to output tokens in the original order. this technique provides good results on an extensive collection of benchmarks. however, we cannot compare with them because we focus on autoencoder architectures only. finally, we mention the work by @xcite , in which the authors list many optimizations that could be applied to the transformers for faster pre-training. they also claim that using larger models with the same runtime leads to better results. we focus instead on the pre-training objective efficiency and on the classification head size. thus we state that those techniques are orthogonal to our work and could be applied along with our alternative pre-training objectives. ## background on pre-training objectives before describing our models, we provide a detailed description of the most common token-level pre-training objectives used in the literature. masked language model (mlm) was proposed by @xcite . in mlm, 15% of the input tokens are replaced with a special mask, e0 e1 e2 em . . . classification @xmath0. . . . . . . . . . . . vocabulary size . . . . . . . . . . . . . . . @xmath1. . . @xmath2. . . . . . . . . . . . e0 e1 e2 em . . . classification @xmath3. . . . . . . . . . . . vocabulary size . . . . . . . . . . . . . . . @xmath4. an improvement of mlm is whole-word-masking, in which masking is applied to every token belonging to a word and not just to independent tokens. the model needs a large classification head for pre-training, as shown in figure 1 . its dimension is proportional to the vocabulary size, and (especially for small models) this constitutes a significant fraction of the entire architecture parameters. in the base architectures, the mlm head constitutes about 20% of the model parameters while for small models, the fraction increases to 30% 1 . the memory footprint of the lm head while training is about 47% for base and 64% for small models 2 . see appendix g for more details. for this reason, a binary classification head (as for td) can provide significant efficiency improvement. moreover, sharing the parameters of the mlm classification head with the embeddings does not reduce the computational cost but leads only to marginally lower memory requirements. in the embedding layer, only a few row vectors corresponding to the actual sentence tokens are updated at every step. on the contrary, mlm's softmax continuously computes gradients for the whole output linear transformation. causal language model (clm) is used to train autoregressive models by predicting the next token in a sequence @xcite . similarly to mlm, it requires a large classification head to output predictions over the whole vocabulary. permutation language modeling (plm) was proposed by @xcite to combine the 1 very often the language modelling head parameters are shared with the embedding layer 2 gradients have to be computed for every token in the vocabulary because of the final softmax layer generative power of autoregressive models with the bidirectional context of autoencoders. this is accomplished by permuting the input tokens and by letting the model use only the left context for the next token prediction. in this way, the model keeps the strengths of autoregressive models while exploiting the whole input sequence for better-contextualized output embedding. token detection (td) was introduced by elec-tra @xcite , which is an architecture composed of a discriminator and a smaller generator network. first, the generator is trained with mlm and finds suitable replacements for the masked tokens, as in bert. then, those candidates are inserted in the original sentence, and the resulting sequence is fed to the discriminator, which classifies whether a token is original or not. td has the advantage of computing the loss on the whole discriminator output and having a minimal memory footprint. however, the whole system is inefficient because of the presence of the generator, which is still mlm-based. ## effective pre-training objectives this section presents our alternative pre-training objectives, which can potentially be applied to a wide range of transformer-based models. random token substitution (rts) like elec-tra, rts trains a model that discriminates between original and substituted tokens. the main difference is that rts replaces 15% of the tokens with random alternatives, thus avoiding using computational resources to train a separate and expensive network. besides, unlike mlm, this approach relies on a smaller classification head that is not proportional to the vocabulary size, see figure 1 . ## datasets this section contains the descriptions of the datasets we used for pre-training and fine-tuning. ## experimental setting in these experiments, we compare the cost and accuracy of our models with the state-of-the-art methods on glue and several as2 benchmarks. finally, we summarize results derived from previous work. then, we provide a description of the training set as well as our pre-training and finetuning experiments and results. ## results we first report the cost of our models in comparison with the state-of-the-art methods in tables 5 and 6 for base and small models, respectively. notice also that rts and c-rts use about half the gpu memory of mlm with base and 1 3 with small models. thus, the training time of the former's to complete all the training steps could have been even shorter by increasing the batch size. however, we keep a constant batch size to make a fair comparison. after that, we carry out a comparison in terms of accuracy on glue datasets as well as question-answering benchmarks. additionally, we compare our models with state-of-the-art results published in previous works, pointing out the cost of training them. test set. we underline the results that are significantly different from the bert-s / mlm baseline after a two-sided t-test with a significance level of 95%. we show the standard deviation after 5 runs with different initialization seeds in rounded brackets. ## discussion and conclusion in this work, we studied several alternative methods to pre-train transformer models. our approaches aim at designing pre-training objectives that (i) match the results of well-known methods using fewer resources or (ii) outperform previous techniques with the same computing budget. this translates into shorter training, lower memory usage, and the possibility of increasing the batch size. among other benefits, more efficient models can reduce carbon footprint and infrastructure costs. notice that this advantage is even higher for smaller models since the mlm classification head is not shrunk proportionally to the architecture size @xcite , thus it increases its weight on the computational complexity for smaller models. moreover, we demonstrate that the mask token is useless and similar or even better results can be achieved using only token substitution. we reiterate that our objectives could be easily applied to many different transformer models; we chose the bert setting due to computational resource constraints and easy reproducibility. in appendix e we provide an overview of negative results. finally, we show that recent models' performance improvements are mostly driven by longer training phases rather than by more refined architectures. we evaluated our approaches on several datasets, such as glue, wikiqa, trec-qa and asnq. the results show that rts and c-rts match the accuracy of mlm in most tasks, requiring a lower amount of computational effort (20% less), while slm outperforms mlm in most tasks. c-rts also shows a lower accuracy in detecting replaced tokens, meaning that the task is harder and better suited for longer pre-training sessions. in addition, we tested our pre-training objectives on smaller transformer architectures. in this scenario, rts and c-rts obtained better performances than mlm by requiring half of its time to be pre-trained. for example, rts outperforms mlm by almost 5 map points on wikiqa. this last finding is potentially helpful for training transformer models from scratch with limited resources. in the future, we plan to explore combinations of our new techniques with efficient architectures such as albert to take advantage of a lighter structure and more effective pre-training objectives.