[논문 리뷰] Few-Bit Backward: Quantized Gradients of Activation Functions for Memory Footprint Reduction

Georgii Novikov, Daniel Bershatsky외 4명|arXiv (Cornell University)|2022. 02. 01.

Advanced Neural Network Applications인용 수 5

한 줄 요약

이 논문은 딥러닝 학습 중 메모리 사용량을 줄이기 위해 점별 활성화 함수의 기울기를 조각별로 일정한 근사값을 사용해 양자화하는 방법을 제안한다. 역전파 과정에서 전체 정밀도 입력 대신 3비트 인덱스만 저장함으로써, 표준 활성화 함수들인 ReLU, GELU, Swish와의 호환성이 있는 드롭인 대체 방법으로, 벤치마크에서 정확도 저하 없이 최대 40%의 메모리 절감을 달성한다. 최적의 양자화 체계는 동적 프로그래밍을 통해 계산된다.

ABSTRACT

Memory footprint is one of the main limiting factors for large neural network training. In backpropagation, one needs to store the input to each operation in the computational graph. Every modern neural network model has quite a few pointwise nonlinearities in its architecture, and such operation induces additional memory costs which -- as we show -- can be significantly reduced by quantization of the gradients. We propose a systematic approach to compute optimal quantization of the retained gradients of the pointwise nonlinear functions with only a few bits per each element. We show that such approximation can be achieved by computing optimal piecewise-constant approximation of the derivative of the activation function, which can be done by dynamic programming. The drop-in replacements are implemented for all popular nonlinearities and can be used in any existing pipeline. We confirm the memory reduction and the same convergence on several open benchmarks.

연구 동기 및 목표

• 딥 네ural 네트워크 학습 중 백프로파게이션 동안 저장되는 활성화 텐서의 메모리 사용량을 줄이는 것.
• 현대 아키텍처에서 점별 비선형 함수의 전체 정밀도 입력을 저장하는 데 발생하는 간과된 메모리 비용을 해결하는 것.
• 데이터에 의존하지 않고 기존 학습 파ip라인을 수정하지 않고도 활성화 함수를 위한 저비트 기울기 저장을 가능하게 하는 드롭인 대체 방법을 개발하는 것.
• 다양한 아키텍처와 작업에서 모델 수렴성과 정확도를 유지하면서도 상당한 메모리 절감을 달성하는 것.
• 재학습이나 아키텍처 변경 없이도 모든 모델, 특히 트랜스포머에 적용 가능한 일반화 가능한 프레임워크를 제공하는 것.

제안 방법

• 기울기의 도함수에 대한 최적의 조각별 일정 근사값을 계산하기 위해 동적 프로그래밍을 사용하여 양자화 오차를 최소화한다.
• 백워드 패스에서 전체 정밀도 활성화 입력 대신 바이너리 인덱스(예: 3비트)만 저장함으로써 텐서 요소당 메모리 사용량을 크게 줄인다.
• 원래 활성화 입력 대신 양자화된 도함수 값을 저장함으로써 손실이 있는 압축을 실현하며 기울기 영향은 극히 미미하게 유지된다.
• PyTorch 파이프라인에 원활하게 통합되는 일반적인 활성화 함수들(ReLU, GELU, Swish, 시그모이드)을 위한 드롭인 대체 구현을 제공한다.
• 데이터 분포나 모델 구조에 의존하지 않고 트랜스포머 및 EfficientNet을 포함한 모든 모델에 일반적으로 적용 가능한 방법이다.
• 최소한의 성능 오버헤드로 효율적인 추론과 학습을 가능하게 하는 CUDA 최적화 커널을 구현한다.

실험 결과

연구 질문

• RQ1 점별 활성화 함수의 기울기를 소수 비트로 양자화함으로써 백프로파게이션 중 메모리 사용량을 줄일 수 있을까? 정확도 저하 없이도 가능한가?
• RQ2 저비트 표현을 위한 최적의 조각별 일정 근사값은 무엇이며, 양자화 오차를 최소화하는가?
• RQ3 3비트 또는 4비트 양자화를 통해 트랜스포머, EfficientNet, GPT-2와 같은 다양한 모델에서 실제로 얼마나 많은 메모리 절감이 가능할까?
• RQ4 제안된 방법은 다양한 모델과 학습 단계(예: 사전학습된 모델의 피니튜닝 포함)에서 드롭인 대체로 적용 가능한가?
• RQ5 데이터에 의존하지 않는 설계를 고려할 때, 이 방법은 다양한 데이터 분포와 모델 아키텍처에서 효과를 유지하는가?

주요 결과

• GPT2와 RoBERTa-large 같은 모델에서 활성화 요소당 3비트 인덱스만 저장함으로써 최대 40%의 메모리 절감을 달성한다.
• 3비트 양자화를 통해 ResNet, EfficientNet, BERT 기반 모델을 포함한 여러 벤치마크에서 평균 20%의 메모리 절감을 기록한다.
• 모든 평가된 모델(ResNet-50, EfficientNet-B7, RoBERTa 포함)에서 전체 정밀도 학습과 동일한 학습 수렴성과 테스트 정확도를 유지한다.
• 기울기 도함수의 최적의 조각별 일정 근사값은 동적 프로그래밍을 통해 계산되며, ReLU, GELU, Swish, 시그모이드를 포함한 모든 테스트된 활성화 함수에 효과적이다.
• 기존 학습 파이프라인과 호환되며, 아키텍처 변경 없이도 학습의 모든 단계(예: 사전학습된 모델의 피니튜닝 포함)에 적용 가능하다.
• CUDA 구현은 낮은 오버헤드를 제공하여 실제 학습 워크플로우에 실용적으로 적용 가능하며 성능 영향을 최소화한다.