Contents
ICLR 2020에서 구글 리서치 팀이 새로운 pre-training 기법을 적용한 language model인 ELECTRA(Efficiently Learning an Encoder that Classifies Token Replacements Accurately)를 발표하였다. BERT 이후의 많은 language model은 MLM task를 통해 pre-training을 하게되는데, 이런 모델들은 학습에 많은 시간과 계산량을 필요하므로 컴퓨팅 리소스가 “많이” 필요하다.
ELECTRA는 모델의 정확도와 더불어 효율성에 포인트를 맞춘 방식의 모델이다. 본 논문에서 학습 효율을 향상하기 위해 새로운 pre-training 방식인 Replaced Token Detection(RTD)라는 방식을 제안하고, 이를 통해 보다 빠르고 효율적으로 학습한다.
결과적으로 ELECTRA는 동일한 크기, 데이터, 컴퓨팅 리소스 조건에서 BERT의 성능을 능가하였다. Small 모델과 Large 모델을 실험하였고, 각각 GPT나 RoBERTa, XLNet 대비 동일 조건 우수한 성능에 도달하였다.
Paper : ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators
Motivation
논문 등장 기준 시점까지의 최신 언어 모델의 학습 방식은 denoising representation 방식으로 학습하는 것을 볼 수 있다. 전체 토큰 중 일부분(15%)를 마스킹하고, 이를 훈련하는 과정에서 원래 입력을 복구하게 된다. 물론 이 방식이 SOTA를 갱신하고 있으므로 좋은 방식임은 틀림없다. 또한, 기존의 순차적으로 토큰을 입력받는 방식에 비해 MLM은 양방향으로 데이터를 본다는 관점(사실 양방향이라기 보다는 문장을 한 번에 전체적으로 본다는 것이 맞겠다.)에서 Masked Language Model은 좋은 모델이다.
그러나 기존 MLM 모델은 다음 세 가지의 문제점들을 가지고 있는데,
- loss가 발생하는 지점은 마스킹된 15% 부분이 전부이다.
마스킹 되는 토큰의 비율이 15%이고, 해당 부분에서만 pred와 origin을 비교하기 때문이다. - 상당한 컴퓨팅 리소스가 필요하다
당연하게도, 15%만 학습에 사용되니 더더더더 많은 양의 데이터가 필요하고 모델이 커지게 되는 것이다. - 학습 때는 [MASK] 토큰을 모델이 참고하여 예측하지만 실제(inference)로는 [MASK] 토큰이 존재하지 않는다.
본 논문은 학습 효율을 향상시키기 위해 Replaced Token Detection (RTD)이라는 새로운 pre-training 태스크를 제안하며, 이를 통해 ELECTRA는 보다 빠르고 효과적으로 학습한다고 주장한다. 또한 실험을 통해 BERT의 성능을 능가함을 보여주고 있으며, Small 모델 세팅에서는 GPU 하나로 4일간 학습한 모델이 계산량이 30배 가량인 GPT를 능가하였다고 보여주고 있다. 또한, Large 모델 세팅에서도 RoBERTa나 XLNet 대비 1/4의 계산량으로 비슷하거나 약간 향상된 성능을 보여주고 있다고 한다.
RTD는 generator를 이용해 실제 입력의 일부 토큰을 그럴싸한 가짜 토큰으로 바꾸고, 각 토큰이 실제 입력에 있는 진짜(original) 토큰인지 generator가 생성해낸 가짜(replaced) 토큰인지 discriminator가 맞히는 이진 분류 문제라고 할 수 있다. RTD Task로 ELECTRA는 입력의 15%가 아니라 100%의 토큰을 모두 학습에 사용하기 때문에 효율적이고 효과적이다. (사실 이게 정상이지 않나…싶다)
위와 같이 다른 방식의 MLM 모델에 비하여 ELECTRA가 빠르게 성능이 향상되는 것을 볼 수 있다. 특히, 동일 조건(모델 사이즈, 데이터, 계산량(FLOPs))에서 더 높은 GLUE 성능을 보인다.
Method
먼저 ELECTRA 학습을 위한 RTD task를 살펴보면, 위 그림과 같이 Generator $G$와 Discriminator $D$로 구성되어 있다. $G$와 $D$는 모두 Transformer의 encoder 부분으로 구성되어 있으며, 뒤에서 설명하겠지만, $G$는 하나의 Transformer, $D$는 두 개의 Transformer로 구성되어있다. $G$는 기존의 MLM 모델과 같은 방식으로 학습된다(아래 step 1~4). $D$는 ELECTRA에서 새롭게 추가된 부분에 해당하는데, 입력 토큰 시퀀스($G$를 통과한 토큰 시퀀스이며, 원본 토큰 시퀀스가 아니라, G가 생성한 Fake 토큰 시퀀스이다. )에서 각 토큰이 원본 토큰과 같은지 아닌지를 binary classification하는 구조이다.
Input부터 output까지의 전체 과정을 step by step으로 설명하면,
- 입력 $\textbf{x} = [x_1, x_2, …, x_n]$를 받아 마스킹할 위치의 집합 $\textbf{m} = [m_1, m_2, …, m_k]$를 결정한다. 이때, 위치 $m_1,\dots, m_k$는 균등분포 \(unif \{ 1,n \}\)를 따른다. (여기서 $k$는 0.15$n$ = 전체 토큰 갯수 $n$개 중 15%)
- 여기까지 공통된 부분이고, 이 입력 $\textbf{x} = [x_1, x_2, …, x_n]$와 $\textbf{m} = [m_1, m_2, …, m_k]$은 $G$와 $D$로 각각 나뉘어 들어가게 되는데 이 부분은 뒤에서 설명하도록 한다.
-
$\textbf{m} = [m_1, m_2, …, m_k]$의 위치에 해당하는 입력 토큰을 $[MASK]$ 토큰으로 치환한다.
\[\textbf{x}^{masked}=\text{REPLACE}(\textbf{x}, \textbf{m}, [MASK])\]위와 같이 $\textbf{x}$를 $\textbf{m}$의 위치에 대하여 $[MASK]$로 바꾼다는 의미의 수식으로 표기한다.
-
$G$에서는 마스킹 된 입력 토큰 $\textbf{x}^{masked}$에 대해서 원래 토큰이 무엇인지를 예측한다.
\[p_G (x_t | \textbf{x}^{masked}) = \exp(e(x_t)^T h_G(\textbf{x}^{masked})_t) / \sum_{x'} \exp(e(x')^T h_G(\textbf{x}^{masked})_t)\]수식은 위와 같이 표기될 수 있다. $\textbf{x}^{masked}$가 주어졌을 때 토큰의 위치 $x_t$의 분포는 $G$를 통과한 Softmax 분포로 표현됨을 알 수 있다. 여기서 $e(\cdot)$은 임베딩을 의미한다. 이후 아래와 같이 MLM Loss를 통해 학습을 진행한다.
\[\mathcal{L}_{\text{MLM}}(\textbf{x}, \theta_G) = \mathbb{E} \left( \sum_{i \in \textbf{m}} -\log p_G (x_i | \textbf{x}^{masked}) \right)\]$G$의 학습 stage는 위의 loss를 구하면서 여기서 끝이 나지만, $D$의 입력 토큰 시퀀스는 앞서 계산한 $p_G (x_t \vert\textbf{x}^{masked})$에서 sampling을 통해 얻게된다.
-
앞서 step2에서 언급 했듯이 준비된 마스킹 위치 $\textbf{m}$에 대하여 $\textbf{x}$를 $G$에서는 $[MASK]$로 바꿨다면, 여기서는 $p_G (x_t \vert\textbf{x}^{masked})$에서 Sampling한 토큰으로 교체한다. 즉,
\[\textbf{x}^{corrupt}=\text{REPLACE}(\textbf{x}, \textbf{m}, \hat{\textbf{x}})\\ \hat{\textbf{x}}\sim p_G (x_i | \textbf{x}^{masked}) \; \text{for} \; i \in \textbf{m}\]으로 표현이 가능하다. 조금 더 설명하자면, $G$에서 구한 softmax의 분포로부터 $x_t$위치의 토큰에 들어갈 가장 높은 확률 값을 지니는 토큰($\hat{\textbf{x}}$)을 샘플링하고, 마스킹하는 대신 해당 토큰으로 치환하는 것이다.
-
이렇게 치환된 토큰 시퀀스 $\textbf{x}^{corrupt}$에 대해서 $D$는 각 토큰이 원래 입력 토큰 시퀀스와 동일한지 예측한다. 여기서 기존 방식과의 차이점은 마스킹 위치의 $\textbf{m}$에 대한 것만 진행하는 것이아니라 전체 FULL 토큰에 대하여 예측을 진행한다. 수식은 아래와 같이 표현 가능하다
\[D(\textbf{x}^{corrupt}, t) = \text{sigmoid}(w^T h_D(\textbf{x}^{corrupt})_t)\]$t$번째 토큰에 대한 예측을 나타낸다.
-
최종적으로는 아래와 같은 Loss로 학습한다.
\[\begin{align} &\mathcal{L}_{Disc} (\textbf{x}, \theta_{D})\\ & =\mathbb{E} \left( \sum_{t=1}^{n} -\mathbb{1}(x_{t}^{corrupt} = x_t) \log D(\textbf{x}^{corrupt}, t) - \mathbb{1}(x_{t}^{corrupt} \neq x_t) \log (1-D(\textbf{x}^{corrupt}, t)) \right) \end{align}\]
빠른 이해를 위해 그림으로 간단히 표현하면 아래와 같다.
| step | Figures | Description |
|---|---|---|
| 1 |  |
입력 토큰 시퀀스를 받아 임베딩한다. |
| 2 |  |
마스킹 후 Generator에 넣는다. |
| 3 |  |
Original input에서 마스킹 위치를 Generator에서 샘플링한 토큰으로 치환한 뒤 Discriminator에 넣어준다. |
| 4 |  |
원본 토큰 시퀀스와 다른지 여부를 모든 토큰에 대하여 분류한다. |
GAN vs ELECTRA
얼핏 보면 Generator와 Discriminator의 등장으로 인해 GAN과 유사하다고 착각하기 쉽지만, 모델 학습 측면에서 매우 다르다. 한 문장으로 요약하자면, GAN은 Fake image를 Discriminator로 분류한 결과가 Generator의 loss에 영향을 미치지만, ELECTRA는 그렇지 않다는 점이다.
앞서 설명했던 Step 1 ~ 7 중 step 4에서 언급했듯이 Generator의 loss는 maximum likelihood로 학습한다.
\[\mathcal{L}_{\text{MLM}}(\textbf{x}, \theta_G) = \mathbb{E} \left( \sum_{i \in \textbf{m}} -\log p_G (x_i | \textbf{x}^{masked}) \right)\]GAN과 달리 ELECTRA의 Generator는 적대적으로 학습하지 않고, Discriminator의 결과와는 상관없이 문맥을 잘 이해하는 방향으로 학습한다는 점이다.
특히, Generator에서 샘플링하는 과정은 미분이 불가능하기에 역전파가 불가능하다. 기존 GAN은 Cost가 발생하면 Generator로 역전파시킬 수킨 뒤, Discriminator의 분류 여부에 따라 Discriinator의 error를 계산하게 된다. 하지만, ELECTRA는 Discriminator에서 발생한 토큰의 loss가 Generator와 상관 없기 때문에($\mathcal{L}_{\text{MLM}}(\textbf{x}, \theta_G)$는 D의 loss와는 무관하다.) 그리고, 샘플링은 역전파가 불가능하기 때문에 Two stage로 학습할 수 밖에 없다.
논문에서는 강화학습을 통한 Adversarial한 학습도 실험했지만, 위와 같은 방식의 two stage 방식이 더 좋은 성능을 냈다고 한다. (Appendix F 참조)
Experiments
다음으로는 Experiments를 살펴보면, 모델의 성능을 끌어올린 3가지 방법을 논문에서는 얘기하고 있다.
Weight sharing
우선 ELECTRA는 Generator와 Discriminator 모두 transformer 인코더 부분을 사용한다. G와 D의 웨이트를 쉐어링하는 것이 파라미터를 줄이고, 어느정도 파라미터를 stabilizing하는 효과가 있다는 연구가 알려져 있었기에(ALBERT: A Lite BERT For Self-Supervised Learning of Language Representations), 논문에서는 이를 사용하였다.
총 3가지를 실험하였는데,
- 임베딩의 가중치만 공유하는 방법 (GLUE scores : 84.3)
- 모든 가중치를 공유하는 방법 (GLUE scores : 84.4)
- 가중치를 공유하지 않는 방법(GLUE scores : 83.5)
을 사용하였고, 2번 “모든 가중치를 공유하는 방법”이 가장 성능이 좋았다고 한다. 하지만, 모든 가중치를 공유하는 것은 G와 D의 크기를 동일하게 맞춰야하는 제약이 있다.
논문에서는 위와 같이 말하고 있다. 요약하면, 임베딩 레이어의 가중치를 공유하지 않는 구조의 경우 G로부터 softmax 샘플링한 토큰만을 입력으로 받는 D는 다소 편향된 학습을 할 것인데 반해, 임베딩 레이어의 가중치를 공유하는 구조는 이로부터 자유로울 수 있다는 점을 이유로 좋은 성능이 나왔을 것이라고 말한다. 위 표에서 첫 번째 Figure에서 볼 수 있듯이 D가 G로부터 편향된 토큰을 샘플링 받는다 하더라도 임베딩 레이어의 input을 통해 극복할 수 있다는 의미이다.
Smaller Generator
제너레이터와 디스크리미네이터의 사이즈 부분에서는 Discriminator 크기 대비 4분의1에서 2분의1 크기의 Generator가 가장 좋은 성능을 보였다고 한다. 논문에서는 제너레이터가 너무 강력하면, Discriminator의 태스크가 어려워진다고 표현하고 있고, Discriminator의 파라미터가 실제 데이터 분포를 분류하는데 쓰여야하는데, 오히려 Generator를 모델링하는데 잘못 사용되는 결과가 나타날 수 있다고 말하고 있다.
앞서 설명한 weight sharing에서 “모든 가중치를 공유하는 것은 G와 D의 크기를 동일하게 맞춰야하는 제약”에 대해 언급했었다. 사실 성능면에서 모든 가중치를 공유하는 것이 가장 좋지만, GLUE scores가 별 차이 나지 않는 “임베딩 가중치만 공유“하는 방식 + G의 사이즈를 D보다 작게하는 방식으로 효율성을 높였다고 논문에서는 말하고 있다.
Results
마지막으로 결과를 보면, 저자들이 제시한 연구의 큰 향상점 중 하나는 pre-training의 효율성 향상이다. 이를 검증하기 위해서 하나의 gPU로 빠르게 학습할 수 있는 small 모델을 만들어 실험을 진행했는데, 공정성을 위해 FLOPs를 동일하게 맞춘 상태에서 BERT와 ELECTRA를 비교한 결과 거의 5%가량의 성능 향상이 있었다고 한다.
심지어는 더 큰 모델인 GPT보다도 좋은 성능을 보이는 것을 확인할 수 있다.
그리고 small모델 뿐만 아니라 라지 모델에서도 실험을 진행했는데 BERT를 기준으로 1배라고 했을때 중간에 5.4배의 사이즈를 지닌 XLNet보다 낮은 4.4배의 계산량으로 ELECTRA는 동등하거나 우수한 성능을 보이는 것을 확인할 수 있다. 큰 성능 차이를 내진 못했지만, ELECTRA가 보다 효율적으로 학습하여 sota에 준하는 성능을 달성하고 있다.