BERT에서 사용되었던 기존의 Masked LM pre-training은 입력 데이터의 일부를 [MASK]로 대체하는 과정을 통해 input corruption을 일으키고, 모델은 original token를 예측하도록 학습된다. BERT 기반의 모델들은 downstream task에서 좋은 성능을 보이지만, 학습하는 과정에서 많은 계산 비용을 필요로 한다.
이 문제를 해결하기 위해 해당 논문은 pre-training을 보다 효율적으로 하기 위해 replaced token detection objective를 사용하여 학습된 ELECTRA(Efficiency Learning an Encoder that Classifies Token Replacements Accurately)를 제안한다. 이 objective는 입력을 마스킹하는 대신, small generator 모델로부터 샘플링된 토큰으로 대체하고, pre-training 모델은 입력이 대체된 토큰인지를 판별하는 discriminative model의 역할을 한다. 기존의 MLM은 모든 입력 토큰을 마스킹할 수 없지만, 제안된 objective는 모든 입력에 걸쳐서 사용될 수 있기 때문에 더 효율적이다. 결론적으로 replaced token detection을 사용하여 pre-train한 모델은 기존의 BERT 기반의 모델보다 뛰어난 성능을 보이고, 특히 작은 모델에 대해서 효과적인 학습 결과를 보인다. 오직 하나의 GPU를 사용해서 학습된 모델이 30배 가량 많은 계산량을 거친 GPT보다 뛰어난 성능을 보였으며, RoBERTa와 XLNet은 1/4의 계산량으로 비슷한 성능에 도달할 수 있었다.
Introduction
최근의 연구 동향을 살펴보면 denoising autoencoders(BERT 기반의 논문, MLM)가 많은 부분을 차지하고 있다. Auto Encoder는 unlabeled 입력의 일부분을 선택하여 마스킹하여 모델이 원래의 입력을 예측하도록 학습시킨다. 하지만 이런 접근 방식은 상당한 계산 비용을 필요로 한다. BERT를 예로 들면, pre-training 과정에서 15%의 토큰을 마스킹하기 때문에 한 번의 학습에서 입력의 15%만 학습에 사용할 수 있는 것으므로, 비효율적이라고 할 수 있다.
따라서 계산 비용 문제를 해결하기 위해 replaced token detection을 제시한다. 이 pre-training objective는 generator에 의해 생성되어 대체된 토큰을 색출하는 discriminator를 학습시킨다. BERT pre-training과 fine-tuning에서 발생하는 불일치를 replaced token detection 방식을 사용함으로써 해결할 수 있다.
[그림 1]은 x축을 pre-training 연산량(FLOP), y축을 GLUE benchmark 성능으로 한 그래프이다. GPU 한개로 4일 동안 학습된 ELECTRA-small은 더 많은 연산량을 수반하는 GPT보다 나은 성능을 보여주며, RoBERTa와 XLNet에 비해 1/4의 연산량과 적은 파라미터량을 가지는 ELECTRA-large 모델은 비슷한 성능을 가진다. 결과적으로, ELECTRA는 discriminator 학습을 통해 적은 계산량을 통해 뛰어난 성능을 가질 수 있게 된다.
Replaced token detection
Replaced token detection은 생성자 G와 판별자 D 두개의 네트워크를 학습시킨다. 판별자 네트워크가 pre-training을 진행하는 주축이 되는 모델이며, 생성자 모델은 입력 모델을 처리하여 판별자에게 입력값을 넘겨준다. [그림 2]는 두개의 모델의 상호 관계를 나타낸다. 생성자 G는 작은 크기의 MLM 모델이고, 원본 입력의 마스킹 처리를 하여 기존 MLM과 동일하게 [MASK] 토큰의 원본 토큰을 예측한다. 생성자 모델은 fine-tuning 과정에서는 사용되지 않는다.
입력 토큰 $x = [x_1, \cdots, x_n]$, 컨텍스트화 된 vector representation $h(x) = [h_1, \cdots, h_n]$, 그리고 위치값 $t$ ($x_t$ = [MASK])가 주어질 때, 생성자는 softmax 레이어를 통해 다음의 probability output을 계산한다:
$$p_G(x_t | x) = exp (e(x_t)^T h_G(x)_t) / \sum_{x'} exp (e(x')^T h_G(x)_t)$$
$e$는 토큰 임베딩을 의미하며, 판별자는 위치 $t$의 입력 토큰이 실제 데이터에서 추출된 'real' 데이터인지를 판별하며, 출력은 sigmoid layer로 계산된다:
$$D(x, t) = sigmoid(w_T h_D(x)_t)$$
입력과 마스킹 위치 정보가 아래와 같이 주어질 때, 마스킹된 입력을 $x^{masked}$으로 나타낸다.
- Input: $x = [x_1, x_2, \cdots, x_n]$
- Masking Position: $m = [m_1, \cdots, m_k]$
- $x^{masked} = REPLACE(x, m, [MASK])$
생성자로부터 출력되어 변형된 입력은 $x^{corrupt}$와 같이 나타내며, 판별자 D의 입력은 다음과 같이 형성된다.
$m_i \sim$ unif{1, $n$} for $i$ = 1 to $k$ $x^{masked} = REPLACE(x,m, [MASK])$
$\hat{x} \sim p_G(x_i | x_{masked})$ for $i \in m$ $x^{corrupt} = REPLACE(x, m, \hat{x})$
Loss function은 다음과 같다.
$$L_{MLM}(x, \theta_G) = \mathbb{E} (\sum_{i \in m} - log p_G(x_i | x^masked))$$
$$L_{Disc(x, \theta_G)} = \mathbb{E} (\sum_{t=1}^n - 1(x_t^{corrupt} = x_t) log D(x_{corrupt}, t) \\- 1(x_t^{corrupt} \neq x_t) log (1 - D(x^{corrupt}, t)))$$
GAN의 생성자/판별자를 기반으로 한 학습은 비슷하지만 몇 가지 다른 점이 있다.
- 생성자 모델이 올바른 토큰을 출력했을 경우, 해당 토큰은 'real'로 여긴다.
- 생성자는 판별자를 속이기 위한 적대적 모델로 학습되는 대신 maximum likelihood로 학습된다.
모델은 다음 loss를 최소화한다:
$$\min_{\theta_G, \theta_D} \sum_{x \in X} L_{MLM}(x, \theta_G) + \lambda L_{Disc}(x, \theta_D)$$
Discussions
ELECTRA는 replaced token detection이라는 pre-training method를 제시하면서 모델이 적은 연산 비용을 요구하면서 좋은 성능을 보일 수 있다는 것을 증명했다. MLM는 마스킹 된 토큰으로 인해 입력의 토큰 연관성을 학습하는데 비효율적이라는 단점이 있다. ELECTRA는 생성자 모델과 판별자 모델을 함께 학습시키기 때문에 판별자 모델의 성능이 낮음에 따라 판별자가 토큰 연관성을 학습하는데 좋지 않은 영향을 미치지 않을 수도 있을 것이라는 의문점을 제기해본다.
논문 참고하실 분은 아래 링크를 참고해주세요.
https://arxiv.org/pdf/2003.10555.pdf
