ELECTRA

Intro

ELECTRA를 살펴보기에 앞서, 기존 BERT 모델이 어떤식으로 작동하는지를 recap해보도록 하자.

 

BERT는 기본적으로 트랜스포머의 인코더 부분만 사용한 모델이다. 또한 추가적으로, MLM과 NSP라는 2가지의 사전학습 과정을 거쳐서 해당 모델이 언어를 이해할 수 있게끔 훈련시킨다.

 

이 중 MLM task의 경우, 다음과 같은 과정을 통해 진행된다.

1. 주어진 corpus를 Tokenizing한다. (Wordpiece tokenizer)
2. 전체 단어의 15% 정도를 무작위로 고르고, 이 중 80%는 MASK 토큰으로 바꾸고, 10%는 임의의 토큰으로 바꾸고, 10%는 아무런 처리를 하지 않는다.(80-10-10 규칙)
3. 3가지의 임베딩 처리를 하고, concat한다.(Token embedding, Segment embedding, Positional embeddgin)
4. MASK처리한 단어를 모델이 최대한 가깝게 예측하게끔 훈련한다.

위와 같이 기존 BERT를 비롯한 많은 language model는 MLM task를 통해 pretraining을 거친다. 

 

하지만, MLM은 크게 2가지 문제점이 존재한다.

첫째. 전체 중 단, 15%만을 이용한다는 점에서 상당한 컴퓨팅 연산을 유발할 수 있게 된다.

둘째. 학습 과정에서는 [MASK]를 모델이 참고하여 예측하는 반면, 실제 과정에서는 해당 토큰이 존재하지 않는 불일치가 존재한다.

 

따라서 본 논문 저자는 "replaced token detection" 방법을 제안한다. 위 방법은 단순히 MASK 토큰 처리를 하는 것이 아니라, Generator을 통해 실제 입력의 토큰을 가짜 토큰으로 바꾸고 Discriminator가 이를 original인지 replaced인지 맞추게 된다.

논문에 따르면 위 과정을 통해 기존의 1/4분의 연산량만으로도 비슷하거나 더 나은 성능을 낼 수 있게 되었다고 한다.

또한, 저자는 단순히 정확성 뿐만 아니라 compute efficiency도 중요하게 고려해야하는 부분이라고 주장한다.

Method

앞서 언급한 것처럼, 위 모델은 크게 Generator와 Discriminator로 구성되어있다.

위 구조를 보면, GAN에서 배웠던 개념이 겹쳐서 보일 수 있다.

GAN

하지만, Generator가 Discriminator에 영향을 받지 않고 원래의 단어를 그대로 예측하는 방향으로 훈련된다는 점에서 adverserial하지 않고, 또한 random variable에 의해 생성되는 과정이 아니라는 점에서 GAN과는 차이가 존재한다. 즉 generator는 likelihood를 maximize하는 방향으로 훈련되는 것이지, Discriminator가 틀리게끔 학습되지 않는다는 것이다.

 

각 과정에 대해서 자세하게 알아보도록 하자.

살펴보기 전에, 몇가지 용어에 대해 정의하고 넘어가도록 하자.

input token : $x = [x_1, \ldots, x_n]$
contextualized vector representations : $h(x) = [h_1, \ldots, h_n]$ (Encoder output)
position : t

Generator

주어진 position t에 대해서, generator가 $x_t$를 generating할 확률은 다음과 같다. (단, e는 token embedding)

$$P_G(x_t|\textbf{x}) = exp(e(x_t)^Th_G(\textbf{x})_t) \text{ / } \sum_{x^\prime} exp(e(x^\prime)^Th_G(\textbf{x})_t$$

 

위 수식 결과를 잘 생각해보면, representation과 $x_i$의 임베딩 사이의 cosine similarity가 얼마냐 유사한지를 기반으로 확률을 추정하고 있다는 것을 쉽게 알 수 있다. 해당 결과를 바탕으로 generator가 MASK 토큰을 위 결과 값이 가장 큰 token값으로 바꾸게 된다.

Discriminator

주어진 position t에 대해서, Discriminator는 주어진 토큰 $x_t$가 generator에 의해 교체된 것인지 아닌지를 판단하게 된다.

즉 결과적으로 Discriminator는 이진 분류 task를 수행하고 있는 것이다.

 

다음 식을 통해 Discriminator는 original인지 replaced인지 판단하게 된다.

$$D(\textbf{x}, t) = sigmoid(w^Th_D(\textbf{x})_t)$$

Summary

앞서 설명한 Generator와 Discriminator를 종합해서 어떻게 진행되는지를 살펴보고, 손실함수를 어떻게 정의하는지를 살펴보도록 하자.

이에 앞서, 용어를 정리하고 넘어가도록 하자.

 

$$ \text{input token} :  \textbf{x} = [x_1, \ldots, x_n]$$

$$m_i \text{ ~ } unif\text{{1, n}} \text{ for } i = 1 \text{ to } k$$

$$ \text{mask out} : \textbf{m} = [m_1, \ldots, m_k]$$

$$\textbf{x}^{masked} = REPLACE(\textbf{x}, \textbf{m}, [MASK])$$

$$\hat{x}_i \text{ ~ } p_G(x_i|\textbf{x}^{masked}) \text{ for } i \in \textbf{ m}$$

$$\textbf{x}^{corrput} = REPLACE(\textbf{x}, \textbf{m}, \bf{\hat{x}})$$

 

위 수식에서 볼 수 있는 것처럼, $m_i$는 uniform distribution에서 뽑은 것이므로 n개의 token 중 k개를 랜덤하게 뽑은 것이다.

따라서 m은 해당 결과들의 집합이다. 

 

loss function은 다음과 같이 정의하였다.

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

$$\mathcal{L}_{Disc}( \textbf{x}, \theta_D) = $$

$$\mathbb{E}(\sum_{t = 1}^n -\mathbb{1}({x_t}^{corrupt} = x_t)logD(\textbf{x}^{corrupt}, t)-\mathbb{1}({x_t}^{corrupt} \ne x_t) log(1 - D(\textbf{x}^{corrupt}, t)))$$

 

위 식의 의미를 잘 이해해보면 Generator의 loss는 log likelihood이고, Discriminator의 loss는 Binary cross entropy loss와 같다는 것을 쉽게 이해할 수 있다.

 

위 2개의 loss를 결합하면 다음과 같다.(단, $\chi$는 large corpus of raw text)

$$\text{min}_{\theta_G, \theta_D} \sum_{\textbf{x} \in \chi} \mathcal{L}_{MLM}(\textbf{x}, \theta_G) + \lambda \mathcal{L}_{Disc}(\textbf{x}, \theta_D)$$

Experiment

Weight sharing에 대한 부분도 논문에서 실험을 진행하였지만, 효과적이지 않았다.

왜냐하면 Weight sharing을 하려면 Generator와 Discriminator의 사이즈를 같게 만들어야 하는데, 이는 필요한 메모리도 그렇고 전체적으로 효과적이지 않기 때문이다. 또한 저자는 Generator의 사이즈를 작게 하고, Generator와 Discriminator의 token embedding만 공유하는 것이 더 효과적임을 발견하였다.

 

실제로 실험을 500k step을 돌려본 결과 아무런 weight tying이 없을 때의 GLUE score는 83.6, token embedding만 했을 때는 84.4, 모든 weight를 sharing했을 경우는 84.4가 나왔다. 

 

저자는 token embedding만 공유하는 것이 효과적인 이유를 Discriminator는 단지 주어진 input에 이미 존재하거나 Generator에 의해 샘플링된 토큰만 업데이트하는 반면, Generator의 softmax는 모든 token embedding을 밀도있게 업데이트하기 때문에 token embedding을 공유하는 것이 효과적이라고 보았다.

또한 위 사진에서 볼 수 있는 것처럼, Generator의 사이즈는 Discriminator의 1/4 ~ 1/2정도일 때 가장 score가 높은 것을 실험적으로 확인하였다.