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

딥러닝 논문리뷰

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

빈이름 2023. 3. 14. 23:39

ELECTRA 모델에 대해 알아 보겠습니다. ELECTRA는 Pre-training 방식의 일종으로 기존의 BERT에서 좀 더 개선된 pre-training 방식을 제시합니다. 한국어로 사전 학습된 모델도 존재합니다.(KoELECTRA, KcELECTRA) BERT에서 어떤 점들이 개선되었을까요?

1. 개요

본 논문에서 지적하는 것은, BERT가 비효율적인 학습 방법이라는 것입니다. BERT의 학습 방식은 MLM pre-training 이죠. 문장의 15%만을 masking하여 그 부분의 단어를 올바르게 예측하는 방식이었습니다. 다르게 말하면 아무리 긴 문장이 들어와도, 모델은 그 문장 중 15%의 단어만을 학습하는 것입니다. ELECTRA는 문장의 일부가 아니라 전체를 모두 학습하는 방법입니다.

2. ELECTRA Pre-training

어떻게 하면 문장의 전체를 학습할 수 있을까요? MLM 방식으로 문장 전체를 masking 하기는 어렵겠죠. 모델이 mask 위치의 토큰을 추론할 수 있도록 멀쩡한 단어들도 존재해야 할 테니까요. ELECTRA는 Generator와 Discriminator 2개의 모델을 사용해 해결합니다.

Generator는 BERT와 같습니다. 문장의 일부를 마스킹하고, 그 위치에 알맞은 토큰을 예측하는 역할을 수행합니다.

Discriminator가 새로 추가된 모델인데, 구조는 BERT와 같이 attention 블럭을 쌓아 사용합니다. 입력으로 Generator가 마스크를 채운 문장을 입력받고, 각 토큰이 Generator가 만들어낸 토큰인지 아닌지를 판별하는 역할을 수행합니다. Discriminator는 문장의 모든 토큰이, 진짜인지 아닌지를 판별하면서 모든 토큰들을 학습할 수 있습니다. 문장 전체를 학습할 수 있는 것이죠.

Generator와 Discriminator라는 단어를 보고 GAN을 떠올릴 수 있을 것 같습니다. 그러나 저자들은 '이는 GAN과는 다른 구조다' 라고 말하며, 차이점들을 말합니다.

1. Generator가 옳게 예측한 토큰을 Discriminator도 이 단어가 '옳다'라고 예측한다.

GAN과 같은 적대적 학습을 한다면, Generator가 mask 위치에 적절한 단어를 만들어 내면, Discriminator가 이것을 '가짜'라고 판별하도록 학습될 것입니다. 그러나 ELECTRA는 mask 위치에 옳은 단어를 만들면 Discriminator도 해당 단어를 '진짜'라고 판별합니다.

2. Generator는 Discriminator를 속이는 적대적 학습이 아니라, mask 위치의 올바른 토큰을 예측하는 maximum likelihood를 통해 학습된다.

Generator는 BERT와 같이 MLM-pretraining을 통해 학습됩니다. Discriminator의 결과와는 관계 없이 독립적으로 crossentropy를 사용해서 학습을 합니다.

3. Generator에 noise vector가 입력으로 제공되지 않는다.

GAN의 Generator는 진짜 같은 결과를 만들어내기 위해 noise vector를 입력으로 받지만, ELECTRA에 그런건 없습니다. 단지 마스킹이 적용된 문장만이 입력으로 제공됩니다.

이를 통해 GAN과의 차이점을 충분히 안 것 같습니다. 그러면 ELECTRA는 어떻게 학습되는 것인지 한 번 정리하고 가겠습니다.

Generator는 BERT와 같은 MLM Pretraining을 통해 학습됩니다. Generator는 마스킹된 위치에 알맞은 단어 토큰을 예측해내야 합니다. Generator의 손실 함수(CrossEntropyLoss)는 아래와 같이 쓸 수 있습니다.

$$ L_{MLM}(x,\theta_G)=E(\sum_{i\in m}-log{p_G}(x_i|x^{masked})) $$

Discriminator는 각 단어 토큰마다 해당 토큰이 Generator가 만들어낸 단어인지 아닌지를 판별하는 이진 분류를 수행합니다. 손실 함수는 BinaryCrossEntropyLoss를 사용해야겠죠. 수식으로 아래와 같이 쓸 수 있습니다.

$$ L_{Disc}(x,\theta_D)=E(\sum_{t=1}^{n}-1(x^{corrupt}_t=x_t)logD(x^{corrupt},t)-1(x^{corrupt}_t\neq x_t)log(1-D(x^{corrupt}, t))) $$

t는 마스킹 된 토큰의 위치를 말합니다. $x^{corrupt}$는 Generator가 예측한 마스크 위치의 토큰입니다.

최종적으로 Loss 함수는 아래와 같이 2개의 Generator와 Discriminator의 Loss를 더하는 식으로 설정하였습니다.

$$ \underset{\theta_G,\theta_D}{min}\sum_{x\in X} L_{MLM}(x,\theta_G)+\lambda L_{Disc}(x,\theta_D) $$

두 Loss를 더하기만 하기 때문에 GAN처럼 Discriminator부터 Generator까지 이어서 역전파가 일어나지 않습니다.

이렇게 사전학습된 모델에서 Generator를 버리고 Discriminator를 사용해 downstream task를 수행합니다.

3. 모델 설정

1. Weight sharing

Generator와 Discriminator 모두 같은 BERT 구조를 사용하기 때문에 두 모델의 크기가 같다면 weight sharing이 가능합니다. Weight tying의 효과를 알아보기 위해 실험한 결과는 아래와 같습니다.

	No weight tying	tying token embedding	tying all weights
GLUE score	83.6	84.3	84.4

토큰 임베딩 레이어의 weight를 공유하는 것은 꽤 큰 성능 향상을 가져왔습니다. 저자들은 그 이유가 Discriminator는 모델에 입력으로 들어온 문장 단어들만 학습을 하지만, Generator는 마스킹된 위치에 들어갈 단어 토큰을 모두 고려해야 하기 때문에 Discriminator의 토큰 임베딩 레이어 학습에 도움을 준 것 같다고 추측했습니다. 반면에 모든 weight를 공유하는건 토큰 임베딩 weight만 공유하는 것과 큰 차이가 없었습니다. 그리고 추후 실험을 통해 Generator는 Discriminator보다 더 작은 모델을 사용하는 것이 효율적이라는 결과를 얻었기 때문에 결과적으로는 토큰임베딩 weight만 공유했습니다.

2. Smaller Generators

만약 Generator와 Discriminator에 같은 크기의 모델을 사용한다면 ELECTRA는 BERT의 2배 크기의 모델을 사용하게 됩니다. 학습의 효율성을 강조하고 있는만큼 이는 ELECTRA의 약점이 될 수도 있겠죠. 그렇기 때문에 실험을 통해 Generator의 크기를 줄이는 방식을 선택합니다. (Discriminator는 downstream task에 사용해야 하기 때문에 BERT와 비교를 위해 크기를 줄이지 않았을 것 같습니다.) Generator의 크기는 Attention block의 수를 줄이는 것보다는 각 레이어의 차원 수를 줄이는 방식을 선택했습니다.

실험 결과, Generator가 Discriminator의 1/2~1/4 정도의 차원 수를 갖도록 크기를 줄이는 것이 제일 효율이 좋은 것을 알 수 있었습니다. 반대로 오히려 Generator의 차원 수가 너무 커질 경우, 오히려 GLUE 점수가 떨어지는데요, 저자들은 이것이 Generator의 성능이 너무 좋아 Discriminator가 충분히 학습하기 어려워진 것이 아닐까 하고 추측합니다.

3. Training Algorithms

Generator와 Discriminator의 효과적인 학습 방법을 알기 위해 3가지 방식을 실험했습니다.

1. Joint training : 두 모델을 초기부터 함께 학습합니다.

2. Two-stage training : Generator만 먼저 일정 step을 따로 학습한 뒤, Generator의 weight를 동결하고 Discriminator만 학습합니다. Discriminator의 weight는 먼저 학습된 Generator의 weight로 초기화하여 학습합니다.

3. Adversarial training : 랜덤으로 텍스트를 마스킹해야 하는 과정 때문에 GAN과 같이 adversarial training을 구현하기는 어렵지만 저자들은 실험을 위해 강화학습을 적용하여 어떻게든 끼워 맞춰서 adversarial training을 수행했다고 합니다.

결과는 위와 같습니다. 실험을 통해 알 수 있는 것은, 3가지 방식 모두 BERT보다 성능이 좋았다는 것입니다. Two-Stage ELECTRA 그래프 중간의 'Switch to Discriminator Loss' 부분까지는 Generator만 학습하기 때문에 BERT와 같은 성능이 나타나지만, Discriminator를 추가하여 ELECTRA 학습을 시작하자, BERT보다 성능이 올라가는 것을 확인할 수 있었습니다.

가장 성능이 좋았던 것은 joint training 방식이었습니다. Adversarial training은 오히려 성능이 떨어졌습니다.

4. 실험

실험과 평가는 GLUE, SQuAD 데이터셋을 사용했습니다. 모델은 small 모델, base 모델, large 모델을 나눠서 실험합니다. 사전학습 데이터는 BERT와 똑같이 BookCorpus와 Wikipedia를 사용했지만, Large 모델은 XLNet과 같이 ClueWeb, CommonCrawl, Gigaword 데이터를 추가해 사전학습을 수행했다고 합니다.

4-1. Small Models

Small model은 아래와 같은 하이퍼 파라미터를 사용했습니다.

max sequence length : 512 → 128
batch size : 256 → 128
hidden dim : 768 → 256
embedding size : 768 → 128

결과는 위와 같았습니다. FLOPs는 학습에 사용되는 컴퓨팅 자원의 수치를 나타내는 점수입니다. 결과를 보면 ELECTRA-Small은 같은 Small 모델인 BERT-Small 보다 GLUE 점수가 5점정도 더 높았고, 심지어 훨씬 큰 크기의 모델인 ELMo, GPT보다도 높은 점수를 받았습니다.

게다가 6시간의 학습만으로도 BERT-Small이 4일 내내 학습한 GLUE 점수와 비슷한 점수를 내는 것을 알 수 있습니다.

Base 크기에서도 ELECTRA는 BERT보다 높은 점수를 냈습니다.

4-2. Large Models

큰 모델에서도 ELECTRA가 효과적이라는 것을 입증하기 위해 Large 모델도 실험을 진행했습니다. 결과는 위 표와 같이 나왔습니다. ELECTRA는 400k step 학습만으로도 RoBERTa, XLNet에 비견할만한 성능을 보였습니다. XLNet과 같이 1.75M step까지 학습한 결과 6개의 task에서 SOTA를 달성했습니다. SQuAD에서도 비슷한 결과가 나왔습니다.

이 결과를 통해 ELECTRA가 기존의 pre-training 방식들보다 더 적은 시간의 학습시간으로 좋은 성능을 낼 수 있다는 것을 알 수 있습니다.

5. ELECTRA의 효과

정확히 ELECTRA의 어느 부분이 효과적인가를 알아봅니다. ELECTRA의 효과를 분석하기 위해 아래와 같은 조건으로 실험을 추가로 진행했습니다.

1. ELECTRA 15% : Discriminator가 문장 전체 토큰이 아니라 Corrupt된 토큰 위치만 학습한 세팅입니다.

2. Replace MLM : Discriminator에서 MLM을 수행하지만, "[MASK]" 토큰을 전혀 사용하지 않고, Generator가 만들어낸 토큰을 마스킹으로 사용합니다.

3. All-Tokens MLM : Replace MLM 방식에서 Discriminator가 마스킹 위치뿐만 아니라, 모든 위치의 토큰을 예측합니다.

3가지 세팅의 결과는 위와 같습니다. 이를 통해 아래와 같은 사실들을 알 수 있습니다.

1. ELECTRA는 문장 전체를 학습한다는 점에서 이득을 많이 본다.

2. "[MASK]" 토큰은 downstream task에서 성능 손해를 가져온다.

추가적으로 모델 크기에 따른 ELECTRA의 효용성도 테스트 해보았습니다. 실험 결과 ELECTRA는 모델이 작을수록 더 많은 이득이 있는 것을 알 수 있었습니다.

6. 결론

ELECTRA는 pre-training 방식에서, 모델의 성능도 좋지만 학습의 효율성도 생각해봐야 한다고 주장합니다. 실제로 pre-training을 수행하는 것은 굉장히 부담되는 작업입니다. 이런 식의 연구가 많이 나와서 더 많은 사람들이 쉽게 pre-training을 수행할 수 있게 되면 좋겠습니다.