[19-2] Multi-head Attention Using PyTorch

1️⃣ 설정

 

필요한 모듈을 import 합니다.

from torch import nn

from torch.nn import functional as F

from tqdm import tqdm


import torch

import math

 

2️⃣ 데이터 전처리

 

Data를 생성합니다.

pad_id = 0

vocab_size = 100


data = [
  [62, 13, 47, 39, 78, 33, 56, 13, 39, 29, 44, 86, 71, 36, 18, 75],
  [60, 96, 51, 32, 90],
  [35, 45, 48, 65, 91, 99, 92, 10, 3, 21, 54],
  [75, 51],
  [66, 88, 98, 47],
  [21, 39, 10, 64, 21],
  [98],
  [77, 65, 51, 77, 19, 15, 35, 19, 23, 97, 50, 46, 53, 42, 45, 91, 66, 3, 43, 10],
  [70, 64, 98, 25, 99, 53, 4, 13, 69, 62, 66, 76, 15, 75, 45, 34],
  [20, 64, 81, 35, 76, 85, 1, 62, 8, 45, 99, 77, 19, 43]
]

padding을 수행합니다.

def padding(data):        # 0을 넣어서 vector들의 length를 맞춰줌

  max_len = len(max(data, key=len))
  
  print(f"Maximum sequence length: {max_len}")
  

  for i, seq in enumerate(tqdm(data)):
  
    if len(seq) < max_len:
    
      data[i] = seq + [pad_id] * (max_len - len(seq))
      

  return data, max_len

data, max_len = padding(data)

>>> data

[[62, 13, 47, 39, 78, 33, 56, 13, 39, 29, 44, 86, 71, 36, 18, 75, 0, 0, 0, 0],
 [60, 96, 51, 32, 90, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
 [35, 45, 48, 65, 91, 99, 92, 10, 3, 21, 54, 0, 0, 0, 0, 0, 0, 0, 0, 0],
 [75, 51, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
 [66, 88, 98, 47, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
 [21, 39, 10, 64, 21, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
 [98, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
 [77,
  65,
  51,
  77,
  19,
  15,
  35,
  19,
  23,
  97,
  50,
  46,
  53,
  42,
  45,
  91,
  66,
  3,
  43,
  10],
 [70, 64, 98, 25, 99, 53, 4, 13, 69, 62, 66, 76, 15, 75, 45, 34, 0, 0, 0, 0],
 [20, 64, 81, 35, 76, 85, 1, 62, 8, 45, 99, 77, 19, 43, 0, 0, 0, 0, 0, 0]]

 

3️⃣ Hyperparameter 세팅 및 embedding

 

Model의 hidden state size $d_{model}$과 head의 개수를 정해줍니다.

d_model = 512  # model의 hidden size

num_heads = 8  # head의 개수

# d_model의 num_heads로 나눠 떨어져야함

vocab_size로 입력을 받아 $d_{model}$ size로 출력하는 embedding을 만들어줍니다.

embedding = nn.Embedding(vocab_size, d_model)


# B: batch size, L: maximum sequence length

batch = torch.LongTensor(data)  # (B, L)

batch_emb = embedding(batch)  # (B, L, d_model)

batch_emb의 shape을 찍어보면 아래와 같습니다.

>>> print(batch_emb.shape)

torch.Size([10, 20, 512])

 

4️⃣ Linear transformation & 여러 head로 나누기

 

Multi-head attention 내에서 쓰이는 linear transformation matrix들을 정의합니다.

w_q = nn.Linear(d_model, d_model)

w_k = nn.Linear(d_model, d_model)

w_v = nn.Linear(d_model, d_model)

w_0 = nn.Linear(d_model, d_model)

q = w_q(batch_emb)  # (B, L, d_model)
k = w_k(batch_emb)  # (B, L, d_model)
v = w_v(batch_emb)  # (B, L, d_model)


>>> print(q.shape)

torch.Size([10, 20, 512])


>>> print(k.shape)

torch.Size([10, 20, 512])


>>> print(v.shape)

torch.Size([10, 20, 512])

Q, k, v를 num_head개의 차원 분할된 여러 vector로 만듭니다.

batch_size = q.shape[0]

d_k = d_model // num_heads      # d_model // num_heads로 나눠줌 => 따라서 나눠 떨어지게 만들어야함


# dimension이 하나의 feature을 나타내기 때문에 multi-headed를 구현할 때 아래 같이 구현이 가능함

# d_model을 num_heads와 d_k로 나눠줌

q = q.view(batch_size, -1, num_heads, d_k)  # (B, L, num_heads, d_k)

k = k.view(batch_size, -1, num_heads, d_k)  # (B, L, num_heads, d_k)

v = v.view(batch_size, -1, num_heads, d_k)  # (B, L, num_heads, d_k)


>>> print(q.shape)

torch.Size([10, 20, 8, 64])


>>> print(k.shape)

torch.Size([10, 20, 8, 64])


>>> print(v.shape)

torch.Size([10, 20, 8, 64])

---

✅ 왜 multi-head를 병렬적으로 처리하는 것이 아닌 $d_{model}$ // num_heads로 나눠주나요?

 

원래 transform의 multi-head attention은 아래 그림과 같이 각각의 W_q, W_k, W_v를 구해서 concat한 후 W_0에서 내적을 하는 방법을 사용합니다.

 

https://www.edwith.org/bcaitech1

하지만 실제 구현된 multi-head attention은 $d_{model}$ / head = $d_k$로 사용합니다.

 

원래 multi-headed attention과 달라 보이지만 결국에는 같은 의미를 나타냅니다.

Hidden vector라 하는 것은 여러 dimension으로 구성되어 있고, 각각의 차원들은 특정 feature를 나타냅니다.

예를 들어, 첫 번째 차원은 category feature, 두 번째 차원은 품사 feature 등을 가질 수 있습니다.

 

즉, 여러 가지 정보를 각기 다른 차원에 가지고 있기 때문에 head를 나눠서 self-attention을 수행하게 되면, 어차피 다른 정보를 focusing 해서 multi-headed attention을 수행하고 있는 것이기 때문에 한 vector를 보더라도 다양한 관점에서 여러 번 수행하는 것과 같은 효과를 냅니다. 따라서 attention을 수행한 후 차원을 다시 합치게 되면, 결국 각기 다른 의미를 focusing 하는 attention을 수행하고 합치는 기존의 multi-headed attention과 의미가 동일해집니다.

 

https://www.edwith.org/bcaitech1

---

이후, (B, L, num_heads, d_k) shape에서 transpose(1, 2)를 통해 num_heads를 밖으로 빼줍니다.

이렇게 (B, num_heads, L, d_k) 구조를 만들면 num_heads에 대해 L x d_k 만큼의 matirx를 가지고 self-attention을 수행하는 것처럼 구조를 만들 수 있습니다.

# head를 밖으로 빼줌(transpose)으로써, 각 head가 L x d_k 만큼의 matirx를 가지고 self-attention을 수행

q = q.transpose(1, 2)  # (B, num_heads, L, d_k)

k = k.transpose(1, 2)  # (B, num_heads, L, d_k)

v = v.transpose(1, 2)  # (B, num_heads, L, d_k)


>>> print(q.shape)

torch.Size([10, 8, 20, 64])


>>> print(k.shape)

torch.Size([10, 8, 20, 64])


>>> print(v.shape)

torch.Size([10, 8, 20, 64])

 

 

 

5️⃣ Scaled dot-product self-attention 구현

 

각 head에서 실행되는 self-attetion 과정입니다.

attn_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)  # (B, num_heads, L, L)

attn_dists = F.softmax(attn_scores, dim=-1)  # (B, num_heads, L, L)


>>> print(attn_dists.shape)

torch.Size([10, 8, 20, 20])

attn_values = torch.matmul(attn_dists, v)  # (B, num_heads, L, d_k)


>>> print(attn_values.shape)

torch.Size([10, 8, 20, 64])

 

6️⃣ 각 head의 결과물 병합

 

각 head의 결과물을 concat하고 동일 차원으로 linear transformation합니다.

attn_values = attn_values.transpose(1, 2)  # (B, L, num_heads, d_k)

attn_values = attn_values.contiguous().view(batch_size, -1, d_model)  # (B, L, d_model)


>>> print(attn_values.shape)

torch.Size([10, 20, 512])

---

✅ contiguous()란?

 

일단 contiguous()를 알기 위해선 stride와 transpose가 어떻게 동작하는지 알아야합니다.

a라는 [5, 10, 20] size의 vector가 존재한다고 가정합니다.

a = torch.zeros(3, 5, 10, 20)


>>> print(a.shape)

torch.Size([3, 5, 10, 20])

 

• stride() : 접근 순서에 대한 가지고 있는 원소의 개수

>>> a.stride()

(1000, 200, 20, 1)

# 3 : 5 * 10 * 20 = 1000

# 5 : 10 * 20 = 200

# 10 : 20 = 20

# 20 : 1 = 1 (마지막 vector에 접근했을 때, 하나의 원소들로만 구성)

 

• transpose() : tensor 모양 변경 

>>> a.transpose(0,1).shape

torch.Size([5, 3, 10, 20])

 

하지만 여기서 중요한 점은 transpose()는 보여지는 형식만 달라지는 것이지 구조가 바뀌는 것은 아닙니다.

즉, transpose를 수행한 후 view를 수행하면 오류가 납니다.

a = a.transpose(0,1)


>>> a.view(-1, 200)

RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.

이 것을 stride로 찍어보면 다음과 같은 구조로 나옵니다.

>>> a.stride()

(200, 1000, 20, 1)

 

즉, 이것이 말하는 의미는 데이터에 접근할 때 접근 순서만 달라진 것 뿐이지, 구조가 바뀌지 않았다는 것을 말하는 것입니다. 원래 stride는 다음과 같이 나와야합니다.

b = torch.zeros(5, 3, 10, 20)


>>> b.stride()

(600, 200, 20, 1)

따라서 transpose를 사용한 후, 데이터 구조를 완전히 바꾸고 싶다면 contiguous()를 사용해야합니다.

a = a.contiguous()


>>> a.shape

torch.Size([5, 3, 10, 20])

contiguous()를 사용하면 view와 stride()가 정상적으로 작동하는 것을 알 수 있습니다.

a = a.view(-1, 200)


>>> a.shape

torch.Size([15, 200])

>>> a.stride()

(600, 200, 20, 1)

---

이제 outputs을 출력하면 아래와 같습니다.

outputs = w_0(attn_values)


>>> print(outputs.shape)

torch.Size([10, 20, 512])

 

7️⃣ 전체 코드

 

위의 과정을 모두 합쳐 하나의 Multi-head attention 모듈을 구현하겠습니다.

class MultiheadAttention(nn.Module):

  def __init__(self):
  
    super(MultiheadAttention, self).__init__()
    

    # Q, K, V learnable matrices
    
    self.w_q = nn.Linear(d_model, d_model)
    
    self.w_k = nn.Linear(d_model, d_model)
    
    self.w_v = nn.Linear(d_model, d_model)
    

    # Linear transformation for concatenated outputs
    
    self.w_0 = nn.Linear(d_model, d_model)


  def forward(self, q, k, v):
  
    batch_size = q.shape[0]
    

    q = self.w_q(q)  # (B, L, d_model)
    
    k = self.w_k(k)  # (B, L, d_model)
    
    v = self.w_v(v)  # (B, L, d_model)
    

    q = q.view(batch_size, -1, num_heads, d_k)  # (B, L, num_heads, d_k)
    
    k = k.view(batch_size, -1, num_heads, d_k)  # (B, L, num_heads, d_k)
    
    v = v.view(batch_size, -1, num_heads, d_k)  # (B, L, num_heads, d_k)
    

    q = q.transpose(1, 2)  # (B, num_heads, L, d_k)
    
    k = k.transpose(1, 2)  # (B, num_heads, L, d_k)
    
    v = v.transpose(1, 2)  # (B, num_heads, L, d_k)
    

    attn_values = self.self_attention(q, k, v)  # (B, num_heads, L, d_k)
    
    attn_values = attn_values.transpose(1, 2).contiguous().view(batch_size, -1, d_model)  # (B, L, num_heads, d_k) => (B, L, d_model)
    

    return self.w_0(attn_values)
    

  def self_attention(self, q, k, v):
  
    attn_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)  # (B, num_heads, L, L)
    
    attn_dists = F.softmax(attn_scores, dim=-1)  # (B, num_heads, L, L)
    

    attn_values = torch.matmul(attn_dists, v)  # (B, num_heads, L, d_k)
    

    return attn_values

multihead_attn = MultiheadAttention()

outputs = multihead_attn(batch_emb, batch_emb, batch_emb)  # (B, L, d_model)