Transformer 아키텍쳐(Vaswani et al., 2017)가 나온 이래 BERT 등으로 발전하면서 NLP 분야에 획기적인 발전이 계속되었지만, 너무 큰 GPU 메모리를 요구하는 구조로 인해 스타트업 직원으로서는 감히 따라가 볼 엄두가 안나는 분야가 되어버렸다.

2020년 새해가 밝아오면서, 16GB짜리 GPU 1장으로도 NLP 연구를 가능하게 한다는 엄청나게 Memory-efficient한 Transformer 구조가 제안되었다는 소식이 전해졌다. 희망찬 소식에 큰 기대를 가지고 바로 그 논문을 들여다본다.

이 논문의 핵심 포인트 3가지

이 논문을 그림 한장으로 요약하면 다음그림과 같다(고 한다.) 지금은 이해하기 어렵다.

LSH Transformer

• 기존 transformer의 dot-product attention을 locality-sensitive hashing attention (이하 LSH attention)으로 바꾸어서 복잡도를 L^2에서 L*logL로 감소시켰다. (L은 sequence 길이)

$$ 복잡도 : O(L^2) => O(LlogL) $$

Reversible Transformer

• 기존 transformer의 standard residual layer를 reversible residual layer 로 변경하면서 N개 layer마다 별도로 사용하던 메모리를 1번만 사용하도록 개선하였다.

Chunked Reversible Transformer

Transformer 내부의 feed forward network의 activation 구간을 chunk 단위로 쪼개서 메모리사용량을 감소

아래는 Illustrated Transformer(http://jalammar.github.io/illustrated-transformer/) 에서 가져온 Transformer 의 기본 구조 그림이다. LSH는 Self-Attention 부분을, Reversible은 Feed Forward 부분을 개선했다고 볼 수 있겠다.