[논문 리뷰] Memory Efficient Optimizers with 4-bit States

Bingrui Li, Jianfei Chen외 1명|arXiv (Cornell University)|2023. 09. 04.

Machine Learning and Data ClassificationComputer Science인용 수 3

한 줄 요약

이 논문은 Adam 유사 최적화기의 1차 및 2차 모멘텀에 대한 새로운 양자화 기법을 도입하여 메모리 효율적인 딥 러닝 훈련을 위한 4비트 최적화기를 제안한다. 더 작은 블록 크기, 2차 모멘텀을 위한 제로포인트 없는 선형 양자화기, 그리고 랭크-1 정규화를 통해 다양한 NLP, 비전 및 지시 훈련 작업에서 8비트 기준선 대비 50% 메모리 감소를 달성하면서도 정밀도가 높은 수준의 수렴 정확도를 확보한다.

ABSTRACT

Optimizer states are a major source of memory consumption for training neural networks, limiting the maximum trainable model within given memory budget. Compressing the optimizer states from 32-bit floating points to lower bitwidth is promising to reduce the training memory footprint, while the current lowest achievable bitwidth is 8-bit. In this work, we push optimizer states bitwidth down to 4-bit through a detailed empirical analysis of first and second moments. Specifically, we find that moments have complicated outlier patterns, that current block-wise quantization cannot accurately approximate. We use a smaller block size and propose to utilize both row-wise and column-wise information for better quantization. We further identify a zero point problem of quantizing the second moment, and solve this problem with a linear quantizer that excludes the zero point. Our 4-bit optimizers are evaluated on a wide variety of benchmarks including natural language understanding, machine translation, image classification, and instruction tuning. On all the tasks our optimizers can achieve comparable accuracy with their full-precision counterparts, while enjoying better memory efficiency.

연구 동기 및 목표

• 고정밀도 최적화기 상태(예: 32비트 Adam 상태)로 인해 발생하는 대규모 신경망 훈련의 메모리 병목 현상을 해결한다.
• 현재 8비트 최소 기준에서 4비트로 최적화기 상태 비트 폭을 줄여 메모리 사용을 추가로 압축한다.
• 기존의 블록 단위 양자화의 한계를 극복하여 1차 및 2차 모멘텀 내 복잡한 아웃라이어 패턴을 포착하지 못하는 문제를 해결한다.
• 2차 모멘텀 양자화에서 발생하는 제로포인트 문제를 해결하여 파rameter 업데이트 방향의 심각한 편차를 방지한다.
• 안정적인 임bedding 레이어에 의존하지 않고도 수렴 속도와 정확도를 유지하는 강력한 4비트 최적화기를 개발한다.

제안 방법

• 1차 모멘텀 내 국소 아웃라이어 패턴을 더 잘 포착하기 위해 더 작은 블록 크기를 도입한다.
• 비제로 값이 0으로 양자화되는 것을 방지하여 업데이트 방향을 왜곡하는 문제를 막기 위해 2차 모멘텀을 위한 제로포인트 없는 선형 양자화기를 제안한다.
• 아웃라이어 구조를 더 효과적으로 모델링하기 위해 2차 모멘텀 근사치를 향상시키기 위해 랭크-1 정규화를 설계한다.
• 4비트 양자화된 1차 모멘텀과 인수분해된 2차 모멘텀을 조합하여 메모리 효율성을 향상시킨 하이브리드 최적화기를 구축한다.
• 각 반복 단위로 최적화기 상태를 압축하고 복원하는 압축 기반 최적화 프레임워크에 제안된 양자화기를 통합한다.
• 실험적 분석을 통해 양자화기 설계를 검증하고 4비트 정밀도에서 손실 없는 수렴을 보장한다.

Figure 1 : Visualization of the first moment in the layers.3.blocks.1.mlp.fc1 layer in a Swin-T model. (a): Magnitude of the first moment. (b): Histogram of the first moment. (c): Moment approximated by B128/DE . (d): Moment approximated by B2048/DE .

실험 결과

연구 질문

• RQ1 모델 수렴 또는 정확도를 저하시키지 않고 최적화기 상태를 4비트로 압축할 수 있는가?
• RQ2 1차 및 2차 모멘텀의 4비트 압축 과정에서 발생하는 주요 양자화 과제는 무엇인가, 특히 아웃라이어 패턴에 대한 영향은?
• RQ3 2차 모멘텀 양자화에서 발생하는 제로포인트 문제의 최적화 안정성에 대한 영향은 무엇이며, 이를 완화할 수 있는가?
• RQ4 제안된 양자화 기법이 안정적인 임베딩 레이어와 같은 아키텍처 수정 없이도 정밀도가 높은 최적화기 성능을 달성할 수 있는가?
• RQ5 다양한 작업에서 기존 8비트 기준선 대비 4비트 최적화기의 메모리 및 훈련 효율성 향상은 어느 정도인가?

주요 결과

• 제안된 4비트 최적화기는 NLP, 기계 번역, 이미지 분류, 지시 훈련을 포함한 모든 평가 벤치마크에서 정밀도가 높은 Adam과 유사한 수렴 속도와 최종 정확도를 달성한다.
• 8비트 최적화기 대비 메모리 소비를 50% 감소시켜 기존 저정밀도 최적화기 중에서 가장 낮은 메모리 프로파일을 확보한다.
• 2차 모멘텀을 위한 제로포인트 없는 선형 양자화기는 큰 파arameter 업데이트 편차를 효과적으로 방지하여 4비트 정밀도에서 안정적인 훈련을 가능하게 한다.
• 랭크-1 정규화는 특히 복잡한 아웃라이어 패턴이 존재할 경우 2차 모멘텀 근사치를 크게 향상시킨다.
• 안정적인 임베딩 레이어가 없더라도 4비트 최적화기는 손실 없는 수렴을 달성하여 다양한 모델 아키텍처에 걸쳐 강건성과 일반화 능력을 입증한다.
• 최적화기 오프로딩을 통해 4비트 최적화기는 더 낮은 정밀도의 상태로 인한 통신 비용 감소로 언어 모델 파인튜닝 시 처리량을 향상시킨다.

Figure 2 : Outlier patterns vary across two first moment tensors. (a): outliers lie in fixed rows (dimension 0). (b): outliers lie in fixed columns (dimension 1).