이 포스트는 개인적으로 공부한 내용을 정리하고 필요한 분들에게 지식을 공유하기 위해 작성되었습니다. 지적하실 내용이 있다면, 언제든 댓글 또는 메일로 알려주시기를 바랍니다.

01. 개요

ICLR 2020에서 Google Brain 팀은 새로운 사전 학습 방법론으로 ELECTRA를 제안했다. ELECTRA는 “Efficiently Learning an Encoder that Classifies Token Replacements Accurately”의 약자로, 사전 학습에서의 효율성을 개선하여 BERT보다 빠르게 사전 학습 할 수 있으면서 downstream task에서도 높은 성능을 보였다.

02. 기존(BERT) 사전 학습 방식의 문제점

언어 모델이 토큰들에 대한 representation을 학습하는 대표적인 방법은 denoising autoencoder 방식이다. 이 방식은 BERT의 Masked Language Modeling(MLM) 방식이 대표적인데, 일반적으로 레이블링이 되어 있지 않은 입력 시퀀스의 일부(대략 15%)를 [MASK] 토큰으로 치환하고, 치환된 [MASK] 토큰의 원래 단어가 무엇이었는지를 맞추는 방식으로 학습한다. Transformer의 self-attention은 양방향의 문맥을 파악할 수 있으므로 이러한 MLM 방식은 매우 효과적인 결과를 낼 수 있다. 하지만 BERT 같은 커다란 언어 모델의 학습이 고작 15% 정도의 [MASK]로 치환된 토큰들만을 이용하여 학습해야 하기 때문에 계산 비용이 매우 크다는 단점이 있다.

ELECTRA의 저자들은 이처럼 적은 수의 [MASK]로 치환된 토큰들만 가지고서 큰 언어 모델을 학습하는 방식이 비효율적이라고 지적하며, MLM 방식의 사전 학습이 아닌, Replaced Token Detection(RTD)이라는 새로운 사전 학습 방식을 제안한다.

03. ELECTRA의 구조

fig01

ELECTRA의 사전 학습 방식은 마치 GAN과 같이 Generator와 Discriminator 네트워크의 학습으로 이루어진다. ELECTRA의 이름에서 알 수 있듯이, Generator와 Discriminator 모두 기본적으로 BERT처럼 Transformer의 Encoder 구조이다. 이때, Discriminator가 바로 ELECTRA 모델에 해당하며, Generator는 사전 학습 시 Discriminator의 입력값을 만들어주는 네트워크에 불과하다. 사전 학습을 마치면 Generator는 떼어내고 Discriminator만을 이용하여 downstream task에 대해 fine tuning 한다.

ELECTRA의 사전 학습 방식이 Generator와 Discriminator 구조로 이루어져 있고 jointly 학습되기 때문에 GAN과 같이 학습될 것이라고 생각할 수 있으나, 엄연히 다르다. GAN에서는 Generator가 Discriminator를 속이기 위해서 adversarial하게 학습을 하지만, ELECTRA의 사전 학습에서는 Generator가 maximum likelihood로 학습한다.

Generator

Generator는 BERT의 사전 학습에서 Next Sentence Prediction(NSP)를 제외한 학습 방법, 즉, MLM만을 적용한 사전 학습 방법과 동일하다.

우선, 입력 토큰 시퀀스 $\mathbf{x}=[x_1, x_2, … , x_n]$이 주어졌을 때, 전체 토큰의 약 15%($k=\lceil 0.15n \rceil$)개의 토큰을 [MASK]로 치환하여 마스킹할 위치의 집합 $\mathbf{m}=[m_1, m_2, …, m_k]$을 만든다. 이때 마스킹 위치는 수식으로 $m_i \sim \text{unif}{1, n} \text{ for } i=1 \text{ to }k$ 로 표현할 수 있다.

마스킹할 위치의 집합 $\mathbf{m}$이 결정되었다면, 해당 위치의 토큰은 [MASK]로 치환하는데, 이를 $\mathbf{x}^{masked}=\text{REPLACE}(\mathbf{x}, \mathbf{m}, \text{[MASK]})$로 표현할 수 있다.

Generator는 이후 BERT의 MLM과 동일한 방식으로 [MASK]로 치환된 토큰의 원래 토큰이 무엇이었는지 예측한다. $t$번째 토큰을 예측한다고 하면, 아래와 같이 softmax로 표현할 수 있다.

$p_G(x_t | \mathbf{x}) = \cfrac{\exp{(e(x_t)^T h_G(\mathbf{x})_t)}}{\sum_x^\prime\exp{(e(x^\prime)^T h_G(\mathbf{x})_t)}}$

이때, $h_G(\mathbf{x})$는 Generator 네트워크를 거쳐 나온 contextualized vector representations를 의미한다. 따라서 $h_G(\mathbf{x})_t$는 $t$번째 [MASK] 토큰의 representation이 되겠다. 또한 $e(\cdot)$는 토큰의 임베딩을 의미한다.

최종적으로 Generator는 BERT의 MLM과 같은 손실 함수를 이용해서 학습한다.

$\mathcal{L}_\text{MLM} (\mathbf{x}, \theta_G) = \mathbb{E}\bigg( \sum_{ i \in \mathbf{m} } -\log{p_G(x_i | \mathbf{x}^{\text{masked}})} \bigg)$

Discriminator

Discriminator는 Generator의 입력 토큰 시퀀스에서 [MASK] 부분을 Generator가 예측한 값으로 대입하여 입력값으로 사용한다. 그 예시는 다음과 같다.

원래 입력 토큰 시퀀스 $\mathbf{x}$
- [ the, chef, cooked, the, meal ]
Generator의 입력 토큰 시퀀스 $\mathbf{x}^{\text{masked}}$
- [ [MASK], chef, [MASK], the, meal ]
  - 원래 입력 토큰 시퀀스의 15%를 마스킹 함
Discriminator의 입력 토큰 시퀀스 $\mathbf{x}^{\text{corrupt}}$
- [ the, chef, ate, the, meal ]
  - the 토큰은 Generator가 정확히 예측하였음
  - ate 토큰은 cooked가 정답이지만, Generator가 확률적으로 plausible 예측하였음

즉, Discriminator의 입력 토큰 시퀀스는 Generator의 예측값을 기반으로 원래 입력 토큰 시퀀스가 재구성되기 때문에, 일부 토큰이 확률적으로 그럴싸할 수는 있지만 원래 토큰과는 다른 토큰으로 구성되어 있을 수 있다(plausible but synthetically generated replacements).

이렇게 Generator의 예측으로 [MASK] 토큰을 대체하여 재구성한 Discriminator의 입력 토큰 시퀀스를 $\mathbf{x}^{\text{corrupt}}=\text{REPLACE}(\mathbf{x}, \mathbf{m}, \mathbf{\hat{x}})$라고 표현할 수 있다. 이때 Generator의 softmax를 통해 예측한 토큰을 $\hat{x}_i \sim p_G(x_i | \mathbf{x}^{\text{masked}}) \text{ for } i \in \mathbf{m}$ 로 표현할 수 있다.

최종적으로 Discriminator는 입력 시퀀스의 각 토큰이 원래 토큰인지(original), 아니면 Generator에 의해 대체된 토큰(replaced)인지를 구분하는 문제로 학습한다. 이는 sigmoid output layer로 구현되며 $t$번째 토큰에 대한 예측값은 $D(\mathbf{x}^{\text{corrupt}}, t) = \text{sigmoid}(w^{T}h_{D}(\mathbf{x}^{\text{corrupt}})_t)$로 표현할 수 있다. 이때, $h_D(\mathbf{x}^{\text{corrupt}})$는 Discriminator 네트워크를 거쳐 나온 contextualized vector representations를 의미한다.

최종적으로 Discriminator는 아래와 같은 크로스 엔트로피 손실 함수를 이용해서 학습한다.

$\mathcal{L}_{\text{Disc}}(\mathbf{x}, \theta_D)=\mathbb{E}\bigg( \sum_{t=1}^{n} -1(\mathbf{x}_t^{\text{corrupt}}=\mathbf{x}_t) \log{D(\mathbf{x}^{\text{corrupt}}, t)} -1(\mathbf{x}_t^{\text{corrupt}} \neq \mathbf{x}_t) \log{(1-D(\mathbf{x}^{\text{corrupt}},t))} \bigg)$

Combined Loss

ELECTRA는 궁극적으로 다음과 같이 Generator의 loss와 Discriminator의 loss를 결합하여 손실 함수를 구성하고, 이 손실 함수가 최소화되도록 사전 학습을 진행한다.

$\min_{\theta_G, \theta_D} \sum_{\mathbf{x} \in \mathcal{X}} \mathcal{L}_\text{MLM}(\mathbf{x}, \theta_G) + \lambda \mathcal{L}_\text{Disc}(\mathbf{x}, \theta_D)$

이때 $\mathcal{X}$는 대규모 raw 텍스트 코퍼스를 의미하며, $\lambda$(논문에서는 $\lambda=50$을 사용함)는 Generator와 Discriminator의 스케일을 맞추는 파라미터이다. 한편, 사전 학습 과정에서 sampling 과정을 거치기 때문에 Discriminator의 loss는 Generator로 back-propagate 되지 않는다.

04. Experiments

만약, Generator와 Discriminator 네트워크 크기가 동일하다면, 두 네트워크의 가중치를 공유해서 사용하는 것이 더욱 효율적일 수 있을 것이라 생각할 수 있다. 이러한 호기심에 대해 ELECTRA 저자들은 두 네트워크가 동일한 크기라고 가정한 후 가중치를 공유하는 실험을 진행해보았다.

(1)가중치를 공유하지 않은 경우와 (2)토큰 임베딩만을 공유한 경우, 그리고 (3)모든 가중치를 공유한 경우 각각에 대하여 GLUE score를 측정한 결과, 각각 83.6, 84.3, 84.4점을 달성했다. 즉, 가중치를 공유하면 성능이 향상된다는 것이다. 그러나, (3)모든 가중치를 공유한 경우에는 (2)토큰 임베딩만을 공유한 경우에 비해 아주 약간의 성능 향상은 있었으나, Generator가 Discriminator와 반드시 동일한 크기만큼 커야 한다는 제약이 있다는 점을 고려하여 단점이 더 크다고 할 수 있다. 따라서 저자들은 (2)토큰 임베딩만을 공유하는 경우를 최적으로 여기고 사용한다.

Smaller Generators

fig02

Generator와 Discriminator의 크기가 동일하면 그만큼 학습할 때의 계산량도 상당할 것이다. ELECTRA의 저자들은 Generator의 크기를 작게하여 계산량을 줄여보았다. 실험 결과, 오히려 Generator의 크기가 Discriminator의 크기보다 작을 때(대략 Discriminator의 1/4 ~ 1/2 크기) 더 높은 GLUE score를 달성할 수 있었다. 저자들은 이 현상에 대해서 너무나 Generator의 크기가 커지면, Discriminator는 실제 데이터의 분포를 모델링 하는 것 보다 Generator를 모델링 하는 데 파라미터를 많이 사용하게되는 단점이 있을 것이라고 해석했다.

Training Algorithms

fig03

ELECTRA의 저자들은 다음과 같이 Generator와 Discriminator의 크기를 동일하게 설정한 후 Two-Stage 방식으로 학습을 해보았다.

우선, Generator만 n스텝 학습한다.
이후, Generator의 가중치로 Discriminator의 가중치를 초기화하고, Generator의 가중치는 고정한 후, Discriminator만 n스텝 학습한다.

또한 GAN과 같이 Generator를 adversarial한 방식으로도 구성하여 실험해보았다. 결과적으로는 다양한 시도을 했지만 성능 향상엔 큰 도움을 얻을 수 없었으며, 기본적인 방식으로 Generator와 Discriminator를 결합하여(jointly) 학습할 때 성능이 가장 좋았다.

Small Models

fig04

저자들은 GPU 한 대에서도 빠르게 ELECTRA를 사전 학습 할 수 있도록 ELECTRA-small 모델을 개발해서 실험해보았다. BERT-Base의 하이퍼파라미터를 사용하되, 시퀀스 길이를 512 토큰에서 128 토큰으로, 배치 사이즈를 256에서 128로, 은닉층의 차원을 768에서 256으로, 임베딩 차원을 768에서 128로 줄여서 실험해보았다. 그럼에도 불구하고 Table 1에서 볼 수 있듯, 놀라울 수준의 GLUE score를 달성하였다.

Large Models

fig05

저자들은 SOTA 모델들과 비교해보기 위해서 BERT-LARGE에 준하는 ELECTRA-LARGE 모델을 개발하여 실험해 보았다. ELECTRA-400K는 RoBERTa 사전 학습의 대략 1/4 수준의 연산만 수행한 모델인데도 불구하고, Table 2에서 알 수 있듯, GLUE dev set에서 RoBERTa와 XLNet에 준하는 성적을 거두었다. RoBERTa의 연산 수준만큼 사전 학습 한 ELECTRA-1.75M 모델은 대부분의 태스크에서 최고점을 기록했다. Table 3에서 알 수 있듯, ELECTRA는 GLUE test set에서도 대부분의 태스크에서 최고점을 달성했다.

fig06

ELECTRA는 SQuAD 데이터셋에서도 매우 우수한 결과를 보였다.

05. 결론

ELECTRA는 언어의 representation을 학습하기 위해서 replaced token detection이라는 새로운 self-supervised task기법을 제시하였다. 이는 기존의 사전 학습 방식들에 비해서 연산을 효율적으로 수행할 수 있으며, 실제로 downstream task에서도 훌륭한 성능을 보였다.

06. 참고 문헌

[1] 원 논문 : ELECTRA: pre-training text encoders as discriminators rather than generators
[2] Scatter Lab의 블로그 : tech.scatterlab.co.kr/electra-review

(NLP) ELECTRA 리뷰 및 설명