Amazon SageMaker에서 지식 증류를 사용하여 소규모 언어 모델 교육

이 기사에서는 언어 모델과 지식 증류에 대해 알아봅니다. Amazon Sagemaker 컴퓨팅 리소스의 도움을 받아 학생 모델에 대한 작업별 지식 추출을 수행합니다.

전체 코드는 이 github 페이지에서 찾을 수 있습니다: https://github.com/gokulsg/BERT-models-complete-code/blob/main/knowledge_distillation.py

최근 몇 년 동안 사전 훈련된 변환기 기반 모델은 여러 NLP 작업에서 최첨단 성능을 발휘했습니다. Self-attention은 병렬 처리를 가능하게 하고 여러 GPU에서 교육을 지원하는 변환기 기반 모델의 핵심 구성 요소입니다. 이러한 변압기 모델은 크게 세 가지 등급으로 분류할 수 있습니다.

인코더 기반 모델 - 변압기 네트워크의 인코더 블록만 사용하는 모델 예: BERT, RoBERTa

디코더 기반 모델 - 변압기 네트워크의 디코더 블록만 사용하는 모델 예: GPT

인코더-디코더 모델 - 인코더 및 디코더 모듈이 모두 있음 예: T5, BART

이 기사에서는 주로 인코더 기반 모델인 BERT에 중점을 둘 것입니다. 이러한 인코더 기반 모델은 일반적으로 두 단계로 학습됩니다. 첫 번째 단계는 MLM(Masked Language Modeling) 목표를 사용하여 모델을 훈련하는 사전 훈련 단계입니다. MLM은 모델이 마스킹된 단어를 예측하기 위해 양방향으로 의미를 학습하도록 합니다. 교육의 두 번째 단계는 미세 조정으로, 감정 분류와 같은 특정 다운스트림 작업에 대해 모델을 교육합니다.

지식 증류는 널리 사용되는 모델 압축 방식 중 하나입니다. 여기에는 거대한 교사 모델에서 작은 학생 모델로 지식을 이전하는 것이 포함됩니다. 지식 증류에는 두 가지 유형이 있습니다.

작업별 지식 증류: 지식 증류는 특정 작업에 대해서만 미세 조정 단계에서 발생합니다.

작업에 구애받지 않는 지식 증류: 지식 증류는 주로 사전 훈련 단계에서 작업에 구애받지 않는 방식으로 발생합니다. 증류 후 이러한 모델은 모든 작업에 대해 미세 조정할 수 있습니다.

작업별 지식 추출을 구현하기 위해 포옹 얼굴 라이브러리를 사용할 것입니다.

student_id = "distilbert-base-uncased"
teacher_id = "textattack/bert-base-uncased-SST-2"

우리는 distilbERT 모델을 학생으로 사용하고 BERT 기본 모델을 교사로 사용하고 있습니다. distilBERT에는 6개의 인코더 레이어만 있는 반면 원래 BERT 기본 모델에는 12개의 인코더 레이어가 있습니다. 따라서 학생 모델은 교사와 비교할 때 인코더 레이어 수가 절반입니다.

from transformers import AutoTokenizer

# init tokenizer
teacher_tokenizer = AutoTokenizer.from_pretrained(teacher_id)
student_tokenizer = AutoTokenizer.from_pretrained(student_id)

# sample input
sample = "Testing tokenizers."

# assert results
assert teacher_tokenizer(sample) == student_tokenizer(sample), "Tokenizers produced different output"

교사 모델과 학생 모델의 토크나이저가 유사한 토큰화 결과를 생성했는지 확인해야 합니다. 그렇지 않으면 학생 모델의 성능에 영향을 미칩니다.

실험을 위해 Stanford Sentiment Treebank(sst-2) 데이터 세트(2클래스 감정 분류 데이터 세트)를 사용할 것입니다.

from datasets import load_dataset

dataset = load_dataset('glue','sst2')

실험을 위한 데이터 세트를 초기화했습니다. 이제 토큰화를 수행해야 합니다.

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(teacher_id)

def process(examples):
    tokenized_inputs = tokenizer(
        examples["sentence"], truncation=True, max_length=512
    )
    return tokenized_inputs

tokenized_datasets = dataset.map(process, batched=True)
tokenized_datasets = tokenized_datasets.rename_column("label","labels")

이제 지식 증류를 수행하기 위한 클래스를 작성합니다.

from transformers import TrainingArguments, Trainer
import torch
import torch.nn as nn
import torch.nn.functional as F

class DistillationTrainingArguments(TrainingArguments):
    def __init__(self, *args, alpha=0.5, temperature=2.0, **kwargs):
        super().__init__(*args, **kwargs)

        self.alpha = alpha
        self.temperature = temperature

class DistillationTrainer(Trainer):
    def __init__(self, *args, teacher_model=None, **kwargs):
        super().__init__(*args, **kwargs)
        self.teacher = teacher_model
        # place teacher on same device as student
        self._move_model_to_device(self.teacher,self.model.device)
        self.teacher.eval()

    def compute_loss(self, model, inputs, return_outputs=False):

        # compute student output
        outputs_student = model(**inputs)
        student_loss=outputs_student.loss
        # compute teacher output
        with torch.no_grad():
          outputs_teacher = self.teacher(**inputs)

        # assert size
        assert outputs_student.logits.size() == outputs_teacher.logits.size()

        # Soften probabilities and compute distillation loss
        loss_function = nn.KLDivLoss(reduction="batchmean")
        loss_logits = (loss_function(
            F.log_softmax(outputs_student.logits / self.args.temperature, dim=-1),
            F.softmax(outputs_teacher.logits / self.args.temperature, dim=-1)) * (self.args.temperature ** 2))
        # Return weighted student loss
        loss = self.args.alpha * student_loss + (1. - self.args.alpha) * loss_logits
        return (loss, outputs_student) if return_outputs else loss

이제 하이퍼파라미터를 지정하고 훈련을 시작할 수 있습니다.

from transformers import AutoModelForSequenceClassification, DataCollatorWithPadding
from huggingface_hub import HfFolder

# create label2id, id2label dicts for nice outputs for the model
labels = tokenized_datasets["train"].features["labels"].names
num_labels = len(labels)
label2id, id2label = dict(), dict()
for i, label in enumerate(labels):
    label2id[label] = str(i)
    id2label[str(i)] = label

# define training args
training_args = DistillationTrainingArguments(
    num_train_epochs=3,
    per_device_train_batch_size=64,
    per_device_eval_batch_size=64,
    learning_rate=5e-5,
    metric_for_best_model="accuracy",
    alpha=0.5,
    temperature=3.0
    )

# define data_collator
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

# define model
teacher_model = AutoModelForSequenceClassification.from_pretrained(
    teacher_id,
    num_labels=num_labels,
    id2label=id2label,
    label2id=label2id,
)

# define student model
student_model = AutoModelForSequenceClassification.from_pretrained(
    student_id,
    num_labels=num_labels,
    id2label=id2label,
    label2id=label2id,
)

trainer = DistillationTrainer(
    student_model,
    training_args,
    teacher_model=teacher_model,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics,
)

trainer.train()
trainer.evaluate()

여기에서는 정확도를 성능 평가 지표로 사용했습니다. AWS sagemaker를 사용한 교육의 경우 기존 코드를 거의 수정하지 않아도 됩니다.

from sagemaker.huggingface import HuggingFace

# hyperparameters, which are passed into the training job #
hyperparameters={
    'teacher_id':'textattack/bert-base-uncased-SST-2',
    'student_id':'distilbert-base-uncased',
    'dataset_id':'glue',
    'dataset_config':'sst2',
    # distillation parameter
    'alpha': 0.5,
    'temparature': 3,
}

# create the Estimator #
huggingface_estimator = HuggingFace(..., hyperparameters=hyperparameters)

# start knowledge distillation training #
huggingface_estimator.fit()

이러한 방식으로 교사 모델보다 몇 배 더 작고 빠른 소형 학생 모델을 생성하여 모바일/저자원 장치에 쉽게 배포할 수 있습니다. 또한 이러한 학생 모델은 교사 모델만큼 잘 수행할 수 있습니다.

다음 글에서는 또 다른 모델 압축 기법인 Quantization에 대해 알아보겠습니다.

참조:

https://www.philschmid.de/knowledge-distillation-bert-transformers