[LLM] 트랜스포머 구조 파헤치기 (4) - 인코더와 디코더

 

 

[LLM] 트랜스포머 구조 파헤치기 (1) - 텍스트를 임베딩 (Embedding)으로 변환하기

[LLM] 트랜스포머 구조 파헤치기 (2) - 어텐션 (Attention) 이해하기

[LLM] 트랜스포머 구조 파헤치기 (3) - 정규화와 피드 포워드 층

 

이제 앞서 배운 내용들을 모두 종합하여 인코더 층을 구현해 보고, 인코더 층을 합쳐 전체 인코더를 코드로 나타내 보자. 이어서 디코더 구조를 살펴보고, 인코더 대비 뭐가 다른 부분인지 살펴보자.

 

 

 

인코더 층

 

class TransformerEncoderLayer(nn.Module):
  def __init__(self, d_model, nhead, dim_feedforward, dropout):
    super().__init__()
    self.attn = MultiheadAttention(d_model, d_model, nhead) # 멀티 헤드 어텐션 클래스
    self.norm1 = nn.LayerNorm(d_model) # 층 정규화
    self.dropout1 = nn.Dropout(dropout) # 드랍아웃
    self.feed_forward = PreLayerNormFeedForward(d_model, dim_feedforward, dropout) # 피드포워드

  def forward(self, src):
    norm_x = self.norm1(src)
    attn_output = self.attn(norm_x, norm_x, norm_x)
    x = src + self.dropout1(attn_output) # 잔차 연결

    # 피드 포워드
    x = self.feed_forward(x)
    return x

 

인코더 한 층에는 멀티 헤드 어텐션, 층 정규화, 드롭아웃, 피드 포워드 층이 포함되어 있다.

 

층 정규화 → 멀티 헤드 어텐션 → 드롭아웃 → 잔차 연결 → 피드 포워드

 

순서대로 진행하면 한 인코더 층의 연산이 모두 완료된다. 위 트랜스포머 구조와 다르다고 생각할 수 있는데, 그 이유는 초기 트랜스포머의 사후 정규화가 아닌 사전 정규화를 사용했고, 멀티 헤드 어텐션과 피드포워드 층 속에 세부적인 층이나 잔차연결, 정규화가 포함된 복잡한 구조로 되어 있기 때문이다.

 

인코더

 

이제 인코더 층까지 모두 포함하여 전체 인코더 구조를 살펴보자.

 

import copy
def get_clones(module, N):
  return nn.ModuleList([copy.deepcopy(module) for i in range(N)])

class TransformerEncoder(nn.Module):
  def __init__(self, encoder_layer, num_layers):
    super().__init__()
    self.layers = get_clones(encoder_layer, num_layers)
    self.num_layers = num_layers
    self.norm = norm

  def forward(self, src):
    output = src
    for mod in self.layers:
        output = mod(output)
    return output

 

한 인코더는 여러 개의 인코더 층으로 구성되어 있다. 따라서 같은 인코더 층을 num_layers 만큼 복사하여 nn.ModuleList에 담아준다. 이어서 num_layers와 정규화 층을 저장해 준다. forward() 메서드에서 반복문을 이용하여 nn.ModuleList에 있는 모든 인코더 층을 입력 데이터가 통과하게 하여 인코더의 기능을 수행하고, 그 결과를 output으로 반환한다.

 

인과적 어텐션 연산

 

이제 이어서 디코더를 보자. 디코더에서는 인코더와 어텐션 연산이 다른데, 그 이유는 텍스트 전체의 의미를 파악하는 것이 아닌, 단어 뒤에 이어질 단어를 생성하는 데에 초점을 두기 때문이다. 뒤에 올 단어를 미리 학습에 사용할 수가 없고, 따라서 뒷부분은 마스킹 처리를 하고 연산을 해야 한다.

 

def compute_attention(querys, keys, values, is_causal=False):
	dim_k = querys.size(-1) # 16
	scores = querys @ keys.transpose(-2, -1) / sqrt(dim_k) # (1, 5, 5)
	if is_causal:
		query_length = querys.size(2)
		key_length = keys.size(2)
		temp_mask = torch.ones(query_length, key_length, dtype=torch.bool).tril(diagonal=0)
		scores = scores.masked_fill(temp_mask == False, float("-inf"))
	weights = F.softmax(scores, dim=-1) # (1, 5, 5)
	return weights @ values # (1, 5, 16)

 

인코더에서의 compute_attention() 메서드와 달라진 부분은 is_casual의 여부가 추가되었다는 점이다. 그 외에 코드는 모두 동일하다. 조건문 내에서 먼저 쿼리와 키의 길이를 가져온다. 쿼리와 키 텐서 차원이 (batch_size, n_heads, sequence_length, head_dim)로 구성되어 있으므로, 세 번째 차원에 해당하는 것이 쿼리와 키의 길이가 된다. 이어서 이 시퀀스의 길이를 바탕으로 텐서를 생성한다. 텐서는 모두 True로 이루어져 있지만, 뒤에 오는 .tril(diagonal=0)으로 인해 주대각선 아래만 True로 되어 있는 하삼각행렬이 생성된다. 이어서 temp_mask에서 False로 되어 있는 부분을 모두 -inf로 채움으로써 미래에 대한 정보를 지운다. 이 부분은 실제 softmax 연산에서 0이 되므로 어텐션 연산에서 제외됨을 알 수 있다.

 

크로스 어텐션이 포함된 디코더 층

 

디코더에서 쓰이는 어텐션 연산에 이어서 디코더 층을 살펴보자.

class TransformerDecoderLayer(nn.Module):
  def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
    super().__init__()
    self.self_attn = MultiheadAttention(d_model, d_model, nhead)
    self.multihead_attn = MultiheadAttention(d_model, d_model, nhead)
    self.feed_forward = PreLayerNormFeedForward(d_model, dim_feedforward, dropout)

    self.norm1 = nn.LayerNorm(d_model)
    self.norm2 = nn.LayerNorm(d_model)
    self.dropout1 = nn.Dropout(dropout)
    self.dropout2 = nn.Dropout(dropout)

  def forward(self, tgt, encoder_output, is_causal=True):
    # 셀프 어텐션 연산
    x = self.norm1(tgt)
    x = x + self.dropout1(self.self_attn(x, x, x, is_causal=is_causal))
    # 크로스 어텐션 연산
    x = self.norm2(x)
    x = x + self.dropout2(self.multihead_attn(x, encoder_output, encoder_output))
    # 피드 포워드 연산
    x = self.feed_forward(x)
    return x

 

대부분의 층 구성은 모두 비슷하지만, forward() 메서드 부분에 특이한 코드 한 줄이 있다. 그것은 바로 크로스 어텐션 연산이다. 크로스 어텐션 연산은 왜 있는 걸까? 그 이유는 코드를 보면 알기 쉽다. 크로스 어텐션에서 쓰이는 멀티 헤드 어텐션 메서드를 보면 쿼리에는 디코더 입력 데이터가, 키와 값에는 인코더의 output 데이터가 들어간다. 두 번째 포스팅에서 다뤘던 쿼리, 키, 값에 대한 설명을 기억하는가? 맞다. 이렇게 쿼리와 키-값 쌍 간의 관계를 학습하여, 디코더가 인코더에서 중요한 정보만 참조하게끔 한다. 이로 인해, 디코더는 출력의 각 단어가 입력의 모든 단어와 문맥적으로 연결되어 번역, 생성 등의 작업에서 보다 정확한 예측을 할 수 있다.

 

디코더

 

import copy
def get_clones(module, N):
  return nn.ModuleList([copy.deepcopy(module) for i in range(N)])

class TransformerDecoder(nn.Module):
  def __init__(self, decoder_layer, num_layers):
    super().__init__()
    self.layers = get_clones(decoder_layer, num_layers)
    self.num_layers = num_layers

  def forward(self, tgt, src):
    output = tgt
    for mod in self.layers:
        output = mod(tgt, src)
    return output

 

그래서 결국 디코더의 구조는 다음과 같다. 인코더와 마찬가지로 디코더 또한 여러 디코더 층으로 이루어져 있으므로 방식은 동일하다.