Encoder: inputs을 Embedding과 Positional Encoding 더한 값을 받는다. 이후 들어온 sequence를 이해하기 위해Self Attention (Multi Head Attention)을 수행, Encoder output은 Key, Value의 역할
Decoder: 문장을 생성하기위해, Causal Attention (Masked Multi - Head Attention) 수행 후 나온 값을 Query, Encoder의 output을 Key, Value를 Cross Attention 수행, 이후 다음 단어 예측 결과 생성
모든 time steps을 병렬적으로 처리하기 때문에, Positional Encoding이 필요하고, 병렬적으로 수행하기 때문에 자연스럽게 Teacher forcing이 적용된다.
Transformer 구조
Multi-Head Attention 구현
Query, Key, Value를 Embedding과 Positional Encoding 더한 input으로 받고 각각 linear를 통해 Q, K, V를 만들고, Head 수만큼 Split진행, 이는 여러 Self-Attention을 앙상블 하는 효과와 같다.
Scaled Dot-Product Attention 의 입력으로 위에서 만든 Q, K, V와 Mask를 준다. 이때, Mask에 따라 Self-Attention, Causal Attention, Cross Attention을 구현할 수 있다.
Self-Attention(Encoder): 입력으로 Q,K,V는 같은 inputs에서 나오고, Mask는 원본 입력에서 Padding 된 부분을 Attention 연산에 영향을 주지 않기 위해 Padding Mask를 생성해 넣어준다.
Causal Attention (Decoder): 입력으로 Q,K,V는 같은 Target에서 나오고, Mask는 Padding Mask & target에서 미래의 입력은 반영하지 않는 Mask를 사용
일반적인 Attention 매커니즘에서는 입력시퀀스의 모든 time step을 함께 계산하여 사용한다. 이는 문장 전체의 문맥을 이해하는 데 유용하지만, 미래의 정보는 아직 알 수 없어야 하기에, 실시간 처리나 예측 작업에서는 문제가 된다. 따라서 학습 과정에서도 미래의 정보는 사용하지 않음으로써 실시간 처리나 예측을 할 수 있게 한다.
Cross Attention(Decoder): 입력으로 Query는 Decoder에서 Causal Attention을 거친 값과, Key와 value는 Encoder에서 self-Attention을 거친 output값을 이용한다. Mask는 Padding Mask만을 이용한다.
Multi-Head Attention 계산수식
class MultiHeadAttention(nn.Module):
def __init__(self, model_dim, n_heads):
# model_dim = 512, num_head = 8
super(MultiHeadAttention, self).__init__()
# 모델 차원이 헤드 수로 나눠떨어져야 함 (모델 차원을 각 헤드로 나누기 위함)
assert model_dim % n_heads == 0
self.model_dim = model_dim
self.n_heads = n_heads
# 각 헤드에서 사용할 key, query, value의 차원 (512 / 8 = 64)
self.head_dim = model_dim // n_heads
# Query, Key, Value를 위한 선형 변환 레이어 정의 (입력 데이터를 임베딩 차원에서 변환)
self.query_linear = nn.Linear(model_dim, model_dim)
self.key_linear = nn.Linear(model_dim, model_dim)
self.value_linear = nn.Linear(model_dim, model_dim)
# 여러 헤드를 결합한 후 출력에 대한 선형 변환
self.output_linear = nn.Linear(model_dim, model_dim)
def scaled_dot_product_attention(self, query, key, value, mask=None):
"""
q,k,v : (batch_size, n_heads, seq_length, head_dim)
attention_scores : (batch_size, n_heads, seq_length, seq_length)
output : (batch_size, n_heads, seq_length, head_dim)
"""
# Query와 Key의 내적을 계산하여 attention score를 구하고, head_dim의 제곱근으로 나누어 스케일링
attention_scores = torch.matmul(query, key.transpose(-1,-2)) / math.sqrt(self.head_dim)
# 마스크가 주어진 경우, 마스크가 0인 부분에 매우 작은 값을 넣어 attention에서 제외
if mask is not None:
attention_scores = attention_scores.masked_fill(mask==0, -1e9)
# Softmax를 통해 각 값에 대한 확률 분포 계산 (dim=-1은 seq_length에 대해 확률화)
attention_probs = torch.softmax(attention_scores, dim= -1)
# Attention 확률과 Value를 곱해 최종 출력을 계산
output = torch.matmul(attention_probs, value)
return output
def split_heads(self, tensor):
"""
(batch_size, seq_length, model_dim) -> (batch_size, seq_length, n_heads, head_dim)
이후 transpose(1,2)로 (batch_size, n_heads, seq_length, head_dim)로 변환
"""
batch_size, seq_length, model_dim = tensor.size()
return tensor.view(batch_size, seq_length, self.n_heads, self.head_dim).transpose(1,2)
def combine_heads(self, tensor):
"""
여러 헤드에서 계산된 결과를 다시 결합
tensor: (batch_size, n_heads, seq_length, head_dim) -> (batch_size, seq_length, model_dim)
"""
batch_size, _, seq_length, head_dim = tensor.size()
# transpose로 n_heads 차원과 seq_length 차원을 바꾼 후 contiguous()로 메모리 재정렬
# (batch_size, seq_length, model_dim)로 변환
return tensor.transpos(1,2).contiguous().view(batch_size, seq_length, self.model_dim)
def forward(self, query, key, value, mask=None):
"""
q,k,v 차원: (batch_size, n_heads, seq_length, head_dim)
attention_output : (batch_size, n_heads, seq_length, head_dim)
combine_heads 후 차원: (batch_size, seq_length, model_dim)
output 차원: (batch_size, seq_length, model_dim)
"""
# query, key, value를 선형 변환 후 헤드별로 분할 (split_heads로 n_heads 차원 추가)
query = self.split_heads(self.query_linear(query))
key = self.split_heads(self.key_linear(key))
value = self.split_heads(self.value_linear(value))
# Scaled Dot-Product Attention 수행
attention_output = self.scaled_dot_product_attention(query, key, value, mask)
# 여러 헤드에서 나온 결과를 결합 후, 최종 선형 변환으로 출력
output = self.output_linear(self.combine_heads(attention_output))
return output
Positional Encoding 구현
transformer는 순차적으로 문장을 처리하는 게 아니라 병렬적으로 처리하기에 Positional Encoding이 필요하다.
위치 벡터를 얻기 위해 조건은 다른 sequence길이라도 같은 위치에 표현되는 위치벡터는 같은 값이면서, 단어 의미정보가 변질되지 않도록 위치 벡터값이 너무 커서도 안된다. 이를 해결하기 위해 Sin과 Cos를 사용한다.
Sin과 Cos는 주기함수 [-1, ~1] 값을 가지는 주기함수이며, 각 time step마다 벡터값들의 위치 차이를 위해 변 갈아 사용하여 주기를 표현한다.
sine and cosine functions of different frequencies
class PositionalEncoding(nn.Module):
def __init__(self, model_dim, max_seq_length):
super(PositionalEncoding, self).__init__()
# 위치 인코딩을 저장할 텐서를 생성
# (max_seq_length, model_dim)
pos_encodings = torch.zeros(max_seq_length, model_dim)
# 각 위치 값이 담긴 텐서를 생성 (각 시퀀스 길이의 위치를 나타냄)
# (max_seq_length, 1)
pos = torch.arange(0, max_seq_length, dtype=torch.float).unsqueeze(1)
# 차원에 따라 변환에 사용할 분모 값 계산
# (model_dim // 2,)
divisor = torch.exp(torch.arange(0, model_dim, 2).float() * -(math.log(10000.0) / model_dim))
# 짝수 인덱스에는 sin 함수를 적용
# 짝수 인덱스에서의 수식: PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
# (max_seq_length, model_dim // 2)
pos_encodings[:, 0::2] = torch.sin(pos * divisor)
# 홀수 인덱스에는 cos 함수를 적용
# 홀수 인덱스에서의 수식: PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
# (max_seq_length, model_dim // 2)
pos_encodings[:, 1::2] = torch.cos(pos * divisor)
# 계산된 위치 인코딩 값을 모델의 버퍼에 등록하여 저장 및 로딩 가능하게 함
# pos_encodings.unsqueeze(0): (1, max_seq_length, model_dim)
self.register_buffer('pos_encodings', pos_encodings.unsqueeze(0))
def forward(self, x):
# 입력 텐서와 위치 인코딩을 더함
# 입력 시퀀스의 길이에 맞추어 위치 인코딩을 슬라이싱하여 적용
"""
입력 텐서 x의 shape: (batch_size, seq_length, model_dim)
위치 인코딩의 shape: (1, max_seq_length, model_dim)
시퀀스 길이(seq_length)에 맞춰 위치 인코딩을 슬라이싱하여 더함
self.pos_encodings[:, :x.size(1)] shape: (1, seq_length, model_dim)
반환 shape: (batch_size, seq_length, model_dim)
"""
return x + self.pos_encodings[:, :x.size(1)]
입력으로 target sequence를 받고, Causal Attention을 수행한다. 이는 미래 단어를 참조하게 못하게 한다. 이 결과를 query로 사용하고, Encoder output을 key와 value를 이용하여 Cross Attention을 수행한다. 이러한 Decoder layer를 6번 거치고, 다음 단어를 예측한다.
class DecoderLayer(nn.Module):
def __init__(self, model_dim, n_heads, fc_dim, dropout):
super(DecoderLayer, self).__init__()
# Causal Attention(self_attention), Cross Attention 와 fc_layer 정의
self.self_attention = MultiHeadAttention(model_dim, n_heads)
self.cross_attention = MultiHeadAttention(model_dim, n_heads)
self.feed_forward = FeedForwardLayer(model_dim, fc_dim)
# Layer Normalization과 Dropout을 정의
self.norm1 = nn.LayerNorm(model_dim)
self.norm2 = nn.LayerNorm(model_dim)
self.norm3 = nn.LayerNorm(model_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, x, encoder_output, src_mask, tgt_mask):
"""
x: (batch_size, seq_length, model_dim) - 디코더의 입력
encoder_output: (batch_size, seq_length, model_dim) - 인코더에서 나온 출력
src_mask: (batch_size, 1, 1, seq_length) - 소스 마스크
tgt_mask: (batch_size, 1, seq_length, seq_length) - 타겟 마스크
"""
# Self-Attention 계산 (디코더에서 타겟 시퀀스가 자기 자신에 대해 어텐션을 수행)
# causal masking을 통해 미래의 단어를 보지 않도록 방지
attention_output = self.self_attention(x, x, x, tgt_mask)
# Self-Attention 결과에 입력 x를 더해 잔차 연결을 하고, 드롭아웃 및 Layer Normalization 적용
x = self.norm1(x + self.dropout(attention_output))
# Cross-Attention 계산 (인코더 출력과 디코더 입력 간의 어텐션)
# 디코더는 인코더에서 얻은 정보에 대해 어텐션을 수행
attention_output = self.cross_attention(x, encoder_output, encoder_output, src_mask)
# Cross-Attention 결과에 잔차 연결을 하고, 드롭아웃 및 Layer Normalization 적용
x = self.norm2(x + self.dropout(attention_output))
# feedforward출력에 잔차 연결을 하고, 드롭아웃 및 Layer Normalization 적용
feedforward_output = self.feed_forward(x)
x = self.norm3(x + self.dropout(feedforward_output))
return x
Transformer 모델 구현
Transformer의 디코더와 병렬 처리
Transformer는 병렬 처리를 통해 각 time step에서 다음 단어를 예측한다. 디코더에서 주어진 타겟 시퀀스가 한꺼번에 디코더 레이어에 입력되고, 이를 기반으로 전체 시퀀스에 대해 한 번에 모든 time step의 예측을 수행한다. 여기서 중요한 부분은 디코더가 미래의 단어를 참조하지 못하게 만드는 것이다(Causal Attention). 이를 통해 현재 sequence의 특정 time step에서 Decoder는 그 이전 time step의 정보만을 활용하여 다음 단어를 예측한다.
구체적으로 forward 과정에서 Decoder는 병렬적으로 모든 time step에 대한 다음 단어의 분포를 계산한다. 각 time step에서의 출력은 self.fc(decoder_output)를 통해 target vocab size 차원의 분포로 변환되고, 이 분포에서 가장 가능성이 높은 단어를 선택한다.
Teacher Forcing
Teacher Forcing은 Decoder의 입력으로 이전 time step에서 실제로 예측된 단어가 아닌, 실제 정답인 단어를 사용하는 방식이다.
Transformer에서는 디코더의 입력으로 target sequence가 그대로 들어가며, 이를 통해 병렬적으로 모든 time step에서 예측을 수행한다. 이로 인해 teacher forcing 효과가 자연스럽게 발생한다.
Transformer 구조
class Transformer(nn.Module):
def __init__(self, src_vocab_size, tgt_vocab_size, model_dim, n_heads, num_layers, fc_dim, max_seq_length, dropout):
super(Transformer, self).__init__()
# 입력과 출력 시퀀스에 대해 각각의 Embedding 레이어와 Positional Encoding 정의
self.encoder_embedding = nn.Embedding(src_vocab_size, model_dim)
self.decoder_embedding = nn.Embedding(tgt_vocab_size, model_dim)
self.positional_encoding = PositionalEncoding(model_dim, max_seq_length)
# Encoder와 Decoder 레이어 정의 (각 num_layers만큼 레이어 생성)
self.encoder_layers = nn.ModuleList([EncoderLayer(model_dim, n_heads, fc_dim, dropout) for _ in range(num_layers)])
self.decoder_layers = nn.ModuleList([DecoderLayer(model_dim, n_heads, fc_dim, dropout) for _ in range(num_layers)])
# Decoder 최종 출력에서 vocab 크기만큼의 차원을 만드는 선형 레이어
self.fc = nn.Linear(model_dim, tgt_vocab_size)
self.dropout = nn.Dropout(dropout)
# 마스크를 생성
def generate_mask(self, src, tgt):
"""
src: (batch, max_seq_len)
tgt: (batch, max_seq_len) # tgt에서는 max_seq_len = max_seq_len - 1 임
src_mask shape: (batch, 1, 1, max_seq_len)
tgt_mask shape: (batch, 1, max_seq_len, 1) # 브로드 캐스팅연산할때 1이 max_seq_len으로 확장된다.
nopeak_mask shape: (1, max_seq_len, max_seq_len)
tgt_mask shape: (batch, 1, max_seq_len, max_seq_len)
"""
# 소스와 타겟에서 패딩(0)이 아닌 위치를 마스크로 생성 (소스 패딩 무시)
# attention 스코어와 연산을 할 수 있게 하기 위해, unsqueeze를 사용하여 차원을 추가 -> 브로드 캐스팅
src_mask = (src != 0).unsqueeze(1).unsqueeze(2).to(device)
tgt_mask = (tgt != 0).unsqueeze(1).unsqueeze(3).to(device)
# 타겟의 시퀀스 길이
seq_length = tgt.size(1)
# nopeak_mask는 디코더의 self-attention에서 미래의 단어를 참조하지 못하게 하는 마스크
# 대각선 위쪽을 0으로 설정하여 미래 정보 참조 금지
nopeak_mask = (1 - torch.triu(torch.ones(1, seq_length, seq_length), diagonal=1)).bool().to(device)
# 패딩 마스크와 nopeak_mask를 결합하여 최종 tgt_mask 생성 -> 패딩과 미래의 단어를 참조하지 못하게 함
tgt_mask = tgt_mask & nopeak_mask
# 소스 마스크와 타겟 마스크를 반환
return src_mask, tgt_mask
# 순전파 정의
def forward(self, src, tgt):
"""
src: (batch_size, max_seq_len)
tgt: (batch_size, max_seq_len)
src_embedded: (batch_size, max_seq_len, model_dim)
tgt_embedded: (batch_size, max_seq_len, model_dim)
encoder_output 초기값: (batch_size, max_seq_len, model_dim)
각 encoder_layer를 통과한 후의 encoder_output: (batch_size, max_seq_len, model_dim)
decoder_output 초기값: (batch_size, max_seq_len, model_dim)
output: (batch_size, max_seq_len, tgt_vocab_size)
"""
# 소스와 타겟 시퀀스에 대한 마스크 생성
src_mask, tgt_mask = self.generate_mask(src, tgt)
# 소스와 타겟에 각각 Embedding과 Positional Encoding을 적용
src_embedded = self.dropout(self.positional_encoding(self.encoder_embedding(src)))
tgt_embedded = self.dropout(self.positional_encoding(self.decoder_embedding(tgt)))
# 인코더 레이어를 순차적으로 통과 (인코더의 각 레이어에 소스 입력과 마스크 전달)
encoder_output = src_embedded # 초기값: 임베딩된 소스 입력
for encoder_layer in self.encoder_layers:
encoder_output = encoder_layer(encoder_output, src_mask)
# 디코더 레이어를 순차적으로 통과 (디코더의 각 레이어에 타겟 입력, 인코더 출력, 마스크 전달)
decoder_output = tgt_embedded # 초기값: 임베딩된 타겟 입력
for decoder_layer in self.decoder_layers:
decoder_output = decoder_layer(decoder_output, encoder_output, src_mask, tgt_mask)
# 디코더 출력에 대해 최종적으로 Linear 레이어를 통해 타겟 어휘 크기만큼의 분포로 변환
output = self.fc(decoder_output)
return output
