[ICLR 2020] Reformer: The Efficient Transformer 논문 정리
ICLR 2020 NLP 논문 읽기 (4)
[paper] [github] [pytorch github]
Abstract만 읽고 정리해보기
- Large Transformer 모델이 여러 task에서 SOTA를 달성했지만, 학습에 너무 큰 비용 → Transformer의 효율성을 향상하는 2가지 기술
- dot-product attention을 locality-sensitive hashing을 사용하는 것으로 대체하여 complexity를 O(L²)에서 O(L log L)로 개선 (L은 sequence length)
- 기존 residual 대신 reversible residual layer를 사용 (activation을 학습 내에 1번만 저장 — 기존 N번: layer 수)
- 기존 Transformer와 비슷한 성능을 보이면서, memory-efficient하고 긴 sequence에 훨씬 빠름
Transformer가 메모리를 잡아 먹는 부분
- Activation이 back-propagation을 위해서 저장이 되어있어야 함 → N개의 layer일 때, N 배 크게 메모리를 잡아 먹음
- 중간의 feed-forward layer의 깊이(depth)가 attention activation의 깊이(depth)보다 훨씬 더 크다 → 메모리 사용의 큰 비중을 차지
- Sequence length가 L일때, attention에 대한 computational/memory complexity는 O(L²) → 64K tokens의 단일 sequence도 메모리를 다 차지할 정도
Reformer 모델의 대처 방안
- Reversible layers → 전체 모델에 대해서 activation을 딱 한번만 저장(copy하여) → N 배 소모 사라진다
- Feed-forward layer들 내 activation을 chunk(나누고 분할)하여 처리 → Feed-forward layer의 depth factor를 없앰 → 그 layer의 메모리 절약
- “Locality-sensitive hashing” 기반의 approximate attention 계산 → complexity를 O(L²) 에서 O(L logL) 로 → Long sequences에 대한 처리 가능하게 함
Locality Sensitive Hashing (LSH) Attention
먼저 LSH와 Angular LSH에 대해 간단하게 짚고 가자.
- LSH (locality-sensitive hashing)은 유사한 입력 항목을 높은 확률의 동일한 bucket으로 hashing하는 알고리즘
- Angular LSH는 LSH 알고리즘의 변형. 각각 고유의 코드를 가지고 미리 정의된 영역으로 분할되어 있는 구(sphere)에 점을 project한 다음, 점을 여러번 랜덤으로 회전하여 점이 속하는 bucket을 정의.
기존 Transformer에서의 self-attention 식은 다음과 같다.
Q, K, V는 [batch size, length, d_model] 의 shape를 가지는데, 앞에서 언급한 main issue는 Dot product 부분의 Q K^T가 [batch size, length, length]의 shape를 가진다는 것이다. softmax(Q K^T)에만 본다면, softmax는 가장 큰 element에 지배적으로 영향을 준다. 이때, 각 query q_i에 대해서 q_i와 가장 가까운 K안의 key들에만 집중하는 것이 어떨까?
예시로 K가 64K의 길이일때, 각 q_i에 대하여 32 또는 64개의 가까운 key들만 고려하면 된다. 그렇다면 nearest neighbor를 어떻게 구할까? → “LSH”
위 그림과 같이 random projection을 한다. b개의 hash를 얻기 위해서 [d_k, b/2] shape를 갖는 random matrix R을 fix한다. 그리고 h(x) = arg max([xR; -xR]) 를 정의하는데, 이를 LSH scheme라 한다. 이 scheme과 hashing attention을 이용하여 LSH attention을 formalize할 수 있다.
Reversible Transformer
- RevNet의 아이디어를 Transformer에 적용
- revnet block 안의 attention과 feed-forward layer를 combine
- F가 attention layer, G가 feed-forward layer가 되는 것
- Reversible Transformer는 각 layer 안에 activation을 저장할 필요가 없다.
Chunking
- feed-forward layer 내 계산은 sequence 내 위치 간 완전히 독립적이기 때문에, 여러 개의 chunk로 split할 수 있다.
- 모든 위치에 대한 작업을 병렬로 처리하고, 한 번에 하나의 chunk로 작동하여 메모리를 줄일 수 있다.
아래는 기존 Transformer와 위의 아이디어를 접목한 Transformer의 변형의 memory and time complexity를 나타낸다.
My Discussion
좋은 성능의 모델이 많아지면서 모델 크기가 엄청나게 커지는 것도 수반되었다. 따라서 모델 경량화도 굉장히 핫한 분야라고 할 수 있는데, Attention is all you need의 Transformer 모델을 메모리를 적게 사용하는 방법을 제시한 Reformer를 한 번 읽어보았다. 사실은 한 달 전에 읽었는데, 기말고사 준비로 너무 바빠서 종강해서야 정리할 기회가 생겼다…
