ViT

Formulation

Transformer block과 MSA(Multihead Self-Attention)의 수식을 과 의미를 살펴보자

Transformer Encoder

• LN : Layer Norm
• MLP : Multi Layer Perceptron
• MSA : Multihead Self Attention

(1) [Patch Embbedings]
$z_0=[x_{class}; x_p^1E;x_p^2;\cdot\cdot\cdot;x_p^NE]+E_{pos}$, $E \in {\mathbb{R}^{(p^2\cdot C) \times D}}$, $E_{pos} \in \mathbb{R}^{N=1)\times D}$

(2) [MLP Block]
$z^{\prime}_l = MSA(LN(z_{l-1}))+z_{l-1}$, $l=1\cdot\cdot\cdot L$

(3) [MLP Block]
$z_l = MLP(LN(z^{\prime}_l))+z^{\prime}_l$ , $l=1\cdot\cdot\cdot L$

(4) [Position Embeddings]
$y=LN(z^0_L)$

Multihead Self Attention(instead of transformer encoder)

(5) $[q, k, v]=zU_{zkv}, U{qkv} \in {\mathbb{R}^{D\times} 3D_h}$

(6) $A = softmax(qk ^\intercal \sqrt{D_h}), A\in\mathbb{R}^{N \times M}$

(7) $SA(z)=Av$

(8) $MSA(z)=[SA_1(z);SA_2(z);\cdots;SA_k(z)]U_{msa}, U_{msa}\in \mathbb{R}^{k \cdot D_h \times D}$

위의 수식들은 MSA으로 입력 시퀀스를 변환하여 입력벡터 $z$로 $qkv$를 생성하고 각 Head에서 self-attention을 적용한 후, 그 결과를 연결하고 다시 변환하여 최종 출력을 생성하는 과정이다.

Explain

(5)입력 벡터 $z$로부터 쿼리(queries), 키(keys), 값(values) 벡터를 생성하는 과정

• $z$ : 입력 시퀀스의 표현으로, 길이$N$의 시퀀스에서 각 위치에서의 벡터$z_i$의 차원은$D$
• $U_{qkv}$ : $z$를 쿼리$q$, 키$k$, 값$v$로 변환하기 위한 가중치 행렬로 크기는$D \times 3D_h$이며, $D_h$는 각 헤드(head)에서의 차원임
• $q$,$k$,$v$는 모두$\mathbb{R}^{N \times D_h}$의 차원을 가짐

(6) 어텐션 가중치$A$를 계산하는 과정

• $qk^\intercal$ : 쿼리와 키의 내적으로,$q$와$k$의 내적을 계산하면 $\mathbb{R}^{N \times M}$ 크기의 행렬이 되고, $N$과$M$은 입력 시퀀스의 길이를 나타냄

(7) 단일 어텐션 헤드에서의 셀프 어텐션(SA)을 계산하는 방법을 설명합니다.

• $A$ : 앞에서 계산된 어텐션 가중치 행렬
• $v$ : 값 벡터
• $Av$ : 어텐션 가중치$A$를 값$v$에 적용하여 셀프 어텐션 출력을 계산하고, 출력은 $\mathbb{R}^{N \times D_h}$의 차원

(8) Multihead Self Attention(MSA)

• $SA_i(z)$ : 각 어텐션 헤드$i$의 셀프 어텐션 출력이고, $k$개의 헤드를 가지므로$k$개의$SA_i(z)$가 있음
• $[SA_1(z); SA_2(z); \cdots; SA_k(z)]$ : $k$개의 어텐션 헤드 출력을 하나로 연결(concatenate)한 것으로, 결과물은 $\mathbb{R}^{N \times k \cdot D_h}$ 차원임
• $U_{msa}$ : 연결된 출력을 최종 출력 $MSA(z)$로 변환하기 위한 가중치 행렬이고, 크기는 $k \cdot D_h \times D$입니다.
• $MSA(z)$는 $\mathbb{R}^{N \times D}$의 차원을 가지는 출력 벡터

Reference
[1] https://github.com/lucidrains/vit-pytorch
[2] https://github.com/huggingface/pytorch-image-models
[3] https://github.com/jankrepl/mildlyoverfitted/tree/master/github_adventures/vision_transformer