ELECTRA. Replaced Token Detection(RTD)
경찰과 도둑
기존 BERT와 같은 모델에서 사용되는 마스크 언어 모델링(MLM)은 입력 문장에서 일부 토큰을 [MASK] 토큰으로 대체한 후, 원래 토큰을 예측하도록 모델을 학습합니다. 이는 추론 작업에 좋은 결과를 보이지만, 효과적인 학습을 위해 많은 계산량이 필요합니다.
리소스와 성능적인 측면에서 이전에 비해 상당히 많이 좋아진 방법인 Replaced Token Detection(RTD)에 대해 소개하겠습니다. 이는 쉽게 말해서 경찰과 도둑에 비유할 수 있습니다. 도둑(Generator)은 기존의 토큰의 일부를 [MASK] 토큰으로 바꾸고 MLM과 같이 다시 원래 토큰을 예측 합니다. 이후 경찰(Discriminator)이 토큰이 원래 토큰인지 새로 생성된 토큰인지 T/F로 맞춥니다.
이 방법은 이미지 생성 모델인 GAN의 방법과 매우 유사하지만 약간의 차이점이 있습니다. 이 방법은 마스킹된 작은 부분집합이 아니라 모든 입력 토큰에 대해 정의되므로 MLM보다 효율적임을 입증했습니다. 결과적으로 ELECTRA는 모델 크기, 데이터, 컴퓨팅 리소스가 동일한 조건에서 기존 BERT의 성능을 능가했습니다. Small 모델인 경우 하나의 GPU를 사용하여 단 4일만 학습한 모델로 계산량이 30배인 GPT를 능가하기도 하였습니다. Large 모델 역시 RoBERTa나 XLNet 대비 1/4의 계산량만으로 비슷한 성능에 도달했습니다.
직접 튜닝해봤을 때에도 엄청난 성능을 보여줬습니다. 링크
Introduction
Efficiently Learning an Encoder that Classifies Token Re-placements Accurately(ELECTRA)
2020 기준 SOTA representation learning(표현 학습) 기법은 일종의 denoising autoencoder(latent vector로 바꾸고 원래 데이터로 복원) 학습이라고 볼 수 있습니다. 주로 입력 시퀀스의 토큰 중 약 15% 정도를 마스킹하고 이를 복원하는 MLM을 통해서 학습을 진행합니다. 이 방법은 하나의 example에 대해서 전체 토큰 중 고작 15%만 학습합니다. 따라서 학습이 오래걸리고 비용이 많이 듭니다. 그리고 학습 때는 [MASK] 토큰을 모델이 참고하여 예측하지만 실제(inference)로는 [MASK] 토큰이 존재하지 않아 성능 저하로 이어질 수 있습니다. 문맥 정보보다 [MASK] 토큰 자체에 더 의존하여 예측을 수행할 수 있기 때문입니다.
위 그래프를 보면 x축은 pre-train에 사용된 계산량 (FLOPs)을 나타내고, y축은 GLUE 벤치마크 점수를 나타냅니다. FLOPs가 크면 당연히 성능도 좋아지겠지만 핵심은 “적은 연산량으로도 좋은 성능을 내는 것”입니다. 전반적으로 RTD pre-train을 사용한 ELECTRA 모델이 MLM pre-train을 사용한 모델보다 동일한 FLOPs에서 더 높은 GLUE 점수를 얻었습니다. 특히 ELECTRA-Small은 BERT-Small이나 GPT와 같이 더 큰 모델이나 더 많은 계산량을 사용한 모델보다도 좋은 성능을 보였습니다. 이 결과는 RTD 사전 학습이 MLM 사전 학습보다 샘플 효율성이 높고, 특히 작은 모델에서 계산량 대비 성능이 뛰어나다는 것을 보여줍니다. 또한, 충분한 계산량이 주어졌을 때 RTD는 대규모 모델에서도 경쟁력 있는 성능을 보여줍니다.
위 결과를 종합해 보면 RTD보다 MLM이 성능과 효율 측면에서 모두 좋다고 할 수 있습니다.
Method
generator G와 discriminator D 두 개의 네트워크가 필요합니다. 두 네트워크는 공통적으로 Transformer 인코더 구조를 가지고 있고, 토큰 시퀀스 x=[x1,x2,…,xn]을 입력으로 받아서 문맥 정보를 반영한 벡터 시퀀스 h(x)=[h1,h2,…,hn] 로 매핑시킵니다. pre-train 후에는 G를 폐기하고, D만 작업에 맞게 fine-tuning합니다.
Generator
G는 BERT의 MLM과 똑같습니다. 이 모델을 학습하려면 우선 입력 x=[x1,x2,…,xn]에 대해서 마스킹할 위치의 집합 m=[m1,m2,…,mk]을 결정합니다. 모든 마스킹 위치는 1과 n 사이의 정수이며, 아래와 같이 수학적으로 표현할 수 있습니다. k는 보통 15%인 0.15n을 사용합니다.
이후 결정한 위치에 있는 입력 토큰을 [MASK]로 치환합니다.
생성자 G가 마스크된 입력 문장 xmasked가 주어졌을 때, [MASK]토큰 중 특정 위치 t에서 토큰 xt를 생성합니다. 여기서 e(⋅)는 임베딩을 의미하고 이는 임베딩 레이어의 가중치를 공유(weight sharing)를 나타냅니다.
e(xt)T hG(xmasked)t 는 토큰 xt의 임베딩 벡터와 생성자 G의 인코더 출력 벡터 사이의 내적을 계산합니다. 이 값은 토큰 xt가 해당 위치 t에 얼마나 적합한지 나타내는 점수로 해석할 수 있습니다. 이후 새롭게 토큰을 생성하고 다음과 같이 나타낼 수 있습니다.
최종적으로 아래와 같은 MLM loss로 학습합니다.
이 과정을 코드로 보면 다음과 같습니다.
def _get_masked_lm_output(self, inputs: pretrain_data.Inputs, model):
# ...
relevant_reprs = pretrain_helpers.gather_positions(
model.get_sequence_output(), inputs.masked_lm_positions)
logits = get_token_logits(
relevant_reprs, model.get_embedding_table(), self._bert_config)
# ...
return get_softmax_output(
logits, inputs.masked_lm_ids, inputs.masked_lm_weights,
self._bert_config.vocab_size)- model.get_sequence_output() 함수를 사용하여 G의 트랜스포머 인코더 출력을 가져옵니다.
- pretrain_helpers.gather_positions 함수를 사용하여 마스킹된 토큰 위치에 해당하는 벡터 표현을 추출합니다.
- get_token_logits 함수를 사용하여 추출된 벡터 표현으로부터 각 토큰에 대한 로짓(확률)을 계산합니다.
- get_softmax_output 함수를 사용하여 로짓을 소프트맥스 함수에 통과시켜 확률 분포를 얻습니다.
Discriminator
D는 ELETRA의 핵심 기능입니다. 입력 토큰 시퀀스에 대해서 각 토큰이 original인지 replaced인지 이진 분류로 학습합니다. 이렇게 하면 마스킹 된 토큰이 아닌 전체 토큰을 학습할 수 있습니다. xcorrupt에 대해서 G가 생성한 새로운 토큰이 원래 입력과 동일한 것인지 생성된 것인지 예측합니다. 시그모이드 함수를 활용하여 단순히 이진 분류를 진행합니다.
다음과 같은 Disc loss로 학습합니다.(cross entropy)
xt가 생성된 토큰이라면 D가 1이라고 추론하고, 생성되지 않은 토큰이라면 0이라고 추론하도록 학습시킵니다.
최종적으로는 다음과 같은 수식으로 두 loss의 합을 최소화 시키도록 학습합니다.
이 과정을 코드로 보면 다음과 같습니다.
def _get_discriminator_output(self, inputs, discriminator, labels):
# ...
weights = tf.cast(inputs.input_mask, tf.float32)
labelsf = tf.cast(labels, tf.float32)
losses = tf.nn.sigmoid_cross_entropy_with_logits(
logits=logits, labels=labelsf) * weights
per_example_loss = (tf.reduce_sum(losses, axis=-1) /
(1e-6 + tf.reduce_sum(weights, axis=-1)))
loss = tf.reduce_sum(losses) / (1e-6 + tf.reduce_sum(weights))
# ...- tf.cast(labels, tf.float32)를 통해 이를 float32 타입으로 변환합니다. 원래 토큰은 1, 대체 토큰은 0으로 표현됩니다.
- weights는 패딩 토큰의 여부를 나타냅니다.(실제 토큰인지 아닌지)
- labels는 각 토큰이 원래 토큰인지 대체 토큰인지 나타내는 텐서입니다.
- tf.nn.sigmoid_cross_entropy_with_logits 함수를 사용하여 로짓(logits)과 레이블(labelsf)을 기반으로 loss를 계산합니다.
- 따라서 loss는 위에서 계산한 loss를 실제 토큰의 갯수로 나눈 값이 됩니다.(1e-6은 0으로 나누는 것을 방지)
GAN과의 차이점
GAN은 생성자가 생성한 토큰은 무조건 “가짜(fake)”로 간주됩니다. 반면 ELECTRA는 생성자가 우연히 원래 토큰과 동일한 토큰을 생성했을 경우, 이를 “진짜(real)”로 간주합니다. 이는 downstream 작업에서 성능 향상에 도움이 되는 것으로 나타났습니다. 사실 이 방법이 “경찰과 도둑”에 더 잘 어울립니다.
또한 GAN은 생성자는 판별자를 속이도록 “적대적(adversarial)”으로 학습됩니다. 즉, 판별자가 생성된 토큰을 진짜로 착각하도록 학습합니다. 반면 ELECTRA는 생성자는 판별자를 속이려고 하지 않고, “최대 우도(maximum likelihood)” 방식으로 학습됩니다. 즉, 마스킹된 토큰의 원래 정체성을 정확하게 예측하도록 학습합니다. 이는 샘플링 과정의 미분 불가능성 문제를 피할 수 있습니다. 논문에서는 강화 학습을 사용하여 생성자를 적대적으로 학습하는 실험을 진행했지만, 최대 우도 학습보다 성능이 좋지 않았다고 말합니다.
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)을 최종 성능으로 사용하였습니다.
MODEL EXTENSIONS
ELECTRA의 성능을 더 향상시키는 몇 가지 기법을 소개하겠습니다.
Weight sharing
Generator와 Discriminator는 모두 Transformer 인코더 구조이기 때문에 두 네트워크의 가중치를 공유하여 학습하는 weight sharing를 통해 성능 향상을 기대할 수 있습니다.
이 모델들을 500k 스텝 동안 학습했습니다.
- 가중치를 연결하지 않았을 때 83.6점
- 토큰 임베딩만 연결했을 때 84.3점
- 모든 가중치를 연결했을 때 84.4점
모든 가중치를 연결했을 때 가장 좋은 성능을 보여주지만(D가 학습하지 못한 부분까지 G가 학습해주기 때문) G와 D의 크기가 같아야 한다는 제약조건이 있습니다. 이후 추가 실험을 위해 임베딩만 연결해서 사용합니다.
Smaller Generators
G와 D의 크기가 같다면 기존 BERT에 비해 단순 계산시 약 2배의 연산량이 필요합니다. 따라서 generator의 크기를 줄이는 연구를 진행합니다. (아래 왼쪽 그림)
다른 하이퍼파라미터는 일정하게 유지하면서 계층 크기를 줄여 모델을 작게 만듭니다. 작은 generator모델이 discriminator크기의 1/4~1/2 크기를사용할 때 가장 잘 작동한다는 것을 발견했습니다. 저자는는 너무 강력한 생성자는 판별자에게 너무 어려운 작업을 제기하여 효과적인 학습을 방해할 수 있다고 추측합니다.
Training Algorithms
또한 효율적인 학습 방법에 대한 탐구를 진행합니다. (위 오른쪽 그림)
- Joint Training: generator와 discriminator를 동시에 학습하는 방법
- Two-stage training: generator를 먼저 학습한 후 discriminator를 generator의 가중치로 초기화 한 후 discriminator를 학습시키는 방법
- Adversarial training: GAN처럼 adversarial training을 모사해서 학습시키는 방법
결론적으로 Joint Training이 가장 좋은 성능을 보여줬습니다. 생성자는 처음에는 약하지만 학습 과정에서 점점 강해지므로, 판별자는 점점 어려워지는 작업을 해결하며 학습할 수 있습니다. Adversarial training가 잘 학습되지 않는 이유는 생성자의 정확도가 너무 낮았기 때문입니다.
Small Models
저자들이 제시한 연구의 핵심 중 하나는 pre-training의 효율성 향상입니다.(ELECTRA의 EL은 Efficiently Learning의 줄임말이다.) 이를 검증하기 위하여 저자들은 하나의 GPU로도 빠르게 학습할 수 있는 수준으로 작은 모델을 만들어보는 실험을 진행했습니다.
Small모델은 Base모델 대비 시퀀스 길이, 배치 크기, 은닉층 크기, 토큰 임베딩 크기를 줄여서 모델 크기를 축소했습니다. 공정한 비교를 위해 학습 FLOPs가 동일하도록 학습시켰습니다.
ELECTRA-Small은 BERT-Small보다 5점 높은 GLUE 점수를 얻었으며, 더 큰 GPT 모델(117M 인것으로 보아 GPT-2로 추측) 보다도 좋은 성능을 보였습니다. ELECTRA-Small은 학습 시간이 짧아도(6시간) 합리적인 성능을 보입니다. 학습 시간이 길어질수록 성능이 향상되며, 4일 학습 시 최고 성능에 도달합니다. 성능과 학습 속도에 대한 향상을 강조합니다.
Large Models
저자들은 Small 모델뿐 아니라 Large 모델로도 실험을 진행하였습니다. ELECTRA-Large의 세팅도 앞선 실험들과 마찬가지로 같은 크기의 BERT-Large에 맞춰서 실험했습니다.
ELECTRA는 RoBERTa 및 XLNet보다 훨씬 적은 fre-train FLOPs(1/4 미만)으로도 비슷한 성능을 달성했습니다. 이는 ELECTRA의 같은 FLOPs 당 효율성을 입증합니다. 또한 ELECTRA-1.75M은 가장 좋은 성능을 보여줬습니다.
다음 표에서는 앙상블과 같은 추가 작업을 했을 때의 성능을 보여줍니다.
ELECTRA는 RoBERTa와 유사한 계산량을 사용하면서도 GLUE 벤치마크 점수(Score)에서 더 높은 점수를 달성했습니다.
또한 SQuAD에서도 같은 연산량 대비 가장 높은 점수를 보여줍니다.
Efficiency Analysis
저자들은 ELECTRA의 성능 향상 요인을 분석하기 위해 다음과 같은 사전 학습 목표들을 비교합니다.
ELECTRA 15%
- ELECTRA 모델과 동일하지만, discriminator 손실은 마스킹된 15%의 토큰에 대해서만 계산됩니다.
- 즉, LDisc 손실 함수의 합산 범위가 모든 토큰(1~n)이 아닌 마스킹된 토큰(i∈m)에 대해서만 계산됩니다.
→ discriminator가 100%를 학습해서 성능이 좋아졌는지 확인
Replace MLM
- MLM과 동일하지만, 마스킹된 토큰을 [MASK] 토큰 대신 생성자 모델에서 샘플링된 토큰으로 대체합니다.
→ [MASK] 토큰 사용으로 인한 pre-train/fine-tune 불일치 문제 극복으로 인해 성능이 좋아졌는지 확인
All-Tokens MLM
- Replace MLM과 유사하게 마스킹된 토큰을 생성자 샘플로 대체합니다. 하지만, 마스킹된 토큰뿐만 아니라 모든 입력 토큰에 대한 정체성을 예측합니다. (이 모델은 BERT와 ELECTRA를 결합한 형태)
- 명시적인 복사 메커니즘을 사용하여 각 토큰에 대한 복사 확률을 출력합니다. 즉, 모델은 입력 토큰을 그대로 복사할지, MLM 소프트맥스 출력을 사용하여 예측할지 결정합니다.
→ ELECTRA의 판별적 접근 방식은 생성적 접근 방식보다 더 효과적일 수 있다.
결론적으로는 ELECTRA는 ELECTRA 15%와 Replace MLM보다 많은 성능 개선이 있는 것으로 보아 이 두 문제를 해결하여 성능이 상당히 좋아진 것을 확인할 수 있습니다.
ELECTRA는 작은 모델에서 더 큰 성능 향상을 보이는데, 이는 ELECTRA가 BERT보다 매개변수 효율성이 높기 때문일 수 있습니다. BERT에 비해 ELECTRA는 모델이 작아도 매우 빠르게 수렴함을 알 수 있습니다.
Conclusion
이 논문은 language representation learning을 위한 새로운 self-supervision 태스크인 Replaced Token Detection을 제안합니다. RTD의 주요 아이디어는 작은 generator가 만들어 낸 질 좋은 negative sample과 입력 토큰을 구별하도록 텍스트 인코더를 학습시키는 것입니다. Masked language modeling에 비해 RTD는 훨씬 효율적(compute-efficient)이고 downstream tasks에 대한 결과 역시 더 좋다는 것을 많은 실험을 통해서 확인할 수 있었습니다. 또한 RTD는 MLM보다 모든 입력 토큰을 활용하여 학습하므로 샘플 효율성이 높습니다.
저자는 향후 NLP pre-train 연구에서는 성능뿐만 아니라 효율성도 고려해야 하며, 계산량 및 매개변수 수를 평가 지표와 함께 보고해야 한다고 주장합니다.
