Hugging Face 팀이 확산 모델 대신 Masked Image Modeling 방식의 aMUSEd 모델을 개발해 추론 단계를 단계 단위로 감소시키고 모델 크기를 800M 파라미터로 축소

Context

기존 텍스트-이미지 생성 모델들은 잠재 확산(Latent Diffusion) 방식을 사용하면서 많은 추론 단계가 필요하고 모델 크기가 커서 추론 레이턴시가 길고 온디바이스 배포가 어려웠다.

Technical Solution

• 확산 모델 대신 Masked Image Modeling(MIM) 방식 도입: VQGAN으로 이미지를 토큰화한 후 코사인 마스킹 스케줄에 따라 마스크된 토큰을 U-ViT 모델이 예측
• CLIP-L/14 텍스트 인코더로 프롬프트 임베딩 생성: T5 대신 사용해 조건부 생성 개선
• 마이크로-컨디셔닝 추가: 이미지 크기, 크롭 정보 등 추가 조건을 U-ViT에 전달 (Stable Diffusion XL 방식 도입)
• 단일 단계 아키텍처 적용: MUSE의 2단계 패치 예측 방식 대신 1단계 구조로 단순화
• Diffusers 라이브러리에 통합: AmusedPipeline으로 텍스트-이미지, 이미지 인페인팅 등 다중 작업 지원

Impact

추론 레이턴시: A100 기준 5~15 추론 단계로 최소 5단계부터 생성 가능 (비교 기준 모델들의 수십~수백 단계 대비) 모델 크기: 텍스트 인코더 및 VQ-GAN 포함 약 800M 파라미터 GPU 메모리: 8비트 Adam 옵티마이저와 float16 정밀도로 11GB 미만 VRAM으로 미세조정 가능, LoRA 사용 시 7GB로 축소

Key Takeaway

확산 모델이 표준인 이미지 생성 분야에서도 언어 모델링의 성숙한 기법인 Masked Image Modeling을 적용하면 추론 효율성과 모델 크기 측면에서 실질적인 이득을 얻을 수 있으며, 제한적인 하드웨어 환경에서의 배포 가능성이 높아진다.