07 트랜스포머 (5) ELECTRA

ELECTRA(Efficiently Learning an Encoder thar Classifies Token Replacements Accurately)는 2020년 구글에서 발표한 트랜스포머 기반의 모델이다.

 

ELECTRA는 GAN과 유사한 방법으로 생성자와 판별자를 사용해 사전 학습을 수행한다.

 

 

1. 사전 학습 방법

생성자와 판별자 모두 트랜스포머 인코더 구조를 따른다. 생성자는 입력 문장의 일부 토큰을 마스크 처리하고 마스크 처리된 토큰이 원래 어떤 토큰이었는지 예측하며 학습한다.

 

반면에 판별자는 입력 토큰이 원본 토큰인지 생성자에 의해 바뀐 토큰인지 구분하는 학습을 수행한다. 이러한 학습 방법을 RTD(Replaced Token Detection)라고 한다.

 

예를 들면 다음과 같다.

 

원본 문장 : ELECTRA 는 RTD 를 사용하는 신경망 입니다.

마스크 처리 : ELECTRA 는 [MASK] 를 사용하는 [MASK] 입니다. <- 생성자의 입력으로 사용

생성자 출력 : ELECTRA 는 MLM 를 사용하는 신경망 입니다. <- 판별자의 입력으로 사용

판별자 출력 : 원본, 원본, 바뀜, 원본, 원본, 원본, 원본

 

생성자는 BERT의 MLM과 동일하다. 판별자는 생성자가 복원한 결괏값을 입력받아 각 토큰이 원본인지 바뀐 것인지 판별한다.

 

GAN과 차이점이 있다. 생성자가 '신경망'을 원본과 동일하게 생성했는데, GAN에서는 원본이 아닌 생성된 것으로 간주하지만 ELECTRA는 원본 토큰으로 간주한다.

 

그리고 GAN은 적대적으로 학습하지만 ELECTRA의 판별자는 단순히 각 토큰이 바뀐 토큰인지 아닌지만 구분하도록 학습한다. 

 

마지막으로 GAN은 완전한 노이즈 벡터를 입력받아 생성하지만, ELECTRA의 생성자는 일부가 마스크 처리된 텍스트를 입력으로 받는다.

 

생성자와 판별자가 둘 다 트랜스포머 인코더 구조를 따르기 때문에, 같은 개수의 계층으로 구성돼 있다면 가중치를 공유할 수 있어 더 빠르게 학습할 수 있다. 그러나 완전히 공유하면 생성자의 성능이 너무 높아져서 판별자가 학습할 수 없게 된다.

 

따라서 ELECTRA에서는 생성자의 크기를 판별자의 1/2에서 1/4 크기로 바꿔 설정하고 모든 가중치를 공유하는 대신 임베딩 계층의 가중치만 공유한다.

 

ELECTRA는 사전 학습이 완료되면 판별자만 사용해 다운스트림 작업을 수행한다. 다운스트림 작업은 BERT와 동일한 방식으로 미세 조정한다.

 

 

2. 모델 실습

허깅페이스 라이브러리의 ELECTRA 모델과 네이버 영화 리뷰 ㄱ마정 분석 데이터세트로 분류 모델을 학습한다.

 

 

네이버 영화 리뷰 데이터세트

import numpy as np
import pandas as pd
from Korpora import Korpora


corpus = Korpora.load("nsmc")
df = pd.DataFrame(corpus.test).sample(20000, random_state=42)
train, valid, test = np.split(
    df.sample(frac=1, random_state=42), [int(0.6 * len(df)), int(0.8 * len(df))]
)

 

 

입력 텐서 생성

import torch
from transformers import ElectraTokenizer
from torch.utils.data import TensorDataset, DataLoader
from torch.utils.data import RandomSampler, SequentialSampler


def make_dataset(data, tokenizer, device):
    tokenized = tokenizer(
        text=data.text.tolist(),
        padding="longest",
        truncation=True,
        return_tensors="pt"
    )
    input_ids = tokenized["input_ids"].to(device)
    attention_mask = tokenized["attention_mask"].to(device)
    labels = torch.tensor(data.label.values, dtype=torch.long).to(device)
    return TensorDataset(input_ids, attention_mask, labels)


def get_datalodader(dataset, sampler, batch_size):
    data_sampler = sampler(dataset)
    dataloader = DataLoader(dataset, sampler=data_sampler, batch_size=batch_size)
    return dataloader


epochs = 5
batch_size = 32
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = ElectraTokenizer.from_pretrained(
    pretrained_model_name_or_path="monologg/koelectra-base-v3-discriminator",
    do_lower_case=False,
)

train_dataset = make_dataset(train, tokenizer, device)
train_dataloader = get_datalodader(train_dataset, RandomSampler, batch_size)

valid_dataset = make_dataset(valid, tokenizer, device)
valid_dataloader = get_datalodader(valid_dataset, SequentialSampler, batch_size)

test_dataset = make_dataset(test, tokenizer, device)
test_dataloader = get_datalodader(test_dataset, SequentialSampler, batch_size)

 

ELECTRA는 허깅페이스에서 다양한 모델이 제공되는데, 영어 텍스트 분류를 위한 ELECTRA 모델과 한국어 텍스트 분류를 위해 만들어진 KoELECTRA가 있다.

 

영어 텍스트 분류는 google/electra-small, google/electra-base, google/electra-large로 적용할 수 있고, 한국어 텍스트 분류는 monologg/koelectra-small-v3, monologg/koelectra-base-v3로 적용할 수 있다.

 

ELECTRA는 판별 모델만을 이용해 다운스트림 작업을 수행하므로 monologg/koelectra-base-discrminator 모델을 불러온다.

 

 

KoELECTRA 모델

from torch import optim
from transformers import ElectraForSequenceClassification


model = ElectraForSequenceClassification.from_pretrained(
    pretrained_model_name_or_path="monologg/koelectra-base-v3-discriminator",
    num_labels=2
).to(device)
optimizer = optim.AdamW(model.parameters(), lr=1e-5, eps=1e-8)

 

for main_name, main_module in model.named_children():
    print(main_name)
    for sub_name, sub_module in main_module.named_children():
        print("└", sub_name)
        for ssub_name, ssub_module in sub_module.named_children():
            print("│  └", ssub_name)
            for sssub_name, sssub_module in ssub_module.named_children():
                print("│  │  └", sssub_name)

electra
└ embeddings
│  └ word_embeddings
│  └ position_embeddings
│  └ token_type_embeddings
│  └ LayerNorm
│  └ dropout
└ encoder
│  └ layer
│  │  └ 0
│  │  └ 1
│  │  └ 2
│  │  └ 3
│  │  └ 4
│  │  └ 5
│  │  └ 6
│  │  └ 7
│  │  └ 8
│  │  └ 9
│  │  └ 10
│  │  └ 11
classifier
└ dense
└ dropout
└ out_proj

 

 

학습

import numpy as np
from torch import nn


def calc_accuracy(preds, labels):
    pred_flat = np.argmax(preds, axis=1).flatten()
    labels_flat = labels.flatten()
    return np.sum(pred_flat == labels_flat) / len(labels_flat)

def train(model, optimizer, dataloader):
    model.train()
    train_loss = 0.0

    for input_ids, attention_mask, labels in dataloader:
        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)

        loss = outputs.loss
        train_loss += loss.item()

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    train_loss = train_loss / len(dataloader)
    return train_loss

def evaluation(model, dataloader):
    with torch.no_grad():
        model.eval()
        criterion = nn.CrossEntropyLoss()
        val_loss, val_accuracy = 0.0, 0.0

        for input_ids, attention_mask, labels in dataloader:
            outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
            logits = outputs.logits

            loss = criterion(logits, labels)
            logits = logits.detach().cpu().numpy()
            label_ids = labels.to("cpu").numpy()
            accuracy = calc_accuracy(logits, label_ids)

            val_loss += loss
            val_accuracy += accuracy

    val_loss = val_loss/len(dataloader)
    val_accuracy = val_accuracy/len(dataloader)
    return val_loss, val_accuracy


best_loss = 10000
for epoch in range(epochs):
    train_loss = train(model, optimizer, train_dataloader)
    val_loss, val_accuracy = evaluation(model, valid_dataloader)
    print(f"Epoch {epoch + 1}: Train Loss: {train_loss:.4f} Val Loss: {val_loss:.4f} Val Accuracy {val_accuracy:.4f}")

    if val_loss < best_loss:
        best_loss = val_loss
        torch.save(model.state_dict(), "../ElectraForSequenceClassification.pt")
        print("Saved the model weights")

Epoch 1: Train Loss: 0.4507 Val Loss: 0.3185 Val Accuracy 0.8662
Saved the model weights
Epoch 2: Train Loss: 0.2800 Val Loss: 0.3033 Val Accuracy 0.8752
Saved the model weights
Epoch 3: Train Loss: 0.2054 Val Loss: 0.3147 Val Accuracy 0.8822
Epoch 4: Train Loss: 0.1486 Val Loss: 0.3597 Val Accuracy 0.8792
Epoch 5: Train Loss: 0.1147 Val Loss: 0.4195 Val Accuracy 0.8728

 

 

평가

model = ElectraForSequenceClassification.from_pretrained(
    pretrained_model_name_or_path="monologg/koelectra-base-v3-discriminator",
    num_labels=2
).to(device)
model.load_state_dict(torch.load("../ElectraForSequenceClassification.pt"))

test_loss, test_accuracy = evaluation(model, test_dataloader)
print(f"Test Loss : {test_loss:.4f}")
print(f"Test Accuracy : {test_accuracy:.4f}")

Test Loss : 0.3118
Test Accuracy : 0.8778

 

 

ELECTRA는 GLUE(General Language Understanding Evaluation) 평가에서 높은 점수를 기록한다. electra-small과 bert-small은 동일한 구조와 가중치 개수를 가지지만 electra가 5% 높은 점수를 기록했다.

 

또한 12시간 학습한 electra-small이 4일 학습한 bert-small 모델보다 더 좋은 GLUE 점수를 기록했다. 이는 ELECTRA가 더 효율적으로 학습한다는 것을 의미한다.