KroQuant: Diffusion Transformers의 효율적인 사후 훈련 양자화(PTQ)를 위한 Kronecker 구조 블록 변환
요약
Diffusion Transformers(DiTs)의 W4A4 양자화 시 발생하는 품질 저하 문제를 해결하기 위해 Kronecker 구조 블록 변환을 사용하는 KroQuant를 제안합니다. 이 방법은 낮은 추론 비용과 높은 양자화 정확도를 동시에 달성하며, 기존 방식보다 효율적인 연산이 가능합니다.
핵심 포인트
- Kronecker 구조를 활용해 DiT의 활성화 함수 이상치 문제를 해결
- SmoothQuant 대비 적은 파라미터와 높은 추론 속도 제공
- MI350 GPU 환경에서 기존 커널 대비 최대 14% 빠른 성능
- PixArt-Σ, FLUX.1 등 주요 모델에서 우수한 이미지 품질 유지
Diffusion Transformers (DiTs)를 W4A4로 사후 훈련 양자화 (Post-training quantization, PTQ)하면 출력 품질이 심각하게 저하되는데, 이는 각 선형 레이어 (linear layer)로 들어가는 활성화 함수 (activations) 값에 4비트 형식이 표현할 수 없는 이상치 (outliers)가 포함되어 있기 때문입니다. 표준적인 해결책은 양자화 전 활성화 함수에는 가역 선형 변환 (invertible linear transform)을 적용하고, 가중치 (weights)에는 그 역변환을 적용하는 것입니다. 블록 사이의 정규화 레이어 (Normalization layers)로 인해 이 변환은 매 디노이징 (denoising) 단계마다 온라인으로 실행되어야 하며, 이로 인해 추론 계산 비용이 설계의 제약 조건이 됩니다. 기존의 옵션들은 양자화 품질과 추론 비용 사이에서 절충을 시도합니다. 채널별 스케일링 (per-channel scaling, SmoothQuant)은 계산 비용은 저렴하지만 채널의 크기에 영향을 주어 양자화 정확도를 해칠 수 있습니다. 고정된 Hadamard 변환은 더 나은 양자화 정확도를 제공하지만, 높은 온라인 비용을 초래하는 큰 블록 크기를 요구합니다. 학습된 전체-$d$ 가역 변환은 가장 잘 보정되지만, 단계별 레이어당 과도한 밀집 $d imes d$ 행렬 곱셈 (GEMM)을 수반합니다. 우리는 활성화 함수의 각 32개 요소 블록에 학습된 Kronecker 구조의 가역 변환을 적용하는 PTQ 방법인 KroQuant를 제안합니다. 이는 채널별 스케일링보다 절반 미만의 파라미터를 저장합니다. 블록 국소 구조 (block-local structure)는 작은 텐서 코어 GEMM으로 실행되며, MI350 GPU에서 KroQuant 양자화 커널은 SmoothQuant 커널보다 최대 14% 더 빠릅니다. 이후 오프라인 LoRaQ 가중치 보정 (weight calibration)을 통해 잔여 가중치별 양자화 오차를 흡수합니다. PixArt-$\Sigma$, SANA, FLUX.1-schnell의 W4A4 (MXFP4e2) 환경에서 KroQuant는 MJHQ-30K 및 SDCI 데이터셋에 대해 SVDQuant 및 LoRaQ보다 FP 참조값에 더 가까운 출력을 생성하는 동시에 이미지 품질을 유지하거나 향상시킵니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 arXiv cs.LG의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기