기술 요약: FPTQuant: LLM 양자화를 위한 함수 보존 변환(Function-Preserving Transforms)

문제 정의

거대 언어 모델(LLM)은 추론 과정에서 상당한 계산 및 에너지 비용을 발생시켜, 자원이 제한된 에지 디바이스로의 배포를 어렵게 만듭니다. 양자화(연속적인 값을 균일한 정수 그리드로 매핑하는 것)는 효율성을 높이는 가장 효과적인 기술이지만, LLM에 대한 나이브한 양자화는 종번 심각한 성능 저하를 초래합니다. 이는 주로 저비트 양자화에 필요한 균일한 분포를 방해하는 "아웃라이어(outliers)"—즉, 큰 크기를 가진 활성화 값(activations)과 가중치(weights)—때문입니다.

근본적인 과제는 범위-정밀도 간의 트레이드오프(range-precision trade-off)에 있습니다:

1. 범위를 늘리면 아웃라이어를 포착할 수 있지만, 0 근처에 밀집된 대다수 값에 대한 정밀도가 손실됩니다.
2. 정밀도를 유지하기 위해 아웃라이어를 클리핑(clipping)하면 정보 손실과 모델 붕괴가 발생합니다.

기존 연구들은 함수 보존 변환(FPTs)—양자화가 없는 상태에서 모델의 출력을 변경하지 않으면서 아웃라이어를 완화하는 스케일링이나 회전과 같은 연산—을 탐구해 왔습니다. 그러나 기존 방법들은 종종 표현력(expressivity)(양자화 오차를 효과적으로 줄이는 능력)과 계산 오버헤드(compute overhead)(추론 비용) 사이에서 트레이드오프를 겪습니다. 많은 표현력이 높은 변환들은 병합이 불가능(non-mergeable)하여, 계산 그래프에 새로운 연산을 추가해야 하며 이는 상당한 레이턴시를 유발합니다.

방법론: FPTQuant

본 논문은 트랜스포머에 특화되어 설계된, 새롭고 가볍고 표현력이 뛰어난 세 가지 FPT를 활용하는 프레임워크인 FPTQuant를 소개합니다. 핵심 철학은 트랜스포머의 정형 연산에 내재된 **등변성(equivariances)과 독립성(independencies)**을 활용하여, 병합 가능하면서도(mergeable) 표현력이 높은 변환을 설계하는 것입니다.

세 가지 주요 변환은 다음과 같습니다:

1. 쿼리 및 키를 위한 병합 가능한 Pre-RoPE 변환 ($T_k, \bar{T}_k$)

• 컨텍스트: 현대 LLM은 회전 위치 임베딩(RoPE)을 사용하는데, 이는 임의의 선형 변환이 RoPE와 교환 법칙이 성립하지 않기 때문에 쿼리($W_q$) 및 키($W_k$) 가중치에 변환을 나이브하게 병합하는 것을 방지합니다.
• 메커니즘: 저자들은 RoPE 이전에 적용되는 한 쌍의 변환을 도입합니다. $T_k$는 키에 적용되며, $\bar{T}_k$(근사 역행렬)는 쿼리에 적용됩니다. 이들은 스케일링된 $2 \times 2$ 회전 행렬로 구성됩니다.
• 특성: RoPE의 블록 대각 구조 덕분에, 이러한 특정 회전/스케일링 행렬은 RoPE 연산과 교환 가능합니다. 이를 통해 훈련 후 이 변환들을 $W_q$ 및 $W_k$ 가중치에 직접 병합할 수 있으며, 추론 오버헤드가 발생하지 않습니다.
• 적용 가능성: RoPE를 사용하는 모델(예: Llama 시리즈) 및 위치 임베딩이 없는 모델에 유효합니다.

2. 병합 가능한 멀티헤드 밸류 변환 ($T_v, \bar{T}_v$)

• 컨텍스트: 어텐션 메커니즘에서 배치 행렬 곱셈(BMM)은 밸류($V$)와 출력 투영($W_o$)에 대해 작동합니다. 차원 $d$(밸류 차원)는 어텐션 확률 계산과 독립적입니다.
• 메커니즘: 가역적인 변환 $T_v$가 밸류 투영 가중치($W_v$)에 적용되고, 그 역행렬 $\bar{T}_v$가 출력 투영 가중치($W_o$)에 적용됩니다.
• 특성: 서로 다른 어텐션 헤드들은 독립적이므로, 각 헤드마다 최적화된 고유한 가역 행렬을 사용할 수 있습니다. 이 변환은 $W_v$와 $W_o$에 병합됩니다.
• 장점: 이는 높은 표현력(헤드당 전체 가역 행렬 제공)을 제공하면서도 병합 가능합니다. 이는 단일 회전을 모든 헤드에 공유하거나 병합이 불가능했던 기존의 회전 기반 방식들과 차별화됩니다.

3. 의사 동적 잔차 스케일링 (Pseudo-Dynamic Residual Scaling, $S_n$)

• 컨텍스트: 트랜스포머의 잔차 연결(residual connections)은 종종 정규화 레이어를 우회하므로, 토큰마다 매우 다른 스케일을 가질 수 있습니다. 이는 후속 블록(예: SwiGLU 출력)에서 아웃라이어를 유발합니다.
• 메커니즘: 저자들은 RMSNorm을 잔차 스트림에도 적용하도록 제안합니다. 잔차 입력을 대상으로 하는 퍼-토큰(per-token) 스케일링 인자 $S_n$(잔차의 제곱평균제곱근 기반)이 적용됩니다. 출력을 보존하기 위해, 어텐션 및 MLP 블록의 출력($Y_n$)은 잔차에 더해지기 전에 동일한 $S_n$으로 재스케일링됩니다.
• 특성: 행렬 곱셈과 편향(bias)이 없는 선형 레이어는 퍼-토큰 스케일링과 교환 가능합니다. 이를 통해 스케일링을 어텐션 및 MLP 블록 내부 깊숙이 적용하여 블록 내 아웃라이어를 줄일 수 있습니다.
• 구현: 스케일은 이전 블록의 정규화된 출력을 사용하여 재귀적으로 계산됩니다. 실제로, LM 헤드가 RMSNorm으로 시작하기 때문에 최종적인 스케일 나눗셈은 불필요합니다.

최적화 전략

FPTQuant는 두 단계의 최적화 과정을 채택합니다:

1. 로컬 최적화: 변환들은 병합된 가중치의 $L_p$ 노름(특히 $L_4$)을 최소화하도록 로컬에서 독립적으로 최적화됩니다. 이는 아웃라이어를 줄이고 안정적인 초기화를 제공합니다.
2. 엔드-투-엔드 (학생-교사) 훈련: 변환이 포함된 양자화된 학생 모델이 젠슨-샤논 발산(Jensen-Shannon Divergence, JSD) 손실을 사용하여 풀 프리시전(full-precision) 교사 모델을 근사하도록 훈련됩니다.
   • 중요한 차이점: 다음 토큰 예측 손실을 사용하는 방법과 달리, 학생-교사 접근 방식은 교사의 전체 확률 분포를 일치시킴으로써 훈련 데이터셋(예: Wikitext)에 대한 과적합을 방 prevented 합니다. 이는 다운스트림 태스크에 대한 일반화 성능을 향상시킵니다.

주요 기여

1. 새로운 FPT: 높은 표현력을 가지면서도 병합 가능한(Pre-RoPE, Value, Residual Scaling) 세 가지 특정 변환을 도입하여 추론 오버헤드를 제거했습니다.
2. 정적 INT4 양자화: FPTQuant가 가중치, 활성화, KV 캐시 모두에 대해 최소한의 오버헤드로 정적 INT4 양자화를 가능하게 함을 입증했으며, 프리필(prefill) 태스크에서 FP16 대비 최대 3.9배의 속도 향상을 달고 달성했습니다.
3. 정확도-속도 트레이드오프: FPTQuant는 FlatQuant나 SpinQuant와 같은 최신(SOTA) 방법들과 대등하거나 이를 능가하는 성능을 보여주면서도, 훨씬 더 빠르며(FlatQuant보다 최대 29% 빠름) 커스텀 커널을 요구하지 않습니다.

실험 결과

저자들은 다양한 모델(Llama-2/3, Ministral, Qwen)에 대해 다양한 양자화 설정(W4A4KV4, W4A8KV8 등)에서 FPTQuant를 평가했습니다.

• 속도: FPTQuant는 FP16 대비 프리필 속도에서 일관되게 2.8배~3.9배의 향상을 달성했습니다. 이는 SpinQuant보다 빠르며, FlatQuant보다 일관되게 빠릅니다(상대적 속도 15~29% 향상).
• 정확도:
  • "어려운" 배포 환경(모든 중간 활성화 값을 양자화하는 경우)에서 FPTQuant는 SpinQuant 및 FlatQuant와 같은 베이스라인보다 우수한 성능을 보입니다.
  • 표준 설정(Linear+KV)에서 FPTQuant는 현재 정확도 측면의 SOTA인 FlatQuant와 대등한 성능을 보입니다.
  • 도전적인 W4A4KV4 설정에서 FPTQuant는 경쟁력 있는 퍼플렉서티(perplexity)와 제로샷 추론 점수를 보여주며, 종종 SpinQuant와 QuaRot을 능가합니다.
• 안정성: 민감도 분석 결과, FPTQuant는 노이즈가 있는 파라미터 업데이트 하에서도 함수를 보존하는 반면, 부동 소수점 역행렬에 의존하는 FlatQuant나 비함수 보존 방식인 OSTQuant는 불안정하거나 성능이 저하됨을 보여줍니다.
• 훈련 비용: FPTQuant 훈련은 효율적입니다. RTN-opt보다는 약간 느리지만, 병합 가능한 변환을 사용하기 때문에 온라인 포워드 패스가 필요한 비병합 방식의 FlatQuant보다는 일반적으로 더 빠릅니다.

의의 및 주장

본 논문은 FPTQuant가 LLM을 에지 배포에 적합하게 만드는 데 있어 중요한 진전임을 주장합니다. 트랜스포머 연산의 수학적 특성(등변성)을 명시적으로 활용함으로써, 저자들은 정확도를 위해 추론 속도를 희생하지 않는 변환을 만들어냈습니다.

• 실질적 영향: FPTQuant는 동적 양자화나 비병합 변환의 무거운 계산적 페널티 없이 더 작은 비트 폭(INT4)을 사용할 수 있게 합니다. 이는 LLM의 메모리 점유율과 에너지 소비를 줄여, 스마트폰 및 에지 디바이스에서의 AI 접근성을 잠재적으로 대중화할 수 있습니다.
• 방법론적 기여: 본 연구는 효과적인 양자화 전략을 설계하는 데 있어 기존 연산의 교환 법칙(commutation properties)을 이해하는 것이 핵심임을 강조합니다. 또한, FPT 훈련 시 다음 토큰 예측 손실을 사용하는 것이 과적합 위험이 있음을 지적하며, 분포 매칭을 통한 학생-교사 훈련을 권장합니다.
• 한계점: 저자들은 이 접근 방식이 다양한 모델 제품군에서 견고하지만, 모든 LLM(예: Mixture-of-Experts)에 동일하게 적용된다는 보장은 없다고 겸허히 언급했습니다. 또한, 현재 하드웨어의 한계(많은 GPU에서 네이티브 INT4 지원 부족)로 인해 벤치마킹이 에뮬레이션에 의존하고 있으며, 완전한 이점은 미래의 하드웨어나 에지 전용 가속기에서 실현될 수 있음을 시사했습니다.