P RETRAINING T EXT E NCODERS WITH A DVERSARIAL M IXTURE OF T RAINING S IGNAL G ENERATORS
Anonymous authors
Paper under double-blind review
A BSTRACT
We present a new framework AMOS that pretrains text encoders with an Adversarial learning curriculum via a Mixture Of Signals from multiple auxiliary generators.
Following ELECTRA-style pretraining, the main encoder is trained as a discriminator to detect replaced tokens generated by auxiliary masked language models (MLMs).
Different from ELECTRA which trains one MLM as the generator, we jointly train multiple MLMs of different sizes to provide training signals at various levels of difﬁculty.
To push the discriminator to learn better with challenging replaced tokens, we learn mixture weights over the auxiliary MLMs’ outputs to maximize the discriminator loss by backpropagating the gradient from the discriminator via Gumbel-Softmax.
For better pretraining efﬁciency, we propose a way to assemble multiple MLMs into one uniﬁed auxiliary model.
AMOS outperforms
ELECTRA and recent state-of-the-art pretrained models by about 1 point on GLUE and SQuAD benchmarks for BERT base-sized models.
We plan to release our pretrained models for future uses.
I NTRODUCTION
Instead of pretraining encoders on texts with random noise ( e.g. , randomly masked tokens in
BERT (Devlin et al., 2019)), ELECTRA-style frameworks employ an auxiliary model to corrupt the original text sequence using its language modeling outputs, and pretrain the main Transformer to detect (Clark et al., 2020) or correct (Meng et al., 2021)
the replaced tokens.
This new self-supervised framework signiﬁcantly improves the efﬁciency and effectiveness of pretraining and leads to strong results in various downstream tasks with ﬁne-tuning (Clark et al., 2020), prompt-based learning (Meng et al., 2021), and zero-shot cross-lingual representations (Chi et al., 2021).
Recent studies revealed that the key to the ELECTRA’s success is its new learning dynamics (Xu et al., 2020; Meng et al., 2021).
By pretraining the auxiliary model jointly with the main Transformer, an implicit learning curriculum is formed as the noise produced by the auxiliary generator becomes more and more plausible during pretraining, posing greater challenges for the discriminator.
The latter has to overcome the difﬁculty by reasoning more deeply using the contexts.
This leads to signiﬁcantly improved sample efﬁciency and effectiveness of ELECTRA-style pretrained models (Clark et al., 2020; Chi et al., 2021; Meng et al., 2021).
On the other hand, such a training dynamic also introduced new challenges in search of the optimal pretraining setting.
First, the conﬁgurations of the auxiliary generator—its depth, width, and masking fraction—require costly trial-and-error pretraining runs, while also have a signiﬁcant impact on the discriminator’s downstream task performance:
A weak auxiliary model does not generate hard enough pretraining signal to push the discriminator, but a too strong one can confuse the discriminator and worsen its downstream task performance (Clark et al., 2020;
Meng et al., 2021).
Second, the sideby-side training of the two models forms a pseudo “GAN-style” (Goodfellow et al., 2014) curriculum which causes difﬁculty to improve or scale: Attempts to make the generator and discriminator learning more interactive ( e.g. , training the generator to maximize the discriminator loss as in actual GAN frameworks) resulted in downgraded performance (Clark et al., 2020).
In this paper, we present a new approach that learns to automatically select pretraining signals and constructs the learning curriculum of ELECTRA-style frameworks.
Our approach, AMOS, “ A dversarial M ixture O f S ignals" for text encoder pretraining, samples a diverse set of pretraining signals from different layers of an auxiliary generator.
The sampling is conducted with GumbelSoftmax gradient estimation, which enables AMOS to backpropagate the (reversed) gradient from the main discriminator to adversarially learn the mixture weights over the signals from the auxiliary generator.
In this way, AMOS constructs a more diverse set of pretraining signals for the main discriminator training while also enables an automatic selection of the signals to form a more effective learning curriculum.
In addition, by sampling from different layers of one generator model,
AMOS ensures minimum computation overhead without training multiple generators.
Our experiments on the GLUE and SQuAD benchmark demonstrate the effectiveness of AMOS.
In the standard BERT Base and RoBERTa Base++ pretraining settings (Devlin et al., 2019), AMOS shows robust advantage across all included language representation tasks.
The improvements are signiﬁcant as it advances the state-of-the-arts by 1+ absolute point on both GLUE scores and SQuAD accuracy.
Our thorough ablation studies also conﬁrm that the source of this effectiveness is the diverse training signals from the mixture of generators and the adversarial learning approach to combine them.
In the rest of this paper, we ﬁrst recap related work (Section 2) and then propose our methods (Section 3).
After that we describe our experimental settings (Section 4) and evaluation results (Section 5).
The last section concludes and discusses potential future research directions.
R ELATED W ORK
Besides the standard auto-regressive language modeling (Radford et al., 2019) and masked language modeling (MLM) (Devlin et al., 2019), many have explored different designs of the pretraining task, for example, preﬁx language modeling (Raffel et al., 2019), permutational language modeling (Yang et al., 2019), and uniﬁed language modeling (Dong et al., 2019).
The advantage of manually constructed pretraining tasks is more often observed in an application-speciﬁc manner, where prior knowledge about the target scenario is introduced (Guu et al., 2020; Roberts et al., 2020;
Jia et al., 2021;
Lu et al., 2021).
To pretrain general purpose language representations, the standard MLM task often offers the best combination of simplicity, robustness, and generalization ability among those manually constructed pretraining signals (Raffel et al., 2019; Lewis et al., 2019).
Clark et al.
(2020) developed a novel framework ELECTRA to pretrain Transformers with modelgenerated signals.
They employed an auxiliary MLM model, the “generator”, to replace the masked positions in the input sequence with tokens sampled from its language model head.
The main
Transformer, the “discriminator”, is pretrained to detect the replaced tokens.
In this way, the main model is trained to the denoise the more challenging noises from the auxiliary language model.
As a result, ELECTRA requires much fewer pretraining data points to reach the performance of MLM pretrained models, and offers better downstream performance when converged.
The unique effectiveness of ELECTRA intrigued many studies to understand its real source of advantage.
Clark et al.
(2020) originally noted the sample efﬁciency comes from ELECTRA’s replaced token detection task being able to leverage all token positions for training.
Later, Xu et al.
(2020) pointed to the more inﬂuential change of the pretraining task complexity for the main
Transformer.
Recently, Meng et al.
(2021) conducted a thorough ablation study on many variants of ELECTRA-style models and revealed the beneﬁts of pretraining with more challenge signals generated by the auxiliary model, as well as the implicit learning curriculum by pretraining two models side-by-side.
One popular line of research to improve ELECTRA-style models is to upgrade the replaced token detection task to more semantically informative ones.
Xu et al.
(2020) and Shen et al.
both proposed the multi-word choice task which pretrains the main model to pick the original token from a small candidate sets.
COCO-LM developed a corrective language modeling task which pretrains the main model to recover the original tokens (Meng et al., 2021).
Another line of research is to improve the training signal construction on the auxiliary side.
In
ELECTRA, Clark et al.
(2020) experimented with adversarially training the auxiliary model using feedback from the discriminator model via reinforcement learning, but observed worse downstream performance.
Hao et al.
(2021) kept the auxiliary model unchanged and augmented the replaced token sampling with predicted signal difﬁculty and smoothed probability via Focal Loss (Lin et al., 2017).
In this work, we continue this line of research and focus on how to more automatically and effectively construct the learning signals for ELECTRA-style frameworks.
M ETHOD
In this section, we ﬁrst recap the preliminary of ELECTRA-style pretraining, the explorations of its learning curriculum, and then present our AMOS framework.
3. P RELIMINARY ON ELECTRA-S TYLE P RETRAINING
In ELECTRA-style pretraining (Clark et al., 2020), two Transformer models are trained side-by-side:
One is trained via standard masked language modeling (MLM) and is used to generate replaced tokens to corrupt the original sequences; the other is trained to denoise the corruptions from the ﬁrst model ( e.g. , via replaced token detection (RTD) (Clark et al., 2020), multi-word choice
(Xu et al., 2020), or corrective language modeling
(Meng et al., 2021)).
The second model is the main model to use in downstream tasks (the “discriminator”).
The ﬁrst one is the auxiliary, often referred to as the training data “generator”.
In this work we built upon the original ELECTRA setup where the main model is trained with RTD.
Generator Training.
Given an original sequence X orig =
[ x orig 1 , . . .
, x orig i , . . .
, x orig n ] , 15% of its tokens are randomly replaced by [MASK] symbols, and the resulting masked sequence X mask =
[ x orig 1 , . . .
, [MASK] i , . . .
, x orig n ] is fed to the generator which is trained via the following loss to predict the original tokens from the vocabulary V at the set of masked positions M :
[Equation]
where { h i } ni =1 are the contextualized representations generated by the Transformer (after the projection head) and { x t } | V | t =1 are the token embeddings.
Discriminator Training.
A replaced sequence X rep is constructed by sampling from the generator’s
MLM probability distribution:
[Equation]
The discriminator takes X rep as input and is trained to distinguish replaced tokens produced by the generator against the original tokens via the binary classiﬁcation loss:
[Equation]
where p RTD ( x rep i = x orig i (cid:12)(cid:12) h i ) = exp( w (cid:62) h i ) / (1 + exp( w (cid:62) h i )) and w is a learnable weight vector.
By training the generator side-by-side with the discriminator, this ELECTRA-Style framework constructs a natural learning curriculum for the main model.
At the beginning, the generator is weak and the replaced tokens are nearly random, making the pretraining task of the main model easy for a head start.
As the generator learns better, the corruptions from it become more challenging, which increase the difﬁculty of the learning task for the main model.
After a series of studies, recent research conﬁrmed that this curriculum learning is the key to ELECTRA-style framework’s effectiveness (Clark et al., 2020; Xu et al., 2020; Meng et al., 2021).
Pretraining with randomly replaced tokens or starting with a converged generator’s harder signals both signiﬁcantly worsen the discriminator effectiveness (Meng et al., 2021).
3. I N S EARCH OF THE O PTIMAL C URRICULUM IN ELECTRA
In general, conﬁguring a learning curriculum for neural model training is tricky (Bengio et al., 2009;
Hacohen & Weinshall, 2019).
Many moving pieces are involved—how to measure the difﬁculty of training signals, and when and how to change the training signal difﬁculty—all require manual trialand-error or careful designs (Graves et al., 2017).
The same challenge persists in ELECTRA-style pretraining models and more speciﬁcally, it is unclear how to choose the size of the generator network and how to control the training dynamics of the generator and discriminator.
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure]
E s t i m a t e d
D e n s i t y
[Figure] [Figure]
(a) Disc
[Figure]
Loss Distribution.
[Figure]
( × 1 e 4)
[Figure]
G e n e r a t o r L o ss
[Figure]
Generator Size.
There are many factors on the auxiliary side that affect the difﬁculty of the generated signals:
Network depth, width, MLM mask fraction, and whether to use dropout when sampling the replaced tokens.
All of these have notable impact on the performance of the pretrained discriminator model.
The general consensus is that the generator should be smaller than the discriminator but all the rest vary setting by setting.
The original ELECTRA framework kept the generator’s Transformer layer depth the same with the discriminator model, enabled dropout in both its training and sampling, and mainly explored using a smaller width ( e.g. , with 1 / 3 hidden dimensions of the main model).
They also empirically found that large-sized discriminators prefer different mask fractions on the generator side.
Later, COCO-LM (Meng et al., 2021) empirically found that it is better to use a shallower generator but of the same width as the main model, while also disabling dropout when sampling replaced tokens.
The depth of the generator, however, is still chosen empirically and the optimal settings vary with the conﬁgurations of the discriminator model.
To illustrate the difﬁculty of training signals by generators with different depths, we train BERT Basesized discriminators with 4 / 6 / 8 -layer generators, and show their RTD loss on replaced tokens by these different-sized generators in Figure 1a.
The average discriminator loss grows with generator depth, indicating the increased signal difﬁculty, and the three distributions overlap with each other.
The current practice is to enumerate different generators and choose the one leading to the best downstream task performance of the discriminator, which is tedious and unsustainable.
Generator Training Dynamics.
As discussed earlier, ELECTRA constructs an implicit learning curriculum with the training process of the generator.
To illustrate this implicit curriculum, in Figure 1b we plot the MLM loss of the generators during pretraining.
The convergence of the MLM models follows a logarithmic schedule: The loss drops sharply in the ﬁrst 10% of pretraining and decreases slowly afterwards.
It is unclear whether such a training process of the generator automatically leads to the optimal curriculum for discriminator training.
Unfortunately, it is challenging to manually explore different generator training dynamics in ELECTRA-style frameworks at it is trained side-by-side with the discriminator.
One idea to improve ELECTRA-style frameworks is to connect the training of the generator with feedbacks from the discriminator ( e.g. , as in GAN-style models (Goodfellow et al., 2014))
so that the generator can be adapted based on the latest discriminator’s state throughout pretraining.
Clark et al.
(2020) experimented with training the generators to produce more difﬁcult signals for the discriminator, using policy gradients from the latter.
However, this adversarial setup makes the training unstable and yields worse results, similar to the instability of GAN-style training in other language modeling tasks (Caccia et al., 2020).
Hao et al.
(2021) introduced a training signal difﬁculty estimator to adjust the sampling of replaced tokens in MLM for more challenging training signals.
It achieved better results than the original ELECTRA but still kept the generator training intact.
3. A DVERSARIAL L EARNING WITH M IXTURE OF G ENERATORS
Intuitively, different-sized generators provide training signals at different levels of difﬁculty and also may advocate the discriminator to capture different linguistic information.
For example, as shown in various probing studies (Tenney et al., 2019a;b), a very shallow MLM model may still make trivial syntactic mistakes, thus the discriminator can focus on learning simple language syntactic knowledge to identify these errors, while the errors made by a deep MLM model may require the discriminator to capture more sophisticated language semantics to detect.
[Table]
d
[Table]
MLM
[Table] [Table]
[Table] [Table]
=</latexit> f ( d 2 )
i
In this work, instead of empirically searching for the best generator setup, we propose the AMOS framework that (1) employs multiple MLM generators to construct a diverse set of pretraining signals, and (2) adversarially learns the mixture of these generator outputs from discriminator feedback to ﬁnd more challenging signals.
These two enable AMOS to automatically compose an effective curriculum.
Figure 2 shows an overview of our AMOS framework.
Multi-Layer MLM Generator.
A straightforward way to obtain a diverse set of signals for replaced token generation is to utilize multiple generator networks of different depths.
However, jointly training K different MLM generators is expensive.
Therefore, we propose to train a single generator with multiple MLM heads at different layers to mimic the effect of using multiple MLM generators.
Speciﬁcally, suppose we aim to use K generators of depths D = { d 1 , d 2 , . . .
, d K } ( d 1 < d 2 < · · · < d
K ), then we will instead train one generator of depth d K , and insert an MLM head to each layer d ∈ D .
The K MLM heads will be jointly trained and the MLM head parameters are shared across the K heads:
[Equation]
where { h ( d ) i } ni =1 are the contextualized representations generated by the Transformer at the d th layer (after the projection head).
During training, gradient backpropagation from upper layers is disabled at each layer d ∈ D so that the generator is partitioned into K blocks where each block is trained via its own MLM objective without being disturbed by other blocks.
This allows each block to act as individual MLM generators, while efﬁciently leverages the representations generated by previous blocks.
Learning Adversarial Mixture Weights.
It would be desirable to progressively adapt the output signals of the multi-layer MLM generator based on the discriminator’s state throughout pretraining.
The generated replaced tokens are thus more informative for discriminator training ( i.e. , avoid generating easy signals which can already be well distinguished by the discriminator).
Given the instability of GAN-style training to language modeling (Caccia et al., 2020; Clark et al., 2020), we maintain the MLM training objective of the generator, and adjust the outputs of the generator by learning mixture weights over the multiple MLM heads’ outputs.
Speciﬁcally, for each masked position i ∈ M , we learn mixture weights γ i over the embeddings h ( d ) i , which is generated by each MLM head, to obtain the combined MLM embedding ¯
h i .
Then the ﬁnal token sampling probability distribution p MLM is computed:
[Equation]
where v is a learnable weight vector;
{ f ( d ) i } ni
=1 are the contextualized representations generated by the Transformer at the d th layer (before the projection head in MLM).
However, directly sampling from the above MLM distribution fails to enable the mixture weight learning to be guided by the discriminator, as the sampling operation is non-differentiable, preventing gradient backpropagation from the discriminator to the generator.
Therefore, we leverage GumbelSoftmax
(Jang et al., 2017) for a continuous approximation to the sampling operation for gradient estimation:
[Equation]
where τ is a temperature hyperparameter (a smaller τ results in a more accurate approximation).
Once we have the gradient estimation backpropagated from the discriminator side, we will update the mixture weights γ to maximize the discriminator loss.
This forms an adversarial learning curriculum as the multiple MLMs’ outputs are always combined to construct challenging replaced tokens given the latest discriminator state.
Such an adversarial learning setup can be easily implemented via a gradient reversal layer (Ganin et al., 2016) that multiplies the gradient backpropagated from the discriminator side by − 1 when updating the mixture weights.
Overall Training.
AMOS jointly trains the multi-layer MLM generator θ GEN , the mixture weight parameter v , and the discriminator θ DISC via the following losses:
[Equation]
where λ is a hyperparameter balancing the weight of the two losses.
Similar to ELECTRA, the generator minimizes L GEN in Eqn.
(2) and the discriminator minimizes L DISC in Eqn.
However, the notable difference from ELECTRA is that AMOS additionally trains the generator’s mixture weights γ (thus v and θ GEN ) to maximize the discriminator loss.
We note that the gradient backpropagated from the discriminator through the Gumbel-Softmax estimation is used to update the mixture weights, but not the MLM embeddings ( i.e. , in Eqn.
(3), γ ( d ) i is updated by the gradient from the discriminator side, while h ( d ) i is not).
This guarantees that the generator’s language modeling ability is still acquired through the MLM task without being disturbed by discriminator signals.
E XPERIMENTAL M ETHODOLOGIES
Pretraining Settings.
We experiment with two standard pretraining settings, base and base++ :
Base is the BERT Base training conﬁguration (Devlin et al., 2019):
Pretraining on Wikipedia and
BookCorpus (Zhu et al., 2015) ( 16 GB of texts) for 256 million samples on 512 token sequences ( 125 K batches with 2048 batch size).
We use the same corpus and 32 , 768 uncased BPE vocabulary (Sennrich et al., 2015) as with TUPE (Ke et al., 2020) and COCO-LM (Meng et al., 2021).
Base++ trains the base size model with larger corpora and/or more training steps.
We add in
OpenWebText (Gokaslan & Cohen, 2019), CC-News (Liu et al., 2019) and STORIES (Trinh & Le, 2018), to a total of 160 GB texts, and train for 4 billion (with 2048 batch size) samples, following recent research Bao et al.
; Liu et al.
; Yang et al.
We follow the prepossessing of UniLMV2 (Bao et al., 2020) and use 64 , 000 cased BPE vocabulary.
Model Architecture.
Our base / base++ discriminator model uses the BERT Base architecture (Devlin et al., 2019): 12 layer Transformer, 768 hidden size, plus T5 relative position encoding (Raffel et al., 2019).
Our generator is an 8 -layer Transformer with the same hidden size, and the MLM heads are inserted at the 4 th, 6 th and 8 th layers (to mimic having three generators).
We disable dropout in the generator following Meng et al.
Downstream Tasks.
We use the tasks included in GLUE (Wang et al., 2018) and SQuAD 2.0 reading comprehension (Rajpurkar et al., 2016).
All models are evaluated with the same standard ﬁne-tuning protocols: Single task learning with vanilla ﬁne-tuning and reporting the median of ﬁve random seeds in GLUE and SQuAD.
Please refer to Appendix A for more details.
Baselines.
We compare with various pretrained models in both settings.
All numbers are from reported results in recent research (more details in Appendix B).
Implementation Details.
Our implementation builds upon the open-source implementation of fairseq Ott et al.
Standard hyperparameters in pretraining and ﬁne-tuning are used.
More details are listed in Appendix C.
[Table]
Results on GLUE and SQuAD 2.0 development set.
All results are single-task, single-model ﬁne-tuning.
We report MNLI matched/mismatched subsets results.
We use Spearman correlation for
STS, Matthews correlation for CoLA, and accuracy for the rest on GLUE.
AVG is the average of the eight tasks on GLUE.
All baseline results are reported by previous research.
Results not available in public reports are marked as “–”.
E VALUATION R ESULTS
We conduct three groups of experiments to evaluate the performance of AMOS, the effect of leveraging signals from multiple generators, and the inﬂuence of its adversarial training.
We also provide some case studies on the constructed pretraining sequences in AMOS.
5. R ESULTS AND A BLATIONS ON GLUE AND SQ U AD
GLUE and SQuAD Results.
Table 1 lists the single-task ﬁne-tuning performance of AMOS and notable baselines that are pretrained under the standard base and base++ setting.
AMOS outperforms all previous state-of-the-art pretraining methods on the overall GLUE score and SQuAD.
AMOS’s improvements are robust, often of large margins, and achieves the new state-of-the-art on multiple tasks under this standard pretraining/ﬁne-tuning setup.
Table 2 presents the ablation studies of AMOS variants on MNLI and SQuAD, the two most stable and representative downstream evaluations.
We organize the ablations into four subgroups to study the effectiveness of Curriculum Learning , Adversarial Training , Multi-Layer MLM , and Training
Signal Diversity in AMOS.
Effect of Curriculum Setup.
Disabling the automatic mixture weights learning in AMOS leads to downgraded downstream performance: Randomly picking the MLM head from the same set of generator layers ( w. random layer ) reduces accuracy on all metrics.
It increases the diversity of pretraining signals but does not conﬁgure then as effectively.
Manually conﬁguring the order of different sets of pretraining signals by manually switching from shallower to deeper layers ( i.e. , 4 to 6 to 8 ) during pretraining ( w. layer switch ) does not lead to better results than random layer selection.
Note that in these manual conﬁgurations we manually switch the signals at 1 / 3 and 2 / 3 of the total pretraining steps.
We have experimented switching at different steps of pretraining but do not observe signiﬁcantly better results.
Manual trials are tedious and expensive in pretraining and underperform automatically learned mixture over multiple training signals.
Effect of Adversarial Setup.
Not performing adversarial learning ( - adv. train ) to learn mixture weights ( i.e. , always use unweighted average signals) hurts the model performance.
However, note that this ablation still beneﬁts from the curriculum learning effect as the generator gradually learns better.
In addition, we also try to backpropagate the adversarial gradient to update the MLM embeddings ( + adv.
MLM ).
Speciﬁcally, in Eqn.
(3), both γ ( d ) i and h ( d ) i are updated by the reversed gradient from the discriminator.
Our observation is that this makes the training unstable, perhaps because it hinders the MLM task of the auxiliary model.
We have to use a very small gradient
[Table]
[Table] [Table]
Example 1. "
Here, let me show you that animal."
He pointed up to (at/at/into) the canopy; hanging from a white branch was (of/stood/stood) a pale, hairless creature (shape/plant/animal).
It was bilaterally symmetrical, with two pairs of tightly (long/ﬁrmly/strongly) folded limbs, but did not appear to have any discernible head.
Example 2.
In 2012 , about 3 , 000 (6/1,800/1,800) villagers remained in Rammun, while there are about 7 , 000 in the Palestinian diaspora, chieﬂy in the United States.
Many in the diaspora have second homes in the (this/the/their) village.
These homes have been troubled by burglaries, therefore some owners have organised night-watches (things/boxes/searches).
multiplier ( e.g. , 0 . 1 ) when updating the MLM embeddings with the discriminator’s backpropagated gradient, which has minimal effects on the model.
Effect of Multi-Layer Generator Setup.
We also experimented with several conﬁgurations of the multi-layer generators: -stop grad.
does not stop the gradient at the 4 th and 6 th layers in the generator of AMOS; w. separate MLM gen. trains three separate generators of 4 , 6 and 8 layers jointly with the discriminator.
Both conﬁgurations result in reduced performance of the discriminator and we keep the simpler setup as in AMOS.
Effect of Diverse Training Signals.
The last ablation compares with using a single generator in
AMOS and without any learned mixtures.
Using a 4 / 6 / 8 -layer generator yields worse results than
AMOS and previous ablations with the multi-MLM generator, especially on SQuAD.
The 12 -layer generator is too strong and makes the pretrained discriminator signiﬁcantly worse.
It is simpler and more effective to grant the model access to a broader set of training signals and automatically learn to leverage them.
5. D IVERSE P RETRAINING S IGNALS FROM D IFFERENT G ENERATORS
To study how generators of different depths provide diverse pretraining signals, we conduct probing experiments on eight NLP tasks covering different linguistic aspects, following (Tenney et al., 2019a;b): part-of-speech (POS)
(Weischedel et al., 2013), constituents (Consts.), dependencies (Deps.), entities, semantic role labeling (SRL), coreference (Coref.), semantic proto-roles (SPR2 (Rudinger et al., 2018)), and relation classiﬁcation (SemEval (Hendrickx et al., 2009)).
The major difference between our setting and Tenney et al.
(2019a;b) is that we do not combine embeddings from multiple layers but directly use the embedding from each MLM layer of AMOS generator as the (frozen) feature to a trainable linear classiﬁer, as we are interested in what information is captured by each MLM layer.
As shown in Table 3, different MLM layers in AMOS generator indeed are good at different tasks–The 6 th layer in the generator has the best performance on POS, constituent labeling, entity recognition and relation classiﬁcation, while the 8 th layer performs the best on the other tasks.
This demonstrates that using different layers in the multi-layer MLM generator is helpful for creating training signals of different levels of difﬁculty and also emphasizing different linguistic aspects.
Note that although the 4 th layer has worse performance than deeper layers across all tasks, it is still useful for providing discriminator pretraining signals, because the discriminator needs to learn from the “mistakes” made by the generator not capturing certain language semantics.
5. E FFECTS OF A DVERSARIAL T RAINING
[Figure]
To better understand the adversarial training dynamics, we plot the mixture weights (averaged over all masked positions) assigned to the 4 th/
th/ 8 th layer MLM during pretraining in Figure 3a.
In the later pretraining stage, the 8 th layer gradually gains more weights with the weights of the other two layers decreasing, showing that the generator indeed tries to create more challenging training signals for the discriminator.
We also compare the training process of AMOS with several generator ablations: (1) A 4 -layer generator; (2) An 8 -layer generator; (3) A 4 / 6 / 8 -multi-layer MLM generator with three layers randomly picked for generating MLM replacements (Random Layer); (4) Switching from shallow layers to deeper layers ( i.e. , 4 to 6 to 8 ) at 1 / 3 and 2 / 3 of the pretraining steps (Layer Switch).
We plot the discriminator accuracy (averaged on replaced tokens) in Figure 3b.
AMOS creates an intuitive learning curriculum that starts simpler than random layer mixing and becomes more challenging as pretraining goes on.
Making “hard” switches suddenly changes the task difﬁculty and may disrupt training.
5. C ASE S TUDY
Table 4 provides concrete cases of replaced tokens generated by different-sized generators.
The “mistakes” made by different generators have different levels of difﬁculty to be detected–The 4 -layer generator sometimes makes simple syntactic mistakes while the replaced tokens given by the 6 / 8 -layer generators are mostly plausible and need to be distinguished based on a deep understanding of the full contexts.
This intuitively conﬁrms our motivation that using multiple MLM generators can provide diverse pretraining signals to compose a more effective learning curriculum.
C ONCLUSIONS AND F UTURE W ORK
In this paper we present AMOS, a new strategy for pretraining ELECTRA-style text encoders using an adversarial mixture of multiple training signal generators.
AMOS constructs the corrupted text sequences by attaching multiple MLM heads to a deeper generator and sampling replaced tokens from their mixed outputs.
The weights of the mixtures are learned to maximize the training signals difﬁculty for the discriminator, by backpropagating the reversed gradient from the discriminator through Gumbel-Softmax.
This upgrades the ELECTRA-style pretraining framework with an automatically learned curriculum that composes more diverse pretraining signals.
Our experiments on the GLUE and SQuAD benchmarks demonstrate the empirical advantages of
AMOS.
Under the standard BERT Base and RoBERTa Base++ pretraining settings, the same Transformer network pretrained with AMOS achieves the new state-of-the-art in nearly all evaluation metrics, with around 1 point gains on GLUE score and SQuAD accuracy.
Our studies and analyses further conﬁrm the source of AMOS’s effectiveness: The more diverse training signals from multiple generators and our adversarial learning design to effectively utilize them throughout pretraining.
Our observations can be viewed as another progress of “data-centric” AI—Effectively constructed training signals from data can lead to signiﬁcant empirical improvements without changing the model itself.
Future work along this direction includes but not limited to: More studies to understand the role of training signals in language model pretraining, explorations of a broader set of training signal sources, and better strategies to leverage different information sources in pretraining.
R EPRODUCIBILITY S TATEMENT
We strive to facilitate the reproducibility of the reported results in this paper, by (1) using the same pretraining and ﬁne-tuning setups with previous research and reporting the median results of multiple ﬁne-tuning runs, (2) providing details about the datasets used in Appendix A, sources of baselines compared in Appendix B, and hyperparameters used in Appendix C and (3) sharing the source code with the reviewers during the reviewing period. (4) reporting the exact pretraining hours of our base models in Appendix E. All experiments in this paper are conducted on 64 A100 GPUs each with 40 GB memory size.
We also plan to release our pretrained models
A D ETAILS OF GLUE T ASKS
[Table]
Below are detailed descriptions of the tasks included in the GLUE benchmark.
The statistics can be found in Table 5.
MNLI: Multi-genre Natural Language Inference (Williams et al., 2018) has 393 K training examples via crowdsourcing.
The task is to predict whether a premise sentence entails, contradicts or is neutral to a given hypothesis sentence.
QQP: Question Pairs (Shankar et al., 2017) has 364 K training examples from the Quora questionanswering website.
The task is to determine whether a given pair of questions asked are equivalent semantically.
QNLI:
Question Natural Language Inference has 108 K training examples collected from the Stanford Question Answering Dataset (SQuAD) (Rajpurkar et al., 2016).
The task is to predict whether a given sentence includes the answer to a given question sentence.
SST-2: Stanford Sentiment Treebank (Socher et al., 2013) has 67 K training examples obtained from movie reviews with manually-annotated sentiment scores.
The tasks is to determine if the sentence contains positive or negative sentiment.
CoLA: Corpus of Linguistic Acceptability (Warstadt et al., 2019) has 8 .
5 K training examples collected from books and journal articles on linguistic theory.
The task is to determine whether a sentence is linguistically acceptable or not.
[Table] [Table] RTE: Recognizing Textual Entailment (Dagan et al., 2005; Haim et al., 2006; Giampiccolo et al., 2007; Bentivogli et al., 2009) has 2 .
5 K training examples from textual entailment challenges.
The task is to predict whether a premise sentence entails a hypothesis sentence or not.
MRPC: Microsoft Research Paraphrase Corpus (Dolan & Brockett, 2005) contains 3 .
7 K training examples from online news sources.
The task is to predict whether two sentences are equivalent semantically.
STS-B: Semantic Textual Similarity (Cer et al., 2017) contains 5 .
8 K training examples collected from multiple sources with human annotations of sentence pair semantic similarity.
The task is to predict how semantically similar two sentences are (with a 1 to 5 scoring scale).
B T HE O RIGINS OF R EPORTED B ASELINE S CORES
The baseline results reported in Table 1 are obtained from the corresponding papers except the following: BERT/RoBERTa from Bao et al.
(2020), ELECTRA from Meng et al.
(2021), XLNet base++ from Bao et al.
When there are different reported scores for the same method, we use the highest of them in our comparisons.
C H YPERPARAMETER S ETTINGS
We use the default/standard values for most hyperparameters for pretraining: The generator pretraining uses the standard 15% masking ratio.
The temperature for Gumbel-Softmax is τ = 0 .
The discriminator loss weight λ
= 50 since the loss of the binary classiﬁcation task is much lower than the
MLM task, which is a 30 , 000 -way classiﬁcation task.
The token embeddings are shared between the generator Transformer and the discriminator Transformer.
Other hyperparameters used in pretraining and ﬁne-tuning are reported in Tables 6 and 7, respectively. [Table]
The same (or equivalent) set of hyperparameters for pretraining and ﬁne-tuning are used for all compared methods.
The reported downstream task results on GLUE/SQuAD are the median of ﬁve runs with the same set of random seeds. D GLUE T EST S ET R ESULTS
We show the GLUE test set scores obtained via private submissions to the GLUE leaderboard in
Table 8.
The baseline results are directly retrieved from the leaderboard or from their original papers.
We use standard single-task ﬁne-tuning to more directly reﬂect the improvements from pretrained models.
The advantage of AMOS over strong baselines holds on the test set: Under both base and base++ settings, AMOS outperforms ELECTRA (Clark et al., 2020) and COCO-LM (Meng et al., 2021) on almost every task.
Still, we would like to note that smaller tasks such as CoLA, RTE, and
MRPC are not stable and leaderboard runs often use more sophisticated ﬁne-tuning method to achieve better performance, for example, continuing training from the checkpoints ﬁne-tuned on MNLI.
E P RETRAINING E FFICIENCY
[Figure]
sets (y-axes) at different pretraining hours on
A100 (40 GB Memory) GPUs.
The ﬁnal training hours and accuracy of COCO-LM, ELECTRA and
RoBERTa (trained under the exact same settings and computing environments) are also shown.
We compare the pretraining efﬁciency of AMOS with COCO-LM (Meng et al., 2021), ELECTRA (Clark et al., 2020) and RoBERTa (Liu et al., 2019) under exactly the same computation environment for base model training.
We show the MNLI-(m/mm) development set accuracy (via standard ﬁne-tuning) of AMOS checkpoints trained for different GPU hours in Figure 4.
For fair comparisons, we train all compared models using the same codebase, pretraining conﬁguration, and computing environments.
While
AMOS takes longer to train than RoBERTa and
ELECTRA, it matches their ﬁnal MNLI performance with signiﬁcantly fewer pretraining hours and achieves much better performance upon convergence.
For example, it achieves RoBERTa’s
MNLI accuracy with two hours of pretraining and outperforms ELECTRA in three hours, more than a 50% reduction in pretraining time.
It also reaches the accuracy of COCO-LM, the recent state-of-the-art in both pretraining accuracy and efﬁciency, only using 60% of COCO-LM’s pretraining hours.
This demonstrates the advantage of the adversarial curriculum learning of AMOS.
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
F M ORE D ETAILED A BLATION R ESULTS
We show more detailed ablation results in Table 9 which extends Table 2 by including results from all
GLUE tasks and showing the standard deviations.
We also include two new ablations: w. ﬁnal mix weights pretrains the discriminator all the way with the ﬁxed mixture weights learned by AMOS at convergence; w. three 4-layer gen. trains the discriminator with mixtures over three 4-layer generators of the multi-layer MLM generator in AMOS.
AMOS has better performance than all other ablation versions on most large tasks (MNLI, QQP,
SST-2 and SQuAD) which are considered more stable and reliable indicators of model effectiveness.
Small tasks (CoLA, RTE, MRPC) have much higher variance than larger tasks, and usually require intermediate task training to yield stable results ( i.e. , starting from checkpoints that are ﬁne-tuned on
MNLI (Clark et al., 2020;
He et al., 2021; Liu et al., 2019)).
The instability of the GLUE small tasks is a widely-observed artifact in pretraining research.
It is standard (and recommended) practice to not rely on these small unstable tasks for ablation studies.
For example, to ensure a correct understanding of research progress, previous studies including
BERT (Devlin et al., 2019), RoBERTa (Liu et al., 2019), ELECTRA (Clark et al., 2020), and
DeBERTa (He et al., 2021)
all perform ablation studies on large GLUE tasks like MNLI.
We include the ablation results on small GLUE tasks in Table 9 only for reference .
Furthermore, we show the standard derivations of the same pretrained checkpoint when ﬁne-tuned on the corresponding tasks with different random seeds.
They further conﬁrm that these smaller tasks have large variance in the ﬁne-tuning stage.
Observing the variance of the pretraining stage ( i.e. , changes in model performance when pretrained with different random seeds) requires running the costly pretraining multiple times which is often infeasible.
We would like to refer to recent studies from Google (Sellam et al., 2021) which reveal that these small GLUE tasks are very unstable, and the same model pretrained with different random initialization can have 2 − 5 points difference in performance upon ﬁne-tuning on these small tasks, while in comparison, observations on MNLI are more reliable.
Since in this paper we only perform single-task ﬁne-tuning without any intermediate task training, we mainly rely on large task results for evaluating the effectiveness of different model components.
The GLUE average score is more of a convenient reference point as used in the pretraining research community.
G P ERFORMANCE S TUDY WITH D IFFERENT N UMBERS OF MLM H EADS
In this experiment, we study the performance of AMOS with different numbers of generator MLM heads K .
In addition to the original AMOS which uses three MLM heads, we also show in Tablethe downstream task performance when two MLM heads (at the 4 th and 8 th layers), four MLM heads (at the 2 nd, 4 th, 6 th and 8 th layers) and eight MLM heads (at each layer) are used.
We note that increasing MLM heads does not necessarily lead to better performance, since having more MLM heads means that each MLM block (partitioned by the stop gradient operators at each MLM layer)
[Table]
will become shallower with weaker MLM learning capacity.
We have also tried inserting the same amount of MLM heads at different layers of the generator, but it does not improve the results.
H D ISCUSSIONS ON U SING G UMBEL -S OFTMAX FOR G RADIENT E STIMATION
In this work
[Table]
we use Gumbel-Softmax to enable gradient approximation of the non-differentiable sampling operation so that we are able to use the gradient backpropagated from the discriminator to train the mixture weights.
There are other possible approaches for discriminator gradient estimation, including REINFORCE (Yu et al., 2017) which has been explored in the ablation studies of
ELECTRA (Clark et al., 2020), or directly operating in the hidden states of the generator instead of in its discrete output space which has been studied in adversarial learning based text generation research (Subramanian et al., 2018; Zhang et al., 2017; Zhao et al., 2018).
We would like to note that as our ﬁrst study on adversarial curriculum for language model pretraining, we prefer a simple framework and standard techniques to demonstrate that such a new direction is promising.
We believe that exploring more sophisticated and advanced realizations for speciﬁc model components in our
AMOS framework ( e.g. , using better gradient estimators than Gumbel-Softmax) will be an interesting future work direction.
I P ROBING E XPERIMENTS WITH D IFFERENT S ETUPS
[Table]
In addition to the probing experiments conducted in Section 5.2, we also show the results under a different setup:
Instead of testing the linguistic information captured by different MLM layers in AMOS generator, we conduct the same tests on separate generators with different depths (layers, 6 layers and 8 layers).
The probing test results on eight tasks are shown in Table 11.
Similar to the results in
Table 3, generators of different depths are good at capturing different types of linguistic information.
Therefore, it is also possible to compose diverse training signals with multiple separate generators with different numbers of layers.
Such ﬁndings may be useful for future studies that explore using multiple different-sized auxiliary models to provide training signals emphasizing different linguistic aspects.
J S OCIETAL I MPACT
There have been concerns about the extensive costs of computing resource required by pretraining language models.
The concerns include whether it is worthwhile to spend such enormous amount of resources in pretraining, and also whether the demanding resource requirements have posed a barrier for most institutions to conduct pretraining research, thus slowing down the overall scientiﬁc development.
One major motivation of this work is to more efﬁciently pretrain language models, including (1) to achieve better model effectiveness with ﬁxed computing resource, and (2) to achieve the same downstream task accuracy with fewer pretraining computes.
As shown in our experimental results, we are able to obtain certain successes in both fronts.
We demonstrate that one can achieve the same pretraining effectiveness of RoBERTa Base in 128 GPU hours on A100 using AMOS (two hours on
A100 GPUs), which is quite affordable for many research institutions.
We hope our observation will inspire more future studies in efﬁcient pretraining and also enable conducting pretraining research in more accessible computing environments.
In addition, we also mainly focus on the base-sized models in our exploratory research which are less costly than large or extra large models, while also being the most commonly used model conﬁguration in the community.
We plan to release our pretrained model checkpoints to the community as well.