[NLP] Transformer Multi-Head Attention 파이썬으로 정리

https://tiabet0929.tistory.com/77

[NLP] Transformer의 Attention Head 파이썬으로 정리

 

이 글을 작성하고 어느덧 두 달이 지나고 해가 바뀌어버렸다. 내 게으름을 탓하며 정리를 마저 하려고 한다.

 

아마도 멀티 헤드 어텐션을 정리하는 이번 포스팅이 지나고 다음 포스팅에서 피드포워드와 최종 부분을 정리하면 길었던 Transformer 정리 글을 마무리하고, 다음 단계로 넘어갈 것 같다. 마음만 굉장히 앞서 있는데 우선 1월에 최대한 힘을 내볼 생각이다.

Transformer의 Multi-Head Attention 부분

 

지난 포스팅에서 Single-Head Attention에 대해서는 다루었으니 이번 포스팅에선 Multi-Head Attention에 대해서만 집중적으로 써보려고 한다.

 

Multi-Head Attention이 Single-Head Attention에 비해 갖는 장점은 명확하다. 더 많은 문맥적 의미를 파악하게 되어 텍스트 이해를 잘 하게 된다. 

https://stackoverflow.com/questions/66244123/why-use-multi-headed-attention-in-transformers

Why use multi-headed attention in Transformers?

위 글의 답변처럼, Transformer에서 Attention Head 하나하나의 역할이 다 다르다고 이해하면 쉽다. Single-Head일 경우, 모델은 단순히 단어와 단어 사이에 어떤 관계가 있는지 학습하게 된다. 이러면 여러 문제가 있을 수 있는데, 그 중 하나가 동음이의어 혹은 다의어 문제다. 한국어든 영어든 전세계 어떤 영어든 한 단어가 여러 의미를 포함하고 있어 문장마다 다르게 사용되는 경우가 굳이 예시를 들 필요가 없을 정도로 빈번하다. Attention Head가 하나라면 많은 문장을 학습한다한들 그 한계가 명확하다. 더 많은 장점은 아래 게시글 참조.

 

https://medium.com/@hassaanidrees7/exploring-multi-head-attention-why-more-heads-are-better-than-one-006a5823372b

Exploring Multi-Head Attention: Why More Heads Are Better Than One

 

나는 AI를 말로 이해하는 것과 코드로 이해하는 영역이 다르다고 생각한다. 사실 처음엔 말로만 이해하는 걸 선호했는데, 이러고 넘어가니 막상 구체적인 질문을 받으면 대답하지 못할 때가 많았다. Multi-Head Attention을 말로 설명하자면, 하나의 헤드는 주어와 목적어의 관계를 학습하고, 또 다른 헤드는 다음 단어에 무엇이 올지 학습하고, 또 다른 헤드는 명사와 동사 사이의 관계를 파악하고, 또 어떤 헤드는 대명사의 역할을 찾는 데에 집중하고.. 이런 식으로 말을 이해하는 역할을 각각의 헤드가 하나씩 맡아 담당함으로써 모델이 문장을 더욱 잘 이해하게 된다고 정리 가능하다. AI를 이렇게 사람에 빗대어 이해하면 쉬운 경우가 많다. 우리도 분업을 하면 효율성이 향상되는 사례를 많이 봤기 때문에 이런 과정은 상당히 합리적으로 보인다.

 

그런데 이렇게 말로만 설명하면 항상 첫번째로 이해되지 않는 부분이 있다. 비전공자 혹은 AI를 잘 모르는 사람들이 많이 하는 질문이다. "아니 그럼 설계자가 사전에 Head마다 뭘 학습할 지 다 설계를 해놓은건가요?" 아니다. 절대절대 아니다. AI는 가중치들로 작동하는 딥러닝 모델이자 하나의 확률계산함수이다. (내가 좋아하는 말) 사전에 뭘 세팅하는 알고리즘이 아니다. "그러면 어떻게 Head마다 뭘 학습하는지 다 다를 수가 있어요?" 이게 바로 가중치와 학습데이터의 역할이다. 다량의 학습데이터들로 학습을 계속 시키면서 가중치들을 업데이트하다보면 자연스럽게 가중치가 저런 식으로 설계되는 것이다. 그래서 사실 윗문단에서 예시로 든 주어와 목적어의 관계나 대명사의 역할을 찾는 데에 집중한다는 것은 사실이 아닐 수도 있다. 어디까지나 예시로 든 것일 뿐, 학습데이터로부터 무엇을 각각의 Head가 학습할지는 학습이 완전히 완료될 때까지는 모른다. 

 

https://tiabet0929.tistory.com/69

[NLP] Transformer의 Attnetion 간단한 정리

예전에 쓴 글에도 Multi-Head Attention에 대해 정리한 적이 있어서, 말로 설명하는 것은 길기도 하고 여기서 줄이기로 한다. 나중에 기회가 되면 직접 저런 이미지들을 만들어서 Head마다 어떻게 단어별로 Attention 점수 차이가 있는지 확인해도 재밌을 것 같다.

 

코드로 구현하는 것과 그를 통해 이해하는 것은 또 다른 영역이다. 그래서 친절한 챗GPT가 생성해준 코드를 살펴보고자 한다.

 

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        
        assert self.head_dim * num_heads == embed_dim, "Embedding dimension must be divisible by number of heads"

        # Linear layers for Q, K, V
        self.q_linear = nn.Linear(embed_dim, embed_dim)
        self.k_linear = nn.Linear(embed_dim, embed_dim)
        self.v_linear = nn.Linear(embed_dim, embed_dim)

        # Output projection
        self.out_linear = nn.Linear(embed_dim, embed_dim)

    def forward(self, query, key, value):
        batch_size = query.size(0)

        # Linear projections for Q, K, V
        Q = self.q_linear(query)  # (batch_size, seq_len, embed_dim)
        K = self.k_linear(key)
        V = self.v_linear(value)

        # Split into multiple heads
        Q = Q.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # (batch_size, num_heads, seq_len, head_dim)
        K = K.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        V = V.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)

        # Scaled Dot-Product Attention
        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)  # (batch_size, num_heads, seq_len, seq_len)
        attn_weights = F.softmax(attn_scores, dim=-1)  # (batch_size, num_heads, seq_len, seq_len)
        attn_output = torch.matmul(attn_weights, V)  # (batch_size, num_heads, seq_len, head_dim)

        # Concatenate heads and project
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.embed_dim)  # (batch_size, seq_len, embed_dim)
        output = self.out_linear(attn_output)  # (batch_size, seq_len, embed_dim)

        return output

# Example usage
batch_size = 2
seq_len = 10
embed_dim = 512
num_heads = 4

query = torch.rand(batch_size, seq_len, embed_dim)
key = torch.rand(batch_size, seq_len, embed_dim)
value = torch.rand(batch_size, seq_len, embed_dim)

multi_head_attention = MultiHeadAttention(embed_dim, num_heads)
output = multi_head_attention(query, key, value)

print("Output shape:", output.shape)  # Should be (batch_size, seq_len, embed_dim)

 

다소 긴 코드가 생성됐다. 이는 학습은 아니고 추론하는 코드이므로 유의하면 된다.

 

Single-Head Attention 때와 특별히 다른 것은 없다.

        # Split into multiple heads
        Q = Q.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # (batch_size, num_heads, seq_len, head_dim)
        K = K.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        V = V.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)

 

이 부분만 짚고 넘어가면 될 것 같다. 모델의 Dimension (이하 줄여서 dim) 을 512, Head의 갯수를 4개로 정의해놓은 상황이다. 그럼 우선 init__  부분에서 head의 dim을 512/4 = 128로 정의한다. 즉, 각각의 Head는 128차원을 갖게 되는 것이다. 다시 말하자면, 모델의 Dimension을 Head의 갯수로 나눈 값을 Head의 dim으로 결정하면서 문장을 이해하는 역할을 나눠갖게 되는 것이다.

 

그럼 여기서 한 가지 의문이 들 수 있는데, 임베딩은 그 각자로 특별한 의미를 갖고 있는데 이렇게 마음대로 잘라버려도 되느냐하는 것이다. 이는 애초에 임베딩을 Query, Key, Value로 각각 변환하는 과정을 간과한 것이다. Wq, Wk, Wv를 통해 거쳐 나온 임베딩 벡터는 더 이상 임베딩이 아니라 쿼리와 키, 밸류이다. 따라서 임의로 차원을 나눈다고 하여 문제될 것은 없으며, 무엇보다 실험적으로 이게 더 좋은 결과를 냈다고 하니 반박할 여지가 없다. 

 

그리고 view 함수를 통해 기존에 3차원이었던 Query, Key, Value를 4차원으로 변환해준다. 위의 예시라면 기존에 Query는 (2, 10, 512)차원이었지만 Multi-Head에서는 (2, 4, 10, 128)로 바뀌는 것이다. 이후 Single-Head에서와똑같이 Attention을 계산해주면 Multi-Head Attention의 완성이다.

 

여기서 Transformer의 우수함을 엿볼 수 있는데, 모든 계산이 병렬적으로 이루어지기 때문에 계산 속도가 어마어마하게 빨라진다는 장점이 있는 것이다. 물론 GPU가 받쳐준다는 가정 하에 얘기지만. 이러한 Transformer의 구조 때문에, AI 대기업들이 더 크고 좋은 GPU에 목숨을 걸고 어떻게든 투자하는 것이라는 사실이  눈에 보인다.

 

다음 포스팅에선 마지막 차례들인 FFNN과 Add&Norm 파트, 마지막 Linear와 Softmax 부분까지 간단하게 정리하도록 하겠다.