[Paper Review] Reformer: The Efficient Transformer 논문 분석
업데이트:
이번 포스팅에서는 Transformer의 architeture를 효율적으로 만든 Reformer(2020)에 대해 살펴본다.
- Paper : REFORMER: THE EFFICIENT TRANSFORMER(2020) Nikita Kitaev, Łukasz Kaiser, Anselm Levskaya
1. Introduction
Transformer model(Vaswani et al, 2017)은 자연어 처리 분야(NLP)에서 널리 사용되는 architecture이다. (Transformer 논문 및 코드)
Transformer의 가장 큰 장점은 확장성이다. 이 모델에서는 아주 큰 데이터셋에 대해서도 학습이 잘되기 때문에 BERT를 비롯한 다양한 모델에 차용이 되었고 계속해서 SoTA의 결과를 내고 있다.
🤔 대규모의 Research에서만 할 수 있는 연구가 아닌가?
transformer의 모델이 점점 커지면서 이 모델을 향한 비판의 목소리도 커지고 있다. Large-scale long-sequence model을 사용하면 좋은 결과를 낼 수 있기는 하지만, 학습을 위해서는 너무나도 많은 resource가 필요하기 때문에 점점 GPU의 싸움이 되고 있는 것이다.
- Ex) 64K의 token을 학습시키는 경우 (embedding size : 1024, batch size : 8)
- Parameters : $64K * 1K * 8 = 0.5B$
- $2GB$의 메모리가 필요
- BERT는 더 많은 corpus를 사용하기 때문에 $17GB$의 memory를 사용한다,,
따라서 Reformer model에서는 Transformer architecture를 구조적으로 수정하여 메모리 효율성과 연산 속도 두마리의 토끼를 잡고자 한다.
✍🏻 Problem & Solution
Reformer에서는 Transformer에서의 비효율성을 이와 같은 방식들로 개선하고자 한다.
(1) N개의 layer를 사용하면 single-layer를 사용했을 때보다 memory를 N배 더 사용한다. (back-propagation과정에서 intermediates들을 저장해야하기 때문)
⇒ Reversible Layers를 사용하자 !
(2) Attention의 depth는 $d_{model}$이지만, Feed-Forward Layer의 depth는 $d_{ff}$ / Feed-Forward Layer를 사용하면 너무나도 많은 memory가 필요
⇒ Feed-Forward layer의 $d_{ff}$를 chuncking하고 activation 함수를 split하여 메모리를 아끼자 !
(3) 기존의 Transformer : 길이가 $L$인 문장을 attention하면 computational & memory complexity가 $O(L^2)$
⇒ Locality-sensitive hashing을 사용하자 ! = complexity = $O(L logL)$
위의 방식에서
(1) 일반 모델 대신 Reversible residuals 를 적용해도 실험결과는 거의 변하지 않으며
(2) Splitting activation은 오직 구현부에만 영향을 줄 뿐, 사실상 transformer의 layer는 동일하다.
(3) 마지막으로 Attention에서 Locality-sensitive hashing 를 사용하면 training 방식에는 영향을 주지만, 결과는 full attention과의 비슷하다.
본 논문에서는 Reformer 모델로 64K의 text task(enwiki8 )와 12K의 image generation task(imagenet-64 ) 실험을 했다. 또한, 본 논문에서는 실험을 통해 Reformer가 Full Transformer와 비슷한 결과를 내면서도 더 빠르고 memory-efficient하다고 말하고 있다.
2. Locality-Sensitive Hashing Attention
Attention in Transformer
Transformer의 standard attention 식은 다음과 같다. Attention 관련한 설명은 이 글 을 참고
\[\operatorname{Attention}(Q, K, V)=\operatorname{softmax}\left(\frac{Q K^{T}}{\sqrt{d_{k}}}\right) V\]또한, transformer에서는 parallel한 계산을 위해 Multi-head attention이라는 mechanism을 사용한다.
이때 Q K V 벡터가 어디서 오는지 살펴보자
multi-head attention에서의 queries, keys, values는 [batch size, length, dmodel]의 $A$ 라는 single tensor에서 3가지의 linear layer를 project함으로써 얻어진다.
- 관련 내용은 이 부분 코드 참고
Shared-QK Transformer
LSH attention의 model에서는 Queries와 Keys가 근본적으로 동일하다. $A$ tensor의 같은 linear layer에서 Q와 K가 나오고, $A$의 또 다른 layer에서 V가 얻어지는 것이다.
이 논문에서는 sharing QK가 Transformer의 성능에 영향을 안주므로 이 방식을 사용하고 있다고 말한다.
Hashing Attention
Transformer의 attention mechanism의 memory사용량을 계산해보자.
- Query, Key, Value : [batch size, seq length, d model]
- $QK^T$ : [batch size, length, length]
- batch size가
1이고, sequence length 가64K, 32-bit float를 사용했다면, $16GB$의 메모리를 사용한다 !!
\[\operatorname{Attention}(Q, K, V)=\operatorname{softmax}\left(\frac{Q K^{T}}{\sqrt{d_{k}}}\right) V\]🤔 $QK^T$의
64K x 64K행렬이 memory에 다 저장될 필요가 있을까?
우리는 위의 식을 통해 attention을 계산한다. 각 query ($q_i)$의 attention 식은 다음과 같다.
\[\operatorname{softmax}\left(\frac{q_{i} K^{I}}{\sqrt{d_{k}}}\right) V\]그런데 사실 $QK^T$의 값을 softmax 할 때 결과값은 largest element에만 영향을 받는다.
즉, 모든 query에 대해 key와 내적해주지 않고 query와 비슷한 (query와 가까운) key에 대해서만 내적을 해줘도 된다.
- 64K의
key가 있을 때, 각query는 모든key에 대해 attention을 하지 않고, 가까운 32 or 64의 key에 대해서만 attention을 해줘도 괜찮음 😉
Locality sensitive hashing
hashing ($x$ → $h(x)$ )에서 가까이에 있는 vector를 높은 확률로 같은 hash에 넣고, 먼 vector들은 다른 hash에 넣는 것을 locality-sensitive라고 부른다.
위에서 가까운 이웃에 대해서만 attention을 해줘도 됨을 확인했다. 따라서 우리는 LSH라는 기법을 사용하여 nearest neighbor를 찾는다.
본 논문에서는 random projection 방식으로 angular locality sensitive hash를 사용했다. (LSH scheme 논문 참고) random으로 $\theta$를 구한 후, 원의 어느 지점에 위치하는 지를 확인하고 그 구간의 index를 붙이는 방식이다.
- We achieve this by employing random projections as follows (see Figure 1). To get b hashes, we first fix a random matrix R of size $[dk, b/2].$ We then define $h(x) = arg max([xR; −xR])$ where $[u; v]$ denotes the concatenation of two vectors. This method is a known LSH scheme (Andoni et al., 2015) and is easy to implement and apply to batches of vectors.
LSH Attention
LSH Attention을 식으로 나타내면 다음과 같다.(이 식에서는 $\sqrt{d_{k}}$ 의 scaling은 생략되어있음)
\[o_{i}=\sum_{j \in \mathcal{P}_{i}} \exp \left(q_{i} \cdot k_{j}-z\left(i, \mathcal{P}_{i}\right)\right) v_{j} \quad \text { where } \mathcal{P}_{i}=\{j: i \geq j\}\]- $\mathcal{P}_{i}$ = $i$ 위치에 있는 query가 attend할 수 있는 set
- $z$ = partition function (ex. softmax에서의 normalizing)
Masking까지 포함하면 다음과 같다.
\[o_{i}=\sum_{j \in \widetilde{\mathcal{P}}_{i}} \exp \left(q_{i} \cdot k_{j}-m\left(j, \mathcal{P}_{i}\right)-z\left(i, \mathcal{P}_{i}\right)\right) v_{j} \quad \text { where } m\left(j, \mathcal{P}_{i}\right)=\left\{\begin{array}{ll}\infty & \text { if } j \notin \mathcal{P}_{i} \\0 & \text { otherwise }\end{array}\right.\]이제부터 LSH Attention 과정에 대해 살펴보자😊
- 각 token의 (Query, Key), Value를 생성한다.
-
shared-QK Transformer : LSH attention에서는 Queries와 keys가 같은 linear layer에 있어야함
-
2. Locality-Sensitive Hashing
- Angular locality sensitive hash -> 비슷한 token들은 같은 bucket에 있을 확률이 큼
- 같은 bucket내에 queries와 keys가 불균형하게 있을 수 있으므로 다음의 식을 만족하도록 Q, K를 뽑는다. \(\mathcal{P}_{i}=\left\{j: h\left(q_{i}\right)=h\left(k_{j}\right)\right\}\)
3. Sort by LSH bucket
- 문장 순서대로 bucket을 sorting한다.
4. Chunk sorted sequence to parallelize
- 고정된 크기로 sequence를 split
5. Attention 적용
- 같은 bucket 내에서 attention
- 이전 chunk에 대해서 attention
Chunking Attention을 하는 이유
- 동일한 bucket에 있는 token들이 많아도 특정 길이 이상 attention을 할 수 없도록 하기 위해
Multi-round LSH Attention
위와 같이 LSH attention을 하면 몇가지 문제가 생길 수 있다.
- Hash bucket의 size가 고르지 않을 수도 있고
- bucket 내의 queries와 keys의 수가 제각각일 수도 있다.
즉, 이렇게 되면 비슷한 item들이 다른 bucket에 떨어질 수 있게 된다.
따라서 이 경우에는 LSH Attention 여러번 하는 방식인 Multi-round LSH Attention 를 사용하여 문제를 해결한다.
\[\mathcal{P}_{i}=\bigcup_{r=1}^{n_{\text {round }}} \mathcal{P}_{i}^{(r)} \quad \text { where } \mathcal{P}_{i}^{(r)}=\left\{j: h^{(r)}\left(q_{i}\right)=h^{(r)}\left(q_{j}\right)\right\}\]Causal masking for shared-QK attention
Transformer의 attention에서는 자기 자신을 포함하여 attention을 할 수 있었다. 다만, shared-QK를 사용하면, 자기자신에 대해 내적을 했을 경우 이 값이 너무나도 커져 전체 attention값에 영향을 미친다.
따라서 자기 자신에 대해서는 attention을 하지 못하도록 조정을 해준 것이 Causal masking for shared-QK attention이다.
3. Reversible Transformer
RevNets
Transformer에 사용되는 Feed-Forward NN에서 Back-Propagation을 하려면 중간결과물을 메모리에 저장해야 한다.
-
Memory use = $b \times l \times d_{ff} \times n_l$
= batch size * dim of layer * dim of feed-forward * numbers of layer
ex) 1 * 64000(dim of layer) * 16(# of layer) * 4000 (dim of feed-forward) * 4 (float32) = 16GB
따라서 Reversible Residual Network(2017) Architecture를 사용해서 intermediate를 메모리에 저장하지 않고, output으로 input을 계산하도록 하였다.
이 방식을 사용하면, layer의 수만큼 메모리를 저장하지 않아도 되므로 Table3의 $n_l$ term이 사라진다.
Chunking
$n_l$ term을 지워도 여전히 메모리는 너무나도 크다 ($d_{ff} = 4K$)
따라서 Feed Forward NN에서의 연산을 병렬적으로 할 수 있도록 chunking을 한다.
\[Y_{2}=\left[Y_{2}^{(1)} ; \ldots ; Y_{2}^{(c)}\right]=\left[X_{2}^{(1)}+\text { FeedForward }\left(Y_{1}^{(1)}\right) ; \ldots ; X_{2}^{(c)}+\text { FeedForward }\left(Y_{1}^{(c)}\right)\right]\]Memory and Time Complexity
4. Experiments
Imagenet64랑enwik8-64k로 실험- $d_{model}=1024, d_{ff}=4096, n_{heads}=8$
- Adafactor optimizer
- evaluation :
WMT 2014Eng-Germ translation task
Effect of sharing QK / reversible layers
LSH attention in Transformer
Large Reformer Models
6. Opinions
재미난 방식으로 Transformer의 mechanism을 풀어간 재미난 model이라는 생각이 든다.
성능 평가를 위한 실험의 지표로 bpd를 쓴게 약간은 아쉽다. 또한, Reformer 자체를 transformer와 비교하는 실험이 약간 부족하다는 생각이 들었다.
그래도 아이디어 자체가 괜찮고 논문도 이해하기 쉽게 쓰여있어서 재미나게 읽었던 논문이다. 관련 논문으로 Longformer, Performer, Sparse attention, Big Bird 등의 비슷한 논문들도 읽어보면 도움이 될 것 같다.
Reference
reformer transformer gif - Google Search
Machine learning 스터디 (18) Neural Network Introduction - README
댓글남기기