Hugging Face Diffusers가 Stable Diffusion 3 Medium을 통합하며 MMDiT 아키텍처와 Rectified Flow Matching으로 텍스트-이미지 생성 파이프라인 제공

Context

이전 Stable Diffusion 모델들은 텍스트 정보를 고정된 표현으로 cross-attention을 통해 레이턴트에 통합했으며, 메모리 제약으로 인해 대규모 GPU(24GB 이상 VRAM)에서만 추론이 가능했다. T5-XXL과 같은 대형 텍스트 인코더 3개를 사용하면서 메모리 압박이 심화되었다.

Technical Solution

• 텍스트-이미지 처리를 단일 시퀀스로 통합: CLIP L/14, OpenCLIP bigG/14, T5-v1.1-XXL의 3개 텍스트 인코더에서 패치 단위로 변환한 이미지 레이턴트와 텍스트 임베딩을 하나의 임베딩 시퀀스로 연결
• Multimodal Diffusion Transformer(MMDiT) 도입: 텍스트와 이미지 시퀀스를 별도의 가중치로 공통 차원성으로 임베딩한 후 attention 연산 직전에 결합하여 양방향 정보 흐름 구현
• Rectified Flow Matching 학습 방식 적용: 데이터와 노이즈 분포를 직선으로 연결하는 조건부 플로우 매칭 목적함수 사용으로 샘플링 스텝 감소 시 성능 유지
• FlowMatchEulerDiscreteScheduler 신규 스케줄러 추가: Rectified flow-matching 공식화와 Euler 방법 스텝 적용, 해상도별 타임스텝 시프트 파라미터로 고해상도 노이즈 스케일링 개선(2B 모델에서 shift=3.0 권장)
• 메모리 최적화 기능 제공: 모델 오프로딩으로 CPU/GPU 간 컴포넌트 전환, torch.compile() 지원으로 fp16 기준 추론 속도 4배 향상, LoRA 기반 DreamBooth 파인튜닝 스크립트 제공

Impact

• torch.compile() 적용 시 A100 80GB 머신에서 fp16 20단계 추론 평균 0.585초로 eager execution 대비 4배 속도 향상

Key Takeaway

텍스트-이미지 생성 모델에서 고정된 텍스트 표현 대신 MMDiT를 통해 이미지 레이턴트와 텍스트를 공통 시퀀스로 처리하면 양방향 정보 흐름을 구현할 수 있으며, Rectified Flow Matching과 FlowMatchEulerDiscreteScheduler 조합으로 샘플링 효율성을 높일 수 있다. 대형 텍스트 인코더의 메모리 부담은 모델 오프로딩과 torch.compile()로 광범위한 GPU 환경에서 추론 가능하게 만들 수 있다.