[NLP — 트랜스포머 ] Multi-head Attention #1

dot product/scaled dot product 결과 시각화 이해하기

multihead attention layer를 좀 더 구체화해보면 Linear 레이어와 Scaled Dot-Product Attention 레이어로 이루어진다.

Q. 트랜스포머의 Attention은 다음의 식으로 Attention 값을 나눠준다는 것에서 “Scaled” Dot-Product Attention이라고 불립니다. 이 Scale 과정은 어떤 의미를 가지나요? (또는, 왜 필요한가요?)

Embedding 차원 수가 깊어지면 깊어질수록 Dot-Product의 값은 커지게 되어 Softmax를 거치고 나면 미분 값이 작아지는 현상이 나타난다. 그 경우를 대비해 Scale 작업이 필요하다. (vanishing gradient 와 유사한 느낌)

차원 수가 깊어짐에 따라 Softmax 값이 작아지는 것을 방지하기 위해 Scale 과정을 포함하였거, 이 과정을 본뜬게 바로 Scaled Dot-Product Attention이다.

Q. embedding 차원은 무엇을 의미할까?

1. 워드를 토크나이징하고나서 한 레이어안에 처리하는 토큰 임베딩
2. embedding의 여러 레이어가 쌓인다고 했을 때 레이어

import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np

def make_dot_product_tensor(shape):
    A = tf.random.uniform(shape, minval=-3, maxval=3)
    B = tf.transpose(tf.random.uniform(shape, minval=-3, maxval=3), [1, 0])

    return tf.tensordot(A, B, axes=1)

length = 30
big_dim = 1024.
small_dim = 10.

big_tensor = make_dot_product_tensor((length, int(big_dim)))
scaled_big_tensor = big_tensor / tf.sqrt(big_dim)
small_tensor = make_dot_product_tensor((length, int(small_dim)))
scaled_small_tensor = small_tensor / tf.sqrt(small_dim)

# print(big_tensor,scaled_big_tensor,small_tensor,scaled_small_tensor)
print(np.min(big_tensor),np.max(big_tensor))
print(np.min(scaled_big_tensor),np.max(scaled_big_tensor))
print(np.min(small_tensor),np.max(small_tensor))
print(np.min(scaled_small_tensor),np.max(scaled_small_tensor))

fig = plt.figure(figsize=(13, 6))

ax1 = fig.add_subplot(141)
ax2 = fig.add_subplot(142)
ax3 = fig.add_subplot(143)
ax4 = fig.add_subplot(144)

ax1.set_title('1) Big Tensor')
ax2.set_title('2) Big Tensor(Scaled)')
ax3.set_title('3) Small Tensor')
ax4.set_title('4) Small Tensor(Scaled)')

ax1.imshow(tf.nn.softmax(big_tensor, axis=-1).numpy(), cmap='cividis')
ax2.imshow(tf.nn.softmax(scaled_big_tensor, axis=-1).numpy(), cmap='cividis')
ax3.imshow(tf.nn.softmax(small_tensor, axis=-1).numpy(), cmap='cividis')
ax4.imshow(tf.nn.softmax(scaled_small_tensor, axis=-1).numpy(), cmap='cividis')

plt.show()

위의 코드를 실행하면 1) dot product 후 softmax처리한 값들과 2) dot product후 scale처 후 softmax 처리한 값들을 시각화한 히트맵 그래프가 출력이 된다.

노드의 설명은 다음과 같다

모두 같은 범위인 [-3, 3]에서 랜덤 Tensor를 만들어 실제 Attention을 하듯 Dot-Product를 하고, Softmax를 취했습니다. 위 히트맵에서 어두운 부분으로는 미분 값(Gradient) 이 흐르기 어렵기 때문에 모델이 넓은 특성을 반영할 수 없게 됩니다. 즉, 히트맵이 선명할수록 모델의 시야가 편협해진다 고 이해하시면 좋아요!

위 시각화를 통해 Embedding의 깊이가 깊을수록 모델의 시야가 편협해지는 문제가 생기고[1, 3], Scale은 그 문제의 훌륭한 해결책임을 확인할 수 있습니다[2, 4]. 심지어 깊이에 무관하게 일정한 결과를 만들어내므로 어떤 경우에도 적용할 수 있는 훌륭한 Attention 기법이 탄생했음을 알 수 있습니다!

좀 나눠서 이해해보자.

선명할수록 모델의 시야가 편협해진다?

softmax의 특성을 돌이켜보면, 해당값이 0과 1일때 미분값이 0으로 학습이 원활하게 되지 않는다. 즉, 이를 미분값이 흐르기 어렵다고 해석할 수 있고, 히트맵에서 0과1의 binary한 값으로 표현된다고 해석할 수 있다. 그래서 히트맵이 선명할수록 모델의 시야가 편협해진다는 것이다.

4개의 그래프 중 가장 선명한 = 가장 모델의 시야가 편협한 그래프는 big tensor에 해당하고 이는 embedding 깊이가 깊은쪽에 해당된다.

Scale은 그 문제(임베딩 레이어가 깊어질 수록 편협해진다)의 훌륭한 해결책이다.

big tensor를 scale 처리하기 전/후와 small tensor의 scale 처리하기 전/후 를 확인해보면 명확하게 알 수 있다(그래프를 위/아래로 비교해보자). scale 처리 후가 잔 컬러라고 표현해야할까? 암튼, 다양한 컬러로 덜 선명한 히트맵이 됨을 알 수 있다.

이는 미분값이 흐른다는 의미이고 해당 컬러들은 softmax 그래프에서 미분값이 0이 아닌 y값이 0과1안의 실수 값이라고 해석 할 수 있다.