ELECTRA에 대해 알아보자

ELECTRA에 대해 알아보자

ELECTRA

발췌 : 스케터랩
기존 MLM 모델의 경우 autoregressive language modeling 학습에 비해 양방향 정보를 고려한다는 점에서 효과적인 학습을 할 수 있지만 몇가지 문제점이 존재한다.
1. 전체 토큰 중 15%에 대해서만 loss가 발생한다. -> 학습 비용이 많이 든다.
2. 학습 때는 [MASK] 토큰을 모델이 참고하여 예측하지만 실제 inference에는 [MASK] 토큰이 존재하지 않음

해당 문제점들을 해결하기 위해 Replaced Token Detection(RTD)라는 새로운 pre-training 태스크를 제안한다.
이는 generator를 이용해 실제 입력의 일부를 그럴싸한 가짜 토큰으로 변경하고, 각 토큰이 실제 입력에 있는 진짜 토큰인지 discriminator가 맞추는 이진 분류 문제이다.

이는 입력 토큰의 15%가 아닌 모든 토큰에 대해서 학습하기 때문에 효율적이면서 효과적이다.

해당 그림에서 볼 수 있듯이 RTD 태스크 학습을 위해서는 generator $G$와 discriminator $D$, 두 개의 네트워크가 필요하다.
두 네트워크는 공통적으로 Transformer 인코더 구조이며, 토큰 시퀀스 $x$를 입력받아 문맥 정보를 반영한 벡터 시퀀스 $h(x)$로 매핑시킨다.

Generator

Generator $G$는 BERT의 MLM과 똑같습니다.

• 입력 $x = [x_1,x_2,\cdots,x_n]$에 대해서 마스킹할 위치의 집합 $m = [m_1,m_2,\cdots,m_n]$을 결정합니다.
  • 모든 마스킹 위치는 1과 $n$사이의 정수이며, 이는 $m_i$ ~ unif {1,$n$} for $i$ = $1$ to $k$
  • 마스킹할 개수 $k$는 보통 $0.15n$을 사용한다. (전체 토큰의 15%)
• 결정한 위치에 있는 입력 토큰을 [MASK]로 치환
  • 이 과정을 $x^{masked}$ = REPLACE($x,m,$[MASK])와 같이 표현
• 마스킹 된 입력 $x^{masked}$에 대해서 generator는 원래 토큰이 무엇인지 예측한다.
• 최종적으로 MLM loss로 학습한다.

Discriminator

Discriminator $D$는 입력 토큰 시퀀스에 대해서 각 토큰이 original인지 replaced인지 이진 분류로 학습한다.

• $G$를 이용하여 마스킹 된 입력 토큰들을 예측 ($G$의 1~3단계)
• $G$에서 마스킹할 위치의 집합 $m$에 해당하는 위치의 토큰을 [MASK]가 아닌 generator의 softmax 분포 $p_G(x_t|$x)에 대해 샘플링한 토큰으로 치환한다.
  • example
    • input : [the,chef,cooked,the,meal]
    • input for $G$ : [[MASK],chef,[MASK],the,meal]
    • input for $D$ : [the,chef,ate,the,meal]
    • 첫 번째 단어는 샘플링 결과 원래 입력 토큰과 동일한 "the"가 나온 것
    • 세 번째 단어는 샘플링 결과 원래 입력 토큰인 "cooked"가 아닌 "ate"가 나온 것
    • $x^{corrupt}$ = REPLACE(x,m,$\hat{x}$), $\hat{x}$ ~ $p_G(x_i|x^{masked})$ for $i \in m$
• 치환된 입력 $x^{corrupt}$에 대해서 $D$는 각 토큰의 치환 여부를 예측한다.

GAN과의 차이점

$G$와 $D$를 사용한다는 점에서 GAN과 유사합니다. 하지만 ELECTRA와 GAN은 다음과 같은 차이점이 존재합니다.
1. $G$가 원래 토큰과 동일한 토큰을 생성했을 때 GAN은 negative sample(fake)로 간주하지만 ELECTRA는 positive sample로 간주합니다.
2. $G$가 $D$를 속이기 위해 adversarial하게 학습한게 아니고 maximum likelihood로 학습한다.
3. $G$의 입력으로 노이즈 벡터를 넣어주지 않는다.

최종적으로 ELECTRA는 대용량 코퍼스에 대해서 $L_{MLM}$과 $\lambda L_{Disc}$의 합을 최소화하도록 학습합니다. ($\lambda$ = 50)
이 때 $\lambda$는 이진 분류인 $L_{Disc}$와 30000개 클래스 분류인 $L_{MLM}$의 스케일을 맞추는 역할을 합니다.
$G$의 샘플링 과정이 역전파가 불가능하기에 $L_{Disc}$는 $G$로 역전파 되지 않으며, 해당 구조로 pre-training 후 $D$만을 취하여 downstream task로 fine-tuning을 진행합니다.

Experiments

ELECTRA 성능 평가를 위해 GLUE 벤치마크와 SQuAD 데이터셋을 사용하였다.
대부분의 실험은 BERT와 동일하게 Wikipedia와 BooksCorpus를 사용하여 pre-training
Large 모델의 경우 XLNet에서 사용한 ClueWeb, CommonCrawl, Gigaword를 사용

Weight sharing

$G$와 $D$ 모두 Trasnformer 인코더 구조이기에 두 네트워크 가중치를 공유하여 학습하는 weight sharing을 적용할 수 있었고 이를 통해 pre-training 효율 향상을 기대할 수 있었다.

GLUE 성능 평가 및 $G$와 $D$의 크기 동일하게 설정

• 공유하지 않음 : 83.5
• 임베딩만 공유 : 84.3
• 모든 가중치 공유 : 84.4

모든 가중치를 공유하는 것이 가장 좋은 성능을 보였으며, 임베딩만 공유하는 방법도 비슷한 성능을 얻을 수 있었다.
이에 대한 논문의 해석은 다음과 같다.

$D$는 입력으로 들어온 토큰만 학습하는 반면, $G$는 출력 레이어에서 softmax를 통해 사전에 있는 모든 토큰에 대해서 밀도 있게 학습이 가능하다. ELECTRA의 경우 $D$만을 취해서 사용하는데, $G$와 임베딩을 공유해서 학습한 경우, $D$는 효과적으로 학습할 수 있기에 좋은 성능을 기록했다.

모든 가중치를 공유하는 방법이 성능이 가장 좋지만 $G$와 $D$의 크기를 반드시 동일하게 맞춰야한다는 제약이 있으며 이는 결국 $D$만을 사용하기에 학습의 효율을 떨어뜨릴 수 있으며, 더 작은 $G$를 사용하는 것이 효과적이라는 실험 결과가 있기에 임베딩만 공유하는 세팅으로 진행하였다고 한다.

Smaller Generators

$G$와 $D$의 크기를 동일하게 가져간다면 일반 MLM 모델에 비해 단순 계산으로 약 2배의 계산량이 필요합니다.
해당 문제를 해결하기 위해 $G$의 크기를 줄이는 실험을 진행합니다.

동일한 500k step을 진행하였기에 size가 작은 $G$는 size가 큰 $G$에 비해 계산량 대비 성능을 손해보았지만 $D$의 크기 대비 1/4 ~ 1/2 크기의 $G$를 사용했을 때 성능이 가장 좋았습니다.
이를 논문은 다음과 같이 해석합니다.

$G$가 너무 강력하면 $D$의 task가 너무 어려워져서..

또한 스캐터랩의 작성자 분은 다음과 같은 해석을 덧붙입니다.

"generator가 강력하면 softmax 분포가 몇몇 토큰에 치우쳐 있을 가능성이 높고 샘플링 결과가 다양하지 않기 때문에 discriminator가 효과적으로 학습할 수 없는 이유도 있지 않을까 싶습니다"

Efficiency Analysis

논문에서는 ELECTRA의 성능이 왜 좋은지에 대해 이해하기 위해 다음과 같은 실험을 세팅했다.

• ELECTRA 15% : ELECTRA의 구조를 유지하되, $L_{Disc}$를 입력 토큰의 15%만으로 만들도록 세팅
• Replace MLM : $D$를 MLM 학습을 하되, [MASK]로 치환하는 것이 아닌 $G$가 만든 토큰으로 치환
• All-Token MLM : Replace MLM 처럼 하되, 일부(15%)만 치환하는 것이 아닌 모든 토큰을 $G$가 생성한 토큰으로 치환

먼저 ELECTRA 15%는 토큰에 대한 학습 효율(15%, 100%) 때문에 성능 차이가 생겼다는 것을 보이기 위한 목적으로 보입니다.
Replace MLM은 pre-training 때만 사용하고 fine-tuning 때는 존재하지 않는 [MASK] 토큰 때문에 생긴 성능차이를 보이기 위한 목적으로 보이며 All-Tokens MLM은 BERT와 ELECTRA를 합친 버전이다.

결과는 다음과 같습니다.

이는 ELECTRA가 학습 효율이 좋으며, [MASK] 토큰에 대한 pre-training과 fine-tuning 간의 불일치 문제도 상당히 완화시킨 것으로 알 수 있다.

또한 히든 레이어의 크기에 따른 BERT와 ELECTRA의 성능 변화 실험을 통해 히든 레이어가 작아질수록 BERT와 ELECTRA의 성능 차이가 커진다는 사실을 알 수 있다.

결론적으로 ELECTRA가 BERT보다 효율적으로 학습을 한다는 것을 알 수 있다.

Conclusion

논문은 language representation learning을 위한 새로운 self-supervision 태스크인 Replaced Token Detection을 제안한다.
제안하는 방법의 주요 아이디어는 작은 generator가 만들어 낸 질 좋은 negative sample과 입력 토큰을 구별하도록 텍스트 인코더를 학습시키는 것이다.
Masked language modeling에 비해, 제안하는 pre-training objective는 훨씬 효율적(compute-efficient)이고 downstream tasks에 대한 결과 역시 더 좋다는 것을 많은 실험을 통해서 확인할 수 있다.
또한 ELECTRA는 타 모델들에 비해서 상대적으로 적은 계산량을 사용하는 경우에 더 효과적이다.

저자는 이 연구를 통해서 연구자들이 적은 컴퓨팅 리소스로도 language model의 pre-training에 대한 많은 연구/개발을 하길 바란다는 말과 pre-training과 관련된 향후 연구가 절대적 성능 지표만큼 계산량과 파라미터 수 등의 효율성도 함께 고려했으면 하는 바람을 남겼습니다.