[논문 리뷰] ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training

Jianfei Chen, Lianmin Zheng외 5명|arXiv (Cornell University)|2021. 04. 29.

Advanced Memory and Neural Computing참고 문헌 43인용 수 8

한 줄 요약

ActNN는 랜덤 양자화를 사용해 활성화 값을 2비트 정밀도로 압축함으로써 메모리 효율적인 훈련 프레임워크를 제안한다. 이는 정확도 손실이 거의 없이 메모리 사용을 크게 줄여주며, 이론적 수렴 보장을 제공하고 이질성 인식 혼합 정밀도 양자화를 적용하여 활성화 메모리 사용을 12배 감소시켜, 최소한의 코드 변경으로 파이토치에서 최대 14배 큰 배치 크기를 허용한다.

ABSTRACT

The increasing size of neural network models has been critical for improvements in their accuracy, but device memory is not growing at the same rate. This creates fundamental challenges for training neural networks within limited memory environments. In this work, we propose ActNN, a memory-efficient training framework that stores randomly quantized activations for back propagation. We prove the convergence of ActNN for general network architectures, and we characterize the impact of quantization on the convergence via an exact expression for the gradient variance. Using our theory, we propose novel mixed-precision quantization strategies that exploit the activation's heterogeneity across feature dimensions, samples, and layers. These techniques can be readily applied to existing dynamic graph frameworks, such as PyTorch, simply by substituting the layers. We evaluate ActNN on mainstream computer vision models for classification, detection, and segmentation tasks. On all these tasks, ActNN compresses the activation to 2 bits on average, with negligible accuracy loss. ActNN reduces the memory footprint of the activation by 12x, and it enables training with a 6.6x to 14x larger batch size.

연구 동기 및 목표

• 제한된 GPU 메모리 증가에도 불구하고 대규모 신경망 훈련에서 증가하는 메모리 병목 현상 해결
• 다양한 아키텍처에 걸쳐 활성화 값 압축 훈련(ActNN)의 일반적인 이론적 수렴 보장 제공
• 특성 차원, 샘플, 레이어 간 활성화 이질성을 활용한 양자화 전략 설계로 기울기 분산 최소화
• 기존 딥러닝 프레임워크(예: 파이토치)에 최소한의 코드 변경으로 쉽게 구현 가능한 ActNN의 실용적 구현 제공
• 시각 작업 전반에서 정확도 저하가 최소한인 고압축 비율(예: 2비트) 달성

제안 방법

• ActNN에서 양자화된 기울기를 정밀도가 전체인 기울기의 불편 추정량으로 간주하여 이론적 수렴 보장
• 기울기 분산에 대한 정확한 표현 유도하여 양자화가 훈련 안정성에 미치는 영향 정량화
• 특성 그룹별 활성화 통계에 기반해 양자화 수준을 적응적으로 조정하는 그룹별 양자화기 도입
• 고정된 메모리 예산 내에서 기울기 분산을 최소화하는 세밀한 혼합 정밀도 양자화 전략 개발
• 파이토치 라이브러리로 ActNN를 구현하여 즉시 사용 가능한 레이어(예: actnn.Conv2d) 제공로 원활한 통합 구현
• 메모리 절감과 훈련 효율성의 균형을 맞추기 위해 실시간으로 양자화 전략 조정 적용

실험 결과

연구 질문

• RQ1 일반적인 신경망 아키텍처에 대해 활성화 값 압축 훈련이 이론적으로 수렴 보장될 수 있는가?
• RQ2 양자화가 기울기 분산에 어떤 영향을 미치며, 이를 정확히 모델링하여 양자화 설계를 이끌 수 있는가?
• RQ3 활성화 이질성을 활용한 혼합 정밀도 양자화 전략이 더 높은 압축 비율과 낮은 정확도 손실을 달성할 수 있는가?
• RQ4 2비트 활성화 압축이 다양한 시각 작업(예: 세그멘테이션, 디텍션)에 효과적으로 적용될 수 있는가?
• RQ5 고정된 메모리 제약 조건 하에서 ActNN는 모델의 깊이, 너비 또는 해상도를 어느 정도까지 확장할 수 있는가?

주요 결과

• ActNN는 이미지 분류, 디텍션, 세그멘테이션 작업 전반에서 평균적으로 활성화 값을 2비트로 압축하면서 정확도 손실이 0.5% 미만이다.
• 동일한 GPU에서 배치 크기를 6.6배에서 14배까지 크게 늘릴 수 있어 훈련 처리량을 크게 향상시킨다.
• 동일한 메모리 예산 하에서 ActNN는 ResNet-152를 최대 6.4배 더 깊게, 3.7배 더 넓게, 또는 3.1배 더 높은 해상도로 확장할 수 있으며, 원래 훈련 처리량의 64%~155%를 유지한다.
• ActNN는 활성화 메모리 사용량을 12배 감소시켜 단일 GPU에서 고해상도 세그멘테이션 및 디텍션 작업을 위한 대배치 훈련을 가능하게 한다.
• 자기지도 학습에서 ActNN는 2비트 활성화를 사용해 1024배치 크기로 ImageNet 상위-1 정확도 72.65%를 달성했으며, 전체 정밀도 기준선과 동일하거나 이를 초월한다.
• 1.25비트 활성화 조차도 ActNN는 수렴하고 합리적인 성능을 내며, 이전 연구에서 4비트 압축에 국한된 것보다 뛰어난 성능을 발휘한다.