ELECTRA(ICLR 2020) 논문 리뷰

 

 

오늘 소개할 논문은 "ELECTRA: PRE-TRAINING TEXT ENCODERS AS DISCRIMINATORS RATHER THAN GENERATORS"이다.(원논문 링크)

 

What is ELECTRA?

---

ELECTRA는 구글 리서치 팀에서 발표한 논문으로 LM의 새로운 pre-training 기법을 제안하였다. ELECTRA는 Efficiently Learning an Encoder that Classifies Token Replacements Accurately의 약자이다.

기존의 MLM(Masked Language Model) 태스트를 통한 pre-training 방식은 많은 연산량을 필요로 한다. 이로 인해 충분한 컴퓨팅 리소스가 없는 연구자들에게 LM을 학습시키기가 점점 어려워지고 있다. 연구팀은 "Pre-training methods should be compute efficient!!!"라는 관점에 주목해 이 같은 연구를 진행하였다.

ELECTRA는 모델의 정확도와 함께 학습의 효율성에도 주목하였다. 학습 효율을 향상시키기 위해 RTD(Replaced Token Detection)라는 새로운 pre-training task를 제안하였다.

결과적으로 ELECTRA는 모델 크기, 데이터, 컴퓨팅 리소스가 동일한 조건에서 기존 BERT의 성능을 능가하였고, small model에서 하나의 GPU로 4일만 학습한 모델로 계산량이 30배나 많은 GPT를 능가하였다.(GLEU score) Large model 역시 RoBBERTa나 XLNet 대비 더 적은 파라미터 수와 1/4의 계산량으로 비슷한 성능에 도달하였다. 즉, RTD task는 기존의 MLM task보다 parameter-efficient하고 compute-efficient한 pre-training task라는 점을 주장한다.

 

Motive

---

기존에 LM을 pre-train 시키는 주된 방법인 MLM은 다음과 같은 문제가 있엇다.

1. 전체 토큰 중 15%에 대해서만 loss가 발생(즉, 하나의 시퀀스에 대해 15%만 [MASK] 토큰으로 masked-out 시켜 15%만 학습함) -> 학습하는데 비용이 많이 듬

2. mismatch 문제 : pre-train 대는 [MASK] 토큰을 모델이 참고해 예측하지만 fine-tuning 때는 [MASK] 토큰이 존재하지 않는다.

이를 해결하기 위해 본 논문에선 RTD라는 새로운 task를 제안한다. RTD 방식에는 입력의 일부 토큰을 그럴싸한 가짜 토큰으로 변환하는 Generator와 각 토큰이 원래 입력에 있던 진짜(original) 토큰인지 generator가 생성한 가짜(replaced) 토큰인지 맞추는 이진 분류 문제를 수행하는 Discriminator가 있다. 이러한 방식의 RTD task는 입력의 15%가 아닌 모든 토큰에 대해 학습하므로 효율적이다. 실제로, ELECTRA는 BERT보다 학습 속도가 빨랐고, downstream task 성능도 좋았다.

위 그래프를 보면 ELECTRA는 다른 모델들에 비해 매우 빠르게 성능이 향상된다.(빠른 수렴 속도) 또한 동일 조건(모델크기, 데이터, 계산량)으로 비교했을 때 더 높은 GLUE score를 낸다. 학습 속도 측면에서도 ELECTRA-Small은 하나의 GPU로 4일이면 학습이 가능하지만 BERT-Small보다 GLUE score가 5 point나 높고, 훨씬 큰 모델인 GPT보다도 성능이 좋다. ELECTRA-Large는 RoBERTa나 XLNet보다 적은 파라미터, 1/4의 계산량으로 비슷한 성능을 보이고, ALBERT를 능가하였다.

즉, ELECTRA는 기존의 language representation learning 기법들보다 compute-efficient하고 paramter-efficient하다.

 

Method

---

RTD task를 위해선 두 개의 네트워크 Generator(G)와 Discriminator(D)가 필요하다. 두 네트워크는 모두 Transformer encoder 구조로 시퀀스 x = [x1, x2, ..., xn]을 입력으로 받아 문맥 정보를 반영한 벡터 시퀀스 h(x) = [h1, h2, ..., hn]로 매핑시킨다.

Generator

G는 BERT의 MLM과 동일하다. 구체적인 학습 메커니즘은 아래와 같다.

1. 입력 x = [x1, ..., xn]에 대해 마스킹할 위치의 집합 m = [m1,...,mk]를 결정(보통 k=0.15n, 전체 토큰의 15% mask-out)

2. 결정한 위치에 있는 토큰을 [MASK]로 치환(마스킹된 입력 x_masked = REPLACE(x, m, [MASK]))

3. x_masked에 대해 G는 원래 토큰이 무엇인지 예측

Discriminator

D는 입력 토큰 시퀀스에 대해 각 토큰이 original인지 replaced인지 이진 분류로 학습한다. 구체적인 메커니즘은 아래와 같다.

1. masked-out된 토큰들을 generator sample로 replace -> corrupted example x_corrupt 생성

2. D가 x_corrupt의 어떤 토큰이 original input과 match 되는지 predict하도록 학습시킴(이진분류)

• original : 이 위치에 해당하는 토큰은 원본 문장의 토큰과 같은 것
• replaced : 이 위치에 해당하는 토큰은 G에 의해 변형된 것

 

GAN과의 차이점

---

G와 D를 학습시키는 방식은 마치 GAN(Goodfellow et al., 2014)과 유사하다. 그러나 ELECTRA는 GAN과의 차이점이 몇 가지 존재한다.

1. Generator가 원래 토큰과 동일한 토큰을 생성했을 때, GAN은 negative sample(fake)로 간주하지만 ELECTRA는 positive sample로 간주한다.

2. Generator는 Discriminator를 속이기 위해 adversarial하게 학습되는 것이 아니고 maximum likelihood로 학습한다.

• Generator에서 샘플링하는 과정 때문에 역전파가 불가
• 따라서, adversarial하게 학습시키는 게 어려움

3. Generator의 입력으로 noise vector를 넣어주지 않는다.

최종적으로 ELECTRA는 대용량 코퍼스 χ에 대해 G loss와 D loss의 합을 최소화하도록 학습시킨다. 이 때, λ는 50을 사용했으며, 이 파라미터는 binary classification인 D loss와 30000 클래스 분류인 G loss의 스케일을 맞추는 역할을 한다. G의 sampling 과정이 있어서 D loss는 G로 역전파되지 않고, 위 구조로 pre-training을 마친 뒤 G는 버리고 D만 취해서 downstream task로 fine-tuning을 진행한다.

 

Experimental Setup

---

- GLUE 벤치마크와 SQuAD 데이터셋을 사용

- BERT와 동일하게 Wikipedia와 BooksCorpus를 사용해 pre-train

- Large model의 경우 XLNet에서 사용한 ClueWeb, CommonCrawl, Gigaword를 사용

- 모델 구조와 하이퍼 파라미터 역시 BERT와 동일하게 세팅

- 일부 평가 데이터셋은 크기가 작아 랜덤 시드의 영향을 많이 받기 때문에 10번의 fine-tuning 결과의 median을 최종 성능으로 사용

 

Experiment #1 : Weight sharing

---

Generator와 Discriminator는 모두 Transformer 인코더 구조이므로 두 네트워크의 가중치를 공유할 수 있다. 각 가중치를 공유했을 때 세팅에 따른 결과는 아래와 같다.(GLUE score)

• 공유하지 않았을 때 : 83.5
• 임베딩 가중치만 공유했을 때 : 84.3
• 모든 가중치를 공유했을 때 : 84.4

결과적으로 모든 가중치를 공유하는 것이 가장 성능이 좋았고, 임베딩 가중치만 공유하는 방법도 거의 비슷한 성능을 보임.(논문에선 아래와 같이 결과를 해석함)

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

결과적으로는 모든 가중치를 공유한 경우 성능이 가장 높았지만, 이 경우 generator와 discriminator의 크기를 맞추어야 한다는 제약이 있다. 이는 결국, discriminator만 사용하는데 generator를 동일한 크기로 맞추며 학습 효율을 저하시키는 결과를 초래한다. 또한, 뒤에 나오겠지만 small generator를 사용하는 것이 더 효과적(cf. Experiment #2)이다. 따라서, 본 논문에선 임베딩 가중치만 공유하는 세팅으로 실험을 진행하였다.

 

Experiment #2 : Smaller Generators

---

Generator와 discriminator의 크기를 동일하게 하면, ELECTRA는 BERT 같은 MLM 모델에 비해 두 배의 계산량이 필요할 것이다. 연구진은 이를 해결하기 위해 generator의 크기(hidden unit #, FFN size, attention head #) 줄여보는 실험을 진행하였다.

Discriminator 크기의 1/4 ~ 1/2 크기의 generator를 사용할 때 가장 좋은 성능(Figure3 Left)을 내는 것을 볼 수 있다. 이 결과를 논문에선 다음과 같이 해석한다.

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

 

Experiment #3 : Training algorithms

---

기본적으로 generator와 discriminator를 jointly 학습시키는 방식을 사용했지만, 아래의 2가지 학습 방식으로도 실험을 진행하였다.

• Two-stage 학습 : generator만 n step 동안 학습 → discriminator를 generator의 학습된 가중치로 초기화 → generator의 가중치는 고정시키고, discriminator만 n step 동안 학습
• Adversarial 학습 : GAN과 같이 adversarial한 방식을 모사하여 training

결과는 두 네트워크를 jointly 학습시키는 방식이 가장 좋은 성능 보임

• Fig3의 Right에서 two-stage 학습 시에 discriminator loss로 바꿨을 때 성능이 급격하게 증가

Adversarial한 방식이 maximum likelihood 학습보다 성능이 낮음

• MLM의 성능이 안 좋아서 : maximum likelihood로 학습한 generator 성능이 더 좋음
• 학습된 generator가 만드는 distribution의 엔트로피가 낮아서 : softmax 분포가 일부 토큰에 쏠려있어 샘플링의 다양성이 떨어짐

 

Experiment #4 : Small Models

---

pre-training의 효율이 얼마나 증대되는 것인지 검증하기 위해 하나의 GPU로도 학습할 수 있는 수준의 작은 모델을 만들어 실험을 진행하였다. ELECTRA-Small은 같은 조건을 가진 BERT-Small보다 5 point 높은 성능, 훨씬 큰 모델인 GPT보다 좋은 성능을 냈다. 또한 학습에 대한 수렴 속도가 빨랐다. 하나의 GPU로 6시간 학습시킨 모델이 4일간 fully-trained시킨 BERT-Small과 비슷한 성능을 보였다. ELECTRA-Base는 성능면에서 BERT-Base와 BERT-Large를 능가하였다.

 

Experiment #5 : Large Models

---

BERT-Large에 맞춰 ELECTRA-Large 모델로도 실험을 진행하였다. Table2는 dev set result, Table3은 test set result이다.모델 옆의 숫자는 pre-training step을 나타낸다. ELECTRA-400K는 RoBERTa-500K와 XLNet의 1/4의 계산량(FLOPs)으로 비슷한 성능을 보였고, 더 많이 학습시킨 ELECTRA-1.75M은 이들을 뛰어넘는 성능을 보이면서도 계산량은 더 적었다.

 

Experiment #6 : Efficiency Analysis(Ablation Study)

---

ELECTRA의 성능이 좋은 이유를 자세히 분석하기 위해 아래와 같은 세팅으로 실험을 진행하였다.

1. ELECTRA 15% : discriminator loss를 입력 토큰의 15%만으로부터 만들도록 세팅하였다. 이는 토큰에 대한 학습 효율을 15% vs. 100%로 비교해 성능 차이를 보기 위한 목적이다.

2. Replace MLM : discriminator를 (RTD 대신) MLM으로 학습 + [MASK]로 치환하지 않고 generator가 만든 토큰으로 치환하였다. 이는 MLM에서 [MASK] 토큰이 pre-training 때만 사용되고 fine-tuning 때는 나타나지 않는 것으로 인한 성능 차이를 보기 위한 목적이다.

3. All-Tokens MLM : Replace MLM처럼 하되, 일부(15%)가 아닌 모든 토큰을 generator가 생성한 토큰으로 치환하였다.

ELECTRA는 ELECTRA 15%와 Replace MLM보다 좋은 성능을 보였다. All-Tokens MLM은 그나마 ELECTRA에 가까운 성능을 보였다. 전반적으로, ELECTRA는 학습 효율이 좋고 [MASK] 토큰에 대한 mismatch 문제도 완화시키는 것을 알 수 있다. 추가적으로, 히든 레이어 크기에 따른 BERT와 ELECTRA의 성능 비교 실험도 진행(Fig 4)하였다.

 

히든 레이어 크기가 작아질수록 BERT와 ELECTRA의 성능 차이가 커지는 것을 볼 수 있다. BERT에 비해 ELECTRA는 모델이 작아도 빠르게 수렴한다.

 

Conclusion

---

1. Language representation learning을 위한 새로운 self-supervised task인 Replaced Token Detection(RTD)을 제안

2. RTD 방식의 pre-training task는 MLM을 사용하는 타모델들에 비해 more compute-efficient하고, 더 좋은 성능을 보임

3. 이를 통해 많은 연구자들이 적은 컴퓨팅 리소스로도 LM을 학습시킬 수 있는데 기여