Learning to Sample Replacements for ELECTRA Pre-Training

TL;DR

This paper enhances ELECTRA pre-training by introducing a hardness prediction mechanism and focal loss for the generator, leading to more effective replacement sampling and improved downstream task performance.

Contribution

It proposes novel sampling and loss techniques to address inefficiencies and biases in ELECTRA's generator-discriminator training process.

Findings

Improved downstream task performance with the new sampling methods.

Reduced training variance of the discriminator.

Mitigated over-confidence bias in generator predictions.

Abstract

ELECTRA pretrains a discriminator to detect replaced tokens, where the replacements are sampled from a generator trained with masked language modeling. Despite the compelling performance, ELECTRA suffers from the following two issues. First, there is no direct feedback loop from discriminator to generator, which renders replacement sampling inefficient. Second, the generator's prediction tends to be over-confident along with training, making replacements biased to correct tokens. In this paper, we propose two methods to improve replacement sampling for ELECTRA pre-training. Specifically, 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 efficient sampling reduces the training variance of the discriminator. Moreover, we propose to use a focal loss for the generator in order to relieve oversampling of correct tokens as replacements. Experimental results show that our method improves ELECTRA pre-training on various downstream tasks.