Depth Registers Unlock W4A4 on SwiGLU: A Reader/Generator Decomposition

이 논문은 SwiGLU 아키텍처를 사용하는 대규모 언어 모델 (LLM) 에서 W4A4(가중치 4 비트, 활성화 4 비트) 양자화 시 발생하는 심각한 성능 저하 문제를 해결하고, 그 원인을 분석한 연구입니다. 저자는 3 억 파라미터 규모의 모델과 50 억 토큰 데이터셋을 사용하여 실험을 진행했으며, **Depth Registers (DR) 와 리지스터 크기 힌지 손실 (sink loss)**을 결합한 새로운 훈련 시간 기법을 제안했습니다.

주요 내용은 다음과 같습니다.

1. 문제 정의 (Problem)

• W4A4 양자화의 실패: SwiGLU 기반 모델 (Llama, Mistral 등) 에 naive round-to-nearest 방식의 W4A4 양자화를 적용하면, 검증 퍼플렉시티 (PPL) 가 FP16 기준 23.6 에서 1727 로 급격히 악화됩니다 (약 70 배 이상).
• 기존 방법의 한계: SmoothQuant, QuaRot, AWQ 와 같은 기존 보정 (calibration) 또는 회전 (rotation) 기반의 사후 (post-training) 방법들은 대부분 퍼플렉시티를 40~60 수준으로만 개선할 뿐, FP16 성능 (23.6) 에 근접하지 못합니다.
• 근본 원인 불명: 기존 연구들은 '어디에 아웃라이어 (outlier) 가 있는지'는 관찰했으나, 어떤 입력 활성화 (input-activation) 지점이 W4A4 오차를 주도하는지를 구조적으로 분리하여 분석하지는 못했습니다.

2. 방법론 및 이론적 배경 (Methodology & Mechanism)

저자는 SwiGLU 블록 내의 5 개의 가변 선형 레이어 (qkv, w1, w3, o_proj, w2) 를 입력 소스에 따라 두 가지로 분류했습니다.

A. 레이어 분류 (Readers vs. Generators)

1. Readers (독자): 잔여 스트림 (residual stream) 의 RMSNorm 출력을 입력으로 받는 레이어들 (qkv, w1, w3).
   • 특징: 잔여 스트림의 크기를 제어하는 개입 (intervention) 이 있다면 이들의 입력 크기도 직접적으로 제어 가능합니다.
2. Generators (생성자): 블록 내부 계산 결과 (어텐션 가중치 합, SwiGLU 이차원 곱셈) 를 입력으로 받는 레이어들 (o_proj, w2).
   • 특징:
     • o_proj: 어텐션 가중치에 의존하므로 잔여 스트림 제어만으로는 제한적인 통제력만 가집니다.
     • w2: silu(w1) ⊙ w3 형태의 **이차원 곱셈 (bilinear product)**을 입력으로 받습니다. 이론적으로 잔여 스트림의 크기 제어만으로는 이 곱셈 결과의 꼬리 (tail) 분포를 제한할 수 없으며, 이는 아웃라이어가 발생하기 쉬운 구조임을 의미합니다.

B. 제안된 기법: Depth Registers + Sink Loss (DR+sink)

• Depth Registers (DR): 잔여 채널 축을 의미 있는 채널 (semantic) 과 레지스터 채널 (register) 로 분할합니다.
• Sink Loss: 레지스터 채널의 최대 크기 (magnitude) 가 특정 임계값을 초과하지 못하도록 힌지 손실 (hinge loss) 을 적용합니다.
• 원리: 이 기법은 훈련 중에 잔여 스트림의 크기를 제어하여, Readers의 입력 크기 분포를 가우시안에 가깝게 만들고 (꼬리 제거), 아웃라이어 에너지를 레지스터 채널로 흡수시킵니다. 이는 사후 보정이나 회전이 아닌 훈련 시간 (training-time) 의 잔여 축 개입입니다.

3. 주요 실험 결과 (Key Results)

3 억 파라미터 모델 (5B 토큰 학습) 에서 다음과 같은 결과를 도출했습니다.

• 성능 개선:
  • Naive W4A4: PPL 1727 (FP16 기준 23.6)
  • DR+sink 적용 후: PPL 119 (약 14 배 개선).
  • SmoothQuant 와 결합 시: PPL 39.9 로 추가 개선.
  • FP16 성능 유지: DR+sink 를 적용해도 FP16 PPL 은 23.6~23.7 로 유지되었으며, 제로샷 능력도 동일하게 유지되었습니다.
• 구조적 분석 (Readers vs. Generators):
  • Readers: DR+sink 적용 후 qkv 의 초과 첨도 (excess kurtosis) 가 35.9 에서 2.85 로 급격히 감소하여 가우시안 분포에 근접했습니다.
  • Generators: w2 의 첨도는 1921 에서 1860 으로 거의 변하지 않았습니다. 이는 잔여 축 개입이 생성자 (Generator) 의 이차원 곱셈 꼬리를 제어하지 못함을 시사합니다.
• 손실 분해 (Damage Decomposition):
  • Readers 만을 FP16 으로 유지하고 나머지를 양자화했을 때, PPL 은 약 31 로 수렴했습니다. 이는 Readers 가 W4A4 오차의 주원인임을 의미합니다.
  • Readers 를 구제 (rescue) 한 후에도 남는 잔여 오차 (Residue) 는 약 0.28 nats (PPL 약 31) 로, DR+sink 유무와 무관하게 모든 체크포인트에서 일정했습니다.
  • 이 잔여 오차는 주로 w2 (MLP 생성자) 에서 기인하며, o_proj 의 영향은 미미했습니다.

4. 후속 검증 및 QuaRot 비교

• Per-Linear QuaRot: DR+sink 와 유사하게 Readers 축에서 25.5 PPL 수준의 성능을 내며 DR+sink 를 대체할 수 있음을 보였습니다.
• Full QuaRot: w2 입력에 대한 온라인 직교 회전 (orthogonal rotation) 을 추가했으나, 성능이 26.0~26.1 PPL 로 오히려 약간 더 나빠졌습니다.
• 결론: 이론적 예측 (Theorem 3.4) 과 일치하게, 직교 회전 (orthogonal rotation) 만으로는 SwiGLU 의 이차원 곱셈으로 인한 꼬리 분포를 제어할 수 없음이 실험적으로 입증되었습니다.

5. 의의 및 기여 (Significance & Contributions)

1. 구조적 분해 (Operator-level Decomposition): SwiGLU 블록의 입력을 '잔여 스트림을 읽는 Readers'와 '블록 내부 생성자 (Generators)'로 명확히 분류하고, W4A4 오차가 어디서 발생하는지 이론적으로 설명했습니다.
2. 새로운 훈련 시간 프로브 (Training-time Probe): DR+sink 를 통해 Readers 의 오차를 효과적으로 제거할 수 있음을 증명했습니다. 이는 배포용 솔루션이라기보다, W4A4 손실의 원인을 규명하기 위한 도구 (probe) 로 제안되었습니다.
3. 이차원 곱셈의 한계 규명: 잔여 축의 크기 제어나 직교 회전만으로는 SwiGLU 의 w2 입력에서 발생하는 이차원 곱셈의 꼬리 (heavy tail) 를 제어할 수 없음을 보여주었습니다. 이는 향후 w2 레이어에 대한 혼합 정밀도 (mixed-precision) 사용이나 양자화 인식 학습 (QAT) 등의 대안이 필요함을 시사합니다.
4. 성능 기록: 단일 시드 (single-seed) 실험에서 Naive W4A4 의 퍼플렉시티를 1727 에서 119 로 대폭 낮추는 성과를 거두었습니다.

요약

이 논문은 SwiGLU 모델의 W4A4 양자화 실패 원인이 잔여 스트림을 읽는 레이어 (Readers) 의 아웃라이어와 **MLP 내부의 이차원 곱셈 (Generators, 특히 w2)**에 있음을 규명했습니다. 훈련 시 Depth Registers를 도입하여 Readers 의 출력을 제어하면 성능을 크게 회복할 수 있지만, w2 의 이차원 곱셈으로 인한 잔여 오차는 기존 사후 회전 기법으로는 해결되지 않음을 보여주었습니다. 이는 차세대 양자화 기법이 단순히 회전을 적용하는 것을 넘어, MLP 내부의 곱셈 연산에 대한 새로운 접근 (예: 정밀도 조절) 이 필요함을 시사합니다.