기술 요약: GEFEN: 최적화된 확률적 옵티마이저 (OPTIMIZED STOCHASTIC OPTIMIZER)

문제 정의

현대 딥러닝은 다양한 아키텍처에서 견고한 성능을 보이는 AdamW 옵티마이저에 크게 의존하고 있습니다. 그러나 AdamW는 모든 파라미터에 대해 1차 및 2차 모멘트 이동 평균을 유지하며, 일반적으로 이는 FP32로 저장됩니다. 이로 인해 옵티마이저 상태 메모리 점유율은 모델 파라미터 크기의 약 2배(파라미터당 약 8바이트)에 달하게 됩니다. 이는 13B 규모의 대규모 모델의 경우, 이 상태값만으로도 약 97 GiB의 메모리를 소비하여, 학습 가능한 모델 크기, 최대 배치 크기, 그리고 메모리가 제한된 하드웨어에서의 학습 처리량(throughput)을 심각하게 제약합니다.

기존의 메모리 효율적인 대안들은 다음과 같은 상당한 트레이드오프를 수반합니다:

• 양자화된 옵티마이저 (예: Adam8bit, Adam4bit): 메모리를 줄여주지만, 수동으로 설계된 최적화되지 않은 코드북이나 고정된 블록 크기에 의존하며(하이퍼파라미터 도입), 범용성이 떨어질 수 있습니다 (예: 임베딩 양자화 회피 등).
• 구조적 옵티마이저 (예: Adam-mini): Hessian의 블록 대각 구조를 모사하는 수동 규칙을 기반으로 2차 모멘트 통계량을 공유합니다. 그러나 이러한 방식은 아키텍처별 메타데이터(예: 어텐션 헤드 수)를 필요로 하며, 왜 Hessian 정렬된 그룹화가 유익한지에 대한 강력한 이론적 근거가 부족합니다.

본 연구의 목표는 AdamW 수준의 성능을 유지하면서도 추가적인 설정, 하이퍼파라미터, 또는 학습률 재조정 없이 메모리 점유율을 (Adam-mini보다 더 크게) 획기적으로 줄이는 옵티마이저를 설계하는 것입니다.

방법론

제안된 방법인 Gefen은 자동 2차 모멘트 공유와 학습된 1차 모멘트 양자화를 통해 두 가지 주요 메커니즘으로 메모리 감소를 달연합니다.

1. 이론적 동기

이 접근 방식은 혼합 Hessian 항목이 큰 파라미터들에 대해, 이들의 제곱 그래디언트 비율이 1에 가깝게 제한된다는 이론적 결과(Theorem 4.1)에 근거합니다. 이는 Hessian 친화도가 높은 파라미터들이 자연스럽게 유사한 제곱 그래디언트 크기를 갖게 됨을 의미하며, 따라서 2차 모멘트 추정치를 공유하기에 이상적인 후보가 됨을 시사합니다. 규모가 큰 모델에서 Hessian을 직접 계산하는 것은 불가능하므로, 저자들은 초기 학습 단계에서의 제곱 그래디언트 구조가 이러한 Hessian 구조와 상관관계가 있음을 관찰했습니다.

2. 자동 블록 분할 (Automatic Block Partitioning)

수동 규칙이나 Hessian 계산에 의존하는 대신, Gefen은 첫 번째 최적화 단계($g_0$)의 제곱 그래디언트를 통해 최적의 블록 구조를 추론합니다.

• 알고리즘: 평탄화된 파라미터 텐서에 대해, 알고리즘은 텐서 크기 $n$의 모든 적절한 약수 $p$를 테스트합니다. 제곱 그래디언트를 크기가 $p$인 블록으로 재형성(reshape)하고, 블록 내 이질성 $E(p)$ (블록 내 제곱 그래디언트의 분산)를 계산합니다.
• 선택 전략: 알고리즘은 블록 크기가 커짐에 따라 $E(p)$가 처음으로 뚜렷하게 감소하는 지점의 주기 $p$를 선택합니다. 이는 메모리 효율성을 위해 더 큰 블록을 선호하면서도 그래디언트 유사성을 보장하기 위함입니다. 만약 의미 있는 감소가 발견되지 않거나 최적의 $p < 8$인 경우, 공유를 하지 않는 기본값인 $p=1$로 설정합니다.
• 사용: 이 단일 분할 체계는 2차 모멘트 공유와 1차 모멘트 양자화 모두에 사용됩니다.

3. 정밀한 동적 계획법(Dynamic Programming) 양자화

1차 모멘트(모멘텀)의 메모리를 줄이기 위해, Gefen은 고정되거나 휴리스틱한 방식이 아닌 학습 가능한 정밀한 양자화 코드북을 채택합니다.

• 코드북 학습: 초기 그래디언트를 사용하여 정규화된 그래디언트 값의 히스토그램을 구성합니다. 그 후, 주어진 분포와 비트 너비에 대해 정확한 최적의 코드북을 찾기 위해 동적 계획법(DP) (Wang & Song, 2011에서 영감을 받음)을 사용하여 최적의 스칼라 양자화 문제를 해결합니다.
• 안정성: "수축(contraction)"(양자화 범위가 시간이 지남에 따라 줄어드는 현상)이나 초기화 민감성 문제를 겪을 수 있는 Lloyd-Max 방식과 달리, Gefen은 코드북의 끝점을 $-1$과$1$로 고정합니다. 이를 통해 학습 과정 전반에 걸쳐 양자화 범위가 유지되도록 합니다.
• 효율성: 코드북은 학습 시작 시 한 번 학습되며(히스토그램 생성을 위한 CUDA 및 DP를 위한 CPU/Numba 사용), 전체 학습 과정 동안 재사용됩니다.

4. 알고리즘 흐름

Gefen은 AdamW의 드롭인 교체(drop-in replacement)로 작동합니다:

1. 초기화: $g_0$를 기반으로 파라미터를 분할하고 양자화 코드북을 학습합니다.
2. 업데이트 단계:
   • 그래디언트를 계산합니다.
   • 저장된 코드북 값을 역양자화하고, 모멘텀을 적용한 뒤, 재양자화하고, 블록별 절대 최댓값을 업데이트하여 1차 모멘트(모멘텀)를 업데이트합니다.
   • 각 블록 내 제곱 그래디언트의 평균을 계산하고 지수 이동 평균을 적용하여 2차 모ment를 업데이트합니다.
   • 편향 수정(bias correction)을 적용하고, 공유된 2차 모멘트와 양자화된 1차 모멘트를 사용하여 파라미터를 업데이트합니다.

핵심 기여

1. 메모리 효율성: Gefen은 AdamW 대비 옵티마이저 메모리 점유율을 약 8배 감소시킵니다 (파라미터 크기의 ~2.0배에서 ~0.25배로 감소, 파라미터 10억 개당 약 6.5 GiB 감소).
2. 자동 적응성: 별도의 아키텍처 메타데이터, 수동 그룹화 규칙, 또는 표준 AdamW 기본값 외의 추가적인 하이퍼파라미터가 필요하지 않습니다. 데이터로부터 블록 구조를 자동으로 추론합니다.
3. 이론적 근거: Hessian 항목과 제곱 그래디언트 비율 사이의 이론적 연결 고리를 제공하여, Hessian 정렬된 파라미터들에 대한 2차 모멘트 통계량 공유의 타당성을 입증합니다.
4. 최적의 양자화: 모멘텀 양자화를 위해 히스토그램 기반의 정밀한 동적 계획법 접근 방식을 도입하여, Lloyd-Max와 같은 반복적 방법의 부적절함과 불안정성을 피합니다.
5. 성능 동등성: 이러한 감소가 모델 품질의 손실을 초래하지 않으며, 다양한 작업에서 AdamW와 대등한 성능을 유지함을 입риди합니다.

실험 결과

저자들은 CNN, ResNet, GPT-2, Llama 3 모델을 대상으로 Gefen을 AdamW, Adam-mini, Adam8bit, Adam4bit, Adafactor, SM3와 비교 평가했습니다.

• 피크 메모리 점유율: Gefen은 비교된 방법들 중 일관되게 가장 낮은 피크 메모리 사용량을 달s했습니다. 1.5B 파라미터 모델의 경우, Gefen은 옵티마이저 상태를 11.2 GiB(AdamW)에서 1.5 GiB로 줄였습니다.
• 성능 (품질): 모든 설정(OpenWebText, C4, ImageNet, MNIST)에서 Gefen은 AdamW와 유사한 검증 손실(validation loss) 및 정확도를 달성했습니다. 특히, 동일한 학습률 스케줄을 유지하면서 여러 지표에서 Adam4bit 및 Adam-mini보다 우수한 성능을 보였습니다.
• 처리량 (Throughput): 감소된 메모리 점유율 덕분에 분산 학습 시 더 큰 마이크로배치(microbatch)를 사용할 수 있었습니다:
  • FSDP (Llama 3-1.5B): Gefen은 마이크로배치 크기를 2로 가능하게 하여 (다른 방식들은 1), 56%의 처리량 증가(tokens/sec)를 실현했습니다.
  • DDP (GPT-2-1B): Gefen은 마이크로배치 크기 2로 학습이 가능했던 반면, AdamW는 마이크로배치 크기 1조차 감당할 수 없었습니다. 이는 Adam-mini 대비 21%의 처리량 개선을 가져왔습니다.
• 안정성: 절제 연구(Ablation study)를 통해 학습된 코드북이 학습 내내 안정적으로 유지됨을 확인하였으며, 이는 초기 단계에 한 번만 학습하기로 한 결정의 타당성을 입증합니다.

의의 및 주장

본 논문은 Gefen을 성능 저하나 수동 튜닝 없이 대규모 학습의 메모리 병목 현상을 해결할 수 있는 실용적인 "드롭인 교체" AdamW로서 포지셔닝합니다.

• 실질적 영향: 옵티마이저 상태를 줄임으로써, Gefen은 고정된 하드웨어에서 더 큰 모델을 학습하거나 더 큰 배치 크기를 사용할 수 있게 하여, 학습 처리량과 샘플 효율성을 직접적으로 향상시킵니다.
• 범용성: 수동 휴리스틱이나 특정 텐서 이름에 의존하는 기존 방식과 달리, Gefen은 아키텍처에 구애받지 않으며 데이터에 따라 자동으로 적응합니다.
• 상호 보완성: 저자들은 Gefen의 옵티마이저 상태 압축이 그래디언트 통신 압축(예: 1-bit Adam) 및 분산 학습 전략(FSDP/DDP)과 직교(orthogonal)하며, 더 큰 이득을 위해 이러한 기술들과 결합될 수 있다고 언급합니다.

저자들은 재현성과 용이한 통합을 강조하며, 융합된(fused) CUDA 커널을 포함한 완전한 Python 구현을 제공합니다.