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

 

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

 
지난 시간에는 텍스트를 어떻게 임베딩으로 변환하는지에 대해 알아보았다. 이번 시간에는 임베딩 벡터를 바탕으로 단어들이 서로를 어떻게 이해하고, 그 속에서 어텐션 메커니즘을 통해 관련성을 어떻게 만들어나가는지에 대해 알아보자.

 

쿼리, 키, 값 (Query, Key, Value)

 
데이터베이스를 다루는 사람이라면 쿼리라는 단어가 익숙할 것이다. 요청, 질문, 등등 다양한 단어로 표현할 수 있지만 왠지 모르게 쿼리라는 말이 더 와닿는 것 같다. 포털 사이트 검색을 생각해 보면, 우리는 쿼리를 던지고, 시스템은 관련 있는 키를 뱉어낸다. 그리고 그 키를 통해 값에 접근할 수 있다. 어텐션의 방식도 마찬가지다! 쿼리에 대한 가장 관련도 높은 키와 그 값을 찾는 방식처럼, 어텐션 또한 단어와 단어 사이의 관계를 계산하고 그 값에 따라 관련이 깊은 단어와 그렇지 않은 단어를 구분한다.
 

셀프 어텐션

 
어텐션 메커니즘을 하기 위해서는 쿼리, 키, 값에 대해 앞서 기술했던 토큰 임베딩을 통해 벡터로 변환해야 한다. 그리고 각 토큰 간 관련도를 계산하기 위해 가중치라는 개념을 도입해야 한다. 이게 갑자기 뭔 뚱딴지같은 소리인가? 사실 단순하다. 가중치라는 건 딥러닝에서 어떤 정확도를 개선하고 싶어 학습을 하며 계속 업데이트되는 녀석이다. 즉, 우리는 쿼리, 키, 값을 두고 학습을 통해 토큰 간 관련도를 개선해 나가겠다는 것이다.
 

head_dim = 16

# 쿼리, 키, 값을 계산하기 위한 변환
weight_q = nn.Linear(embedding_dim, head_dim)
weight_k = nn.Linear(embedding_dim, head_dim)
weight_v = nn.Linear(embedding_dim, head_dim)
# 변환 수행
querys = weight_q(input_embeddings) # (1, 5, 16)
keys = weight_k(input_embeddings) # (1, 5, 16)
values = weight_v(input_embeddings) # (1, 5, 16)

 
Pytorch가 있어서 참 다행이다. 우리는 그저 차원을 설정해 주고 각 단어 임베딩 벡터를 쿼리, 키, 값으로 변환하여 어텐션 계산의 준비 단계로 들어갈 수 있다. nn.Linear 는 완전 연결층을 만드는 Pytorch의 클래스이다.
 

from math import sqrt
import torch.nn.functional as F

def compute_attention(querys, keys, values, is_causal=False):
	dim_k = querys.size(-1) # 16
	scores = querys @ keys.transpose(-2, -1) / sqrt(dim_k)
	weights = F.softmax(scores, dim=-1)
	return weights @ values

 
이제 어텐션 계산을 어떻게 하는지 보자. 일반적으로 스케일 점곱 방식을 사용하는데, compute_attention의 코드를 한 줄씩 살펴보자. 이해를 위해 예시를 들어본다.
 

ex) query, key, value 텐서의 차원 (batch_size, sequence_length, head_dim) = (2, 5, 16)이라고 하자.

 

1. 먼저 쿼리의 마지막 차원을 가져온다. 16이 해당된다.

 

2. 이어서 scores는 유사도 점수를 의미한다. 문제는 쿼리와 키가 차원이 똑같기 때문에 행렬 곱 연산이 안된다는 것인데, 행렬 곱 연산을 위해 키의 마지막 두 차원을 바꿔줬다. 2차원 행렬 곱에서 앞선 행렬의 두 번째 차원과 뒷행렬의 첫 번째 차원이 같아야 하는 것과 같은 원리다. 따라서 (2, 5, 16)과 (2, 16, 5)가 곱해져서 scores의 차원은 (2, 5, 5)가 된다. (배치 크기는 행렬 곱 연산에 이용되지 않는다)

 

3. 계산된 score에 softmax 함수를 적용하여 각 단어가 다른 단어들과의 유사도를 합해 1이 되는 확률 분포로 변환된다. 이 값이 가중치이고, 차원은 여전히 (2, 5, 5)이다.

 

4. 마지막으로 가중치에 값(Value)을 내적해 준다. 값은 (2, 5, 16)이므로 연산 결과 다시 (2, 5, 16)의 차원으로 반환이 완료된다.

 
결국 어텐션 연산이란 것은 쿼리와 키의 스케일 점곱, 가중치 변환, 그리고 그 결과를 값 (value)과 내적 하는 것이라 보면 된다.
 

class AttentionHead(nn.Module):
  def __init__(self, token_embed_dim, head_dim, is_causal=False):
    super().__init__()
    self.is_causal = is_causal
    self.weight_q = nn.Linear(token_embed_dim, head_dim) # 쿼리 벡터 생성을 위한 선형 층
    self.weight_k = nn.Linear(token_embed_dim, head_dim) # 키 벡터 생성을 위한 선형 층
    self.weight_v = nn.Linear(token_embed_dim, head_dim) # 값 벡터 생성을 위한 선형 층

  def forward(self, querys, keys, values):
    outputs = compute_attention(
        self.weight_q(querys),  # 쿼리 벡터
        self.weight_k(keys),    # 키 벡터
        self.weight_v(values),  # 값 벡터
        is_causal=self.is_causal
    )
    return outputs

 
이를 클래스로 나타내면 위와 같다. 우리는 쿼리, 키, 값을 임베딩 벡터로 변환한 후, compute_attention 함수에 집어넣어 어텐션 연산을 수행하게 되는 것이다. is_causal은 인과적 어텐션을 적용할지 여부를 결정하는 것인데 이는 추후 다룰 디코딩 부분에서 다시 설명하겠다.
 

멀티 헤드 어텐션

 
한 번에 하나의 어텐션 연산을 수행하는 것보다는 여러 어텐션 연산을 동시에 하는 것이 성능을 더 높일 수 있을 것이다. 왜냐하면, 토큰 사이의 관계를 여러 측면에서 고려할 수 있기 때문이다. 따라서 기존의 셀프 어텐션 방식에서 (헤드를 여러 개 둠으로써) 차원을 확장시켜 계산 후에 다시 원래 차원으로 변환시켜 주는, 여러 헤드를 이용하는 멀티 헤드 어텐션을 코드로 이해해 보자.
 

self.n_head = n_head
self.weight_q = nn.Linear(token_embed_dim, d_model)
self.weight_k = nn.Linear(token_embed_dim, d_model)
self.weight_v = nn.Linear(token_embed_dim, d_model)
self.concat_linear = nn.Linear(d_model, d_model)

 
우리는 먼저 d_model의 차원으로 확장하는 선형층을 선언해 준다. 그리고 각 헤드의 출력을 병합한 후에 최종 출력을 생성하는 선형 변환 층 concat_linear도 선언한다.
 

querys = self.weight_q(querys).view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
keys = self.weight_k(keys).view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
values = self.weight_v(values).view(B, T, self.n_head, C // self.n_head).transpose(1, 2)

 
앞서 셀프 어텐션에서 살펴봤었던 차원은 (2, 5, 16)이었다. 만약 헤드가 4개라면 선형층 통과 후에 차원으로 (2, 5, 4, 4)로 변경해 주고, 이 상태에서 두 번째, 세 번째 차원을 바꿔준다. (2, 4, 5, 4)가 되는 것이다.

 

행렬 곱 연산은 어텐션 연산 함수에 구현해 뒀으니 그 문제는 아닌 거 같고, 그렇다면 왜 차원을 바꿔줄까? 먼저 행렬 곱은 마지막 두 차원, 그러니까 세번째 네번째 차원에 대해서만 수행된다. 우리가 헤드의 개수로 차원을 나누었다는 것은 각 헤드별로 독립적인 연산을 가능하게 하기 위해서다. 따라서 각 헤드의 차원 4를 앞으로 배치함으로써 연산에 필요한 차원을 뒤쪽으로 옮기는 것이다.

 

요약하자면, 지금 단계에서 하는 transpose는 헤드별로 독립적인 연산을 하기 위해서이고, 어텐션 연산에서 마지막 두 차원을 transpose하는 것은 그래야 차원이 맞아 행렬 곱 연산이 가능하기 때문이다.

 

attention = compute_attention(querys, keys, values, self.is_causal)

 
이렇게 변환된 키, 쿼리, 값에 어텐션 연산을 적용해 주고,
 

output = attention.transpose(1, 2).contiguous().view(B, T, C)
output = self.concat_linear(output)

 
나온 결과의 차원을 다시 원래대로 되돌린 후에 concat_linear를 통해 다시 학습 가능한 가중치를 적용한다. 그런데 output은 이미 (2, 5, 16) 차원인데 굳이 concat_linear를 한번 더 사용하는 이유가 뭔지 궁금할 것이다. 이는 모델이 concat_linear 가중치를 학습하면서, 각 헤드에서 추출한 정보를 효과적으로 결합하여 최종 출력을 조정할 수 있기 때문이란다. 결국 하고 싶은 말은 학습 한 번 더 하면 좋다는 소리다.
 

class MultiheadAttention(nn.Module):
  def __init__(self, token_embed_dim, d_model, n_head, is_causal=False):
    super().__init__()
    self.n_head = n_head
    self.is_causal = is_causal
    self.weight_q = nn.Linear(token_embed_dim, d_model)
    self.weight_k = nn.Linear(token_embed_dim, d_model)
    self.weight_v = nn.Linear(token_embed_dim, d_model)
    self.concat_linear = nn.Linear(d_model, d_model)

  def forward(self, querys, keys, values):
    B, T, C = querys.size()
    querys = self.weight_q(querys).view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
    keys = self.weight_k(keys).view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
    values = self.weight_v(values).view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
    attention = compute_attention(querys, keys, values, self.is_causal)
    output = attention.transpose(1, 2).contiguous().view(B, T, C)
    output = self.concat_linear(output)
    return output

n_head = 4
mh_attention = MultiheadAttention(embedding_dim, embedding_dim, n_head)
after_attention_embeddings = mh_attention(input_embeddings, input_embeddings, input_embeddings)
after_attention_embeddings.shape

 
다음 글에서는 층 정규화와 피드 포워드 층에 대해 살펴보겠다. 이번 글도 마찬가지로 설명이 부족한 부분이 있다면 조금씩 살을 붙여보도록 하겠다.

 

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

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