꼼꼼하고 이해하기 쉬운 ELECTRA 논문 리뷰

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

이주홍 | 2020년 05월 08일

이번 ICLR 2020에서 구글 리서치 팀은 새로운 pre-training 기법을 적용한 language model, ELECTRA (Efficiently Learning an Encoder that Classifies Token Replacements Accurately)를 발표하였습니다. 기존 BERT를 비롯한 많은 language model들은 입력을 마스크 토큰으로 치환하고 이를 치환 전의 원본 토큰으로 복원하는 masked language modeling 태스크를 통해 pre-training을 합니다. 하지만 이런 모델들은 학습 시 상당히 많은 계산량을 필요로 하기 때문에 충분한 컴퓨팅 리소스가 없는 영세한 연구자들은 language model의 pre-training을 연구하기 점점 더 어려워지고 있습니다(ㅠㅠ). ELECTRA는 모델의 정확도와 함께 학습의 효율성에도 주목하고 있습니다. 본 논문은 학습 효율을 향상시키기 위해 Replaced Token Detection (RTD)이라는 새로운 pre-training 태스크를 제안하며, 이를 통해 ELECTRA는 보다 빠르고 효과적으로 학습합니다. 결과적으로 ELECTRA는 모델 크기, 데이터, 컴퓨팅 리소스가 동일한 조건에서 기존 BERT의 성능을 능가했습니다. 이러한 결과는 Small 모델 세팅에서 더욱 두드러졌는데요, 하나의 GPU를 사용하여 단 4일만 학습한 모델로 계산량이 30배인 GPT를 능가하기도 하였습니다. Large 모델 역시 RoBERTa나 XLNet 대비 1/4의 계산량만으로 비슷한 성능에 도달했습니다.

효율적인 학습으로 단 며칠만에 기존 language model들을 뛰어넘는 성능을 낼 수 있는 ELECTRA를 좀 더 꼼꼼하고 이해하기 쉽게 알아보겠습니다!

Table of Contents

  1. Introduction
  2. Method
    1. Generator
    2. Discriminator
    3. GAN과의 차이점
  3. Experiments
    1. Experimental Setup
    2. Model Extensions
      1. Weight sharing
      2. Smaller Generators
      3. Training Algorithms
    3. Small Models
    4. Large Models
    5. Efficiency Analysis
  4. Conclusion
  5. References

Introduction

현재 state-of-the-art representation learning 기법은 일종의 denoising autoencoder 학습이라고 볼 수 있습니다. 주로 입력 시퀀스의 토큰 중 약 15% 정도를 마스킹하고 이를 복원하는 masked language modeling (MLM)이라는 태스크를 통해서 학습을 하죠. MLM은 기존의 autoregressive language modeling 학습에 비해 양방향 정보를 고려한다는 점에서 효과적인 학습을 할 수 있지만 몇 가지 문제점들이 있습니다.

  1. 전체 토큰 중 15%에 대해서만 loss가 발생한다. (= 하나의 example에 대해서 고작 15%만 학습함)
  2. (그래서) 학습하는 데 비용이 많이 든다.
  3. 학습 때는 $[MASK]$ 토큰을 모델이 참고하여 예측하지만 실제(inference)로는 $[MASK]$ 토큰이 존재하지 않는다.

논문은 이런 문제를 해결하기 위해 Replaced Token Detection (RTD)이라는 새로운 pre-training 태스크를 제안합니다. RTD는 generator를 이용해 실제 입력의 일부 토큰을 그럴싸한 가짜 토큰으로 바꾸고, 각 토큰이 실제 입력에 있는 진짜(original) 토큰인지 generator가 생성해낸 가짜(replaced) 토큰인지 discriminator가 맞히는 이진 분류 문제입니다. (Figure 2를 보시면 직관적으로 이해하실 수 있습니다.) ELECTRA는 RTD 태스크로 입력의 15%가 아닌 모든 토큰에 대해서 학습하기 때문에 상당히 효율적이면서도 효과적입니다. 일례로 ELECTRA는 BERT (Devlin et al., 2019)보다 훨씬 빠르게 학습할 수 있으며 downstream task에서의 성능도 더 좋습니다.

Figure 1을 통해서 ELECTRA의 성능과 효율성을 간단히 살펴보겠습니다. ELECTRA가 다른 방식들에 비해 매우 빠르게 성능이 향상되는 것을 볼 수 있으며, 동일한 조건 (모델 크기, 데이터, 계산량)으로 비교했을 때 모든 학습 과정에서 더 높은 GLUE 성능을 보인다는 것을 알 수 있습니다. 또한 ELECTRA-Small의 경우에는 하나의 GPU로 단 4일이면 학습이 가능한데, 이는 BERT-Large 대비 파라미터 수는 1/20, 계산량은 1/135에 해당하는 수치입니다. 그럼에도 불구하고 ELECTRA-Small은 BERT-Small보다 GLUE 성능이 5 point나 높고, 심지어는 훨씬 큰 모델인 GPT (Radford et al., 2018)보다도 성능이 높습니다. Small 모델뿐만 아니라 Large 모델 역시 좋은 성능을 보입니다. ELECTRA-Large는 RoBERTa (Liu et al., 2019)나 XLNet (Yang et al., 2019)보다 더 적은 파라미터, 1/4의 계산량으로 학습했지만 이들과 비슷한 성능을 보였습니다. GLUE에서는 ALBERT (Lan et al., 2019)도 능가하였고 SQuAD 2.0은 state-of-the-art 성능을 기록하였습니다. 위 결과들을 종합적으로 보았을 때, ELECTRA는 새롭게 제안하는 RTD 태스크를 통해 기존 representation learning 기법들보다 더욱 효율적으로(compute-efficient and parameter-efficient) 학습했다고 할 수 있습니다.

Method

이번 챕터에서는 ELECTRA 학습을 위한 전체 모델 구조와 RTD 태스크를 자세히 알아보겠습니다. 먼저, Figure 2에서 볼 수 있듯이 RTD 태스크를 통해 학습하기 위해서 generator $G$와 discriminator $D$, 두 개의 네트워크가 필요합니다. 두 네트워크는 공통적으로 Transformer 인코더 구조 (Vaswani et al., 2017)이며, 토큰 시퀀스 $\textbf{x} = [x_1, x_2, …, x_n]$을 입력으로 받아서 문맥 정보를 반영한 벡터 시퀀스 $h(\textbf{x}) = [h_1, h_2, …, h_n]$로 매핑시킵니다. Generator와 discriminator, 각 네트워크의 구체적인 학습 방법과 특징은 세부 챕터에서 살펴보겠습니다.

Generator

Generator $G$는 한마디로 BERT의 MLM과 똑같습니다. 구체적으로 학습 매커니즘은 다음과 같습니다.

  1. 입력 $\textbf{x} = [x_1, x_2, …, x_n]$에 대해서 마스킹할 위치의 집합 $\textbf{m} = [m_1, m_2, …, m_k]$을 결정합니다.
    • 모든 마스킹 위치는 1과 $n$ 사이의 정수이며, 아래와 같이 수학적으로 표현할 수 있습니다.
      • $m_{i} \sim \text{unif { 1, n}} \; \text{for} \; i = 1 \; \text{to} \; k$
    • 마스킹할 개수 $k$는 보통 $0.15n$을 사용합니다. (전체 토큰의 15%)
  2. 결정한 위치에 있는 입력 토큰을 $[MASK]$로 치환합니다.
    • 이 과정을 $\textbf{x}^{masked} = \text{REPLACE}(\textbf{x}, \textbf{m}, [MASK])$와 같이 표현합니다.
  3. 마스킹 된 입력 $\textbf{x}^{masked}$에 대해서 generator는 아래와 같이 원래 토큰이 무엇인지 예측합니다.

    • 이런 과정을 수학적으로 표현하면 아래와 같습니다. ($t$ 번째 토큰에 대한 예측)
    • $e(\cdot)$는 임베딩을 의미합니다. 위의 식은 language model의 출력 레이어과 임베딩 레이어의 가중치를 공유(weight sharing)하겠다는 의미합니다.
  4. 최종적으로 아래와 같은 MLM loss로 학습합니다.

Discriminator

Discriminator $D$는 입력 토큰 시퀀스에 대해서 각 토큰이 original인지 replaced인지 이진 분류로 학습합니다. 구체적인 학습 매커니즘은 다음과 같습니다.

  1. Generator $G$를 이용해서 마스킹 된 입력 토큰들을 예측합니다. (Generator의 학습 매커니즘 1~3단계)
  2. Generator $G$에서 마스킹할 위치의 집합 $\textbf{m}$에 해당하는 위치의 토큰을 $[MASK]$가 아닌 generator의 softmax 분포 $p_G(x_t | \textbf{x})$에 대해 샘플링한 토큰으로 치환(corrupt)합니다.

    • Original input : [the, chef, cooked, the, meal]
    • Input for generator : [$[MASK]$, chef, $[MASK]$, the, meal]
    • Input for discriminator : [the, chef, ate, the, meal]
    • 첫 번째 단어는 샘플링했는데 원래 입력 토큰과 동일하게 “the”가 나온 것입니다.
    • 세 번째 단어는 샘플링했는데 원래 입력 토큰인 “cooked”가 아니고 “ate”가 나온 것입니다.
    • 이 치환 과정은 수학적으로 다음과 같습니다.
  3. 치환된 입력 $\textbf{x}^{corrupt}$에 대해서 discriminator는 아래와 같이 각 토큰이 원래 입력과 동일한지 치환된 것인지 예측합니다.

    • Target classes (2)
      • original: 이 위치에 해당하는 토큰은 원본 문장의 토큰과 같은 것
      • replaced: 이 위치에 해당하는 토큰은 generator $G$에 의해서 변형된 것
    • 이런 과정을 수학적으로 표현하면 아래와 같습니다. ($t$ 번째 토큰에 대한 예측)
  4. 최종적으로 아래와 같은 loss로 학습합니다.

GAN과의 차이점

Generator와 discriminator가 나오니 GAN (Goodfellow et al., 2014)이 가장 먼저 떠오르는데요, ELECTRA의 training objective는 GAN과 몇 가지 차이점이 있습니다.

  1. Generator가 원래 토큰과 동일한 토큰을 생성했을 때, GAN은 negetive sample (fake)로 간주하지만 ELECTRA는 positive sample로 간주하는 점.
  2. Generator가 discriminator를 속이기 위해 adversarial하게 학습하는 게 아니고 maximum likelihood로 학습한다는 점.
    • Generator에서 샘플링하는 과정 때문에 역전파가 불가능하고, 따라서 adversarial하게 generator를 학습하는 게 어려움.
    • 그래서 강화 학습으로 이를 구현해보았지만 maximum likelihood로 학습시키는 것보다 성능이 좋지 않았음 (논문의 Appendix F 참조).
  3. Generator의 입력으로 노이즈 벡터를 넣어주지 않는 점.

최종적으로 ELECTRA는 대용량 코퍼스 $\mathcal{X}$에 대해서 generator loss와 discriminator loss의 합을 최소화하도록 학습합니다. 이때 $\lambda$는 50을 썼으며 (논문의 Appendix A 참조), 이 파라미터는 이진 분류인 discriminator loss와 30000 클래스 분류인 generator loss의 스케일을 맞추는 역할을 합니다.

앞에서 설명했듯이 샘플링 과정이 있기 때문에 discriminator loss는 generator로 역전파 되지 않으며, 위의 구조로 pre-training을 마친 뒤에 generator는 버리고 discriminator만 취해서 downstream task으로 fine-tuning을 진행합니다.

Experiments

Experimental Setup

ELECTRA의 성능 평가를 위해 General Langauage Understanding Evaluation (GLUE) 벤치마크와 Stanford Question Answering (SQuAD) 데이터셋을 사용하였습니다. 대부분의 실험은 BERT와 동일하게 Wikipedia와 BooksCorpus를 사용해서 pre-training을 시켰으며, Large 모델의 경우에는 XLNet에서 사용한 ClueWeb, CommonCrawl, Gigaword를 사용하였습니다.

모델의 구조와 대부분의 하이퍼 파라미터 역시 BERT와 동일하게 세팅하였습니다. 몇몇 평가 데이터셋은 크기가 작아서 랜덤 시드의 영향을 많이 받을 수 있기 때문에 10번의 fine-tuning 결과의 중간값(median)을 최종 성능으로 사용하였습니다. 학습에 대한 더 자세한 사항과 하이퍼 파라미터 값은 Appendix를 참조하시기 바랍니다.

Model Extensions

ELECTRA의 성능을 더 향상시키는 몇 가지 기법과 실험 결과를 알아보겠습니다. 실험은 BERT-Base와 동일한 크기의 모델과 데이터를 사용하였습니다.

Weight sharing

Generator와 discriminator는 모두 Transformer 인코더 구조이기 때문에 두 네트워크의 가중치를 공유하여 학습하는 weight sharing 기법을 써볼 수 있고, 이로써 pre-training의 효율 향상을 기대할 수 있습니다. 동일한 구조의 두 네트워크가 주어졌을 때, (1) 임베딩의 가중치만 공유하고 그 외의 가중치는 따로 학습시키거나 (2) 모든 가중치를 서로 공유하는 방법을 생각해 볼 수 있습니다. 논문에서는 각 weight sharing 세팅에 대해 500k 스텝만큼 학습시키고 GLUE로 성능을 평가했습니다. 이 실험에서 저자들은 generator와 discriminator의 크기를 동일하게 설정하였습니다.

결과적으로 모든 가중치를 공유하는 것이 가장 좋은 성능을 보였으며, 임베딩만 공유하는 방법도 거의 비슷하게 좋은 성능을 기록하였습니다. 이런 결과를 논문은 다음과 같이 해석합니다.

Discriminator는 입력으로 들어온 토큰만 학습하는 반면, generator는 출력 레이어에서 softmax를 통해 사전에 있는 모든 토큰에 대해서 밀도 있게 학습할 수 있다. ELECTRA는 결국 discriminator만을 취해서 사용하는데, 이때 generator와 임베딩을 공유해서 학습한 경우의 discriminator는 훨씬 효과적으로 학습했을 것이고 결과적으로 좋은 성능을 기록한 것으로 보인다.

성능 자체는 모든 가중치를 공유하는 방법이 약간 더 좋긴 했지만 generator와 discriminator의 크기를 반드시 동일하게 맞춰야 한다는 제약이 따릅니다. 결국 discriminator만 사용하는데 generator를 이와 동일한 크기로 가져가는 것은 학습의 효율을 떨어뜨릴 수 있으며, 이후 챕터(Smaller Generators)에서 자세히 살펴보겠지만 더 작은 generator를 사용하는 것이 상당히 효과적이기 때문에 논문에서 소개하는 이후 실험은 모두 임베딩만 공유하는 세팅으로 진행하였다고 합니다.

Smaller Generators

Generator와 discriminator의 크기를 동일하게 가져간다면, ELECTRA를 학습하기 위해서는 BERT와 같은 일반 MLM 모델에 비해 단순 계산으로 거의 두 배의 계산량이 필요합니다. 이 문제를 완화하기 위해서 generator의 크기를 줄여보는 실험을 진행합니다. 다른 하이퍼 파라미터는 그대로 두고 레이어의 크기만을 줄이는 것이죠. 여기서 레이어의 크기란 히든 레이어의 크기, FFN의 크기, 어텐션 헤드의 수를 의미합니다. 논문에서는 크기가 작은 generator뿐만 아니라 unigram generator로도 실험했습니다. Unigram generator는 간단한 베이스라인으로, 학습 코퍼스에 등장하는 unigram의 분포를 기반으로 샘플링하는 generator입니다.

실험은 모두 동일하게 500K 스텝 동안 학습시켰으며 결과는 Figure 3의 왼쪽과 같습니다. 모두 동일한 스텝만큼을 학습했기 때문에 작은 모델은 똑같은 계산량, 시간만큼 학습하면 더 많은 스텝을 돌 것이고, 결과적으로 작은 모델 입장에서는 계산량 대비 성능을 손해본 셈입니다. 그럼에도 불구하고 discriminator의 크기 대비 1/4 - 1/2 크기의 generator를 사용했을 때 가장 좋은 성능을 보입니다. 이런 실험 결과를 논문은 다음과 같이 해석합니다.

Generator가 너무 강력하면 discriminator의 태스크가 너무 어려워져서 이런 현상이 나타날 수 있다. 게다가 discriminator의 파라미터를 실제 데이터 분포가 아닌 generator를 모델링하는 데 사용할 수도 있다.

개인적인 생각을 덧붙이자면, generator가 강력하면 softmax 분포가 몇몇 토큰에 치우쳐 있을 가능성이 높고 샘플링 결과가 다양하지 않기 때문에 discriminator가 효과적으로 학습할 수 없는 이유도 있지 않을까 싶습니다.

Training Algorithms

이번에는 ELECTRA를 효과적으로 학습시킬 수 있는 알고리즘을 알아볼텐데요. 기본적으로 앞서 Method 챕터에서 소개한 training objective는 generator와 discriminator를 jointly 학습시키는 방식입니다. 이와 함께 다음 2가지 학습 방식으로 실험을 진행하였습니다.

결과는 Figure 3의 오른쪽과 같으며, generator와 discriminator를 jointly 학습시키는 방식이 가장 좋았습니다. 좀 더 자세히 결과를 살펴보면 위의 Two-stage 학습에서 discriminative objective로 바꿨을 때, 성능이 쭉 오른 것을 볼 수 있습니다. 또한 Adversarial 학습이 maximum likelihood 기반의 학습보다 성능이 낮다는 것도 알 수 있는데요, 이런 현상의 원인은 다음과 같습니다. (Caccia et al., 2018)

  1. MLM 성능이 안 좋아서: MLM 성능은 58% 밖에 안됨. Maximum likelihood로 학습한 generator는 65%였음.
  2. 학습된 generator가 만드는 분포의 엔트로피가 낮아서: Softmax 분포는 하나의 토큰에 확률이 쏠려있고, 이러면 샘플링할 때 다양성이 많이 떨어짐.

Small Models

저자들이 제시한 연구의 큰 향상점 중 하나는 pre-training의 효율성 향상입니다. 이를 검증하기 위하여 저자들은 하나의 GPU로도 빠르게 학습할 수 있는 수준으로 작은 모델을 만들어보는 실험을 진행하였습니다.

ELECTRA-Small 모델의 하이퍼 파라미터는 위의 표와 같으며, 크기가 대조군인 BERT-Base에 비해서 매우 작습니다. 공정한 비교를 위해서 학습 FLOPs가 동일하도록 BERT-Small은 1.5M 스텝, ELECTRA-Small은 1M 스텝 동안 학습시켰습니다. BERT와 함께 ELMo (Peters et al., 2018)와 GPT도 비교하였습니다.

결과는 위의 Table 1과 같습니다. ELECTRA-Small은 BERT-Small보다 무려 5 포인트나 높은 성능을 보였고, 심지어는 훨씬 큰 모델인 GPT보다도 좋은 성능을 보였습니다. (추가적인 실험 결과는 논문의 Appendix D 참조) 또한 수렴 속도가 매우 빠른데요, 하나의 GPU로 단 6시간 만에 꽤 괜찮은 성능을 보여줍니다. Base 크기의 경우에도 ELECTRA-Base는 BERT-Base를 능가할 뿐 아니라 심지어 BERT-Large보다도 더 좋은 성능을 기록하였습니다.

Large Models

저자들은 Small 모델뿐 아니라 Large 모델로도 실험을 진행하였습니다. ELECTRA-Large의 세팅도 앞선 실험들과 마찬가지로 같은 크기의 BERT-Large에 맞춰서 실험하였고 학습 데이터로는 XLNet에서 사용한 데이터를 사용하였습니다. 결과는 아래의 Table 2 (dev), Table 3 (test)와 같았습니다. (모델명 옆에 숫자는 학습 스텝을 의미)

ELECTRA-400K는 RoBERTa(-500K)나 XLNet의 단 1/4의 계산량(FLOPs)만으로 이들과 필적할 만한 성능을 보입니다. 그리고 더 많이 학습시킨 ELECTRA-1.75M은 이들을 뛰어넘는 성능을 보였고, 이 역시도 계산량은 두 모델보다 작습니다. GLUE뿐만 아니라, SQuAD에서도 마찬가지로 ELECTRA는 가장 좋은 성능을 보입니다. (논문의 Table 4 참고)

Efficiency Analysis

논문은 ELECTRA가 왜 성능이 좋은지 좀 더 자세히 이해하기 위해서 일종의 ablation study의 성격으로 다음 실험을 세팅합니다.

ELECTRA 15%는 토큰에 대한 학습 효율(15% vs 100%) 때문에 성능 차이가 생겼다는 것을 보이기 위한 목적, Replace MLM는 pre-training 때만 사용하고 fine-tuning 때는 없는 $[MASK]$ 토큰 때문에 생긴 성능 차이를 보이기 위한 목적으로 보입니다. All-Tokens MLM은 BERT와 ELECTRA를 합친 버전입니다. 논문은 이 세팅의 성능을 좀 더 개선하기 위해 sigmoid 레이어를 통해 입력 토큰을 복사할지 결정하는 확률 $D$을 뽑는 매커니즘을 도입했다고 하는데요, 결과적으로 모델의 softmax 분포는 $D * \text{input-token-distribution} + (1-D) * \text{MLM-output-distribution}$와 같은 형태로 만든 것으로 보입니다.

결과는 위의 Table 5와 같으며, ELECTRA는 ELECTRA 15%Replace MLM 보다 훨씬 좋은 성능을 보였고 All-Tokens MLM은 그나마 ELECTRA에 가까운 성능을 보이며 BERT와 ELELCTRA의 성능 차이를 많이 줄였습니다. 전반적으로 결과를 봤을 때, ELECTRA가 학습 효율도 굉장히 좋고 $[MASK]$ 토큰에 대한 pre-training과 fine-tuning 간의 불일치 문제도 상당히 완화시킨 것을 알 수 있습니다.

추가적으로 히든 레이어 크기에 따른 BERT와 ELECTRA의 성능 변화 실험도 있었는데요. Figure 4에서 볼 수 있듯이 히든 레이어의 크기가 작아질수록 BERT와 ELECTRA의 성능 차이는 커진다는 사실을 알 수 있으며, BERT에 비해 ELECTRA는 모델이 작아도 매우 빠르게 수렴함을 알 수 있습니다.

앞서 소개한 많은 실험을 통해서 ELECTRA가 BERT보다 효율적으로(parameter-efficient) 학습한다고 결론지을 수 있습니다.

Conclusion

논문은 language representation learning을 위한 새로운 self-supervision 태스크인 Replaced Token Detection을 제안합니다. 제안하는 방법의 주요 아이디어는 작은 generator가 만들어 낸 질 좋은 negative sample과 입력 토큰을 구별하도록 텍스트 인코더를 학습시키는 것입니다. Masked language modeling에 비해, 제안하는 pre-training objective는 훨씬 효율적(compute-efficient)이고 downstream tasks에 대한 결과 역시 더 좋다는 것을 많은 실험을 통해서 확인할 수 있었습니다. 또한 ELECTRA는 타 모델들에 비해서 상대적으로 적은 계산량을 사용하는 경우에 더 효과적이었습니다.

저자는 이 연구를 통해서 연구자들이 적은 컴퓨팅 리소스로도 language model의 pre-training에 대한 많은 연구/개발을 하길 바란다는 말과 pre-training과 관련된 향후 연구가 절대적 성능 지표만큼 계산량과 파라미터 수 등의 효율성도 함께 고려했으면 하는 바람을 남겼습니다.

References