Reformer - 트랜스포머(Transformer)를 발전시킨 Re(trans)former

트랜스포머를 통해서 자연어처리는 한단계, 아니 몇단계 성장했지만, 그 한계와 단점도 존재합니다. 가장 대표적인 것은 다음과 같습니다.

 

1. 어텐션 기반 모델이 갖는 메모리 문제

2. 트랜스포머 (기반) 모델이 정말 강인한가?

 

이 문제를 자세히 알아보기위해 먼저 트랜스포머를 살펴보도록 합시다.

 

ref: https://towardsdatascience.com/illustrating-the-reformer-393575ac6ba0

 

Reformer의 저자는 트랜스포머의 메모리 문제를 지적하였습니다. 그 부분들을 위 그림에서는 안경으로 표시했습니다. (위 레퍼런스에 자세히 설명되어있습니다.)

 

1. 검정 안경: attention 연산

• 길이 L의 시퀀스에 대한 어텐션 연산은 시간과 메모리 복잡도에서 모두 O(L²)이다. 시퀀스의 길이가 6만 4천이라고 한다면, 40억 정도 됩니다...

2. 빨간 안경: 많은 레이어

• 역전파를 위해 각 레이어의 활성화를 저장해야하므로, N 레이어의 모델은 하나의 레이어의 모델보다 N배 메모리를 더 소비합니다. BERT large 모델은 24 레이어를 갖고 있습니다.

3. 초록 안경: feed-forward 레이어의 깊이

• 중간의 feed-forward 레이어의 깊이는 대부분 어텐션 활성화의 깊이보다 매우 크다.

Reformer 모델은 Transformer의 위의 세가지 메모리 문제를 해결합니다. 이 메모리 문제를 해결하는데 사용하는 기술은 아래 두 가지를 결합한 것입니다.

 

1. LSH (Locality-Sensitive-Hashing)

• 긴 시퀀스에 집중(attention)하는 복잡도를 줄임

2. RevNet(Reversible Residual Network)

• 사용 가능한 메모리를 보다 효율적으로 사용

1. LSH (Locality-Sensitive-Hashing)을 활용한 NN

위의 검정 안경에 해당하는 O(L²)의 메모리 문제를 되돌아 보았을 때, 어텐션 연산은 전체 행렬을 계산하고 저장하는데 이것이 필요 없다는 점을 논문의 저자는 꼬집습니다. LSH는 고차원 데이터셋에서 가장 가까운 이웃(Nearest Neighbors; NN)을 효율적이고 근사적으로 검색할 수 있는 알고리즘입니다. LSH의 주요 아이디어는 두 지점 p와 q에 대해 q가 p에 가까우면 충분한 확률로 hash(q) == hash(p)를 가질 수 있습니다. 이를 달성하는 가장 간단한 방법은 임의의 초평면(hyperplane)으로 절단 공간을 유지하고 각 포인트의 해시 코드에 부호(pᵀH)를 추가하는 것입니다. 

 

ref: https://towardsdatascience.com/illustrating-the-reformer-393575ac6ba0

 

위의 그림과 같이 초평면(h)를 이용해 공간을 나누고, 그에 해당하는 부호를 추가하여 해시코드를 만듭니다. 그리고 해시 코드를 기준으로 포인트를 버킷으로 분할합니다. 위의 예시에서 a와 b는 같은 해시값을 갖는 것을 볼 수 있습니다. 따라서 동일한 bucket에 속하게됩니다. 그러면 a와 가장 가까운 이웃을 찾기 위해서는 전체 점들과 비교하는 것이 아니라, 하나의 버킷 안에서 즉, 같은 초 평면으로 나뉜 공간 내에서만 찾으면 됩니다. 즉, NN을 하는 공간이 좁아진 것이죠. 이를 통해 복잡도가 O(L²)에서 O(L log L)로 바뀝니다.

 

ref: REFORMER: THE EFFICIENT TRANSFORMER

 

이제 LSH를 어텐션에 적용하는 것을 생각해봅시다. Q와 K 행렬의 모든 벡터에 대해서 어텐션을 계산하는 대신에 LSH 해시 기법을 사용합니다. Q와 K 행렬의 LSH 해시를 만듭니다. 그리고, 같은 버킷 안에있는 q(Q의 부분)와 k(K의 부분)에 대해서만 어텐션을 계산합니다. 이 과정을 한번 하게 되면, 유사한 항목이 다른 버킷에 놓일 수 있기 때문에 문제가 생길 수 있습니다. 따라서 위의 과정을 여러번 반복하는 Multi-round로 진행해야합니다. 이 과정도 illustrating the reformer 글에 잘 나타나있습니다. (정말 이만한 글이 없습니다. 꼭 보셔야합니다!)

 

ref: https://towardsdatascience.com/illustrating-the-reformer-393575ac6ba0

 

위 그림과 같이 먼저 초평면(h)를 이용해 공간을 나눕니다. 그리고 그것을 버킷으로 만들어 정렬합니다. (색으로 구분) 시퀀스를 청크 단위로 나누고, 버킷 내에서의 어텐션 관계를 계산합니다.

 

2. Reversible Transformer (Residual Layers)

이번에는 빨간, 초록 안경에 해당하는 많은 수의 인코더 및 디코더의 레이어와 feed-forward 레이어의 깊이 문제를 해결하는 RevNet(Reversibvle residual Network)입니다. ResNet은 역전파동안 그래디언트를 계산하기 위해 각 레이어의 활성화를 메모리에 저장해야하므로 메모리측면에서 병목 현상이 발생합니다. 

 

이 문제를 해결하기 위한 RevNet은 가역적(Reversible)인 블록으로 구성됩니다. 가격적 블록에서 각 레이어의 활성화는 후속 레이어의 활성화에서 정확하게 재구성될 수 있으므로 활성화를 메모리에 저장하지 않고 역전파를 수행할 수 있어서 메모리 소비 문제를 해결합니다. 

 

ref: https://towardsdatascience.com/illustrating-the-reformer-393575ac6ba0

 

Reformer는 RevNet 블록 내부의 어텐션 레이어와 feed-forward 레이어를 결합합니다. 즉, 위의 그림에서 F는 어텐션 레이어, G는 feed-forward 레이어로 만들어주어 Transformer의 형태로 만들어줍니다. 

 

Y₁ = X₁ + Attention(X₂),
Y₂ = X₂+ FeedForward(Y₁)

 

residual 대신에 reversible residual 레이어를 사용하여, 학습 과정 동안에 N번이 아니라 한번 활성화를 저장하여 메모리 낭비를 막습니다.

 

3. 결론

LSH와 RevNet을 적용한 것이 성능에 어떠한 영향을 미쳤는지 결과적으로 살펴봅시다. 

 

먼저 LSH 어텐션의 경우, 해시 라운드의 수에 다라 성능이 나뉘었습니다. 해싱을 안쓰고 전체 어텐션을 계산하는 것이 성능이 가장 좋았고, 해시 라운드 수가 많을 수록 성능이 유사하게 나타났습니다.

ref: REFORMER: THE EFFICIENT TRANSFORMER

 

RevNet을 적용시킨 것은 성능에서 차이가 거의 없는 것을 볼 수 있습니다.