[논문리뷰] Sparse Transformer

Sparse Transformer 간단 정리. 

arxiv.org/pdf/1904.10509.pdf

서론

Transformer는 Sequence Input 의 context 정보를 잘 내포하도록 학습할 수 있는 모델인데, 
차지하는 메모리 용량이 매우 크다. attention layer의 memory complexity는 무려 O(n^2). 
input sequence의 모든 token들이 다른 token들에 대해서 attention하는 값을 구하게 되기 때문이다. 

이 memory complexity를 낮추려는 거의 초기의 시도! 

sparse tranformer는 sparse factorization을 통해서 이 memory complexity를 낮추었다. 

Key Idea 

Factorized Self Attention 

일반적인 attention mechanism 은 다음과 같이 나타낼 수 있다. 

Fully self attention autoregressive model은 Si = {j : j ≤ i} 를 가정하고 모든 element가 previous position element에 attend 할 수 있도록 한다. 

Factorized self attention은 대신에,, p개의 self attention head를 가지고 m번째 head가 제한된 indices에만 attend 할 수 있게하는것이다. A (m) i ⊂ {j : j ≤ i} and lets Si = A (m) 여기서 subset A를 잘 고르는 게 중요하다고 한다. 그러면 memory complexity O(n^2)를 줄일 수 있다. n * (줄어든 subset 의 memory complexity) 가 되겠지 

다음 그림을 참고하자 

Row가 output column이 input이다. 
각 output은 제한된 개수의  input token에만 attention을 볼 수 있고, strieded 와 fix의 경우 그 구현이 조금 다르다 (연한 하늘색깔을 보자)

수식으로 설명이 되어있는 부분은 여기이다 (arxiv.org/pdf/2009.06732.pdf) 참고함 

1. Local attention head  

2. Strided Attention head

이렇게 된다... 

파라미터 개수는 Query, Key, Value 계산해 줘야 하는 부분은 똑같으니까 그대로고, 

Memory complexity만 attention 계산을 줄이니까 반으로 줄어서 O(Nlogn)이 된다. 

하지만 이거 구현하려면 custom kernel을 사용해야 해서 reproduceblity가 떨어진다구 한당 ㅎㅎ ;;