AdamW 4비트 옵티마이저 상태 양자화를 프리컨디셔너 공간 반올림으로 개선한다

무엇인가

AdamW는 파라미터마다 FP32 1차·2차 모멘트를 유지하므로 파라미터당 8바이트가 옵티마이저 상태로 상주한다. 대규모 학습에서 이는 무시할 수 없는 상주 메모리이며, 그래서 블록 단위 동적 양자화로 8비트·4비트 옵티마이저 상태를 저장하는 연구가 이어져 왔다. 문제는 양자화 오차가 한 번으로 끝나지 않는다는 점이다. 모멘트는 지수이동평균 재귀로 다음 스텝으로 되먹임되고, 그 오차가 다시 적응적 업데이트를 교란한다. Li 등(2023)은 작은 양의 2차 모멘트를 0으로 매핑하면 업데이트가 지나치게 커지는 zero-point failure를 보고했고, TorchAO의 4비트 AdamW를 비롯한 후속 구현은 0을 제외한 양수 코드북으로 이를 피해 왔다. 이 논문은 그 회피책을 그대로 두지 않고, 반올림을 어느 좌표에서 하느냐는 관점에서 4비트 AdamW를 다시 설계한다.

어떻게 동작하나

저자들의 핵심 통찰은 반올림 공간(rounding space), 즉 양자화기가 인접한 재구성 레벨 중 하나를 고르는 좌표다. 기존 반올림은 2차 모멘트 상태 공간에서 값을 비교해 가까운 레벨을 고른다. 그런데 AdamW는 2차 모멘트를 1/(sqrt(v/(1-β2^t)) + ε) 꼴의 프리컨디셔너로 비선형 변환한 뒤에야 업데이트를 스케일링한다. 그래서 상태 공간에서 정확하거나 평균적으로 무편향인 반올림도 결과 업데이트는 왜곡할 수 있다. 반올림 결정을 프리컨디셔너 공간에서 내리면 0을 포함한 코드북을 그대로 유지하면서도 해로운 0 선택 확률을 낮출 수 있다는 것이 이 논문의 출발점이다.

무엇과 다른가

논문은 이를 두 개의 명제로 정식화한다. 명제 1은 다음 스텝 프리컨디셔너 편향을 비교한다. 0 ≤ a < ṽ_t < b인 인접 레벨에서 Update-SR(프리컨디셔너 공간 확률적 반올림)과 State-SR(상태 공간 확률적 반올림)을 비교하면 r_{t+1}(b) < E[r_{t+1}(Q_U(ṽ_t))] ≤ r_{t+1}(ṽ_t) ≤ E[r_{t+1}(Q_S(ṽ_t))]가 성립한다. 또 0과 b 중 하나를 고르는 양자화기에 대해, 프리컨디셔너 오차가 유계가 될 필요충분조건은 0으로 반올림할 확률이 O(ε)인 것이고, 그 확률이 Θ(1)이면 오차가 Θ(ε⁻¹) 속도로 발산한다. State-SR은 2차 모멘트 값에 무편향이지만 0으로 반올림할 확률이 1 - ṽ_t/b = Θ(1)이라 발산하고, State-RTN은 셀 아래 절반에서 결정론적으로 0을 고른다. 반면 Update-SR이 0을 고를 확률은 ε(1 - sqrt(ṽ_t/b)) / (sqrt(ṽ_t/(1-β2^t)) + ε)로 Θ(ε)에 불과해 오차가 유계다. Update-RTN은 원점의 Θ(ε²) 근방에서만 0을 고른다. 명제 2는 스칼라 이차함수 μw²/2에서의 동역학을 본다. 같은 코드북 {0, b}에서 State-RTN은 0이 아닌 안정적인 2-사이클로 수렴하는 반면 Update-RTN은 최소점으로 기하급수적으로 수렴하고, 최적점 근방에서 State-SR은 국소적으로 확장적, Update-SR은 국소적으로 수축적이다.

어떻게 쓰나

이 결과가 두 가지 4비트 구성으로 이어진다. ZIP-SR은 0을 포함한 Dyn4 2차 모멘트 코드북을 유지한 채 Update-SR로 반올림 확률을 계산한다. ZE-EDEN은 반대로 0을 제외한 Dyn4-NZ 코드북을 쓰되, 양의 양자화 바닥이 만드는 프리컨디셔너 왜곡을 EDEN 블록 보정으로 완화한다. 보정 계수는 c_τ = ||x||² / max{<x, Q(x)>, τs²} (τ = 10⁻¹²)로, 양자화된 블록의 스케일을 다시 조정해 작은 2차 모멘트의 과대평가와 그에 따른 큰 프리컨디셔너 항 억제를 줄인다. 두 구성 모두 1차 모멘트에는 NF4를 쓰고, 학습 마지막 10%에서 LM-head의 1차 모멘트만 확률적 반올림으로 전환한다.

전제와 한계

1차 모멘트 코드북 선택도 다시 검토한다. TorchAO는 SDyn4를 쓰지만, 블록 정규화된 1차 모멘트 분포에서 평균제곱 재구성 오차는 FP4 1.739×10⁻³, SDyn4 2.495×10⁻³, NF4 1.056×10⁻³로 NF4가 가장 작다. GPT-style 834M 단일 시드 사전학습에서 최종 격차는 ZE-EDEN이 SDyn4 0.0289, FP4 0.0272, NF4 0.0183, ZIP-SR이 0.0239, 0.0214, 0.0140이었다. 컴포넌트 기여를 분해하면 2차 모멘트를 Lin4-NZ에서 Dyn4-NZ로 바꾸는 것만으로는 격차가 오히려 커지고, 1차 모멘트를 SDyn4에서 NF4로 바꾸는 것이 가장 큰 단계적 감소를 준다. 그 위에 EDEN 보정(3a)과 Dyn4 Update-SR(3b)이 각각 TorchAO 대비 28.2%, 43.2% 낮은 격차를 만든다. 다만 3b는 2차 모멘트 형식과 반올림 규칙을 동시에 바꾼 단일 실행이라 Update-SR 자체의 효과를 분리하지는 못한다고 저자들은 밝힌다. 본 사전학습은 GPT-style 162M·405M·1.4B·2.7B와 Llama-style 130M·350M·1.1B에서 3쌍 시드로 돌렸고, 두 방법 모두 평가한 모든 모델 크기에서 TorchAO의 평균 검증 손실 격차를 줄였다. GPT-style 2.7B에서 TorchAO는 학습률 쿨다운 후반에 급격히 발산하는데, 발산 전인 학습 토큰 95% 체크포인트에서 평균 쌍 격차는 TorchAO +0.0754(표본표준편차 0.0034), ZE-EDEN +0.0471(0.0018), ZIP-SR +0.0286(0.0024)로 각각 37.6%, 62.1% 감소했다. 논문이 보고한 최대 감소율 70.1%는 1.4B에서 나온 값이며, 발산한 2.7B 종점에서는 백분율을 계산하지 않는다.

전체 파라미터 지도학습 파인튜닝은 Tulu-3 데이터로 Qwen3-8B-Base와 Llama-3.2-3B에서 5쌍 시드로 수행하고 MMLU, GSM8K, HumanEval, IFEval로 평가했다. 두 방법 모두 TorchAO 4비트 AdamW보다 낮은 평균 검증 손실을 냈고, ZE-EDEN은 32비트 AdamW와의 평균 쌍 검증 손실 격차가 Qwen3-8B에서 사실상 0, Llama-3.2-3B에서 약간 음수였다. 다운스트림 점수는 32비트 AdamW에 근접하지만 과제별로 일관된 우위는 없었다. 검증 손실은 ZE-EDEN이 두 SFT 모델 모두에서 더 낮았고, 사전학습 비교에서는 ZIP-SR이 대체로 더 작은 격차를 보였다.

저장 비용은 두 방법 모두 TorchAO 4비트와 사실상 같다. 블록 크기 128, 모멘트마다 블록당 FP32 스케일 하나를 두면 적격 텐서에서 파라미터당 1.0625바이트로, FP32 모멘트 저장의 8바이트 대비 7.53배 압축이다. 2.7B 사전학습 모델에서 약 19.2GB, 3B~8B 전체 파라미터 SFT 모델에서 약 21~56GB의 상주 모멘트 상태를 덜 쓴다는 계산이다. 정규화 레이어의 모멘트 텐서만 적격성 기준 밖으로 남아 FP32를 유지한다. 실무 관점에서 이 논문이 주는 실용적 메시지는 두 가지다. 첫째, 이미 TorchAO 4비트 AdamW를 쓰고 있다면 코드북을 갈아엎지 않고 반올림 좌표만 프리컨디셔너 공간으로 옮기는 것만으로 32비트와의 격차를 줄일 수 있다. 둘째, 2.7B급 이상 사전학습에서 1차 모멘트, 특히 LM-head의 결정론적 반올림이 후반 학습률 쿨다운 구간에서 발산을 유발할 수 있으므로, 그 부분만 확률적 반올림으로 바꾸는 저비용 처방을 검토할 가치가 있다. 구현은 github.com/nubank/adamw4bit에 공개돼 있다.

저자들이 명시한 한계는 분명하다. 이론은 조건부 1스텝 프리컨디셔너 편향과 장난감 스칼라 이차 모델의 동역학만 비교할 뿐, 양자화된 AdamW의 일반적인 종단간 수렴을 보장하지 않는다. LM-head 확률적 반올림이 관찰된 후반 회귀를 완화하기는 하지만, 방향성 1차 모멘트 편향이 왜 생기는지, 그것이 학습률 감쇠와 어떤 관계인지는 규명되지 않았다. 4비트 아래로 내려가는 양자화와 다른 적응적 옵티마이저로의 확장도 향후 과제로 남겨 두었다.