cs231n Assignment 3: Q2-2 (Multi-head Transformer 구현)

내 풀이 링크: https://github.com/lionkingchuchu/cs231n.git

GitHub - lionkingchuchu/cs231n: cs231n Spring 2022 Assignment

이번 과제는 Transformer을 사용한 Captioning 구현이다. Transformer는 지난번에 배운 self attention에서 데이터를 나누어서 attention layer를 적용시키는 multihead self attention을 사용한다. [1706.03762] Attention Is All You Need (arxiv.org) 을 참고하면 과제에 도움을 얻을 수 있다. 사진은 모두 논문에서 나온 사진이다.

Multi-Head Attention

위는 Transformer 모델의 전반적인 architecture이다. Multihead Attention, 지난 시간에 배운 Positional Encoding을 사용하고 잘 보면 MultiHead Attention layer 의 출력 데이터와 이전 데이터를 Add&Norm 하는 것을 볼 수 있다. 이렇게 이전 데이터와 layer 이후 데이터를 이어주는 것을 residual connection 이라고 한다. Norm은 Layer Norm을 이용했다고 한다. 또한 Feed Forward Network에서도 residual connection과 layer norm을 사용했다. Feed Forward Network로는 2개의 Linear layer 사이에 ReLU layer를 끼운 layer를 사용했다고 한다 ( FFN(x) = max(0, xW1 + b1)W2 + b2 ). 밑에 architecture에는 보이지 않지만 매 sublayer마다 ratio 0.1 의 residual dropout 을 사용해서 regularize 했다고 한다.

그럼 본격적으로 Transformer를 이용한 Image Captioning 과제를 해보겠다. 데이터로는 저번처럼 COCO datasets를 사용한다. 처음으로 MultiHeadAttention의 init 함수이다. 

MultiHeadAttention init

key, query, value의 FCLayer를 한개의 Linear layer로 사용하여 미리 선언해 둔것을 볼 수 있다. proj Linear layer은 multiple head attention에서 나온 여러 head의 attention 값들을 Linear transform 해주기 위한 layer이다. attn_drop은 dropout을 위한 layer이고 n_head는 multihead attention에서 사용할 head의 개수, head_dim 은 head의 개수에 따른 dimension 이 될 것이다.

 

다음은 MultiHeadAttention의 forward 구현인데, 과제에서 친절하게 과정의 수식에 대해 설명이 되어있다.

MultiHead Scaled Dot-Prodcut Attention 수식 설명

forward layer 구현이다. 먼저 key, query, value를 각각 위에서의 Linear layer를 통과시켜 주고, head의 수에 맞게 key, query, value의 E 차원을 H로 나누어 reshape 해준다. 다음으로 query와 key(2,3 dim Transpose) 를 matmul 해주어 xqxk (N, H, S, T) 를 얻는다. key의 2,3 dim을 바꿔주는 이유는 N, H 차원은 데이터 번호와 Head 번호의 차원이므로, 같은 N과 H차원에서 일어나는 계산이므로 0, 1차원은 바꾸면 안된다. 그리고 scaling을 위해 sqrt(E/H)를 나눠준다. 여기서 attn_mask가 있다면 저번 시간에 배운 attention에서 sequece의 order를 반영시켜주기 위해서는 query가 미래의 key를 반영하지 못하게 attn_mask의 값을 -inf 로 바꾸어주어 softmax를 통과시켰을때 0의 값이 나올 수 있도록 한다.

 

이제 softmax layer를 통과시켜 주고, dropout layer를 통과시켜 주고, 이제 value와 곱해주어 최종 out을 얻는다. 이때 최종 out의 차원은 (N, H, S, E//H)가 된다. Linear Transformation을 위한 self.proj (E, E) layer를 통과시키기 위해서, out을 concatenate 해주어 차원을 (N, S, E) 로 바꿔주어야 한다. H 와 E//H 차원이 서로 떨어져 있으므로 H와 S차원을 swap 해주고 (N, S, H, E//H), out을 이제 (N, S, E)로 reshape 해주고 self.proj layer를 통과시켜주면 결과를 얻을 수 있다.

잘 구현했다면 e-3 미만의 오차를 얻을 수 있다.

다음으로 Positional Encodning layer 구현이다. 여기서도 과제에서 친절하게 Positonal Encoding layer의 P를 어떻게 만들어야 할 지 수식을 통해 설명해 준다. 아래의 식을 통해 P에 각 행(i), 열(j)마다 다른 값을 저장하는 방식으로 Positional Encoding 정보가 담긴 P를 input (X) 에 더해주어서 기존의 X에 위치 정보를 추가시켜 준다.

Positional Encoding 수식

Positional Encoding P 구현

먼저 row vector를 torch.arange 로 만들어 주고, col 은 TODO의 힌트처럼 홀수, 짝수에 따라 (0, 0, 2, 2, 4, 4..) 이런 방식으로 저장된 벡터를 만들어준다. 그리고 col에 10000 ^ (- (j/d)) 함수를 취해주고, row 와 col을 곱해주기 위해 col은 col vector로 만들어준다. 그리고 pe 행렬 전체에 에 row * col을 할당 해주면 (i * 10000 ^ (j/d)) 까지 반영이 되었을 것이다. 이제 짝수 col에는 sin함수, 홀수 col에는 cos 함수를 반영시켜 저장해 주면 Postional Encoding 정보가 담긴 P 행렬을 만들 수 있다. 

forward 구현 (X + P) -> dropout

이제 위에서 구한 P를 통해 forward 함수 구현이다. x에 위에서 구한 pe를 더하는데 x의 input sequence (S) 가 max length보다 짧을 수 있으므로 [:S]를 통해 input sequence까지 잘라서 차원을 맞춰준 후 더해준다.

 

Inline Question 1

Several key design decisions were made in designing the scaled dot product attention we introduced above. Explain why the following choices were beneficial:

1. Using multiple attention heads as opposed to one.
2. Dividing by √d/h before applying the softmax function. Recall that d is the feature dimension and h is the number of heads.
3. Adding a linear transformation to the output of the attention operation.

Only one or two sentences per choice is necessary, but be sure to be specific in addressing what would have happened without each given implementation detail, why such a situation would be suboptimal, and how the proposed implementation improves the situation.

Your Answer: Using multiple attention heads helps the attention to reflect data position and positions' representation better. Dividng by sqrt(d/h) scales the attention before applying softmax. By scaling, we can prevent some output of softmax value to be extremely small, which will cause vanishing gradient problem. Adding a linear transformation to the output of attention makes model to infer about those attentions, especially when we use multiple attention, each of head output might not fit with other heads.

 

다음 문제는 Transformer에서 사용한 다음의 방법들이 왜 좋은지 설명하는 것이다.

1. 하나의 head 대신 여러개의 attention head를 사용하는 것

- 하나의 head대신 여러개의 attention head를 사용하는 것은 attention이 데이터의 embed dimension 마다의 위치 정보와 위치 정보가 표현하는 바를 더 반영할 수 있다. 

2. softmax 함수 이전에 sqrt(d/h) 를 나누는 것

- softmax함수 이전의 sqrt(d/h)를 전체 값에 나누어 주어 각 값을 scaling 하는 이유는 각 값들의 차이를 줄여주기 위해서이다. 최대값과 최소값의 차이가 클 때, 특히 softmax 함수를 지나치면 작은 값들은 0에 거의 가까워져 vanishing gradient 문제가 생길 수 있다.

3. attention output에 Linear transformation을 추가하는 것

- attention 결과에 linear transformation을 적용하는 이유는 신경망을 통해 각 attention에 대해 모델이 추론할 수 있게 하기 위해서다. 특히 multiple attention을 사용하는 과정에서 결과로 각 head의 값을 단순히 연결하였는데, 이렇게 연결만 하면 각 head의 출력들이 다른 head의 출력들 값과 잘 맞지 않을 수 있어 linear transformation으로 결합 시켜준다.

 

다음은 layer 구현이 끝났으므로 구현한 layer를 사용해 transformer 모델 CaptioningTransformer 를 구현해본다. 먼저 init 함수이다.

CaptioningTransformer init함수

word_to_idx 는 각 숫자가 뜻하는 word를 저장, input_dim은 입력할 사진의 features의 dimension, wordvec_dim은 caption 데이터 word vector의 dimension, num_heads는 사용할 head 개수, num_layers는 사용할 transformer layer 개수, max_length는 sampling에서 생성할 output caption의 최대 길이가 될 것이다. visual_projection은 사진의 features를 wordvec_dim으로 바꾸어주고, embedding은 caption을 wordvec_dim으로 바꾸어 줄 것이고 positional encoding은 caption을 통해 얻은 wordvec_dim에 transformer를 사용해 주기 전에 위치 정보를 추가해 줄 것이다. 다음은 decode_layer와 transformer 함수인데, 아래의  TransformerDecoderLayer를 보며 이해해 보자.

TransformerDecoderLayer init함수

self_attn, multihead_attn은 은 우리가 아까 구현한 MultiHeadAttention을 사용한다. 아래의 linear, dropout, norm들은 forward 함수에서 사용할 layer를 먼저 구현한 것이다. 그러면 forward 함수는 어떻게 생겼을까?

TransformerDecoderLayer forward 함수 / Transformer Decoder Architecture와 비교

forward 함수이다. tgt는 decoder layer의 sequence, memory 는 encoder의 마지막 layer의 sequence이다. 먼저 tgt (decoder sequece)를 multihead self attention -> dropout -> residual connection + layernorm 시켜준다. 그리고 memory (encoder sequence) 를 key, value로 tgt를 다시 multihead self attention -> dropout -> residual connection + layernorm 시켜준다. 그리고 Feed Forward network (linear -> ReLU -> dropout -> linear) -> dropout ->  residual connection + layernorm 시켜주어 결과를 출력한다. 옆의 transformer architecture의 decoder 부분과 비교하며 보자. 아래는 위의 TransformerDecoderLayer를 사용한 CaptioningTransformer의 forward 함수이다.

먼저 이미지 features를 처리하는 encoder sequnece 는 visual_projection layer를 통해 구현해 준다. 이따 decoder 에서 사용할 것이므로 decoder sequence와 같은 dimension (N, 1, D)로 reshape 해준다. 그러면 encoder sequence 완성이다.

 

다음으로 decoder에서 사용할 caption을 embedding -> positional encoding으로 위치 정보를 추가해주어 decoder sequence를 만들어 준다. 여기서 tgt_mask는 전에 multihead attention layer을 만들떄 사용한 위치 정보를 반영하기 위해 미래 key값을 보는 query의 정보를 지워주는 mask이다. torch.tril 함수를 사용해 ones 행렬에서 아래의 삼각형 값들을 날려서 구현한다.

 

encoder가 features를 이용하고, decoder가 caption을 이용하는 이유는 우리가 사진을 통해 caption을 만드는 모델을 만들기 때문이다. train과정에서는 features, caption이 모두 주어지지만 sample 과정에서는 features 이 input으로 주어졌을 때 caption을 output으로 만들어야 한다.

 

그렇게 decoder sequence, encoder sequence, tgt mask로 transformer() (TransformerDecoderLayer의 forward함수) 를 통해 transforemr decoder를 통과시켜 주면 out (N, T, D), out을 마지막에 Linear layer (self.out (D, V))를 통과시켜주면 scores (N, T, V)를 얻을 수 있게 된다.

 

model이 잘 만들어졌는지 평가하기 위해 Train data를 overfit 시켜본다.

잘 overfit 되는 것을 알 수 있다.

다음은 위에서 구현한 transformation captioning을 사용해 실제 데이터들을 sampling 해본 결과이다. train data는 잘 작동하고, val data는 항상 잘 작동하지 않지만 지난번의 RNN captioning보다는 좀 더 문법적으로 말이 되는 결과를 보여준다. 두번째 사진은 <a open dog with a stuffed and red pizza> 로 아마 가운데 비닐같은 물체를 개로 오해한 것 같고, 빨간색과 누런색 침대 이불을 피자로 오해한 것 같다. 세번째 사진은 <a man is with a bite of his food in a hand> 라고 결과가 나왔는데, 사진속 인물의 성별이 여자인 것을 제외하면 손에 음식을 든 것을 잘 인식하는 엄청나게 정확한 결과를 보여주는 것을 볼 수 있다.