[논문 리뷰] Sub-Linear Memory: How to Make Performers SLiM

Valerii Likhosherstov, Krzysztof Choromański외 3명|arXiv (Cornell University)|2020. 12. 21.

Topic Modeling참고 문헌 31인용 수 5

한 줄 요약

이 논문은 메모리 효율적인 역전파 알고리즘을 제안하여, 시퀀스를 조각으로 처리함으로써 O(1) 메모리 사용을 달성하고, C=1일 때 하위선형 메모리 스케일링을 가능하게 한다. 이 방법은 근사치 없이 정확한 기울기를 계산하여, 스마트폰이나 이전 세대 GPU와 같은 저메모리 장치에서의 훈련 또는 미세조정에 완전한 후행 호환성을 제공한다.

ABSTRACT

The Transformer architecture has revolutionized deep learning on sequential data, becoming ubiquitous in state-of-the-art solutions for a wide variety of applications. Yet vanilla Transformers are notoriously resource-expensive, requiring $O(L^2)$ in serial time and memory as functions of input length $L$. Recent works proposed various linear self-attention mechanisms, scaling only as $O(L)$ for serial computation. We perform a thorough analysis of recent Transformer mechanisms with linear self-attention, Performers, in terms of overall computational complexity. We observe a remarkable computational flexibility: forward and backward propagation can be performed with no approximations using sublinear memory as a function of $L$ (in addition to negligible storage for the input sequence), at a cost of greater time complexity in the parallel setting. In the extreme case, a Performer consumes only $O(1)$ memory during training, and still requires $O(L)$ time. This discovered time-memory tradeoff can be used for training or, due to complete backward-compatibility, for fine-tuning on a low-memory device, e.g. a smartphone or an earlier-generation GPU, thus contributing towards decentralized and democratized deep learning.

연구 동기 및 목표

• 장시퀀스 트랜스포머의 높은 메모리 비용을 해결하기 위해, 특히 저메모리 장치에서의 훈련을 위한 대안을 모색한다.
• 기본 트랜스포머에서 O(L²) 메모리 복잡도로 인해 배치 크기가 감소하는 제한을 극복한다.
• 스마트폰이나 이전 세대 GPU와 같은 제약된 메모리가 있는 장치에서 Performer 모델의 훈련 또는 미세조정을 가능하게 한다.
• 메모리 소비를 줄이면서도 후행 호환성과 정확한 기울기 계산을 유지한다.
• 시간-메모리 트레이드오프를 탐색하여 사용자가 메모리 사용량을 조절할 수 있도록 하되, 병렬 계산 시간이 증가하는 대가를 지ay한다.

제안 방법

• C 크기의 조각으로 입력 시퀀스를 처리하며, C는 메모리-시간 트레이드오프를 제어한다.
• 알고리즘 1을 사용하여 순차적으로 각 조각별로 전방 및 역방향 전파를 수행하며, 손실 및 중간 버퍼만 저장한다.
• 역방향 전파 중에 기호적 표현 Φ^(n)을 구성하여 근사치 없이 정확한 기울기를 계산한다.
• 각 조각 내에서 전방 및 역방향 전파 간에 드롭아웃 마스크에 동일한 난수 시드를 재사용하여 확률적 계산을 지원한다.
• 기울기 누적 및 가역 레이어와 같은 기존 최적화 기법과 통합하여 추가로 메모리 소비를 줄인다.
• 정확한 기울기 계산을 통해 후행 호환성을 확보하여, 훈련 및 미세조정 모두에 적합한 방법이다.

실험 결과

연구 질문

• RQ1 근사 오차 없이 O(L) 이하로 Performer 모델의 메모리 소비를 줄일 수 있는가?
• RQ2 장시퀀스 훈련에서 메모리 사용과 병렬 계산 시간 사이의 트레이드오프는 어떠한가?
• RQ3 제안된 방법은 저메모리 장치에서 모델 정확도를 유지하면서도 미세조정에 사용될 수 있는가?
• RQ4 기존 딥러닝 최적화 기법(예: 기울기 누적, 가역 레이어)과의 호환성은 어떻게 되는가?
• RQ5 스라이스 기반 역전파가 전체 배치 훈련과 비교해 기울기 정확도를 유지하는가?

주요 결과

• C=1일 때 제안된 알고리즘이 O(1) 메모리 사용을 달성하며, 오직 O(L) 시간 복잡도만을 가지므로, 최소한의 메모리가 있는 장치에서도 훈련이 가능하다.
• 복사 작업(L=512)에서, 메모리 사용량은 전체 훈련 시 0.0449 GB에서 C=64일 때 0.0374 GB로 감소하였고, 시간은 0.0474초에서 0.1228초로 증가하였다.
• Penn Treebank(L=1024)에서, 메모리 사용량은 전체 훈련 시 0.300 GB에서 C=256일 때 0.231 GB로 감소하였고, 시간은 0.1377초에서 0.3060초로 증가하였다.
• Enwik8(L=4096)에서, 메모리 사용량은 전체 훈련 시 1.513 GB에서 C=1366일 때 0.909 GB로 감소하였고, 시간은 0.4598초에서 0.8654초로 증가하였다.
• 실험 결과, 모든 방법 간 성능이 구분되지 않아, 알고리즘이 정확한 기울기를 계산하고 모델 정확도를 유지함을 확인하였다.
• 이 방법은 완전히 후행 호환성이 있으며, 저메모리 장치에서의 미세조정을 지원하여 탈중앙화되고 민주화된 딥러닝을 가능하게 한다.