마이크로스케일 양자화를 위한 학습 가능한 어파인 변환 LATMiX

초록

LATMiX는 MX(마이크로스케일) 양자화 환경에서 활성값의 이상치를 완화하기 위해 일반적인 가역 어파인 변환을 학습한다. 이론적 오류 상한을 도출하고, LU·QR 분해 기반 파라미터화와 왜곡 최소화 손실을 이용해 변환을 최적화한다. 실험 결과, 다양한 크기의 LLM에 대해 4‑bit 이하 MX 양자화 시 기존 회전·하다마드 기반 방법보다 일관된 정확도 향상을 달성한다.

상세 분석

본 논문은 MX 양자화가 블록 단위 스케일링을 사용함에 따라, 변환 행렬이 블록 구조와 활성값 분포를 동시에 고려해야 함을 이론적으로 증명한다. 정의 3.1‑3.2와 정리 3.3을 통해 변환 후 평균제곱오차(MSE)는 두 요인, 즉 변환 행렬의 역스펙트럼 노름 ‖A⁻¹‖₂와 각 블록 내 최대 절대값의 기대값 평균 ∑₁ᴺᴮMᵢ에 의해 상한이 결정된다고 제시한다. 여기서 Mᵢ는 변환된 활성값의 서브가우시안 노름과 블록 크기 B에 대한 로그항에 의존한다. 따라서 ‖A⁻¹‖₂를 작게 하면 (즉, A의 최소 특이값을 크게 하면) 변환이 안정적이지만, 동시에 블록 내 값들의 분산이 커져 Mᵢ가 증가할 위험이 있다. 반대로 블록‑다이아고날 회전은 각 블록을 독립적으로 최적화해 Mᵢ를 감소시키지만, 전체 텐서 차원 간의 상호작용을 차단해 ‖A⁻¹‖₂가 커지는 부작용을 낳는다.

이러한 트레이드오프를 해결하기 위해 저자들은 전체 가역 어파인 변환 A·x+v 를 허용하고, v를 -Aμ 로 설정해 평균값을 0으로 맞춤으로써 Mᵢ의 첫 번째 항을 제거한다. 변환 파라미터는 LU와 QR 분해를 통해 자유롭게 학습되며, 볼륨 보존 정규화(‖det(A)‖≈1)와 디스틸레이션 손실(양자화된 모델 출력과 FP 모델 출력 간 L₂ 차이)으로 최적화된다. 학습된 변환은 선형 레이어와 결합해 사후에 가중치에 흡수될 수 있으므로 추론 시 추가 연산 비용이 거의 없으며, 편향이 존재할 경우 완전히 무손실로 통합된다.

실험에서는 Llama‑3.2‑1B, Llama‑2‑7B, Mistral‑7B 등 다양한 모델에 대해 MX‑FP4·INT4 양자화를 적용하고, BIG‑Bench, GSM‑8K, WikiText‑2 등 7개의 제로‑샷 벤치마크에서 평균 정확도가 0.5~1.2%p 상승함을 보고한다. 특히 블록 크기가 커질수록(예: B=64) 기존 회전·하다마드 방식은 급격히 성능이 떨어지는 반면, LATMiX는 블록 간 에너지 재분배를 통해 오류를 균등화하고 전체 퍼플렉시티를 낮춘다. 수치 실험(그림 2)에서도 학습된 어파인 변환이 블록‑별 MSE를 고르게 감소시키며, 전체 MSE는 가장 낮은 수준을 기록한다.

결과적으로, 본 연구는 MX 양자화가 가진 블록‑스케일링 특성을 고려한 변환 설계가 필요함을 이론·실험적으로 입증하고, 일반 선형·편향 변환을 학습함으로써 기존 제한적 회전 기반 방법을 뛰어넘는 성능 향상을 제공한다는 점에서 LLM 양자화 분야에 중요한 전진을 이룬다.