aMUSEd: 효율적인 텍스트-이미지 생성

Hugging Face가 aMUSEd를 소개했습니다. 이는 효율적인 비확산 텍스트-이미지 모델입니다. Google의 MUSE를 오픈 재현한 aMUSEd는 Masked Image Modeling (MIM)을 활용하여 기존의 latent diffusion 모델에 비해 더 빠르고 해석 가능한 대안을 제공합니다.

Masked Image Modeling 아키텍처

aMUSEd는 Masked Image Modeling (MIM) 접근 방식을 활용하여 latent diffusion에 비해 추론 단계 수를 줄이고 해석 가능성을 높입니다. 모델 아키텍처는 세 가지 주요 구성 요소로 이루어져 있습니다: 토큰화를 위한 VQGAN, 프롬프트 임베딩을 위한 CLIP-L/14 텍스트 인코더, 그리고 마스크된 패치를 예측하는 U-ViT 모델.

학습 과정

학습 중 모델은 다음 단계들을 수행합니다:

  1. Tokenization: 입력 이미지가 VQGAN을 사용해 이미지 토큰으로 변환됩니다.
  2. Masking: 이미지 토큰이 코사인 마스킹 스케줄에 따라 마스크됩니다.
  3. Prediction: 마스크된 토큰은 CLIP-L/14 텍스트 인코더의 프롬프트 임베딩에 조건화된 뒤 U-ViT 모델에 전달되어 마스크된 패치를 예측합니다.

추론 과정

추론 시 모델은 반복적인 과정을 통해 이미지를 생성합니다:

  1. Embedding: 입력 프롬프트가 CLIP-L/14 텍스트 인코더를 사용해 임베딩됩니다.
  2. Iterative Prediction: 무작위로 마스크된 토큰으로 시작합니다. U-ViT 모델이 마스크된 토큰을 예측하고, 마스크 스케줄 및 단계 수 N에 따라 가장 자신 있는 예측의 일부만 유지됩니다. 나머지 토큰은 다시 마스크되어 U-ViT 모델에 다시 전달됩니다.
  3. Decoding: 최종 출력은 VQGAN 디코더를 통해 최종 이미지를 생성합니다.

MUSE와의 주요 차이점

aMUSEd는 원래 MUSE 모델과 세 가지 주요 측면에서 다릅니다:

  • Single-Stage Prediction: 최종 마스크된 패치를 예측하기 위해 두 단계 접근 방식을 사용하지 않습니다.
  • Text Conditioning: 텍스트 임베딩을 계산할 때 T5 대신 CLIP L/14를 사용합니다.
  • Micro-conditioning: SDXL을 따라 이미지 크기 및 크롭과 같은 추가 조건을 포함합니다.

기능 및 성능

aMUSEd는 효율성과 속도를 위해 설계되었습니다. 텍스트 인코더와 VQ-GAN을 포함해 약 8억 개의 파라미터를 가지고 있어 높은 성능을 발휘합니다.

추론 속도 및 지연 시간

A100 GPU에서 수행된 벤치마크 결과, aMUSEd는 다른 모델에 비해 추론 지연 시간이 크게 낮음을 보여줍니다. 이러한 효율성은 온디바이스 애플리케이션에 적합한 후보가 됩니다.

제로샷 인페인팅

사전 학습 목표 덕분에 aMUSEd는 제로샷 이미지 인페인팅이 가능하며, 이는 SDXL과 같은 모델에는 없는 기능입니다.

스타일 전이

MUSE의 기능을 이어받아 aMUSEd는 단일 이미지를 사용한 스타일 전이 능력을 보여주며, 이는 개인화 및 스타일별 이미지 생성에 잠재력을 제공합니다.

파인튜닝 및 접근성

aMUSEd는 OpenRAIL 라이선스로 출시되어 상업적으로 친화적입니다. Hugging Face는 커스텀 데이터셋에 대한 파인튜닝을 위한 학습 스크립트를 제공합니다.

파인튜닝을 위한 하드웨어 요구사항

  • Standard Fine-Tuning: 8-bit Adam 옵티마이저와 float16 정밀도를 사용하면 파인튜닝에 11GB 미만의 GPU VRAM이 필요합니다.
  • LoRA: Low-Rank Adaptation (LoRA)을 사용하면 메모리 요구사항이 7GB VRAM으로 더욱 감소합니다.

제한 사항

aMUSEd는 원시 이미지 품질 면에서 최첨단은 아니지만, 커뮤니티가 이미지 생성에 대한 MIM의 잠재력을 탐구하도록 장려하기 위해 출시되었습니다. 저자들은 MIM 프레임워크의 여러 장점을 강조합니다. 여기에는 추론 효율성, 온디바이스 사용을 위한 작은 모델 크기, 그리고 비용이 많이 드는 파인튜닝 없이 작업 전환을 수행할 수 있는 능력이 포함됩니다.

Sources