LATMiX: Learnable Affine Transformations for Microscaling Quantization of LLMs

LLM 의 마이크로스케일링 양자화를 위한 학습 가능 아핀 변환 (LATMiX) 에 대한 상세 기술 요약입니다.

1. 문제 제기

대형 언어 모델 (LLM) 은 메모리 및 연산 비용을 절감하기 위해 increasingly 학습 후 양자화 (PTQ) 를 통해 배포되고 있습니다. 주목받고 있는 특정 데이터 포맷은 마이크로스케일링 (MX) 으로, 이는 텐서를 블록 단위로 분할하여 각 블록에 고유한 스케일링 인자를 할당함으로써 전역 스케일링보다 동적 범위를 더 효율적으로 처리합니다.

그러나 MX 양자화를 기존 아웃라이어 완화 기법 (저비트 양자화에 필수적임) 과 결합하는 것은 상당한 도전을 제기합니다:

• 활성화 아웃라이어: LLM 활성화에는 양자화 오차를 지배하는 극단적인 값 (아웃라이어) 이 포함되어 있습니다.
• 기존 방법의 한계: 기존 접근법은 활성화 에너지를 재분배하고 아웃라이어를 억제하기 위해 가역 변환 (보통 회전 또는 해다마르 변환) 을 사용합니다.
  • 전역 변환: MX 양자화 전에 활성화에 전역 회전을 적용하면 MX 의 블록 단위 스케일링 구조와 전역 변환이 충돌하여 심각한 성능 저하가 발생합니다.
  • 블록 대각 변환: 최근 연구는 각 MX 블록 내에서 독립적으로 변환을 적용하는 (블록 대각 행렬) 방식을 제안합니다. 이는 불일치를 피하지만, 변환을 작고 고립된 부분 공간으로 제한하여 텐서 전체에 걸쳐 활성화 질량을 효과적으로 재분배하지 못하게 합니다. 결과적으로 지배적인 아웃라이어가 충분히 확산되지 않아 저비트 영역 (예: 4 비트) 에서 최적이지 않은 정확도를 초래합니다.

2. 방법론: LATMiX

저자들은 MX 데이터 포맷에 최적화된 학습 가능한 풀랭크 가역 아핀 변환으로 아웃라이어 감소를 일반화하는 LATMiX(Microscaling 을 위한 학습 가능 아핀 변환) 를 제안합니다.

A. 이론적 분석

저자들은 MX 양자화 하에서 양자화 오차에 대한 이론적 상한을 유도했습니다. 분석 결과, 오차는 두 가지 경쟁 요인에 의존함이 밝혀졌습니다:

1. 역변환의 스펙트럼 노름 ($||A^{-1}||_\sigma$): 변환이 오차를 얼마나 증폭시키는지에 대한 척도.
2. 블록당 기대 최대 활성화: 각 MX 블록 내 절대값의 최댓값들의 평균.

본 논문은 기존 연구에서 사용된 블록 대각 변환이 이러한 요인들을 전역적으로 균형 있게 맞추지 못하는 특수한 경우에 불과하다고 주장합니다. 대신, MX 블록 구조와 전역 특징 분포를 모두 명시적으로 고려하기 위해 풀 (비블록 대각) 아핀 변환이 필요합니다.

B. LATMiX 아키텍처

LATMiX 는 모델의 선형 레이어에 접혀 추론 오버헤드가 제로인 일련의 가역 아핀 변환 ($T(x) = Ax + v$) 을 학습합니다.

• 변환 위치:
  • $T_1$: 잔여 경로 (입력 임베딩) 에 전역적으로 적용.
  • $T_2$: 멀티헤드 어텐션 (MHA) 의 값 벡터와 피드포워드 네트워크 (FFN) 내부에 로컬하게 적용.
  • $T_3$: 기존 베이스라인과의 일관성을 위해 유지되는 온라인 블록 - 해다마르 변환.
• 매개변수화: 가역성을 보장하고 표준 딥러닝 최적화를 가능하게 하기 위해, 저자들은 행렬 $A$에 대해 두 가지 매개변수화를 제안합니다:
  1. LU 분해: $A = P(L + \text{diag}(s))$, 여기서 $L$은 하단 단위 삼각행렬입니다.
  2. QR 분해: $A = Q(R + \text{diag}(s))$, 여기서 $Q$는 (반대칭 행렬의 행렬 지수를 통해) 직교 행렬이고 $R$은 상단 삼각행렬입니다.
  • 부피 보존 정규화가 적용되어 행렬식이 1 에 가깝게 유지되도록 하여 수치적 불안정성을 방지하면서도 개별 방향의 확장 또는 수축을 허용합니다.
• 최적화 목적 함수:
  • 증류 손실: 직교 행렬을 요구하는 엄격한 계산 불변성을 강제하는 대신, LATMiX 는 풀 정밀도 교사 모델과 변환된 양자화 학생 모델 간의 클루백 - 라이블러 (KL) 발산 손실을 사용합니다. 이를 통해 변환이 데이터 분포에 더 잘 적합하도록 직교성에서 벗어나게 할 수 있습니다.
  • 정규화: 변환의 부피를 1 에 가깝게 유지하는 항.

C. 접기 (Folding) 메커니즘

핵심 기여는 학습된 아핀 변환 ($A$ 및 $v$) 을 LLM 의 가중치 행렬과 편향에 직접 접을 (fold) 수 있는 능력입니다.

• 역변환은 어텐션 메커니즘의 입력 프로젝션 (K, Q, V) 에 접힙니다.
• 순방향 변환은 어텐션 및 FFN 레이어의 출력 프로젝션에 접힙니다.
• 이로 인해 추론 시 모델은 추가적인 계산 비용 없이 표준 양자화 모델과 정확히 동일한 행동을 합니다.

3. 주요 기여

1. 이론적 통찰: MX 포맷에서의 양자화 오차가 전역 특징 분포와 블록 구조 모두에 의존함을 보여주는 이론적 상한을 제시하여, 블록 대각 변환이 최적의 오차 감소를 위해 불충분함을 증명했습니다.
2. LATMiX 알고리즘: LU/QR 매개변수화와 증류 손실을 사용하여 (회전으로 제한되지 않는) 일반적인 가역 아핀 변환을 학습하는 방법을 도입했습니다.
3. 오버헤드 없는 배포: 이러한 복잡한 학습된 변환이 기존 선형 레이어에 수학적으로 접혀 표준 양자화 모델의 추론 속도와 메모리 풋프린트를 유지함을 입증했습니다.
4. 종합적 평가: 다양한 벤치마크에서 여러 모델 크기 (1B~32B) 와 아키텍처 (Llama, Qwen) 에 대해 방법을 검증했습니다.

4. 실험 결과

저자들은 LATMiX 를 7 가지 제로샷 추론 벤치마크(ARC, HellaSwag, WinoGrande, PIQA, BoolQ, OpenBookQA) 와 WikiText2 퍼플렉시티 작업에서 평가했습니다.

• 성능 향상: LATMiX 는 다음을 포함한 강력한 베이스라인보다 일관되게 우수한 성과를 보였습니다:
  • 표준 MX 양자화 (RTN, GPTQ).
  • 회전 기반 방법 (QuaRot, SpinQuant).
  • 블록 대각 회전 방법 (MR-GPTQ, OSTQuant).
  • 기타 변환 기반 방법 (FlatQuant).
• 정확도 회복: MXFP4(4 비트) 양자화에서 LATMiX 는 최상의 베이스라인 대비 평균 정확도 회복에서 최대 4% 향상을 달성했습니다. 예를 들어, Llama-3.2-1B 에서 LATMiX-LU 는 93.07% 인 MR-GPTQ 대비 **93.88%**의 회복률을 달성했습니다.
• 강건성: 이 방법은 다양한 모델 크기 (1B~32B) 와 양자화 입자 (MXFP4 및 MXINT4) 에서 일관된 개선을 보여주었습니다.
• 애블레이션 연구:
  • 풀 vs 블록: 풀 (텐서 단위) 아핀 변환이 블록 대각 변환보다 훨씬 우수했습니다.
  • 아핀 vs 직교: 편향을 포함한 일반적인 아핀 변환을 학습하는 것이 직교 (회전) 행렬만 학습하는 것보다 우수했습니다.
  • 손실 함수: KL 증류 손실이 MSE 또는 교차 엔트로피 손실보다 제로샷 일반화에 더 효과임이 입증되었습니다.

5. 의의

• 격차 해소: LATMiX 는 (전역 변환을 통한) 공격적인 아웃라이어 억제 필요성과 현대 하드웨어가 지원하는 MX 포맷의 구조적 제약 사이의 격차를 성공적으로 해소합니다.
• 하드웨어 효율성: 변환을 접음으로써, LATMiX 는 MX 양자화가 제공하는 추론 효율성 향상을 희생하지 않고도 강력하고 학습된 수학적 연산을 사용할 수 있게 합니다.
• 미래 방향: 이 연구는 변환을 회전이나 블록 대각 행렬로 제한하는 것이 최적이지 않음을 시사합니다. 차세대 AI 하드웨어의 특정 데이터 포맷을 존중하는 더 표현력 있고 데이터 의존적인 변환을 학습할 수 있는 문을 엽니다.

요약하자면, LATMiX 는 활성화 분포와 MX 블록 구조 간의 상호작용을 이론적으로 분석하고 유연한 아핀 변환을 학습함으로써, 추론 오버헤드를 발생시키지 않고도 극도로 낮은 비트 폭 (4 비트) 에서 LLM 의 손실 없는 양자화를 달성할 수 있음을 보여줍니다.