https://arxiv.org/abs/2406.02540 ICLR 2025
ViDiT-Q: Efficient and Accurate Quantization of Diffusion Transformers for Image and Video Generation
Diffusion transformers have demonstrated remarkable performance in visual generation tasks, such as generating realistic images or videos based on textual instructions. However, larger model sizes and multi-frame processing for video generation lead to inc
arxiv.org
Abstract
기존의 diffusion transformer 양자화 방식은 text-to-image 및 text-to-video와 같이 복잡한 multi-modal task에서 성능저하를 유발함. 본 논문에서는 diffusion transformer(DiT)에 특화된 Post-Training Quantization (PTQ) 방식인 ViDiT-Q를 제안함.
ViDiT-Q는 W8A8(Weight 8bit, Activation 8bit)와 W4A8(Weight 4bit, Activation 8bit) 환경에서 정밀도 저하를 최소화하며 1.7배의 지연시간(latency) 감소를 달성함.
Introdution
Problem Definition
DiT 기반의 이미지 및 비디오 생성은 우수한 성능을 보이지만, 메모리 사용량 및 계산 비용이 높아 edge device에 부적합함
이러한 한계를 개선하기 위해 Q-DiT, PTQ4DiT와 같은 기존의 DiT Quantization 기법은 coarse quantization을 진행함.
하지만 이러한 방식들은 낮은 bit-width(높은 압축률)에서 시각적 품질, 시간정보 고려, 텍스트 정합도 등 여러 지표에서 심각한 성능저하를 보임.
Proposed Methods
본 논문의 ViDiT-Q는 DiT에 특화된 3가지 기법을 도입하여 위 한계를 해결함.
1) Fine-grained Grouping & Dynamic Quantization
Token 및 Condition 단위의 세분화된 그룹으로 quantization 수행, timestep이나 조건(condition)에 따라 동적으로 양자화 파라미터(scale, zero-point) 계산
→ 하나의 큰 그룹 내 값들의 분산이 큰 경우 발생하는 rounding error를 줄여 quantization error 최소화
2) Static-Dynamic Channel Balancing
시간에 따라 변하는 channel 간 불균형을 줄이기 위해 정적인 scale factor와 동적인 rotation을 결합
→ timestep마다 달라지는 channel 별 분포차이를 안정적으로 보정하여 일관된 quantization 품질을 유지
3) Metric-Decoupled Mixed Precision
시각품질(visual quality), 텍스트 정합도(text alignment), 시간적 일관성(temporal consistency) 등 생성 품질 지표별로 민감한 layer를 구분하고, 각 그룹에 맞춤형 bit-width를 할당하여 중요 layer의 정보를 보존
→ 낮은 비트환경(W4A8) 환경에서도 다양한 품질 항목을 모두 유지할 수 있도록 정밀도를 분산 조절
Contributions
1) DiT Quantization에서 발생하는 quantization error를 체계적으로 분석함
2) DiT 모델에 최적화된 quantization기법인 ViDiT-Q 제안함
3) 다양한 DiT 모델 (text-to-image, text-to-video)에 적용해 성능을 검증하였으며, GPU 커널 구현을 통해 메모리 절감 및 latency 를 단축
Related Work
Diffusion Transformers for Image and Video Generation
기존에는 U-Net 기반의 CNN구조가 diffusion model의 주류였음. 하지만 최근에는 Transformer 기반의 DiT 구조가 각광받고 있음. 하지만, Video Generation task에서는 multi-frame을 처리하기 때문에 시공간적 복잡성이 매우 증가하며 이에 따른 메모리사용량과 계산 비용이 급증함. 따라서, 경량화 필요성이 커짐
Evaluation Metrics for Image and Video Generation
Image Generation에서는 다음과 같은 지표들을 사용함
- FID(Frechet Inception Distance) : 생성 이미지와 실제 이미지의 feature 분포 차이
- IS (Inception Score) : 생성 이미지의 다양성과 정합성
- CLIP Score : 텍스트-이미지 정합도
- ImageReward : 인간 선호도 기반의 평가
Video Generation에서는 다음과 같은 지표들을 사용함
- FVD(Frechet Video Distance) : FID의 비디오 확장
- CLIPSIM, CLIP-Temp : 텍스트-비디오 정합성, 시간적 정합성 평가
- Flow-score : 모션 품질
- DOVER : 종합적인 비디오 품질
※ 단일 함수(ex. MSE)로는 전체적인 생성 품질을 측정하기에 부족하며, 다양한 지표 기반의 정밀한 평가가 필요하기 때문에
ViDiT-Q에서는 metric-decoupled mixed precision을 도입함.
Model Quantization
PTQ (Post-Training Quantization)는 별도의 재학습 없이 모델을 압축할 수 있는 실용적인 양자화 방법임.
하지만, Transformer와 Diffusion 모델은 구조적 특성상 quantization 시 문제 발생 가능성이 높음.
Transformer에 대한 기존 Quantization 연구들은 다음과 같이 진행됨.
- SmoothQuant : 채널별로 scaling을 적용해 weight 및 activation의 quantization 난이도 조정
- Quarot : 회전(orthogonal rotation)을 적용해 outlier 제거 및 distribution 평준화
Diffusion에 대한 기존 Quantization 연구들은 다음과 같이 진행됨.
- Q-Diffusion, PTQ4DM : timestep 단위로 activation 수집 후 quantization parameter 설정
- Q-DiT : 채널 간 imbalance를 해결하기 위해 channel 간 grouping을 통 quantization parameter 설정
- PTQ4DiT : timestep의 변화를 고려해 고정된 마스크 방식의 channel balancing 도입
기존 방식들은 대부분 static한 quantization parameter를 사요하거나 timestep의 변화에 따라 성능이 저하됨.
특히, 다양한 품질 요소가 동시에 요구되는 상황(이미지 및 비디오 생성)에서는 성능저하가 심각함.
이를 해결하기 위해 본 논문에서는 ViDiT-Q를 제안함.
Preliminaries
Quantization Error Analysis
quantization error의 주요원인은 clipping error와 rounding error로 구분할 수 있음.
clipping error는 min/max 범위를 벗어나서 생기는 손실이며, rounding error는 scale이 클수록 라운딩 오차가 커지면서 발생함.
따라서, 그룹 내 데이터 변동이 크면 scale 값이 커지고, rounding error가 급증하여 품질 저하를 유발함.
Unique Challenges for DiT and Visual Generation
본 논문에서 분석한 DiT의 고유한 데이터 분포 특성은 다음과 같음
- Token-wise Variation
Visual Tokne의 경우 공간 및 시간 축을 따라 token 간 큰 분산 발생 (특히 비디오에서 심각)
ex. 앞프레임과 뒷프레임의 feature 차이
- Condition-wise Variation
Classifier-free guidance로 인해 conditional / unconditional forward 결과가 다름.
(Prompt의 적용 유뮤에 따라 activation값이 매우 크게 달라짐)
이에 따라, 같은 timestep이어도 condition 유무에 따라 activation이 다르게 분포 - Timestep-wise Variation
timestep에 따라 동일한 layer의 feature 값이 상당히 다르게 나타남
고정된 quantization parameter를 timestep 전체에 적용하는 것은 부적절
- Channel-wise Variation
weight나 activation 모두에서 channel 간 분산이 존재하며, 특히 activation은 timestep에 따라 동적으로 분산이 변함 (time-varing imbalance라고 함)
Methods
Fine-grained Grouping and Dynamic Quantization
데이터의 분포가 넓은 상태에서 quantization group을 설정하는 것은 quantization error를 발생시키는 요인임.특히, group-size가 커질수록 큰 variation을 포함하는데 이러한 접근은 효율적인 하드웨어 구현을 위해서는 더 세분화된 grouping을 사용해야함을 시사함. (Group-size가 작을수록 Quantization 성능 증가)
하지만, fine-grained(세분화된) quantization은 하드웨어에서 연산이 매우 비효율적임.
따라서, ViDiT-Q에서는 다음과 같은 quantization을 진행함.
activation : token-wise grouping
(각 토큰의 activation vector 전체를 하나의 group으로 사용, group size = 채널 수)
weight : token-wise or channel-wise grouping
(Linear weight matrix의 row 또는 column 단위를 group으로 설정, group size = head/channel 수)
해당 방식을 통해 세밀한 group-wise Quantization을 진행함.
또한, 기존의 PTQ방식은 소량의 calibaration 데이터를 활용하여 timestep에 관련없이 모두 한개의 scaling factor로 scaling하는 Static Quantization을 진행함.
하지만 ViDiT-Q 저자들은 Timestep-wise variation을 효과적으로 처리하기 위해서 timestep마다의 scaling을 하는 Dynamic Quantization을 진행함.
Static-Dynamic Channel Balancing
아래 그림과 같이 기존의 채널 scaling 방식은 고정된 scale factor($\alpha$)를 사용함. 하지만 시간에 따른 동적 변화(timestep variance)에 적응하지 못함. 또한, 기존의 rotation방식은 outler가 여전히 남아 있다는 문제가 존재함.
Channel Scaling과정을 구체적으로 보면, 기존의 Smooth-Quant는 다음 수식과 같이 동작함.
여기서 $X$는 activation을 의미하고, $W$는 weight를 의미함. $Y$는 activation과 weight의 projection을 마친 결과값을 의미함.
위 수식에서는 $X$는 줄이고, $W$는 키우는 연산을 시도함.
이러한 연산을 시도하는 이유는 양자화 오류(quantizaiton error)를 줄이기 위함인데, 구체적으로는 activation과 weight 간 불균형한 크기를 채널 단위의 scaling으로 보정하기 위함임.
Smooth Quant : 큰 값은 줄이고, 작은 값은 키워서 곱하기 → 표현범위는 줄이면서 원본과 유사한 결과값 도출 → 양자화 오류를 줄이는 핵심 방안 ☞ 자세한 개념은 논문에서 확인하기
Xiao, Guangxuan, et al. "Smoothquant: Accurate and efficient post-training quantization for large language models." International Conference on Machine Learning. PMLR, 2023.
$X$와 $W$는 scale factor인 $s$와의 연산을 통해 조정되는데, 이때 activation에 해당하는 $X$는 무조건 값을 줄이고 (역행렬 연산을 통해 나눔), weight에 해당하는 $W$는 곱셈연산을 통해 값을 증폭시키는 구조를 가짐. 이러한 구조를 통해 rotation이 되어 표현범위를 줄일 수 있음.
※ 하지만, ViDiT-Q는 이러한 기존 quantization 방식의 한계에 대해 지적함.
구체적으로는, $X$와 $W$값 중 $X$가 무조건 값이 크고, $W$가 무조건 값이 작을것이라고 가정하고 위와 같이 고정된 scaling 연산이 진행되는데, 이러한 방식은 양자화 관점에서 더 큰 오류를 만들게 하는 구조임을 지적함.
따라서, ViDiT-Q에서는 다음과 같은 두가지 케이스에 따라 채널들을 구분하였다고 함.
1) Timestep에 따라 activation값의 변동이 매우 큰 경우 (Dependent Channel, Dynamic Scaling이 효과적)
2) Timestep에 관계 없이 activation값이 매우 크거나, 혹은 매우 작은 경우 (Independent Channel, Static Scaling이 효과적)
저자들은 이러한 채널들을 모두 효과적으로 처리하기 위해 Static 방식과 Dynamic 방식을 모두 활용하는 방식을 제안함.
구체적으로, Static방식에서는 Scaling을 진행 할때 고정된 scaling factor를 사용하는 것이 아닌, Trainable한 scale table을 사용해 activation을 정규화함.
또한, Dynamic 방식에서는 Rotation이라는 개념을 도입하여 Trained Scale Table에서 현재 timestep에 맞는 scale 값을 선택하여 적용함.
결과적으로 ViDiT-Q에서는 Trainable scale table을 정적(static)인 방식, Rotation을 동적(dynamic)인 방식으로 정의하여 결합하였으며, 이를 통해 시간축까지 고려한 scaling이 가능하다고 주장함.
Rotation : 크기가 매우 큰 Vector값의 표현범위를 줄이기 위한 개념
☞ 자세한 개념은 선형대수학에서 찾아보기
Metric-Decoupled Mixed Precision
어떤 레이어는 quantization error가 작아서 전체 성능에 큰 영향을 주기 때문에 모든 레이어가 quantization에 동등하게 민감한 것은 아님.
따라서, 기존에 많이 활용되던 MSE 기반의 민감도 측정 방식은 시각적인 품질이나 시간 일관성과 같은 복합적인 측면의 반영이 불가능함.
따라서, 다양한 지표를 고려한 민감도 분석이 필요함
ViDiT-Q에서는 다음과 같은 단계를 거쳐 레이어마다 적절한 비트수를 할당하는 과정을 거침
1) 그룹별 parameter 예산 분배
모델의 모든 레이어를 기능별로 그룹화 하는 과정으로, FFN, Self-Attention과 같은 레이어를 그룹화
2) Metric-Decoupled 민감도 분석
각 그룹별로 중요 기능에 대한 지표를 선택함
- Visual Quality → FFN, Spatial Attn
- Textual Alignment → Cross-Attn
- Temporal Consistency → Temporal Attn
각 지표들에 대한 평균 Sensitivity값을 계산
3) 비트 수 할당
Sensitiviy값을 기반으로 민감한 레이어에 더 높은 비트를 할당하며, 가장 영향력이 큰 레이어에 대해서는 높은 bit-width를 할당
반면, 영향력이 작은 레이어에 대해서는 낮은 bit-width(높은 압축)을 할당하여 Quantization 진행
Experiements
Image Generation에서는 DiT, Video Generation에서는 VideoDiT 모델을 실험에 활용함.
Image Generation 실험은 다음과 같이 진행함
ViDiT-Q는 4-bit weight까지 성공적으로 quantize했으며, 기존 모델들보다 높은 압축률에도 불구하고 더 좋은 성능을 보임.
Video Generation은 다음과 같이 실험함.
Video Generation에서도 마찬가지로, 높은 압축률에도 불구하고 대체적으로 좋은 성능을 도출함을 확인할 수 있음.
자원사용량에 대한 실험도 Figure 7과 같이 진행하였으며, 압축을 진행하지 않은 Full-Precision 모델에 비해 자원활용을 획기적으로 줄임.
Table 2는 Video Generation에 대한 ablation study로 method에 제시한 3가지 방식을 모두 적용하였을 때 가장 우수한 성능이 나옴을 검증함. 특히, VAQ에서 획기적으로 높은 성능을 기록함.
Channel Balancing 기법에 대한 효과를 정성적으로 평가하기 위해 Figure 8과 같은 시각화 또한 진행함.
아무런 balancing도 적용하지 않은 Original Weight같은 경우 채널간의 스케일 차이가 매우 크며, incoherence 또한 40.75로 가장 높음. 이는 양자화가 매우 불리한 상태임을 의미함.
반면 ViDiT-Q에서 제안한 Static-Dynamic Channel Balanced의 경우 5.02의 incoherence를 보여 양자화에 가장 유리한 상태임을 분석하였음.
Incoherence : 채널 간 분포의 불균형 정도를 의미함. 값이 클수록 양자화에 취약하며, 일반적으로 스케일이 너무 크거나, 너무 작을때 incoherence가 큼
Conclusion & Limitation
ViDiT-Q는 DiT의 양자화 한계를 해결하기 위해 설계된 기법으로 세가지 method를 결합함.
1) Fine-graiend, Dynamic Quantization
2) Static-Dynamic Channel Balancing
3) Metric-Decoupled Mixed Precision
한계점으로는 Mixed Precision에 대한 전략은 아직까지 개선 여지가 있으며, 특히 layer sensitivity 측정 방식이나 비트수 할당방식에 대한 고도화가 필요하다고 함.
또한, activation에 대한 비트수를 더욱 압축하는 연구가 필요하다고 함.
Review
Strengths
- Timestep에 따라 달라지는 activation분포를 적절히 분석하고 이에 동적인 방식을 활용해 해결함
- 기존 연구에서는 단순히 MSE만 평가지표로 사용하지만 ViDiT-Q에서는 다양한 지표를 통해 철저한 검증을 진행함
- incoherence의 감소를 보여 양자화 오류를 줄일 수 있다는 것을 보였음
- 시간축과 공간축을 동시에 고려해 처음으로 성공적인 DiT Quantzation
Weakness
- Method에서 수식이 정리되어 있지 않음 (기존연구의 수식에 대한 한계만 설명하고 제안방식의 수식이 없음)
- Rotation 내용과 관련한 ablation study가 없음
- 전체적인 Method 흐름에 대한 figure는 존재하지만, 상세한 개념을 이해하기 위한 figure가 부족함.
