언어모델

언어 모델(Language model)은 주어진 단어들의 시퀀스에 대해 임베딩 표현(Embedding representation) 또는 벡터(Embedding vector)를 만들어 줍니다. 언어 모델을 통해 만든 임베딩 표현으로 주제 찾기, 감정 분류, 개체명 인식과 같은 분류 문제나 기계 번역, 요약 등 텍스트 생성 문제를 풀 수 있습니다.

현재 언어 모델 중 가장 유명한 것은 GPT입니다. GPT는 Generative Pretrained Transformer의 약자로, 딥러닝 기반의 자연어 처리 모델입니다. GPT 모델은 대화형 인공지능 챗봇인 ChatGPT 및 자동 코드 완성 인공지능인 Github Copilot 등에 사용되고 있습니다. 또한, Copilot은 MS의 워드, 엑셀, 파워포인트에 내장되어 문서나 프레젠테이션을 자동으로 완성할 수 있다고 합니다.

이런 정교한 언어 생성이 가능한 이유는 매우 큰 모델에 수많은 데이터를 학습한 결과입니다. 전 버전인 GPT-3 모델은 약 1750억 개의 파라미터가 존재한다고 하고, 현재 버전인 GPT-4 모델은 그보다 훨씬 많은 파라미터가 존재할 것이라고 추측하고 있습니다.

이렇게 챗봇과 문서 자동완성과 같은 작업은 거대한 모델에 수많은 데이터를 사전학습을 하고 각 작업에 맞게 미세조정(Fine-tuning)을 한 결과로 나옵니다. 사전훈련된 모델의 장점은 다양한 분야의 지식이 포함되어 있어 범용성이 높고, 목표 작업으로 미세조정을 하는데 작은 데이터세트로도 가능하다는 점입니다. 하지만 단점으로는 특정 도메인에 맞게 맞춤화하는 것이 어렵고 모델의 크기가 커서 리소스가 많이 필요하다는 점이 있습니다. 그렇기 때문에, 도메인에 맞는 작은 모델이 필요하다면 직접 모델을 구축해야 합니다.

그래서 이번 글에서는 딥러닝 기반 자연어 처리 분야에서 대표적인 기술인 Masked Language Model(이하 MLM)과 Causal Language Model(이하 CLM)에 대한 원리를 알아보고 PyTorch 라이브러리로 직접 모델을 구축하는 방법을 소개하겠습니다.

텍스트 데이터세트와 토크나이저

토크나이저

모델에 대해 알아보기 전에 우선 데이터를 마련하겠습니다. 데이터세트는 wikitext-2의 로우(Raw) 데이터를 사용하겠습니다. wikitext-2는 https://blog.salesforceairesearch.com/the-wikitext-long-term-dependency-language-modeling-dataset에서 다운로드할 수 있습니다.

텍스트 데이터를 모델의 입력으로 사용하기 위해서는 우선 토크나이저를 만들어야 합니다. 토크나이저는 단어들의 통계수치로 우선순위를 구해, 어휘사전을 만들고 텍스트를 토큰으로 파싱(분해) 합니다. 이번 예시에서는 토크나이저로 Huggingface tokenizers의 CharBPETokenizer 구현체를 사용하겠습니다.

아래의 코드로 토크나이저를 생성합니다.

from tokenizers import CharBPETokenizer
 
tokenizer = CharBPETokenizer(suffix='', lowercase=True)

special_tokens = ['','','','','']

vocab_size = 36000
min_frequency = 2
 
data_path = '/data/wikitext-2-raw'
tokenizer.train(files=data_path + '/wiki.train.raw',
				vocab_size=vocab_size,
				min_frequency=min_frequency,
				special_tokens=special_tokens,
				suffix='')

tokenizer.save('tokenizer.json')

특수 토큰(Special tokens)은 5개로 다음 표의 용도로 쓰입니다.

token	token_id	설명
	0	서로 다른 길이의 문장을 처리하기 위해 짧은 문장을 긴 문장의 길이와 맞추기 위해 로 패딩합니다.
	1	토크나이저가 모르는 단어를 만나면 unknown으로 처리하기 위한, 처리용 토큰입니다.
	2	문장의 시작을 알리는 토큰입니다. BOS(Begin of sentence), CLS(Classification) 토큰으로도 사용됩니다.
	3	문장의 끝을 알리는 토큰입니다. EOS(End of sentence), SEP(Seperator) 토큰으로도 사용됩니다.
	4	MLM 학습 시 쓰이는 토큰입니다. 토큰을 마스킹해서 이 토큰을 맞추는 문제를 풀 때 사용됩니다.

데이터세트

아래는 데이터세트를 구현하겠습니다. PyTorch의 torch.utils.data.Dataset을 상속한 클래스를 만듭니다. 이 클래스는 텍스트 파일을 불러와 한 줄씩 읽어 토크나이저로 인코딩을 한 토큰 시퀀스를 담고 있습니다. 이 데이터세트는 인덱스로 접근한다면 토큰 시퀀스 하나를 리턴합니다.

학습 효율을 향상하기 위해 버퍼를 사용해 토큰들을 채우고, 시퀀스 최대 길이로 자르면서 토큰 시퀀스를 만들겠습니다.

# dataset.py
from torch.utils.data import Dataset
from tokenizers import Tokenizer

class MyDataset(Dataset):
	def __init__(self, text_path, tokenizer_path, seq_length):
		super().__init__()
		self.tokenizer = Tokenizer.from_file(tokenizer_path)
		self.pad_token = self.tokenizer.encode("").ids[0]
		self.unk_token = self.tokenizer.encode("").ids[0]
		self.bos_token = self.tokenizer.encode("").ids[0]
		self.eos_token = self.tokenizer.encode("").ids[0]
		self.mask_token = self.tokenizer.encode("").ids[0]
		self.input_ids = []
		buffer = []
		with open(text_path, "r") as f:
			for text in f.readlines():
				buffer.extend(self.tokenizer.encode(text).ids)

				# eos, bos 토큰을 붙이기 위해 seq_length-2 만큼 자른다.  
				while len(buffer) >= seq_length - 2:
					input_id = (
						[self.bos_token] + buffer[: seq_length - 2] + [self.eos_token]
					)
					self.input_ids.append(input_id)
					buffer = buffer[seq_length - 2 :]

	def __len__(self):
		return len(self.input_ids)
		
	def __getitem__(self, idx):
		return self.input_ids[idx]

이렇게 텍스트 데이터를 가지고 토크나이저를 만들고 학습에 쓸 데이터세트를 구성하는 코드를 만들어 보았습니다.

다음으로는 모델을 구성하고 Masked Language Model과 Causal Language Model 학습하는 방법에 대해 알아보겠습니다.

트랜스포머 인코더

MLM이나 CLM은 정확하게는 모델을 사전학습하는 메커니즘입니다. 여기서 학습시키는 모델은 Attention Is All You Need 논문(https://arxiv.org/abs/1706.03762)에서 소개된 트랜스포머 모델입니다. 트랜스포머는 크게 인코더와 디코더로 이루어져 있는데, 여기서는 인코더를 학습해 보겠습니다.

트랜스포머 인코더는 아래와 같은 요소로 구성되어 있습니다.

Embedding Layer
Transformer Encoder
- Multi Head Self-Attention Layer
- Feed-Forward Layer

트랜스포머 인코더의 각 요소에 대해 간략히 알아보고 PyTorch로 구현하는 방법을 알아보겠습니다.

임베딩 레이어

트랜스포머의 임베딩 레이어의 특징은 위치 정보를 포함하고 있다는 점입니다. 토큰 시퀀스의 인덱스를 임베딩한 것을 토큰 임베딩에 더하는 방법으로 위치 정보를 추가합니다. 다시 말해, 토큰 시퀀스와 토큰 시퀀스의 인덱스를 각각 nn.Embedding()을 거쳐 더하는 방법으로 임베딩 레이어를 구성합니다.

원래 논문에서는 Positional Encoding으로 위치 정보를 sin, cos 값을 이용해 인코딩을 하는 방법을 제시했지만, 여기에서는 더 많이 쓰이는 방법인 학습 가능한 위치 정보 임베딩을 사용하겠습니다.

# model.py  
import torch  
from torch import nn 

class Embedding(nn.Module):
	def __init__(self, config: Config):
		super().__init__()

		self.position_embeddings = nn.Embedding( 
			config.max_position_embeddings, config.hidden_size
		) self.word_embeddings = nn.Embedding( 
			config.vocab_size, config.hidden_size, padding_idx=0
		)
		self.dropout = nn.Dropout(config.hidden_dropout_prob)
		
	def forward(self, input_ids: torch.Tensor):
		seq_len = input_ids.size(1)
		position_ids = (
			torch.arange(seq_len, dtype=torch.long).unsqueeze(0).to(input_ids.device)
		)
		position_embeddings = self.position_embeddings(position_ids)
		word_embeddings = self.word_embeddings(input_ids)
		embeddings = word_embeddings + position_embeddings
		embeddings = self.dropout(embeddings)
		return embeddings

여기서 초기화 변수로 Config를 파라미터로 받는데, Config 클래스는 모델의 하이퍼파라미터를 담고 있습니다. 자세한 내용은 아래에서 설명드리겠습니다.

트랜스포머 인코더

트랜스포머 인코더는 여러 개의 트랜스포머 인코더 레이어를 쌓은 형태이고, 트랜스포머 인코더 레이어는 멀티 헤드 어텐션 레이어와 피드 포워드 레이어로 구성되어 있습니다.

(출처: https://arxiv.org/abs/1706.03762)

이제, 멀티 헤드 어텐션과 피드 포워드에 대해서 알아보겠습니다.

멀티 헤드 어텐션 레이어(Multi Head Attention Layer)

(출처: https://arxiv.org/abs/1706.03762)

멀티 헤드 어텐션은 여러 개의 어텐션 헤드로 나누어지는데, 이 각각의 어텐션 헤드는 주어진 시퀀스에서 특정한 정보를 추출합니다. h개의 어텐션 헤드가 있다면(오른쪽 그림에서 h개의 어텐션 헤드), 토큰 임베딩의 각 토큰은 h개의 서로 다른 정보로 분석됩니다. 이는 이미지 처리에서 합성곱 신경망의 필터와 비슷합니다. 이것은 하나의 필터가 원을 처리하고, 다른 필터는 사각형을 처리하는 경우와 비슷합니다. 하나의 어텐션 헤드는 가중치 행렬을 계산 후 토큰 임베딩의 각 토큰마다 가중 평균을 구합니다. 여기서 가중치 행렬은 어텐션 점수(Attention score)로도 불립니다. 이런 과정을 거치면 토큰 임베딩의 각 토큰 사이에 유사도 점수가 계산됩니다. 연관성이 큰 토큰의 점수는 크고, 연관성이 작다면 점수가 작습니다. 이렇게 강조할 부분의 점수를 크게 함으로써 어텐션 메커니즘이 적용이 되고, 같은 토큰 임베딩끼리 점수를 계산해서 Self Attention이라고 이름이 붙여졌습니다.

어텐션 헤드의 구현은 토큰 임베딩을 가지고 완전 연결 층(Fully connected layer)을 통해 쿼리 벡터(Q), 키 벡터(K), 밸류 벡터(V)를 만들고 계산합니다. 쿼리 벡터와 키 벡터의 점 곱(Dot product 또는 MatMul)으로 가중치 행렬을 만들고, 이 행렬을 정규화하고 소프트맥스 함수를 적용해 열의 합이 1이 되게 만듭니다. 그다음 가중치 행렬과 밸류 벡터를 곱해서 최종 결과를 만듭니다. 한편, 다음 토큰을 예측하는 CLM 학습을 할 때에는 추가작업이 필요합니다. 현재 토큰에서 다음 토큰들을 참고하지 못하도록 가중치 행렬의 소프트맥스 함수 전에 다음 토큰들의 위치에 음의 무한대의 값으로 처리합니다. 이 작업을 마스킹이라 하며, 논문에서 Masked Multi Head Attention이라고 부르고 있습니다. 이 어텐션 헤드를 여러 개 쌓으면 멀티 헤드 어텐션이 됩니다.

아래는 멀티 헤드 셀프 어텐션 레이어의 예시 코드입니다.

# 예시 코드  
class AttentionHead(nn.Module):  
	def __init__(self, hidden_size, num_attention_heads):
		super().__init__()
		self.attention_head_size = hidden_size // num_attention_heads
		self.q = nn.Linear(hidden_size, self.attention_head_size)
		self.k = nn.Linear(hidden_size, self.attention_head_size)
		self.v = nn.Linear(hidden_size, self.attention_head_size)
	
	def forward(self, hidden_state, mask=None):
	attention_outputs = self.scaled_dot_product_attention( 
		self.q(hidden_state), self.k(hidden_state), self.v(hidden_state), mask
		)
		return attention_outputs
		
	def scaled_dot_product_attention(self, query, key, value, mask=None):
		dim_k = query.size(-1)
		scores = torch.bmm(query, key.transpose(1, 2)) / sqrt(dim_k)
		if mask is not None:
			scores = scores.masked_fill(mask == 0, float("-inf"))
		weights = F.softmax(scores, dim=-1)
		return weights.bmm(value)

class MultiHeadAttention(nn.Module):  
	def __init__(self, hidden_size, num_attention_heads):
		super().__init__() 
		self.attention_heads = nn.ModuleList(
			[
				AttentionHead(hidden_size, num_attention_heads)
				for _ in range(num_attention_heads)
			]
		)
		self.output_linear = nn.Linear(hidden_size, hidden_size) 

	def forward(self, hidden_state, mask=None):
		x = torch.cat([h(hidden_state, mask) for h in self.attention_heads], dim=-1)
		x = self.output_linear(x)
		return x

피드 포워드(Feed-Forward Layer) 레이어

피드 포워드는 간단하게 완전 연결층 2개로 구성되어 있습니다. 논문에서는 첫 번째 층의 크기를 임베딩 크기의 네 배로 설정하고, 활성화 함수로 GELU를 사용하고 있습니다.

아래는 피드 포워드 레이어의 예시 코드입니다.

# 예시 코드  
class FeedForward(nn.Module):  
	def __init__(self, hidden_size, intermediate_size):
		super().__init__()
			self.linear1 = nn.Linear(hidden_size, intermediate_size)
				self.gelu = nn.GELU()
			self.linear2 = nn.Linear(intermediate_size, hidden_size)
			self.dropout = nn.Dropout(0.1)
	def forward(self, x):
		x = self.linear1(x)
		x = self.gelu(x)
		x = self.linear2(x)
		x = self.dropout(x)
		return x

인코더

위의 MultiHeadAttention과 FeedForward를 조합하면 트랜스포머 인코더를 만들 수 있습니다. 위의 코드는 설명을 위해서 예시를 드렸고, 실제 학습에서는 PyTorch에서 구현한 torch.nn.TransformerEncoder와 torch.nn.TransformerEncoderLayer를 사용하겠습니다. torch.nn.TransformerEncoderLayer는 MultiHeadAttention과 FeedForward를 연결시켜 놓은 구현체이고, torch.nn.TransformerEncoder는 torch.nn.TransformerEncoderLayer와 레이어 수를 입력으로 받아 인코더를 구성하는 구현체입니다. 인코더를 만들기 전에 인코더의 하이퍼파라미터부터 먼저 설명드리겠습니다. 아래의 코드는 설정 클래스를 구현한 것이고, 표는 하이퍼파라미터의 설명과 BERT, GPT2에서 사용한 기본값들입니다.

# model.py  
class Config:  
	def __init__(
		self, 
		vocab_size=10000, 
		hidden_size=512, 
		num_hidden_layers=4, 
		num_attention_heads=4, 
		intermediate_size=2048, 
		max_position_embeddings=128, 
		layer_norm_eps=1e-12, 
		hidden_dropout_prob=0.1, 
		initializer_range=0.02, 
		is_causal=False, 
	): 
	self.vocab_size = vocab_size 
	self.hidden_size = hidden_size 
	self.num_hidden_layers = num_hidden_layers 
	self.num_attention_heads = num_attention_heads 
	self.intermediate_size = intermediate_size 
	self.max_position_embeddings = max_position_embeddings 
	self.layer_norm_eps = layer_norm_eps 
	self.hidden_dropout_prob = hidden_dropout_prob 
	self.initializer_range = initializer_range 
	self.is_causal = is_causal

설정	설명	BERT 기본값	GPT2 기본값
vocab_size	어휘 크기	30522	50257
hidden_size	임베딩과 hidden states의 차원	768	768
num_hidden_layers	인코더 레이어 수	12	12
intermediate_size	피드 포워드의 중간 차원. 보통 hidden_size X 4로 설정	3072	hidden_size * 4 고정
max_position_embeddings	시퀀스의 최대 길이	512	1024
layer_norm_eps	LayerNorm의 epsilon 값	1e-12	1e-5
hidden_dropout_prob	드롭아웃 확률	0.1	0.1
initializer_range	모델 가중치의 초기화 값 구간	0.02	0.02
is_causal	CLM 학습인지 체크. True라면 Attention Mask 추가

아래는 인코더 구현체입니다. 인코더에서는 Embedding과 torch.nn.TransformerEncoderLayer, torch.nn.TransformerEncoder를 선언하고 초기화합니다. 여기서 하나 주의해야 할 부분은 마스킹을 처리하는 부분입니다. 마스킹 관련해서 여러 문서에서 혼용해서 쓰고 있으므로 트랜스포머에서 쓰이는 마스킹의 용어 및 torch.nn.TransformerEncoder의 파라미터 정리부터 하겠습니다.

Mask	설명	TransformerEncoder Parameters	차원
Padding Mask	시퀀스의 길이가 서로 다를 때, 길이를 맞추기 위해 패딩 토큰의 위치를 알려주는 마스킹	src_key_padding_mask	Batch size X Seq length
Attention Mask	미래(혹은 다음) 단어를 참조하지 못하도록 마스킹	mask	Seq length X Seq length

torch.nn.TransformerEncoder의 마스킹 처리는 FloatTensor가 들어오면 어텐션 레이어의 가중치 행렬에 더하고, BoolTensor가 들어오면 True인 부분을 무시합니다. 나중에 설명할 CLM 학습할 때는 config.is_causal을 true로 설정해 상삼각행렬의 1인 부분(대각은 0)을 float(”-inf”)로 채우고 가중치 행렬에 더하게 합니다.

# model.py
import torch
from torch import nn
import torch.nn.functional as F

class Encoder(nn.Module):
	def __init__(self, config: Config):
		super().__init__()
		self.config = config
		self.embeddings = Embedding(config)
		layers = nn.TransformerEncoderLayer(
			d_model = config.hidden_size,
			nhead = config.num_attention_heads,
			dim_feedforward = config.intermediate_size,
			dropout = config.hidden_dropout_prob,
			activation = F.gelu,
			layer_norm_eps = config.layer_norm_eps,
			batch_first = True,
		)
		self.encoder = nn.TransformerEncoder(
			encoder_layer = layers, num_layers = config.num_hidden_layers
		)
		
	def forward(
		self,
		input_ids,
		attn_mask = None,
		padding_mask = None,
	):
		if self.config.is_causal and attn_mask is None:
			size = input_ids.shape[1]
			device = input_ids.device
			attn_mask = torch.triu(
				torch.ones(size, size) * float("-inf"), diagonal = 1
			).to(device)
			
		x = self.embeddings(input_ids)
		x = self.encoder(x, mask = attn_mask, src_key_padding_mask = padding_mask)
		return x

여기까지 인코더에 대한 설명을 드렸고, 그 세부적인 구현에 대해 알아보았습니다. 다음 섹션에서는 이 인코더를 MLM과 CLM 두 방법으로 학습하는 과정을 알아보겠습니다.

Masked Language Model

MLM 구현

MLM(Masked Language Model)은 시퀀스의 일부 토큰을 마스킹(Masking)하고, 해당 마스킹된 위치에서 원래 토큰을 맞추는 작업을 수행합니다. 잡음(Noise)이 섞인 시퀀스를 원래대로 복원하는 Denoising Autoencoder로도 볼 수 있습니다. 대표적인 모델은 BERT가 있고, 주로 텍스트 분류에서 강점을 보입니다.

MLM은 아래와 같은 과정으로 훈련합니다.

훈련 텍스트 시퀀스에서 일부 토큰을 임의로 선택하여 마스킹합니다. 이때 보통 15%의 토큰을 마스킹합니다.
마스킹을 할 토큰의 80%는 라는 특수 토큰이 들어갑니다. 이 토큰은 모델이 해당 위치에 있는 단어를 추론할 때 사용됩니다. 10%는 랜덤 토큰이 들어가고, 나머지 10%는 그대로 둡니다.
마스킹된 문장을 모델에 입력하고, 마스킹된 위치에 있는 단어를 추론하도록 합니다.
추론한 단어와 정답 단어 간의 오차를 계산하여 모델을 학습시킵니다.

이렇게 모델에게 마스킹된 토큰을 찾도록 훈련시키면, 모델이 토큰들 사이의 상관관계 및 문장의 문맥을 해석하는 능력을 기르게 할 수 있습니다.

MLM을 학습시키기 위해서는 위의 인코더 모델에 학습용 헤드를 붙여야 합니다. 이 헤드는 Huggingface Transformers의 BERT 구현을 참고했습니다. 이 헤드에서는 인코더의 hidden states를 받아 vocab_size 길이를 가진 벡터를 결과로 냅니다. 이 결과와 마스킹된 위치의 원래 단어와 오차를 계산해 학습을 하게 됩니다. 또, 레이어의 가중치를 초기화하는 메서드를 구현해 모델 초기화를 할 수 있도록 합니다. 이 과정은 아래와 같이 구현했습니다.

# model.py  
# https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/modeling_bert.py#L704  
class MLMHead(nn.Module):
	def __init__(self, config: Config):
		super().__init__()
		self.linear = nn.Linear(config.hidden_size, config.hidden_size)
		self.gelu = nn.GELU()
		self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
		self.decoder = nn.Linear(config.hidden_size, config.vocab_size)
		self.bias = nn.Parameter(torch.zeros(config.vocab_size))
		self.decoder.bias = self.bias

	def forward(self, x):
		x = self.linear(x)
		x = self.gelu(x)
		x = self.layer_norm(x)
		x = self.decoder(x)
		return x

class MaskedLanguageModel(nn.Module):
	def __init__(self, config: Config):
		super().__init__()
		self.config = config
		self.encoder: Encoder = Encoder(config)
		self.mlm_head = MLMHead(config)
		self.apply(self._init_weights)
		
# https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/modeling_bert.py#L748 
	def _init_weights(self, module):
		if isinstance(module, nn.Linear):
			module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
			if module.bias is not None:
				module.bias.data.zero_()
		elif isinstance(module, nn.Embedding):
			module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
			if module.padding_idx is not None:
				module.weight.data[module.padding_idx].zero_()
		elif isinstance(module, nn.LayerNorm):
			module.bias.data.zero_()
			module.weight.data.fill_(1.0)
			
	def forward(self, input_ids, padding_mask=None):
		x = self.encoder(input_ids, padding_mask=padding_mask)
		x = self.mlm_head(x)
		return x

다음은 토큰 시퀀스를 마스킹을 하고 레이블을 생성하기 위해 torch.utils.data.DataLoader의 collate_fn을 구현해 보겠습니다. collate_fn은 메서드를 입력으로 받으며, 데이터 샘플을 미니배치로 만들 때 다양한 길이의 시퀀스를 패딩 할 때 주로 쓰입니다. callate_fn을 Callable 한 클래스로 구현하면 추가 입력을 받을 수 있는데, 패딩뿐만 아니라 마스킹까지 처리하겠습니다. collate_fn 메서드는 입력으로 Dataset의 요소들의 리스트로 받습니다. MyDataset에서 토큰 시퀀스를 리턴하므로 collate_fn은 토큰 시퀀스들의 리스트를 입력으로 받습니다. 이 토큰 시퀀스를 학습하기 위해 텐서로 변환하고 길이 패딩, 패딩 마스크, 토큰 마스킹과 레이블을 생성해, 최종 데이터세트를 만듭니다.

# dataset.py
class MLMCollator:
	def __init__(self, tokenizer=None, special_token_cnt=5):
		self.tokenizer = tokenizer
		self.special_token_cnt = special_token_cnt
		
	def __call__(self, batch):
		input_ids = [torch.tensor(x) for x in batch]
		padding_mask = [get_padding_mask(x) for x in input_ids]
		input_ids = pad_sequence(input_ids, batch_first=True, padding_value=0)
		input_ids, labels = get_mask_tokens(
			input_ids, self.tokenizer, special_token_cnt=self.special_token_cnt
		)
		padding_mask = pad_sequence(padding_mask, batch_first=True, padding_value=True)
		return {
			"input_ids": input_ids,
			"labels": labels,
			"padding_mask": padding_mask,
		}

# https://github.com/huggingface/transformers/blob/main/src/transformers/data/data_collator.py#L748
def get_mask_tokens(input_ids, tokenizer, mlm_probability=0.15, special_token_cnt=5):
	input_ids = input_ids.clone()
	labels = input_ids.clone()

	probability_matrix = torch.full(labels.shape, mlm_probability)
	special_token_mask = input_ids < special_token_cnt
	probability_matrix.masked_fill_(special_token_mask, value=0.0) 
	masked_indices = torch.bernoulli(probability_matrix).bool()
	labels[~masked_indices] = -100

	indices_replaced = (
		torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices
	)
	input_ids[indices_replaced] = tokenizer.encode("").ids[0]

	indices_random = (
		torch.bernoulli(torch.full(labels.shape, 0.5)).bool()
		& masked_indices
		& ~indices_replaced
	)

	random_words = torch.randint(
		tokenizer.get_vocab_size(), labels.shape, dtype=torch.long
	)
	input_ids[indices_random] = random_words[indices_random]

	return input_ids, labels
	
def get_padding_mask(input_id):
	return torch.zeros(input_id.shape).bool()

학습

학습 코드에서는 데이터세트를 불러오고 모델을 선언합니다. 여기서 모델은 작게 설정하겠습니다. train 메서드는 train_dataloader를 돌면서 오차 계산을 하고 모델의 파라미터를 갱신합니다. 학습이 되고 있는지 확인하기 위해 간단한 메트릭을 설정하고 학습을 진행하면 됩니다.

from dataset import MyDataset, MLMCollator  
from model import Config, MaskedLanguageModel  
import torch  
from torch.utils.data import DataLoader  
from torch import nn

vocab_size = 36000  
batch_size = 256  
seq_length = 128

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  
path = "/data/wikitext-2-raw"  
train_dataset = MyDataset(path + "/wiki.train.raw", "tokenizer.json", seq_length)  
collate_fn = MLMCollator(train_dataset.tokenizer)  
train_dataloader = DataLoader(train_dataset, batch_size=64, collate_fn=collate_fn)  

config = Config(
	vocab_size=vocab_size,
	hidden_size=256,
	num_hidden_layers=4,
	num_attention_heads=4,
	intermediate_size=1024,
	max_position_embeddings=seq_length,
)
model = MaskedLanguageModel(config).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.0005)
criterion = nn.CrossEntropyLoss()

def train(model: nn.Module, epoch):
	model.train()

	total_loss = 0.0
	total_iter = 0
	
	for batch in train_dataloader:
		input_ids = batch["input_ids"].to(device)
		labels = batch["labels"].to(device)
		padding_mask = batch["padding_mask"].to(device)

		logits = model(input_ids, padding_mask=padding_mask)
		loss = criterion(logits.view(-1, vocab_size), labels.view(-1))
		loss.backward() optimizer.step()
		optimizer.zero_grad()

		total_loss += loss.item()
		total_iter += 1
		
	mean_loss = total_loss / total_iter
	print(f"epoch {epoch+1} : loss {mean_loss:1.4f}")

# 학습  
for epoch in range(10):  
	train(model, epoch)

Causal Language Model

Causal Language Model(CLM)은 시퀀스에서 이전 토큰들을 기반으로 다음 토큰을 찾는 작업을 수행합니다. 이렇게 인과적인 관계를 찾기 때문에 Causal이라 불립니다. 이 방법은 다음에 나올 토큰은 이전 토큰들에만 영향을 받는다고 가정하고 동작합니다. 이전 토큰들이 주어졌을 때 다음에 나올 토큰 중 가장 높은 확률을 선택하기 때문에 CLM으로 생성한 텍스트는 자연스럽다는 장점이 있습니다. 대표적인 모델은 GPT가 있습니다.

CLM 구현

MLM과 같이 CLM도 인코더에 붙일 헤드를 먼저 만들어야 합니다. BERT와 다르게 GPT2의 헤드는 간단하게 완전 연결층 하나로 되어 있습니다. 그리고 여기에서도 MLM처럼 모델을 초기화하는 함수를 적용해 줍니다.

# https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py#L958  
class CausalLanguageModel(nn.Module):
	def __init__(self, config: Config):
		super().__init__()
		self.config = config
		self.encoder: Encoder = Encoder(config)
		self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
		self.clm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
		self.apply(self._init_weights)
		
	# https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py#L457 
	def _init_weights(self, module):
		if isinstance(module, nn.Linear):
			module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
			if module.bias is not None:
				module.bias.data.zero_()
		elif isinstance(module, nn.Embedding):
			module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
			if module.padding_idx is not None:
				module.weight.data[module.padding_idx].zero_()
		elif isinstance(module, nn.LayerNorm):
			module.bias.data.zero_()
			module.weight.data.fill_(1.0)
	def forward(self, input_ids, attn_mask=None, padding_mask=None):
		x = self.encoder(input_ids, attn_mask=attn_mask, padding_mask=padding_mask)
		x = self.layer_norm(x)
		x = self.clm_head(x)
		return x

CLM의 collate_fn 함수는 MLM보다 간단합니다. 마스킹 없이 현재 토큰 시퀀스를 한 칸 앞으로 당겨서 레이블을 만들기만 하면 됩니다. torch.roll()을 사용해 텐서를 한 칸 앞으로 당기고 마지막 토큰()의 다음 예측을 막기 위해 -100으로 설정합니다.

class CLMCollator:
def __init__(self, tokenizer=None):
	self.tokenizer = tokenizer
	
def __call__(self, batch):
	input_ids = [torch.tensor(x) for x in batch]
	padding_mask = [get_padding_mask(x) for x in input_ids]
	
	input_ids = pad_sequence(input_ids, batch_first=True, padding_value=0)
	labels = input_ids.clone()
	labels = torch.roll(labels, -1, -1)
	labels[:-1] = -100
	padding_mask = pad_sequence(padding_mask, batch_first=True, padding_value=True)
	return {"input_ids": input_ids, "labels": labels, "padding_mask": padding_mask}

학습

MLM의 학습 코드에서 일부분을 수정하면 바로 CLM을 학습할 수 있습니다. Config에서 is_causal=True로 설정하고 collate_fn를 CLMCollator로 model을 CausalLanguageModel로 변경합니다.

from dataset import MyDataset, CLMCollator
from model import Config, CausalLanguageModel
...
collate_fn = CLMCollator(train_dataset.tokenizer)  
...
config = Config(
	vocab_size=vocab_size,
	hidden_size=256,
	num_hidden_layers=4,
	num_attention_heads=4,
	intermediate_size=1024,
	max_position_embeddings=seq_length,
	is_causal=True,
)
model = CausalLanguageModel(config).to(device)
...

맺음

이상으로 트랜스포머 인코더 모델에 MLM 및 CLM을 직접 학습시키는 방법을 알아보았습니다. 데이터를 충분히 모으고 모델의 크기를 조절해 위의 코드로 실험을 한다면 사용하기에 적당한 모델이 나옵니다. 보통 사전학습된 모델을 미세조정(Fine-tuning) 하는데 학습과 추론이 느리기 때문에, 모델의 결과를 늦게 확인할 수 밖에는 없습니다. 사전학습된 모델은 범용성을 고려한 매우 큰 파라미터를 가지는 모델이기 때문입니다.

하지만 도메인에 맞는 적당한 크기의 모델을 만들어 미세조정을 한다면, 성능은 비슷하거나 더 좋게 나올 수 있으면서 리소스는 오히려 줄어들어, 실험을 하는데 생산성이 향상됩니다.

이번 글을 통해 독자분들이 언어모델에 대한 이해가 높아지고 직접 모델을 구축하는데 도움이 되셨으면 좋겠습니다.

감사합니다.

Reference

Transformer : https://arxiv.org/abs/1706.03762
BERT : https://arxiv.org/abs/1810.04805
GPT2 : https://github.com/openai/gpt-2
PyTorch Transformer Tutorial : https://pytorch.org/tutorials/beginner/transformer_tutorial.html
Natural Language Processing with Transformers Book : https://transformersbook.com/
Huggingface Transformers BERT 구현 : https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/modeling_bert.py
Huggingface Transformers MLM 구현 : https://github.com/huggingface/transformers/blob/main/src/transformers/data/data_collator.py
Huggingface Transformers GPT2 구현 : https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py

언어모델의 원리와 만들기