ALBERT( A Lite BERT)는 구글에서 발표한 언어 모델로, BERT 모델의 파라미터 수를 대폭 줄이고 효율적인 학습을 가능하게 하는 기법을 사용한다.
ALBERT은 더 작은 모델 크기와 더 높은 정확도를 동시에 달성하는 것을 목표로 하고 있다.
ALBERT의 구조는 BERT 모델과 비슷하지만, 몇 가지 중요한 변경 사항이 있는데 아래와 같은 방법으로 파라미터 수를 줄이게 된다.
1. Parameter Sharing
ALBERT에서는 입력 임베딩과 출력 임베딩의 차원 수를 줄이는 것이 첫 번째 단계이다. 이를 위해, ALBERT은 두 개의 임베딩 행렬을 공유하는데, 이는 BERT와 달리 입력과 출력의 차원이 다르기 때문이다. 이 방법으로 인해 ALBERT는 BERT보다 18배 더 작은 파라미터 수로도 높은 정확도를 달성할 수 있다.
2. Cross-layer Parameter Sharing
두 번째 파라미터 공유 방법은 cross-layer parameter sharing 이다. 이 방법은 BERT의 각 레이어마다 존재하는 fully-connected layer를 하나의 레이어에서 모든 레이어로 공유한다. 이는 BERT의 fully-connected layer를 대체하고 모델의 크기를 2배로 늘리는 것을 방지할 수 있다.
3. Sentence Order Prediction(SOP) Loss
ALBERT은 BERT와 달리 Next Sentence Prediction(NSP) 태스크를 수행하지 않는다. 대신, Sentence Order Prediction(SOP) Loss를 사용한다. SOP Loss는 입력 문장에서 문장의 순서를 랜덤으로 바꿔서 이전 문장과 다음 문장을 맞추는 태스크로 이를 통해, ALBERT은 NSP 태스크보다 훨씬 더 효과적인 학습을 가능하게 한다.
이러한 변경 사항으로 ALBERT은 BERT보다 10배 이상 작은 파라미터 수로도 높은 정확도를 보인다.
다만 pre-training과 fine-tuning 프로세스는 BERT 와 동일한 프로세스를 사용하고 있다.
아래는 ALBERT 모델을 단순하게 구현해본 샘플 코드이다.
import torch
import torch.nn as nn
class MultiHeadAttention(nn.Module):
def __init__(self, embedding_size, num_heads):
super(MultiHeadAttention, self).__init__()
self.embedding_size = embedding_size
self.num_heads = num_heads
self.query_fc = nn.Linear(embedding_size, embedding_size)
self.key_fc = nn.Linear(embedding_size, embedding_size)
self.value_fc = nn.Linear(embedding_size, embedding_size)
self.softmax = nn.Softmax(dim=2)
self.dropout = nn.Dropout(0.1)
self.output_fc = nn.Linear(embedding_size, embedding_size)
def forward(self, x, mask=None):
batch_size = x.size(0)
q = self.query_fc(x).view(batch_size, -1, self.num_heads, self.embedding_size//self.num_heads).transpose(1, 2)
k = self.key_fc(x).view(batch_size, -1, self.num_heads, self.embedding_size//self.num_heads).transpose(1, 2)
v = self.value_fc(x).view(batch_size, -1, self.num_heads, self.embedding_size//self.num_heads).transpose(1, 2)
scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.embedding_size, dtype=torch.float32))
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
scores = self.softmax(scores)
scores = self.dropout(scores)
weighted_values = torch.matmul(scores, v)
weighted_values = weighted_values.transpose(1, 2).contiguous().view(batch_size, -1, self.embedding_size)
output = self.output_fc(weighted_values)
return output
class ALBERT(nn.Module):
def __init__(self, vocab_size, embedding_size, num_heads, num_layers):
super(ALBERT, self).__init__()
self.token_embedding = nn.Embedding(vocab_size, embedding_size)
self.position_embedding = nn.Embedding(512, embedding_size)
self.dropout = nn.Dropout(0.1)
self.layers = nn.ModuleList()
for i in range(num_layers):
layer = nn.TransformerEncoderLayer(embedding_size, num_heads, dim_feedforward=4*embedding_size)
self.layers.append(layer)
self.layernorm = nn.LayerNorm(embedding_size)
def forward(self, x, mask=None):
positions = torch.arange(x.size(1), device=x.device).expand(x.size(0), x.size(1)).contiguous()
x = self.token_embedding(x) + self.position_embedding(positions)
x = self.dropout(x)
for layer in self.layers:
x = layer(x)
x = self.layernorm(x)
return x
위 코드에서는 MultiHeadAttention 클래스와 ALBERT 클래스를 정의한다. ALBERT 클래스 생성자에서는 입력 단어 집합의 크기, 임베딩 크기, 어텐션 헤드 수, 레이어 수를 인자로 받는다.
클래스 내부에서는 먼저 입력 토큰과 위치 정보를 임베딩한 후, 인코더 레이어를 num_layers개 만큼 쌓아서 인코딩을 수행한다. 인코딩 결과는 LayerNorm을 거쳐 출력되며, 인코더 레이어는 TransformerEncoderLayer 클래스를 사용하여 구현한다. 이 클래스는 self-attention, multi-head attention 및 position-wise feedforward 네트워크를 포함하고 있다.
댓글