Elastic Diffusion Transformer

1. 연구 배경 및 문제 정의 (Problem)

• 배경: 확산 모델 (Diffusion Models), 특히 Diffusion Transformer (DiT) 는 이미지, 비디오, 3D 자산 등 다양한 모달리티에서 뛰어난 생성 능력을 입증했습니다.
• 문제점:
  • 높은 계산 비용: 대규모 모델 크기와 많은 디노이징 (denoising) 단계로 인해 추론 비용이 매우 높고 지연 시간 (latency) 이 길어 실용적 배포에 제한이 있습니다.
  • 기존 가속화 방법의 한계: 가지치기 (Pruning) 나 증류 (Distillation) 와 같은 기존 가속화 기법들은 고정된 계산 용량을 기반으로 합니다. 이는 모든 샘플과 모든 디노이징 단계에 동일한 구조를 적용하므로, 생성 품질을 희생하지 않으면서 효율성을 극대화하는 데 한계가 있습니다.
  • 샘플 의존적 희소성 (Sample-dependent Sparsity) 의 간과: 생성 과정에서 각 블록 (Block) 의 중요성이나 필요한 계산량은 입력 샘플의 난이도와 내용에 따라 크게 달라집니다. 그러나 기존 고정 구조는 이러한 차이를 반영하지 못해 불필요한 계산을 수행하거나 중요한 계산을 생략하는 비효율이 발생합니다.

2. 제안 방법: Elastic Diffusion Transformer (E-DiT)

저자들은 DiT 의 생성 과정에 존재하는 샘플 의존적 희소성을 활용하여, 생성 품질을 유지하면서 효율성을 극대화하는 적응형 가속화 프레임워크인 E-DiT를 제안합니다.

핵심 메커니즘

E-DiT 는 각 DiT 블록에 경량화된 **라우터 (Router)**를 부착하여 입력 잠재 변수 (latent) 와 디노이징 타임스텝에 기반하여 동적으로 계산을 조절합니다.

1. 적응형 블록 스킵 (Adaptive Block Skipping):

   • 라우터가 각 블록의 기여도를 예측하여, 기여도가 미미한 블록은 완전히 건너뜁니다 (Skip).
   • 이는 샘플마다 다른 블록 세트를 스킵할 수 있게 하여 불필요한 계산을 제거합니다.

2. 적응형 MLP 폭 축소 (Adaptive MLP Width Reduction):

   • 스킵되지 않는 블록 내부에서도 라우터가 샘플의 복잡도에 따라 MLP 의 활성화 폭 (Width) 을 동적으로 조절합니다.
   • 복잡한 샘플에는 넓은 폭을, 간단한 샘플에는 좁은 폭을 사용하여 계산량을 최적화합니다.

3. 블록 단위 캐싱 (Block-wise Feature Caching):

   • 훈련 불필요 (Training-free): 라우터의 예측을 활용하여 인접한 디노이징 단계 간의 시간적 중복성을 제거합니다.
   • 경계선 영역 (Borderline Region) 활용: 스킵할 수는 없지만 기여도가 낮은 (라우터 확률이 임계값 근처인) 블록의 경우, 이전 단계에서 계산된 잔차 (residual) 를 재사용하여 계산을 생략합니다. 이는 추가적인 학습 없이 추론 속도를 더 높여줍니다.

학습 전략

• 공동 최적화 (Joint Optimization): 생성 품질을 유지하는 성능 손실 (Performance Loss) 과 계산 효율성을 장려하는 효율성 손실 (Efficiency Loss) 을 결합하여 엔드 - 투 - 엔드 (End-to-End) 로 학습합니다.
• 초기화 전략: 학습 초기에는 모든 블록을 활성화하고 전체 폭을 사용하는 'Full-capacity' 상태로 시작하여 학습 안정성을 확보한 후, 점차 적응형 가속을 학습하도록 합니다.

3. 주요 기여 (Key Contributions)

1. 샘플 적응형 가속 프레임워크: 고정된 구조가 아닌, 입력 콘텐츠의 난이도에 따라 계산 자원을 동적으로 할당하는 E-DiT 를 제안했습니다.
2. 세 가지 상호 보완적 구성 요소: 블록 스킵, MLP 폭 축소, 블록 단위 캐싱을 통합하여 기존 방법보다 유연하고 강력한 가속화를 달성했습니다.
3. 광범위한 검증: 2D 이미지 생성 (Qwen-Image, FLUX) 과 3D 자산 생성 (Hunyuan3D-3.0) 등 다양한 모달리티와 모델 아키텍처에서 유효성을 입증했습니다.
4. 훈련 불필요한 캐싱 메커니즘: 라우터 예측을 기반으로 한 단순하고 효과적인 캐싱 전략을 도입하여 추가 학습 없이도 추론 속도를 추가로 향상시켰습니다.

4. 실험 결과 (Results)

• 속도 향상: Qwen-Image, FLUX, Hunyuan3D-3.0 모델에서 약 2 배 (2×) 의 추론 속도 향상을 달성했습니다.
• 품질 유지: 가속화에도 불구하고 생성 품질 (GenEval, DPG, ULIP 등) 의 손실은 미미하거나 거의 없었습니다.
  • 예시 (Qwen-Image): Base 모델 (2431ms) 대비 E-DiT-turbo 는 1283ms 로 약 1.9 배 빠르면서 GenEval 점수는 0.853 (Base: 0.870) 으로 유사한 수준을 유지했습니다.
  • 예시 (Hunyuan3D-3.0): 5012ms 에서 2587ms 로 약 2 배 빨라졌으며, ULIP-I 점수는 오히려 소폭 상승 (0.1446 → 0.1473) 하였습니다.
• 비교 우위: 기존 가지치기 기반 방법 (TinyFusion, PPCL 등) 이나 동적 네트워크 방법 (DyDiT 등) 보다 더 낮은 지연 시간과 더 나은 품질 트레이드오프를 보였습니다.

5. 의의 및 결론 (Significance)

• 실용성: 대규모 생성 모델의 배포 장벽인 높은 계산 비용을 획기적으로 낮추어, 실시간 또는 저사양 환경에서의 생성 모델 활용을 가능하게 합니다.
• 일반성: 특정 모델 아키텍처에 국한되지 않고 다양한 DiT 백본에 적용 가능한 범용 프레임워크를 제시했습니다.
• 미래 방향: 생성 과정의 '희소성'을 체계적으로 활용하는 새로운 패러다임을 제시하여, 향후 더 효율적인 생성 모델 설계에 중요한 통찰을 제공합니다.

요약하자면, E-DiT는 "모든 샘플이 같은 계산이 필요한 것은 아니다"라는 통찰에서 출발하여, 동적 라우팅과 캐싱을 통해 불필요한 계산을 제거함으로써 높은 품질을 유지하면서 2 배의 속도 향상을 이룬 획기적인 가속화 기술입니다.