ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators

Inc Lomin's avatar
Oct 31, 2021
ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators

Introduction

Motivation

BERT와 같은 언어 모델들은 Masked Language Modeling task를 이용하여 pre-training을 하는 것이 일반적입니다.
  • Masked Language Modeling(MLM) : 언어 모델의 input token들 중 일부(보통 15% 정도)를 mask시키고 이 토큰들을 예측해내도록 하는 task입니다. Language Model들은 MLM task를 통해 문맥 정보를 파악(정확히는 양방향의 정보를 고려)하는 능력을 기르게 됩니다.
하지만 MLM Task는 example의 15%에 대해서만 학습할 수 있으므로 상당한 computing cost가 필요합니다. 또한, pre-training 시에는 [MASK] 토큰을 모델이 참고하여 예측하지만 fine-tuning 때에는 [MASK] 토큰이 존재하지 않는다는 문제점이 있습니다.
이러한 문제들을 해결하기 위해 ELECTRA는 Replaced Token Detection이라는 새로운 pre-training task를 제안하여 효율적인 pre-training을 수행하도록 하였습니다.
notion image
 

Contributions

  1. Replaced Token Detection이라는 새로운 pre-training task 제안 : 적은 파라미터와 적은 계산량으로도 SOTA(State-Of-the-Art)에 준하는 성능이나 새로운 SOTA 성능을 달성
      • ELECTRA-large 모델의 경우, RoBERTa, XLNet보다 더 적은 파라미터와 1/4의 계산량으로 학습했지만 이들과 비슷한 성능을 보임
      • GLUE(General Language Understanding Evaluation) benchmark와 SQuAD(Stanford Question Answering Dataset) 2.0에서 SOTA 성능을 달성
  1. Language Model을 Generator가 아닌 Discriminator로 보도록 하는 것을 제안 : 학습 과정에서 input token들의 효율적인 사용이 가능
    1. This corruption procedure solves a mismatch in BERT (although not in XLNet) where the network sees artificial [MASK] tokens during pre-training but not when being fine-tuned on downstream tasks. A key advantage of our discriminative task is that the model learns from all input tokens instead of just the small masked-out subset, making it more computationally efficient.
 

Proposed Method

Replaced Token Detection(RTD)

ELECTRA에서는 pre-training을 위해 RTD라는 task를 도입하였습니다.
RTD task를 통해 학습을 하기 위해서는 다음의 그림과 같이 generator 와 discriminator 의 두 개의 네트워크가 필요합니다. 이 때, 두 개의 네트워크는 모두 Transformer 인코더의 구조이며, 토큰 시퀀스 을 입력으로 받아 문맥 정보를 반영한 벡터 시퀀스 을 출력(입력에 대해 출력을 매핑)합니다.
 
구체적으로, RTD task에서 generator 네트워크는 BERT의 MLM task와 동일하게 단순히 masking된 토큰들(아래의 그림에서 [MASK] 토큰들)을 원래의 token으로 복원합니다.
다음으로 discriminator 네트워크는 generator의 output token sequence들을 입력 받아 이들이 original인지, replaced인지를 예측합니다.
notion image
 

Generator

generator 의 구체적인 학습 매커니즘은 다음과 같습니다.
  1. 입력 토큰 시퀀스 에 대해서 masking할 위치에 해당하는 를 랜덤으로 결정합니다.(이 때, masking 비율은 BERT와 동일하게 15% 비율을 유지)
  1. 결정한 위치에 있는 입력 토큰들을 [MASK]로 치환합니다.
  1. masking된 입력 에 대해서 generator는 원래의 토큰이 무엇인지를 예측합니다.(아래의 식에서 는 masking된 토큰의 위치를 의미. 또한, 는 임베딩을 의미. 즉, 아래의 식은 모델의 출력 레이어와 임베딩 레이어의 weight sharing을 하겠다는 의미)
    1. notion image
  1. 학습을 위한 loss function은 기존의 MLM loss (Cross Entropy loss)를 사용합니다.
    1. notion image
 

Discriminator

discriminator 의 구체적인 학습 매커니즘은 다음과 같습니다.
  1. generator 를 통해 예측된 output token sequence(에 의해 masking token들이 치환된 것으로 이해)를 입력으로 받습니다.(이를 라 합니다.)
  1. discriminator 는 치환된 입력 에 대해서 각 토큰이 원래 입력과 동일(original)한지 치환(replaced)된 것인지를 예측한다.(아래의 수식은 번째 토큰에 대해 예측한것)
  1. 학습을 위한 loss function은 Binary Cross Entropy loss를 사용합니다.
    1. notion image
 
최종적으로, ELECTRA에서는 generator는 버리고 discriminator만을 downstream task에 사용합니다.
 

Differences between GAN and ELECTRA

  • generator가 원래 토큰과 동일한 토큰을 생성했을 때, GAN의 discriminator는 negative sample(fake)로 간주하지만 ELECTRA의 discriminator는 positive sample(real)로 간주합니다. — GAN이 학습이 완료된 상태에서를 의미합니다.
  • generator가 discriminator를 속이기 위해 adversarial하게 학습이 진행되는 것이 아니라, maximum likelihood를 통해 학습이 진행됩니다. 따라서, 최종적으로 ELECTRA의 loss function은 다음과 같습니다. — Reinforcement Learning, Adversarial Learning을 구현하여 학습시켜봤지만 성능이 좋지 않게 나왔다고 합니다.
    • notion image
  • 일반적인 GAN에서는 generator의 입력으로 noise vector를 넣어주지만 ELECTRA의 generator에는 noise vector가 들어가지 않습니다.
 

Experiments

Experimental Setup

  • 논문에서는 GLUE benchmark와 SQuAD 데이터셋과 같은 downstream task들에 대한 fine-tuning을 수행하여 성능 평가를 합니다.
  • pre-training을 위해서는 BERT와 동일하게 Wikipedia, BooksCorpus 말뭉치들을 사용했습니다.
  • Large model의 경우, ClueWeb, CommonCrawl, Gigaword와 같은 말뭉치들을 사용했습니다.
  • 몇몇 evaluation dataset들은 크기가 작아서 random seed의 영향을 받을 수 있으므로 10개의 다른 checkpoint들의 평균 성능으로 최종 성능을 측정했습니다.
 

Weight Sharing

논문에서는 generator와 discriminator 네트워크 간의 weight를 공유하는 weight sharing 방법에 대해 평가합니다. GLUE benchmark에 대해 다른 weight sharing 방법을 적용하여 평가한 성능은 다음과 같습니다.
  • weight sharing 사용 x : 83.5
  • embedding의 weight만을 공유 : 84.3
  • 모든 weight를 공유 : 84.4
결론적으로 generator와 discriminator 간의 모든 weight를 sharing을 하는 것이 가장 좋은 성능을 보였지만 이는 generator와 discriminator의 크기를 동일하게 맞추어야 한다는 제약이 따릅니다.
논문의 Smaller Generators 부분에서 다루고 있듯이 generator의 크기를 줄이는 것이 상당히 효과적이었기 때문에 이후의 모든 실험들을 embedding의 weight만을 공유하도록 셋팅하였다고 합니다.
 

Smaller Generators

generator와 discriminator의 크기를 동일하게 하는 것은 기본 BERT를 학습 시키는 것과 비교하여 2배의 계산량을 필요로 합니다.(Transformer encoder 두개를 학습시키는 것과 동일하기에)
따라서, 논문에서는 generator의 크기를 줄여보는 실험을 진행했습니다.
아래의 그림에서 왼쪽에 있는 차트는 discriminator의 크기를 고정시켜두고 generator의 크기를 변경해가며 실험을 한 결과를 보여줍니다. 차트에서 확인할 수 있듯이 generator는 discriminator의 크기의 1/2에서 1/4정도로 지정하는 것이 가장 좋다고 합니다.(논문에서는 generator가 너무 strong할 경우, discriminator의 학습에 어려움을 야기하는 것으로 추측함)
notion image
 

Training Algorithms

논문에서는 ELECTRA 학습을 위해 다음과 같이 다른 algorithm들의 실험도 수행했습니다.
  • Adversarial Training : GAN과 같이 adversarial training 방법으로 학습
  • Two-stage Training : generator만 으로 step 학습 —> discriminator를 generator의 학습된 가중치로 초기화하고 discriminator만 step 학습
위의 그림에서 오른쪽에 있는 차트는 해당 실험에 대한 결과를 보여줍니다.
최종적으로는 generator와 discriminator를 joint하게 학습시키는 방법이 성능이 가장 좋았다고 합니다.
 

Experimental Results

마지막으로 논문에서는 pre-training의 효율성 향상이라는 목표를 성공적으로 이루어 냈음을 다양한 크기의 모델들과 비교하며 보여줍니다.
먼저, small size의 ELECTRA는 훨씬 큰 모델인 GPT보다도 좋은 성능을 기록하였으며, base size의 ELECTRA는 동일한 크기인 BERT-Base보다도 좋은 성능을 기록했습니다.
notion image
다음으로 large size의 모델들에 대한 비교입니다.
ELECTRA-400K는 RoBERTa 모델의 계산량에 1/4에 해당하는 학습만으로도 RoBERTa나 XLNet에 필적할만한 성능을 보여줍니다. 또한, ELECTRA-1.75M의 경우, 이들을 뛰어넘는 성능을 보여주었습니다. 이 또한 RoBERTa-500K와 XLNet에 비해 적은 계산량으로 달성했습니다.
notion image
notion image
 

Conclusion

NLP 분야에서 pre-trained language model의 중요성이 커지는 가운데, ELECTRA는 pre-training을 효율적으로 할 수 있는 방법론을 제안한 논문입니다.
Text Recognition 또는 Document Understanding에 Language Model을 접목시키려는 노력이 커지고 있으므로 저희의 독자적인 Language Model(또는 Vision Transformer)을 만드는 것은 언젠가는 해야할 일이라고 생각했습니다.
따라서, 이러한 공부를 미리 해보는 것도 좋을거 같습니다.
 
 
 
 
Share article