Transformer 대규모화의 돌파구 될까? 고효율 Reformer의 등장

3 개의 요점

✔️ Attention 계산량을 O(n^2)에서 O(n log n)로 크게 감소

✔️ 액티베이션 등의 메모리 사용량을 대폭 절감

✔️ 속도와 메모리 양쪽 모두의 구현 효율을 크게 개선하면서도 Transformer의 성능을 유지

Reformer : The Efficient Transformer

written by Nikita Kitaev ,  Łukasz Kaiser ,  Anselm Levskaya

(Submitted on 13 Jan 2020, ICLR 2020)

subjects : Machine Learning (cs.LG) ; Computation and Language (cs.CL); Machine Learning (stat.ML)

Reformer: The Efficient Transformer

 

대규모 응용이 뒷받침하는 Transformer의 기록적인 성과

Transformer를 사용한 연구는 계속해서 state-of-the-art 결과를 만들어 내고 있습니다만, 그 응용은 점점 커지고 있습니다. 그리고 이 대규모화 흐름은 큰 연구기관 외에는 연구 자체가 어려워지고 있는 것을 의미합니다. 현재 그 문제로 논란이 되고 있는 것 같습니다.

 

그럼 어느 정도의 규모인 것일까요?

● Transformer층(이하 Trm) 하나의 파라미터수 0.5B(= 500,000,000)를 유지하는데, 2GB의 메모리가 필요(32 비트 부동 소수점).

● 토큰 길이가 64,000이고, Embedding 크기가 1,024, 배치 사이즈 8일 때 Activation(순 전파의 계산 결과)는 동일하게 64K x 1K x 8 = 0.5B에서 2GB.

이 크기는 실제로 계산 시 치명적이라고 할 수 있는 문제입니다.

● Trm이 12층에 도달하면 Activation은 2GB x 12 = 24GB. 학습 시에는 역전파 될 때까지 유지가 필요합니다.

● Attention 계산은 토큰 길이 (L)에 대해 계산량과 메모리량 모두 O(L^2), 즉 배치 사이즈 1도 L = 64KB인 경우 64KB^2 x 4 (32비트 부동 소수점) = 16GB.

이 규모가 되면 몇 개의 GPU정도의 구성으로는 전혀 처리할 수 없습니다.

그럼, 대규모 용도의 예를 살펴보겠습니다.

 

대규모 프로젝트의 예

논문 "Mesh-Tensor Flow : Deep Learning for Supercomputers" 1)

Google의 기계학습용 칩 TPU를 메쉬 처리하는 연구 논문에서는 기계 번역을 예제로 하고 있습니다. Trm층의 피드 포워드 262,144차원, Attention 헤드를 256으로 한 Transformer를 계산하는 데 TPU를 512메쉬로 했다고 합니다. 이는 매우 막대하며, Google정도의 규모가 아니면 불가능합니다.

 

Google's scalable supercomputers for machine learning, Cloud TPU Pods, are now publicly available in beta

Google’s scalable supercomputers for machine learning, Cloud TPU Pods, are now publicly available in beta | Google Cloud Blog

 

논문 "Image Transformer" 2)

이미지에 Transformer를 응용한 연구 사례로 어중간한 이미지로부터 나머지 부분을 생성하는 것과 초해 상의 응용이 있습니다.

 

초해상은 입력 8 × 8 픽셀부터 출력 32 × 32 픽셀로, 이 규모조차 계산량에 고심을 하고 'Local Self-Attention'이라는 방식으로 계산량을 떨어뜨리고 있을 정도입니다.

 

 

논문 "Music transformer : Generating music with long-term structure" 3)

이 음악의 생성은 1분의 시간을 10밀리 초마다 구분하여 6,000개의 데이터로 취급하는 것으로, 감량 데이터로 2,000차원까지 떨어뜨리고 있습니다.

 

이미지로서 문장의 경우에는 2,000단어를 하나의 데이터로 하는 것을 2,000음으로 하나의 곡 데이터로 취급하는 것이 됩니다. 그럼에도 불구하고 2,000차원이라면 8헤드의 경우 8.5GB의 메모리가 소비되어 취급되지 않기 때문에 "Relative Positional Self-Attention"으로 동적 메모리 소비량을 줄일 수 있습니다. 이 방식을 사용하면 2,000차원이라는 긴 표현을 취급할 수 있는, 좋은 논문입니다만, 더 많은 데이터 량을 사용하고 싶지만, 병목현상이 발생하는 현상도 동시에 보여주고 있습니다.

 

 

논문에서 인용한 이 Figure 3에서 "Relative Positional Self-Attention"를 사용하는 것으로(우측), 카덴차(음악 용어)가 프레이즈의 끝에 나타날 수 있다는 것입니다. 즉 보통의 악곡에 잘 나타날 것 같은 진행이 생성되도록 개선된 것입니다.

 

 

여담이지만이 위 그림을 보는 법을 설명드리겠습니다. 이것은 '피아노 롤'이라는 것입니다. 세로축은 건반의 위치, 가로축은 시간입니다. 소리를 내는 시간을 막대로 표현합니다. 아래 그림에서는 2개의 소리가 울리고 있는 듯합니다. 낮은음은 한 소절마다 소리가 진행되고, 높은음은 멜로디를 연주하고 있는 것 같습니다.

 

Attention을 근사치 계산으로 줄여 리소스화 ~ LSH Attention

키 포인트 중 하나는 Attention의 처리를 빠르게 · 적게 만드는 것입니다.

 

이를 실현하는 경량 리소스화의 근사 계산을 수행하기 위해 'Locality-sensitive hashing'을 사용했습니다. 1999년부터 진화를 계속한 알고리즘(LSH Algorithm and Implementation(E2 LSH) - MIT)으로, 유사한 입력 데이터가 동일한 버킷(해시 테이블의 각 요소)에 들어가는 성질을 이용하고 있습니다.

Locality Sensitive Hashing (LSH) Home Page

(참고로 이 LSH는 머신러닝 알고리즘의 일종인 k근접 법에도 사용되는 것입니다.)

 

비슷한 데이터를 정리하여 계산량을 줄이는 방법의 포인트는 Attention에 포함된 SoftMax에 있습니다.

 

SoftMax에 주어진 경로에서 MatMul(Q, K)의 크기가 매우 커질 것입니다만, SoftMax 자체의 연산 결과는 입력값 중 큰 값이 거의 지배적입니다. 그러므로 이 큰 값 이외의 연산을 생략 하자는 아이디어입니다.

 

그리고 그 큰값으로 되는 것은, Q와 K가 가까운 때입니다. 그래서 이 가까운 데이터를 선택하는데 LSH을 사용합니다. 이렇게 하면 원래의 계산은

아래의 계산식처럼 됩니다.

여기서 i는 Q를 요소마다 다루고 있는 의미에서 z(·, ·)는 원래의 SoftMax의 분모라고 파악해 주십시오.

 

j ∈ P_i 가 포인트로, P_i는 q_i에 가까운 k_j만의 j 집합입니다. 즉, q_i에 가깝지 않은 k_j의 계산은 생략될 것입니다.

※ 매우 간략히 설명했습니다. 정확한 내용은 논문의 계산식을 참조하십시오.

 

논문의 그림을 확인해 보겠습니다.

[Reformer 2020]에서 Figure 2

 

왼쪽 그림에 주목합시다. 위에서 'LSH bucketing'을 하는 것으로 비슷한 데이터가 모입니다. 같은 색상이 동일한 버킷을 나타내고 있고, 그들 사이에서만 계산되는 모습이 그려져 있습니다. 오른쪽 그림은 대각선 상에 비슷한 데이터만 모여 그들 사이에서만 계산이 이뤄지는 것을 알 수 있습니다. 모든 요소 간에 계산하는 것이 아니라 계산할 요소를 선택합니다.

 

이제 계산량과 메모리 모두 O(n^2)에서 O(n log n)이 됩니다.

 

Multi-round LSH attention

실제로 해시가 동일한 것만 아니라 근접한 n_r 개(round)의 버킷까지 계산 대상으로 하게 되면, 해시가 잘 맞지 않는 경우도 커버할 수 있습니다. round값을 늘림으로써 계산량과 메모리량은 거의 증가하지 않고도 근사 계산의 질을 높일 수 있습니다.

 

알고리즘의 개선은 놀라움, O(n^2)에서 O(n log n)로의 임팩트

어느 정도의 개선 임팩트가 있을까요?

예를 들면 속도에 관해서는 CPU에 비하여 GPU로 계산할 경우 대략 10배 정도 빠르게 처리할 수 있습니다. 반대로 말하면, GPU로 한다고 쳐도 고작 10배 밖에 안된다고 말할 수 있습니다.

n	O (n^2)	O (n log n)
10	100	33
30	900	147
100	10,000	664
1,000	1,000,000	9,966
10,000	100,000,000	132,877

이에 대해 O(n^2)와 O(n log n)의 주요 관계를 보면, n=100에서 이미 O(n log n) 쪽이 O(n^2)에 비해 10,000 / 664 = 약 15 배입니다. Transformer로 문장의 단어 수 100개 정도의 규모는 보통이라고 생각됩니다만, 이미 CPU대 GPU에 필적하는, 혹은 그 이상의 효과가 있습니다.

 

만약 n=1,000 단어를 취급하고자 하면 O(n log n)의 9,966에 O(n^2)는 1,000,000이므로 약 100배의 차이가 있습니다. 이것은 이미 하드웨어의 개선규모를 넘어서고 있고, 알고리즘으로만 가능한 영역의 개선입니다.

 

Algorithms, 4th edition 강의 슬라이드로 부터

O(n log n)은 linearithmic로, 거의 linear O(n)와 동일합니다.

14AnalysisOfAlgorithms-2x2.pdf

4.42MB

 

리버시블화 ・조각화로 인한 메모리 절감

지금까지 진행한 경량 리소스화는 Attention 주변이었습니다. 남은 것은 역전파에 대비해 보존해 둘 Activation(계산 결과)와 FeedForward 계산을 위한 메모리량 감소입니다.

리버시블(Reversible)화로 유지 메모리량 감축

계산 결과를 '재현'하는 선행연구에 ResNets 5)를 기본으로 한 RevNets 6) 이 있습니다.

일반적으로 신경망의 학습에서는 일단 순 전파를 네트워크의 입력부터 순서대로 출력을 향해 계산해가면 각처에서 계산 결과를 유지합니다. 이들은 그 역 전파의 계산에 사용됩니다. 그런데 RevNets에서는 곳곳에서 계산을 보유할 필요 없이 출력 측에서 순차적으로 재계산하여 재현할 수 있는 것입니다.

 

Trm 층에 포함되는 Attention과 FeedForward의 블록은 각각 ResNets과 마찬가지로 skip connection 합니다. 그래서 RevNets을 응용하여 Reversible Transformer를 구성하고 있습니다. 따라서 리버시블 가능해지고 중간 계산결과를 유지할 필요가 없어졌습니다.

 

조각화로 FeedForward의 계산 시 메모리 절감

피드 포워드(FeedForward)도 큰 사이즈로 사용되기 때문에 이 계산에 사용되는 메모리의 양을 줄일 필요가 있습니다. 이것은 간단하며, 각 요소의 계산에 상관없기 때문에 작게 나누어 계산하려는(조각화) 것입니다.

 

실험

실험 내용은 논문에서 찾을 수 없었기 때문에 기사 "Google & UC Berkeley 'Reformer'Runs 64K Sequences on One GPU (medium)"를 참조했습니다.

Google & UC Berkeley ‘Reformer’ Runs 64K Sequences on One GPU

• imagenet64 이미지 생성 작업으로, 입력 토큰 길이 12K
• enwik8 텍스트 작업으로, 입력 토큰 길이 64K

둘 다 매우 긴 입력 크기입니다. 위의 작업을 통해 동일한 설정을 사용한 일반 Transformer와 비교하였습니다. Reformer를 경량 리소스로 하는 변경점 하나하나에 대해 Transformer와의 비교를 통해 성능에 영향이 없는 것을 확인한 후, LSH Attention에서의 근사 영향은 아래 그래프와 같습니다.

 

보통의 'full attention'에 비해, round 8까지 올리면 거의 동일한 성능이 나오고 있습니다. 즉, 8개의 해시값과 일치하는 것을 계산 대상으로 하는 것으로, 성능을 떨어뜨리지 않고 계산량을 줄일 수 있는 것 같습니다. round를 올리면 성능이 떨어지지 않는 만큼, 이 근처는 이용 가능한 계산 리소스와의 균형이 필요하다는 논문의 댓글이 있습니다.

 

다음은 큰 Reformer에서의 실험 결과입니다. 20층까지의 Reformer도 GPU하나를 사용하여 긴 입력 데이터의 빠른 학습이 실현된 것입니다.

 

여기서는 오른쪽 그래프에 주목합시다. 'full attention'은 입력 시퀀스 길이가 커짐에 따라 소요시간이 기하급수적으로 급등(※log스케일)하는 것에 비하여 LSH attemtion은 일정합니다. 이는 섬뜩한 무서움이 있습니다.

 

정확도는 어떤가요? 이는 논문 전반에서 언급되고 있으며 Transformer와 필적하는 것으로 확인되고 있습니다.

 

 

마지막으로 계산량과 메모리량의 비교표를 보면 Reformer는 l(입력 데이터의 길이), nl(층의 수), d_ff(피드 포워드의 수)에 대하여 메모리량이 개선되었고 계산량은 l가 하나 줄어든 것을 알 수 있습니다.

 

 

정리

지금까지 소개한 바와 같이 핵심 개선기법 자체는 이 논문에서 고안된 것이 아닙니다. 그것들을 도구로 사용하여 Transformer의 문제점을 해결하고 응용의 활로를 열고 있는 것에 큰 의미가 있습니다. Reformer라면, TPU를 512개 메쉬 가능 등, 큰 규모가 아니더라도 이미지, 영상 등의 큰 입력을 취급할 수 있게 될 가능성이 있으며, AI의 민주화에 큰 의미가 있습니다.

 

알고리즘 개선에 의한 임팩트, 어쩌면 기계학습 시대이기 때문에 중요성을 더해가고 있는지도 모릅니다. ICLR2020는 대부분 완벽한 점수 '8, 8, 6'로 채택되어 향후의 연구에 큰 영향이 있을 것이라고 인정하고 있습니다. 이는 정말 주목할 만한 일로써 향후 발전에 매우 기대가 높습니다.

reformer_the_efficient_transformer.pdf

0.55MB

References

1.	[Mesh-TensorFlow 2018] Shazeer et al., 2018, "Mesh-TensorFlow : Deep Learning for Supercomputers"
2.	[Image Transformer 2018] Parmar et al., 2018, "Image transformer"
3.	[Music transformer 2018] Huang et al., 2018, "Music transformer : Generating music with long-term structure"
4.	[Transformer 2017] Vaswani et al., 2017, "Attention is all you need"
5.	[ResNets 2015] He et al., 2015, "Deep Residual Learning for Image Recognition"
6.	[RevNets 2017] Gomez et al., 2017, "The reversible residual network : Backpropagation without storing activations"

Mesh-TensorFlow_Deep Learning for Supercomputers.pdf

0.39MB

Image Transformer.pdf

1.06MB

Music Transformer_GENERATING MUSIC WITH LONG-TERM STRUCTURE.pdf

0.96MB

Attention Is All You Need.pdf

2.10MB

Deep Residual Learning for Image Recognition.pdf

0.78MB

The Reversible Residual Network_Backpropagation Without Storing Activations.pdf

0.52MB