[논문 리뷰] SANA: Efficient High-Resolution Image Synthesis with Linear Diffusion Transformers

Enze Xie, Jun Song Chen외 9명|arXiv (Cornell University)|2024. 10. 14.

Advanced Vision and Imaging인용 수 4

한 줄 요약

Sana는 32배의 깊은 압축 오토인코더, DiT 아키텍처에서의 선형 어텐션, 지문 기반 지도 학습을 통한 디코더 전용 LLM(Gemma)을 활용한 텍스트 인코딩, 그리고 추론 단계를 줄이는 플로우 기반 DPM 솔버를 통해 기존 최고 수준의 성능을 낼 뿐 아니라 초고속으로 4096×4096 해상도의 이미지를 생성하는 매우 효율적인 텍스트-이미지 디퓨전 프레임워크를 소개한다. Sana-0.6B는 FLUX와 같은 SOTA 모델보다 100배 빠른 추론 속도를 기록하며, 크기가 20배 작고 16GB 랩탑 GPU에서도 배포 가능한 특징을 지닌다.

ABSTRACT

We introduce Sana, a text-to-image framework that can efficiently generate images up to 4096$ imes$4096 resolution. Sana can synthesize high-resolution, high-quality images with strong text-image alignment at a remarkably fast speed, deployable on laptop GPU. Core designs include: (1) Deep compression autoencoder: unlike traditional AEs, which compress images only 8$ imes$, we trained an AE that can compress images 32$ imes$, effectively reducing the number of latent tokens. (2) Linear DiT: we replace all vanilla attention in DiT with linear attention, which is more efficient at high resolutions without sacrificing quality. (3) Decoder-only text encoder: we replaced T5 with modern decoder-only small LLM as the text encoder and designed complex human instruction with in-context learning to enhance the image-text alignment. (4) Efficient training and sampling: we propose Flow-DPM-Solver to reduce sampling steps, with efficient caption labeling and selection to accelerate convergence. As a result, Sana-0.6B is very competitive with modern giant diffusion model (e.g. Flux-12B), being 20 times smaller and 100+ times faster in measured throughput. Moreover, Sana-0.6B can be deployed on a 16GB laptop GPU, taking less than 1 second to generate a 1024$ imes$1024 resolution image. Sana enables content creation at low cost. Code and model will be publicly released.

연구 동기 및 목표

• 기존 고해상도 텍스트-이미지 디퓨전 모델의 높은 계산 비용과 느린 추론 속도 문제, 특히 4K 해상도에서의 성능을 해결한다.
• 랩탑 등 소비자용 하드웨어에서 고해상도 4K 이미지를 실시간으로 효율적으로 추론할 수 있도록 한다.
• 이미지 품질이나 텍스트-이미지 정렬 성능을 저하시키지 않은 채 모델 크기와 추론 지연 시간을 줄인다.
• 기존의 텍스트 인코더(예: CLIP, T5)의 한계를 극복하기 위해 디코더 전용 LLM의 지시 따르기 및 추론 능력을 활용한다.
• 초고해상도 이미지 합성에 최소한의 하드웨어 요구 조건으로도 적용 가능한 확장 가능한 종단 간 훈련 및 추론 파이프라인을 개발한다.

제안 방법

• 이미지의 공간적 차원을 32배로 압축하는 깊은 압축 오토인코더(AE-F32)를 제안하여 기존 8배 압축 오토인코더 대비 잠재 토큰 수를 16배 줄였다.
• DiT 아키텍처 내 모든 표준 자기어텐션 모듈을 선형 어텐션 메커니즘으로 대체하여 고해상도 생성에서 계산 복잡도를 O(N²)에서 O(N)으로 감소시켰다.
• 지역적 토큰 집계를 가능하게 하여 위치 인코딩이 필요 없게 하는 Mix-FFN을 MLP 블록에 통합한 3×3 깊이 분리형 컨볼루션을 도입했다.
• T5를 작은 디코더 전용 LLM(Gemma)으로 대체하고, 복잡한 인간 지시(CHI) 프롬프팅과 인라인 학습을 통해 지시 따르기 능력과 이미지-텍스트 정렬 성능을 향상시켰다.
• 클립스코어 기반 동적 캡션 선택을 통한 다중 VLM 캡션 파이프라인을 설계하여 훈련 데이터의 품질과 수렴 성능을 향상시켰다.
• 기존 28–50단계에서 14–20단계로 추론 단계를 줄이는 플로우 기반 샘플링 방법인 Flow-DPM-Solver를 제안하여 추론 속도를 향상시키고 품질을 유지하거나 향상시켰다.

실험 결과

연구 질문

• RQ1 작은 아키텍처를 사용하여 고성능과 저지연 시간을 동시에 달성할 수 있는 텍스트-이미지 디퓨전 모델이 4K(4096×4096) 해상도에서 이미지를 생성할 수 있는가?
• RQ2 표준 자기어텐션을 선형 어텐션으로 대체해도 품질 저하 없이 고해상도 효율적 추론을 가능하게 할 수 있는가?
• RQ3 적절한 지시 프롬프팅을 통해 작은 디코더 전용 LLM(Gemma 등)이 지도 학습된 T5나 CLIP 인코더를 능가할 수 있는가?
• RQ4 고도로 개선된 오토인코딩, 효율적인 어텐션, 최적화된 샘플링 기법의 조합이 SOTA 모델 대비 추론 시간을 100배 이상 단축시킬 수 있는가?
• RQ5 16GB GPU 메모리가 탑재된 랩탑 GPU에서 1초 이내로 1024×1024 해상도 이미지를 추론할 수 있는 고해상도 디퓨전 모델을 구현하는 것은 가능한가?

주요 결과

• Sana-0.6B는 16GB 랩탑 GPU에서 1024×1024 이미지를 1초 이내로 생성하여 실시간 추론 능력을 입증한다.
• Sana-0.6B는 A100 GPU에서 4096×4096 해상도 이미지 생성 시 FLUX 대비 106배 빠른 추론 속도를 기록하며, 지연 시간을 469초에서 9.6초로 감소시켰다.
• Sana-0.6B는 1024×1024 해상도 생성에서 FID 5.81, CLIP 스코어 28.36을 기록하여 SDXL 및 PixArt-Σ와 같은 모델들을 모두 초월하는 성능을 보였다.
• Gemma-2B를 텍스트 인코더로 사용할 경우 FID 5.9, CLIP 스코어 26.8을 기록하여 유사하거나 더 낮은 지연 시간에서 T5-Large 및 T5-XXL 수준의 성능을 달성하거나 초월했다.
• Flow-DPM-Solver는 추론 단계를 28–50단계에서 14–20단계로 줄여 추론 속도를 향상시키면서도 생성 품질을 유지하거나 향상시켰다.
• Mix-FFN를 활용한 NoPE(위치 인코딩 없음) 설계는 위치 인코딩 없이도 효과적인 모델링을 가능하게 하여 아키텍처를 단순화하고 훈련 안정성을 향상시켰다.