[논문 리뷰] Compressing Gradient Optimizers via Count-Sketches

Ryan Spring, Anastasios Kyrillidis외 2명|arXiv (Cornell University)|2019. 02. 01.

Advanced Neural Network Applications참고 문헌 32인용 수 5

한 줄 요약

이 논문은 Adam, 모멘타임, 아다그레드와 같은 1차 최적화 방법의 보조 변수를 Count-Sketch 데이터 구조를 사용하여 압축하여 대규모 딥러닝 모델에서 메모리 사용량을 줄이는 방법을 제안한다. 임bedding 및 소프트맥스 레이어의 희소성 특성을 활용함으로써 수렴 속도와 모델 성능를 유지하면서도 메모리 사용량을 최대 25% 감소시키고, 4950만 개의 클래스를 가진 극단적 분류 작업에서 학습 속도를 38% 향상시킨다.

ABSTRACT

Many popular first-order optimization methods (e.g., Momentum, AdaGrad, Adam) accelerate the convergence rate of deep learning models. However, these algorithms require auxiliary parameters, which cost additional memory proportional to the number of parameters in the model. The problem is becoming more severe as deep learning models continue to grow larger in order to learn from complex, large-scale datasets. Our proposed solution is to maintain a linear sketch to compress the auxiliary variables. We demonstrate that our technique has the same performance as the full-sized baseline, while using significantly less space for the auxiliary variables. Theoretically, we prove that count-sketch optimization maintains the SGD convergence rate, while gracefully reducing memory usage for large-models. On the large-scale 1-Billion Word dataset, we save 25% of the memory used during training (8.6 GB instead of 11.7 GB) by compressing the Adam optimizer in the Embedding and Softmax layers with negligible accuracy and performance loss. For an Amazon extreme classification task with over 49.5 million classes, we also reduce the training time by 38%, by increasing the mini-batch size 3.5x using our count-sketch optimizer.

연구 동기 및 목표

• Adam과 모멘타임과 같은 1차 최적화 방법에서 모델 크기에 비례하여 증가하는 보조 파rameter로 인한 증가하는 메모리 오버헤드를 해결하기 위해.
• 딥러닝에서 최적화 방법 상태의 메모리 소비를 줄여 더 큰 배치 크기와 더 표현력 있는 모델을 가능하게 하기 위해.
• 임베딩 및 소프트맥스 레이어의 구조적 희소성을 활용하여 성능 저하 없이 최적화 방법 상태를 효율적으로 압축하기 위해.
• 완전 정밀도 최적화 방법 상태에 대한 이론적으로 탄탄한 메모리 효율적인 대안을 제공하여 SGD 수렴 속도를 유지하기 위해.
• 최대 규모의 언어 모델링 및 극단적 분류 작업에서 실용적인 이점을 입증하기 위해, 정확도 손실를 최소화한다.

제안 방법

• 임베딩 및 소프트맥스 레이어의 2차 최적화 방법 상태(예: Adam의 속도 및 모멘타임)를 압축된 표현으로 유지하기 위해 Count-Sketch 데이터 구조를 사용한다.
• 전체 최적화 방법 상태 대신 작은 고정 크기의 스케치([3, 266, 1024] 등)만 저장함으로써 2차 모멘트 추정치에 대해 최대 99%의 메모리 감소를 달성한다.
• 기울기 업데이트를 버킷에 매핑하기 위한 무작위 해시 함수를 사용하여 스케치를 점진적으로 업데이트함으로써 상수 시간 연산을 가능하게 한다.
• 파rameter 업데이트 중에 스케치를 쿼리하여 전체 최적화 방법 상태를 근사화하며, 이론적 경계를 통해 오차를 통제한다.
• 해시 함수 수와 버킷 수를 조정하여 스케치 크기와 정확도 사이에 부드러운 트레이드오프를 허용한다.
• 표준 최적화 방법 상태를 압축된 변형으로 대체하기 위해 스케치 기반 최적화 방법을 직접 학습 파이프라인에 통합한다.

실험 결과

연구 질문

• RQ1 Count-Sketch를 사용하여 딥러닝에서 보조 최적화 방법 변수를 압축할 수 있는가? 이로 인해 수렴 또는 성능 저하가 발생하지 않는가?
• RQ2 대규모 모델에서 스케치 기반 최적화 방법의 메모리 사용량과 학습 시간은 표준 Adam과 비교해 어떻게 되는가?
• RQ3 압축된 최적화 방법 상태를 사용할 경우 배치 크기를 얼마나 늘릴 수 있으며, 이는 학습 속도에 어떤 영향을 미치는가?
• RQ4 이론적 분석을 통해 스케치 기반 최적화 방법은 표준 1차 최적화 방법과 동일한 수렴 속도를 유지하는가?
• RQ5 4950만 개 이상의 클래스를 가진 극도로 대규모인 분류 작업에 이 방법을 효과적으로 적용할 수 있는가?

주요 결과

• 10억 단어 언어 모델링 작업에서 Count-Sketch 최적화 방법은 정확도나 성능 손실 없이 메모리 사용량을 25% 감소시켰다(11.7GB에서 8.6GB로).
• 4950만 개의 클래스를 가진 극단적 분류 작업에서, Count-Sketch 최적화 방법을 사용해 미니배치 크기를 3.5배 증가시켜 학습 시간을 38% 단축시켰다.
• Count-Sketch 최적화 방법은 Recall@100이 0.6889를 기록했고, 표준 Adam 기준 0.6881과 비교해 거의 차이가 없었으며, 성능 저하가 거의 없음을 보여주었다.
• 1% 크기의 스케치를 사용하여 2차 모멘트 추정치를 압축함으로써 모델당 메모리 비용을 4GB에서 2.6GB로 35% 감소시켰다.
• 이론적 분석을 통해 Count-Sketch 최적화 방법이 완전 정밀도 기준선과 동일한 SGD 수렴 속도를 유지함을 확인했다.
• 메모리 사용량과 근사 오차 사이에 실용적인 트레이드오프를 가능하게 하였으며, 계산 오버헤드도 최소화하였다.