Collaborative Few-Step Distillation and Low-Bit Quantization for Wan2.2 Dual-Expert Video Diffusion Models

기술 요약: Wan2.2 듀얼 전문가 비디오 확산 모델을 위한 협력적 적은 단계 증류 및 저비트 양자화

1. 문제 정의

Wan2.2-T2V-A14B와 같은 대규모 비디오 확산 모델은 높은 시각적 품질을 달もの하지만, 다음과 같은 두 가지 주요 병목 현상으로 인해 배포에 상당한 장벽이 존재한다:

1. 계산 비용: 반복적인 디노이징(denoising) 과정은 많은 수의 샘플링 단계를 요구하며, 이는 추론 시 지연 시간(latency) 측면에서 비용을 높인다.
2. 메모리 점유율: 이 모델들은 거대한 상주 파라미터 수(듀얼 전문가 아키텍처에서 약 14B 개의 활성 파라미터를 포함하여 총 약 27B)를 보유하고 있어, 많은 배포 시나리오의 메모리 예산을 초과한다.

기존의 압축 전략은 이러한 모델에 직접 적용될 때 종종 실패한다. 단순히 재학습 없이 샘플링 단계만을 줄이면 품질이 저하되며, 원래의 긴 단계 궤적(long-step trajectory)을 양자화하면 추론 중 활성화 분포 불일치(activation-distribution mismatch)가 발생하여 시간적 오류 전파와 비디오 충실도 저하를 유발할 수 있다. 또한, Wan2.2는 듀얼 전문가 MoE(Mixture of Experts) 아키텍처(고노이즈 전문가 및 저노이즈 전문가)를 사용하므로, 모델을 밀집된 트렁크(dense trunk)로 취급하는 표준 양자화 접근 방식은 복잡성을 더한다.

2. 방법론

본 논문은 **적은 단계 증류(few-step distillation)**와 **저비트 양자화(low-bit quantization)**를 공동 설계(co-designed) 방식으로 통합한 배포 지향적 압축 파이프라인을 제안한다. 이 파이프라인은 모델 적응 $\rightarrow$ 적은 단계 증류 $\rightarrow$ 저비트 양자화 $\rightarrow$ 캘리브레이션 구축 $\rightarrow$ 생성 $\rightarrow$ 평가의 특정 순서를 따른다.

주요 기술 구성 요소:

• 듀얼 전문가 인지 아키텍처 처리:

  • 파이프라인은 고노이즈 전문가가 글로벌 레이아웃/모션을 담당하고 저노이즈 전문가가 텍스처/디테일을 정교화하는 Wan2.2의 특정 라우팅 메커니즘을 준수한다.
  • 분기별 캘리브레이션(Branch-wise Calibration): 활성화 값을 풀링하는 대신, 이 방법은 고노이즈 및 저노이즈 브랜치에 대해 별도로 통계량을 포착하고 캘리브레이션한다. 이는 양자화기가 서로 호환되지 않는 활성화 분포를 혼합하는 것을 방지한다.
  • 입구 레이어 보호(Entrance-Layer Protection): 초기 레이어와 라우팅에 민도 민감한 블록은 더 높은 정밀도를 유지하거나 더 공격적으로 캘리브레이션된다. 이는 초기 고노이즈 단계(레이아웃/모션 초기화)에서의 오류가 전체 시퀀스로 전파될 수 있는 반면, 후기 오류는 국소적인 데 영향을 미치기 때문에 매우 중요하다.

• 분포 매칭을 통한 적은 단계 증류:

  • 이 방법은 단순히 샘플러를 자르는 것이 아니라, 더 짧은 디노이징 궤적을 학습하기 위해 **분포 매칭 증류(Distribution-Matching Distillation, DMD)**를 채택한다.
  • 학생 모델은 줄어든 단계에 걸쳐 교사 모델의 분포를 일치시키도록 최적화되며, 고노이즈 단계와 저노이즈 단계의 구별된 역할을 존중하는 단계별 목적 함수($L_{distill} = \lambda_1 L_{high} + \lambda_2 L_{low}$)를 갖는다.

• 공동 설계된 양자화:

  • 궤적 정렬 캘리브레이션(Trajectory-Aligned Calibration): 핵심 혁신은 양자화 캘리브레이션을 원래의 긴 단계 모델이 아닌 증류된 적은 단계 학생 모델에서 수행한다는 것이다. 이는 양자화기가 실제 추론 중에 마주하게 될 활성화 통계량을 관찰하도록 보장하여 분포 불일치를 줄인다.
  • HiF4 (HiFloat4) 표현: 파이프라인은 저비트 추론을 위해 설계된 4비트 블록 부동 소수점 형식인 HiF4를 활용한다. 이 형식은 3단계 스케일링 계층(Global E6M2, Level-2 E1, Level-3 E1)과 S1P2 부호 있는 진폭을 특징으로 한다. 이 설계는 디노이징 단계 전반의 이질적인 활성화 범위를 처리하기 위해 넓은 글로벌 동적 범위(약 $2^{-50}$ ~ $2^{18}$)를 제공하는 동시에, 마이크로 지수(micro-exponents)를 통해 국소적 디테일을 보존한다.
  • 그래뉼러리티(Granularity): 가중치는 채널별(per-channel)로 양자화되며, 활성화 값은 프롬프트 및 프레임 수준의 변화를 추적하기 위해 토큰별(per-token)로 양자화된다.

• 엔지니어링 및 평가:

  • 메모리 관리: 단일 GPU(예: 80GB H100)에서의 OOM(Out-Of-Memory) 오류를 방지하기 위해, 캘리브레이션과 교체는 레이어별로 수행되며 보조 버퍼는 필요에 따라 로드된다.
  • 폐쇄 루프 평가: 시스템은 피사체 일관성(SC), 미적 품질(AQ), 이미징 품질(IQ), 전체적 일관성(OC), 모션 부드러움(MS)의 다섯 가지 특정 차원을 평가하기 위해 VBench를 사용한다. 평가 프로토콜은 베이스라인에 대해 동일한 프롬프트, 해상도 및 샘플링 스케줄을 사용하여 공정한 비교를 보장한다.

3. 주요 기여

1. 증류와 양자화의 공동 설계: 본 논문은 양자화 캘리브레이션을 (원래의 긴 단계 모델이 아닌) 증류된 적은 단계 궤적에서 수행하는 것이 활성화 분포 불일치를 크게 줄여, 압축된 모델이 풀 정밀도 학생 모델에 근접한 품질을 유지할 수 있음을 입증한다.
2. MoE 인지 압축 전략: Wan2.2의 듀얼 전문가 구조에 특화된 분기별 캘리브레이션 및 입구 레이어 보호 기법을 도입하여, 고노이즈 및 저노이즈 활성화 통계가 혼합되는 것을 방지한다.
3. 비디오 확산을 위한 HiF4 통합: HiF4의 계층적 스케일링 적용이 비디오 확산 모델에 효과적임을 보여주며, 고노이즈 단계에서의 오버플로를 방지하기 위한 넓은 글로벌 범위와 저노이즈 단계에서의 텍스처 보존을 위한 국소적 정밀도 사이의 균형을 맞춘다.
4. 단계 수에 대한 경험적 검증: 4, 8, 20, 40단계 설정을 체계적으로 평가하여 품질과 효율성 사이의 최적의 절충점을 식별한다.

4. 실험 결과

실험은 캘리브레이션을 위해 OpenS2V 데이터셋을, 평가를 위해 VBench를 사용하여 Wan2.2-T2V-A14B 모델에 대해 수행되었다.

• 베이스라인 대비 성능:

  • 4단계에서, 양자화된 모델은 동일 단계의 풀 정밀도 모델에 비해 작은 오차를 유지하지만, 원래의 긴 단계 베이스라인에는 뒤처진다.
  • 8단계에서, 양자화된 모델(평균 점수: 0.7002)은 원래의 풀 정밀도 베이스라인(0.6976)을 능가하기 시작한다.
  • 20단계에서, 양자화된 모델은 최고의 성능(평균 점수: 0.7074)을 달성하며, 원래의 풀 정밀도 베이스라인과 동일 단계의 풀 정밀도 증류 모델(0.7051)을 모두 앞선다.
  • 40단계에서는 성능이 정체되며, 이는 궤적 길이를 더 늘려도 이득이 줄어듦을 나타낸다.

• 지표 분석 (20단계 설정):

  • 양자화된 모델은 피사체 일관성(0.9629 vs 0.9495), 미적 품질, 이미징 품질, 모션 부드러움 측에서 원래의 베이스라인보다 개선된 성능을 보였다.
  • 전체적 일관성은 동일 단계 모델과 대등한 수준을 유지했다.
  • 결과는 분포 정렬 캘리브레이션이 시간적 안정성을 희생하지 않으면서 저비트 손실을 성공적으로 회복했음을 나타낸다.

• 오류 민감도:

  • 분석 결과, 단일 프레임 섭동(perturbation)은 주로 이미징 및 미적 품질에 영향을 미치는 반면, 단계 전반에 걸쳐 누적된 오류는 피사체 드리프트(Subject Drift) 및 모션 지터(Motion Jitter)로 나타난다. 제안된 방법의 입구 레이어 보호는 후자를 효과적으로 완화한다.

5. 의의 및 주장

본 논문은 제안된 파이프라인이 제한된 GPU 메모리와 지연 시간 예산 하에서 고품질 비디오 확산 모델을 배포할 수 있는 실질적인 경로를 제공한다고 주장한다.

• 품질-효율성 절충: 저자들은 20단계 설정을 최적의 구성으로 식별하였으며, 이는 기존 샘플러보다 훨씬 적은 단계를 요구하면서도 시간적 및 시각적 품질을 회복하기에 충분한 디노이징 용량을 제공한다.
• 전략적 전환: 본 연구는 양자화가 "손실이 발생하는 사후 처리(lossy afterthought)"가 아니라 가속 전략의 필수적인 부분이어야 한다고 주장한다. 증류와 양자화를 공동 설계함으로써, 이 방법은 저비트 변환을 추론 파이프라인의 안정적인 구성 요소로 만든다.
• 배포 준비성: 입구 레이어 보호 및 분기별 캘리브레이션의 사용은 듀얼 전문가 비디오 모델의 특정 실패 모드(예: 피사체 드리프트 및 모션 불안정성)를 해결하여, 압축된 모델을 대화형 생성 및 자원이 제한된 배치 프로세싱에 적합하게 만든다.

결론적으로, 현재의 파이프라인은 사후 학습 양자화(post-training quantization)에 집중하고 있지만, 이 프레임워크는 특정 지표의 실패(예: 피사체 일관성 저하)가 캘리브레이션 및 단계 수를 정교화하기 위한 직접적인 디버깅 신호를 제공하는 검증 루프를 구축한다.