[논문 리뷰] GACT: Activation Compressed Training for Generic Network Architectures

Xiaoxuan Liu, Lianmin Zheng외 10명|arXiv (Cornell University)|2022. 06. 22.

Advanced Neural Network Applications인용 수 5

한 줄 요약

GACT는 활성화 압축 학습(Activation Compressed Training, ACT)을 위한 일반적 프레임워크로, 수학적 수렴 보장을 갖는 동적 활성화 압축을 통해 다양한 신경망 아키텍처—합성곱, 트랜스포머, 그래프 네트워크—의 메모리 효율적인 학습을 가능하게 한다. 이는 활성화 메모리를 최대 8.1배까지 감소시켜 정확도 손실이 거의 없는 조건에서 최대 24.7배 큰 배치 크기를 허용하며, 기존 모델 및 메모리 절약 기법과의 원활한 통합을 위한 파이토치 라이브러리로 구현되어 있다.

ABSTRACT

Training large neural network (NN) models requires extensive memory resources, and Activation Compressed Training (ACT) is a promising approach to reduce training memory footprint. This paper presents GACT, an ACT framework to support a broad range of machine learning tasks for generic NN architectures with limited domain knowledge. By analyzing a linearized version of ACT's approximate gradient, we prove the convergence of GACT without prior knowledge on operator type or model architecture. To make training stable, we propose an algorithm that decides the compression ratio for each tensor by estimating its impact on the gradient at run time. We implement GACT as a PyTorch library that readily applies to any NN architecture. GACT reduces the activation memory for convolutional NNs, transformers, and graph NNs by up to 8.1x, enabling training with a 4.2x to 24.7x larger batch size, with negligible accuracy loss. We implement GACT as a PyTorch library at https://github.com/LiuXiaoxuanPKU/GACT-ICML.

연구 동기 및 목표

• 아키텍처나 연산자에 특화되지 않은 일반적인 활성화 압축 프레임워크를 개발함으로써 다양한 신경망 아키텍처의 학습을 가능하게 하기.
• 연산자 유형이나 모델 구조에 대한 사전 지식 없이도 활성화 압축 학습의 수렴성을 증명하기.
• 기울기 영향도 추정을 기반으로 런타임 시에 각 텐서에 대해 최적의 압축 비율을 선택하는 적응형 알고리즘 설계하기.
• 기존 학습 파이프라인 및 기타 메모리 절약 기법(예: 기울기 체크포인팅, 스왑)과의 원활한 통합을 위한 파이토치 라이브러리로 GACT를 구현하기.

제안 방법

• 아키텍처 간 일반화를 위해 ACT를 연산자 기반의 계산 그래프로 수식화하여 레이어 특화 가정을 피하기.
• 선형화 근사법을 사용해 ACT의 확률적 기울기를 분석하고, 편향이 없으며 구조화된 분산을 가지는 것을 증명함으로써 수치적 추정이 가능함을 보장하기.
• 다양한 압축 전략 하에서 기울기 분산을 예측할 수 있는 수치 알고리즘 개발을 통해 최적의 전략 선택 가능하게 하기.
• 각 텐서의 기울기 업데이트에 대한 민감도를 기반으로 최적의 압축 비율을 자동으로 결정하기 위해 정수계획법 근사를 활용하기.
• 사용자 정의 연산자 지원 및 기존 메모리 효율 기법과의 통합을 가능하게 하는 여러 최적화 수준을 갖춘 파이토치 라이브러리로 GACT 구현하기.
• 학습 중 각 텐서의 기울기 분산에 대한 영향도를 추정함으로써 런타임 시 압축 비율의 적응적 조정 가능하게 하기.

실험 결과

연구 질문

• RQ1 아키텍처에 대한 특수한 가정 없이도 일반적인 신경망 아키텍처에서 활성화 압축 학습이 수렴 보장 가능하게 만들 수 있는가?
• RQ2 각 텐서에 대해 압축 비율을 어떻게 적응적으로 선택할 수 있을까? 이는 학습 안정성과 정확도 유지에 기여하는가?
• RQ3 사용자 정의 연산자 포함한 임의의 연산자를 지원하는 일반 목적의 ACT 프레임워크를 커스터마이제이션 없이 구축할 수 있는가?
• RQ4 일반적인 딥러닝 모델에서 메모리 감소, 학습 속도, 정확도 손실 간의 상호 교환 관계는 어떠한가?
• RQ5 GACT는 기존 메모리 효율 기법(예: 기울기 체크포인팅, 스왑)과 어떻게 통합되고 향상되는가?

주요 결과

• GACT는 합성곱 네트워크, 트랜스포머, 그래프 신경망에서 활성화 메모리를 최대 8.1배까지 감소시키며, 정확도 손실은 거의 없음.
• 동일한 GPU 메모리 예산 하에서 GACT는 최대 24.7배 큰 배치 크기로 학습이 가능해져 학습 처리량을 크게 향상시킴.
• ResNet-152에서 GACT는 16GB GPU 메모리 조건에서 깊이 7.0배, 넓이 3.6배, 해상도 3.0배까지 확장 가능함.
• BERT-large에서는 동일한 메모리 제약 하에서 깊이 2.0배 또는 넓이 1.6배까지 확장 가능함.
• GACT와 기울기 체크포인팅을 조합하면 배치 크기 288일 때 피크 활성화 메모리를 최대 5.4배 감소시키며 정확도 손실 없음.
• GACT와 스왑을 조합하면 압축된 텐서의 CPU-GPU 전송 속도 향상으로 최대 2.3배 빠른 학습 속도를 기록하며, 프리패칭은 약 7%의 추가 속도 향상과 최소한의 메모리 오버헤드를 제공함.