Learning to Sample Replacements for ELECTRA Pre-Training

Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021, pages 4495–4506 August 1–6, 2021. ©2021 Association for Computational Linguistics 4495 Learning to Sample Replacements for ELECTRA Pre-Training Yaru Hao†,‡∗, Li Dong‡, Hangbo Bao‡, Ke Xu†, Furu Wei‡ †Beihang University ‡Microsoft Research {haoyaru@,kexu@nlsde.}buaa.edu.cn {lidong1,t-habao,fuwei}@microsoft.com Abstract ELECTRA (Clark et al., 2020a) pretrains a dis- criminator to detect replaced tokens, where the replacements are sampled from a gener- ator trained with masked language modeling. Despite the compelling performance, ELEC- TRA suffers from the following two issues. First, there is no direct feedback loop from dis- criminator to generator, which renders replace- ment sampling inefficient. Second, the gen- erator’s prediction tends to be over-confident along with training, making replacements bi- ased to correct tokens. In this paper, we propose two methods to improve replacement sampling for ELECTRA pre-training. Specif- ically, we augment sampling with a hardness prediction mechanism, so that the generator can encourage the discriminator to learn what it has not acquired. We also prove that the ef- ficient sampling reduces the training variance of the discriminator. Moreover, we propose to use a focal loss for the generator in order to relieve oversampling correct tokens as re- placements. Experimental results show that our method improves ELECTRA pre-training on various downstream tasks. Our code and pre-trained models will be released at: https: //github.com/YRdddream/electra-hp 1 Introduction One of the most successful language model pre- training tasks is masked language modeling (MLM; Devlin et al. 2019). First, we randomly mask some input tokens in a sentence. Then the encoder learns to recover the masked tokens given the corrupted input. ELECTRA (Clark et al., 2020a) argues that MLM only produces supervision signals at a small proportion of positions (usually 15%), and uses the replaced token detection task as an alternative. Specifically, ELECTRA contains a generator and a discriminator. The generator is a masked language ∗Contribution during internship at Microsoft Research. model, which substitutes masks with the tokens sampled from its MLM predictions. The discrimi- nator learns to distinguish which tokens have been replaced or kept the same. Experimental results on downstream tasks show that ELECTRA can largely improve sample efficiency. Despite achieving compelling performance, it is usually difficult to balance the training pace be- tween the generator and the discriminator. Along with pre-training, the generator is expected to sam- ple more hard replacements for the detection task in a curriculum manner, while the discriminator learns to identify the corrupted positions. Although the two components are designed to compete with each other, there is no explicit feedback loop from the discriminator to the generator, rendering the learn- ing games independent. The absence of feedback results in sub-efficient learning, because many re- placed tokens have been successfully trained while the generator does not know how to effectively sample replacements. In addition, a well trained generator tends to achieve reasonably good MLM accuracy, where many sampled replacements are correct tokens. In order to relieve the issue of oversampling correct tokens, ELECTRA explored tweaking the mask probability larger, raising the sampling …

Content truncated to 500 words