프리컨디셔너 공간에서의 반올림: 4비트 AdamW 옵티마이저 상태 양자화의 재설계

AdamW는 대규모 모델 학습의 기본 옵티마이저이지만, 그 비용은 과소평가되기 쉽다. 가중치와 그래디언트 외에도 각 파라미터는 1차 모멘트와 2차 모멘트라는 두 가지 상태를 가진다. 32비트로 저장하면 이 상태가 차지하는 상주 메모리는 가중치와 비슷하거나 더 크다. 옵티마이저 상태를 4비트로 압축하는 것은 학습 메모리를 줄이는 가장 직접적인 방법 중 하나다. TorchAO는 이미 4비트 AdamW를 제공하지만, 32비트 AdamW와의 사이에는 측정 가능한 검증 손실 격차가 남아 있다. 2026년 10월 8일 arXiv cs.LG에 공개된 Hanyang Li 등 다섯 저자의 이 논문은 코드북의 모양이나 블록 크기를 조정하지 않는다. 대신 더 기본적인 질문을 던진다. 양자화기가 인접한 두 재구성 수준 중 하나를 고를 때, 어느 좌표에서 골라야 하는가. 저자들은 이 좌표를 '반올림 공간'이라 부른다. 이것이 중요한 이유는 양자화 오차가 제자리에 머물지 않기 때문이다. 오차는 모멘트의 점화식을 따라 전파되어 이후 모든 적응적 업데이트를 교란한다. 1차 모멘트에서는 업데이트가 상태에 대해 거의 선형이므로 상태 공간에서의 반올림은 큰 문제가 아니다. 2차 모멘트는 다르다. Adam은 2차 모멘트의 제곱근으로 나누며, 이 역수가 프리컨디셔너다. 논문은 0에 인접한 양자화 셀을 국소적으로 분석하여, 상태의 평균 오차가 작다고 해서 다음 단계 프리컨디셔너의 평균 오차도 작다고 할 수 없음을 보인다. 상태가 0에서 조금만 벗어나도 가파른 역수 함수를 거치며 오차가 프리컨디셔너 쪽에서 증폭된다. 이어 1차원 이차 문제 구성은, 상태 공간 반올림과 프리컨디셔너 공간 반올림이 단순한 수치 차이가 아니라 질적으로 다른 최적화 동역학을 낳는다는 것을 보여준다.

이 관찰에서 첫 번째 방법 ZIP-SR, 즉 0을 포함하는 프리컨디셔너 공간 확률적 반올림이 나온다. 2차 모멘트 코드북에 0을 남기고, 확률적 반올림의 확률을 상태 공간이 아닌 프리컨디셔너 공간에서 계산한다. 확률적 반올림은 기댓값이 불편이며, 결정적인 체계적 오차를 평균 0의 잡음으로 바꾼다. 공간을 제대로 고르면 이 불편성이 업데이트 크기를 실제로 결정하는 양 위에 놓인다. 두 번째 방법 ZE-EDEN은 보완적인 길을 택한다. 2차 모멘트 코드북에서 0을 제외하므로 양자화된 값에는 양의 하한이 생기고, 0 근처에서 역수가 발산하지 않는다. 그 하한 자체가 프리컨디셔너를 왜곡하기 때문에, EDEN 방식의 보정으로 양자화된 2차 모멘트 블록의 크기를 다시 조정하여 왜곡을 상쇄한다. 두 구성 모두 1차 모멘트에는 4비트 NormalFloat(NF4)를 쓰고, 학습의 마지막 10% 동안 언어 모델 출력 헤드의 1차 모멘트에 표적화된 확률적 반올림을 적용한다.

실험은 GPT와 Llama 스타일의 사전학습을 다루며 모델 규모는 1억 3천만에서 27억 파라미터에 이른다. 평가한 모든 규모에서 두 방법은 TorchAO 4비트 AdamW와 32비트 AdamW 사이의 평균 검증 손실 격차를 줄였고, 보고된 최대 감소폭은 70%에 달했다. 전체 파라미터 지도 미세조정에서는 두 레시피 모두 TorchAO보다 낮은 검증 손실을 기록했고, 하위 작업에서는 32비트 AdamW에 가까운 수준을 유지했다. 한 번의 큰 승리보다 일관성이 더 중요하다. 모든 규모에서 개선이 나타난다면 그것은 한 번의 운이 아니라 작동하는 메커니즘을 시사한다. 우리의 판단으로 이 논문의 핵심 기여는 설계 공간을 다시 그은 데 있다. 기존 논의는 코드북 설계와 블록 분할에 집중했다. 이 논문은 세 번째 축, 즉 반올림이 일어나는 좌표를 더한다. 이 관점은 저장된 상태에 비선형 변환을 가하는 다른 옵티마이저에도 옮겨질 가능성이 크다. 그래도 신중함이 필요하다. 초록에는 절감된 메모리나 처리량 수치가 없고, 가장 큰 모델이 27억 파라미터이며, 더 큰 규모와 더 긴 학습에서도 이득이 유지되는지는 독립적인 재현이 답해야 한다. 학습 메모리가 부족한 팀이라면 운영 레시피를 바로 바꾸기보다, ZIP-SR과 ZE-EDEN을 자체 규모에서 TorchAO와 비교하는 소규모 시험의 후보로 다루는 것이 합리적이다.