4. Reversible Transformer

Illustrating the Reformer를 보며 정리.

RevNet, Reversible residual Network

트랜스포머에서 인코더와 디코더 레이어를 여러개를 쌓을 때, Residual Network에서 역전파를 위해 그래디언트 값들을 저장하고 있다. 이때 저장하고 있는 값들이 매우 많아, 큰 트랜스포머 모델을 사용할때 메모리 부족 현상이 쉽게 나타난다. 이러한 문제를 RevNet을 이용해 메모리를 계산 문제로 바꿔서 해결을 시도한다.

각각의 피드포워드레이어와 인코더, 디코더 블록의 경우 Residual Network를 사용하는데, 메모리 부족을 reversible block으로 구성된 reversible residual network (RevNet)로 해결한다. RevNet의 경우 아래와 같이 ResNet을 재구성한다.

Reversible Transformer

리포머는 RevNet의 아이디어를 트랜스포머에 적용했다. 어텐션과 피드포워드 네트워크를 아래와 같이 바꾸었다.
$Y_1 = X_1 + Attention(X_2)$
$Y_2 = X_2 + FeedForward(Y_1)$

위와 같이 RevNet으로 변경 하면서, 기존에 학습 하면서 사용했던 메모리를 레이어 N개 일때, $1/N$으로 줄이게 되었다.

Reference

Illustrating the Reformer