사전학습된 언어모델은 특정한 수의 Vocab이라고 하는 Token의 리스트를 가집니다. 토크나이저에 따라 다르겠지만, 언어모델을 훈련시키기 위한 데이터가 Vocab에 없다고 하더라도 subword로 분할하여 새로운 단어를 이해할 수 있습니다.

그런데 만약, 새로운 도메인의 데이터로 사전학습된 모델을 fine-tuning 한다면 어떨까요? 전문영역에서 사용하는 특정한 용어의 경우에는 subword로 분절해서 살펴보는 경우에는 본래의 의미를 잃어버리는 경우가 생길 수 있습니다.

따라서 도메인에서 사용되는 특정한 용어가 매우 중요하거나 성능에 큰 영향을 미칠 수 있는 경우에는 Vocab에 새로운 단어나 용어를 추가해줄 필요가 있습니다.

John Hewitt의 블로그 글(글 아래 링크 참고)은 이를 위한 가이드입니다.

간단하게 요약하면, 기존의 Vocab으로 훈련된 Embedding 벡터의 평균과 분산을 활용하여 새로 추가되는 단어에 대한 임베딩 값을 미리 설정하는 것입니다.

아래 코드에서 살펴볼 수 있듯이 기존의 embedding layer의 파라미터값의 평균과 분산을 활용하여 정규분포 값을 도출하고, 이 값을 새로 추가된 embedding layer의 초깃값으로 넣어주는 직관적인 방식입니다.

성능에 대해서는 검토해봐야겠지만, 직관적이고 이해하기 쉬워서 쉽게 활용할 수 있을 것 같습니다.

import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel

# 토크나이저와 모델 인스턴스 생성
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')

# 새 단어 추가
tokenizer.add_tokens(['Aragorn', 'Frodo', 'Lothlorien'])
model.resize_token_embeddings(len(tokenizer))

# 임베딩의 평균과 분산 계산
params = model.state_dict()
embeddings = params['transformer.wte.weight']
pre_expansion_embeddings = embeddings[:-3,:]
# 평균
mu = torch.mean(pre_expansion_embeddings, dim=0)
n = pre_expansion_embeddings.size()[0]
# 분산 (공분산)
sigma = ((pre_expansion_embeddings - mu).T @ (pre_expansion_embeddings - mu)) / n
# 다변량 정규 분포
dist = torch.distributions.multivariate_normal.MultivariateNormal(
	mu, 
	covariance_matrix=1e-5*sigma # 스케일링 팩터
)

# 새 임베딩 샘플링 및 적용
new_embeddings = torch.stack(tuple((dist.sample() for _ in range(3))), dim=0)
embeddings[-3:,:] = new_embeddings
params['transformer.wte.weight'][-3:,:] = new_embeddings
model.load_state_dict(params)

Reference

https://nlp.stanford.edu/~johnhew/vocab-expansion.html