[Review] Mimetic Initialization of Self-Attention Layers

한창 attetion layer의 가중치 초기화를 어떻게 하면 효과적으로 할 수 있을까 

고민하던 중 읽었던 논문을 리뷰한다

 

서론과 선행연구는 스킵하도록하고 핵심만 적겠다

 

뭐하는 논문인가

Pretrained model의 self-attention weight를 [모방]해보겠다는 것이다

이것이 무슨 의미냐면, 학습된 모델을 직접적으로 쓰지않고, 흉내내는 것만으로 성능을 높일 수 있다는 말이다.

 

어떻게 하는가

 

이것은 이미지처리 분야에서 ViT를 통해 학습된 attention layer의 가중치들이다.

a와 b의 Query, Key, Value, Projection의 가중치들이 보이는데 

왜 저렇게 둘 씩 곱한 것을 보이냐고하면, self attention 계산식에서 알 수 있다.

 

$ Attn =  Softmax(X_{Q} W_{Q} W_{K}^{T} X_{K}) X_{V} W_{V} W_{proj} $

input X만 빼면 위와 같은 방식으로 묶을 수 있고, 이걸 시각화했더니,

그림처럼 대각성분들이 두드러진 Diagonal 비슷하게 나온다는 것이다. 

특히 $W_{Q}W_{K}$는 대각성분들이 양수이고, $W_{V}W_{proj}$은 음수라는 걸 기억해놓자

이걸 행렬 트릭으로 모방시키며 초기화하는 것이 전부다.

 

 

이렇게 한다. Z는 normal distribution matrix이고이 뒤에 Identical한 행렬을 더해줌으로

대각성분들을 더 두드러지게 만드는 것이다.

 

두 가중치 행렬에 대한 semi-diagonal matrix를 만들었으면 

각 파라미터를 초기화하기 위해 SVD를 한다.

이때, $W_{Q}$와 $W_{K}$는 Reduced SVD를 하는데, k차원까지 굳이 줄이는 이유는

multi head attention의 각 헤드의 k차원에 해당하는 가중치만 뽑으려는 것으로 추측된다

 

어떻게 됐는가

CIFAR 데이터셋으로 ViT 성능을 비교했더니 3~7% 이상 높아진 것을 볼 수 있다.

또한 모델이 파라미터가 커질수록<밑으로 내려갈수록> 성능 향상이 더 뚜렷하다

Pretrained 모델을 쓰기 위해서 모델 파라미터도 맞춰줘야하는 등의 세팅을 매우 유사하게 맞춰야하는데

다음과 같이 가중치를 모방하는 방법으로도 빠르게 Optimal을 찾을 수 있다는 것을 보여줬다

 

그 외의 Task

 

NLP의 GPT모델에서도 다음과 같이 Semi diagonal 한 attettion matrix가 나왔다.

그래서 적용을 해보았더니 성능이 기존보다 크게 개선되진 않았다.

저자는 이유를 이렇게 추측했는데, CV보다 NLP에서 다뤄지는 파라미터의 수가 더 많아서 

상대적으로 초기화되는 가중치들이 적기때문이라고 한다.

 

추천시스템에서는?

추천시스템, 특히 트랜스포머 기반이라면 GPT나 BERT와 같이 self attention을 사용한다

적용한다면 BERT4Rec과 같이 파라미터가 GPT보다 적고 퍼포먼스를 낼 수 있는 모델이라 적합할 듯하다.

그런데 BERT4Rec도 NLP와 마찬가지로 임베딩 레이어도 가지고 있어서,

비전분야 같이 퍼포먼스가 나올지는 미지수이다.

위 연구에서도 NLP 환경 더 튜닝된 채로 실험을 해야한다고 언급한 바 있는데,

추천시스템에서 이 방법을 적용할 때는

먼저 Pretrained model이 모방할만한 레이어를 가지고 있는지,

이를 행렬 조작을 통해 모방할 수 있는지, 

모방된 레이어들이 전체 파라미터 중 의미있는 파이를 차지하고 있는지를

확인해서 진행해야 할 것 같다

 

저자 : Asher Trockman, J. Zico Kolter

출판 : 2023.05.16