[AI 논문 리뷰] WURSTCHEN: AN EFFICIENT ARCHITECTURE FOR LARGE-SCALE TEXT-TO-IMAGE DIFFUSION MODELS

1. 논문이 나온 배경

최첨단 확산 모델들은 이미지 합성 분야에서 놀라운 성과를 달성하며 사실적인 이미지 생성에 근접했습니다. 하지만 이러한 기반 모델들은 매우 높은 계산 비용을 요구한다는 심각한 단점을 가지고 있었습니다. 예를 들어, Stable Diffusion 1.4는 학습에 150,000 GPU 시간을, SD 2.1은 200,000 GPU 시간을 사용했습니다. 더 경제적인 text-to-image 모델들도 존재했지만, 이들은 낮은 해상도와 전반적인 미적 특성 측면에서 이미지 품질이 떨어진다는 문제가 있었습니다.

 

이러한 불일치의 핵심 딜레마는 해상도를 높이면 시각적 복잡성과 계산 비용도 함께 증가하여 이미지 합성이 더 비싸지고 데이터 집약적으로 변한다는 점이었습니다. Encoder 기반 Latent Diffusion Models(LDMs)가 압축된 잠재 공간에서 작동함으로써 이 문제를 부분적으로 해결했지만, 결국 인코더-디코더 모델이 품질 저하 없이 이미지를 압축할 수 있는 정도에 의해 제한되었습니다.

2. 어떤 것을 개선할 것인지

연구진은 계산 요구사항을 획기적으로 줄이면서도 경쟁력 있는 성능을 유지하는 것을 목표로 했습니다. 구체적으로는 다음을 개선하고자 했습니다:

• 학습 효율성: SD 2.1 대비 약 8배 감소된 GPU 시간(24,602 GPU 시간)으로 유사한 품질의 모델을 학습하고자 했습니다.
• 추론 속도: 기존 최첨단 모델들보다 2배 이상 빠른 추론 속도를 달성하여 비용과 탄소 발자국을 크게 줄이고자 했습니다.
• 데이터 효율성: 더 적은 학습 데이터로도 최첨단 결과를 달성하고자 했습니다.
• 이미지 품질: 계산 비용을 줄이면서도 사용자 선호도 연구를 통해 텍스트 조건부 이미지 생성의 품질을 개선하고자 했습니다.

3. 어떤 원리로 개선된 것인지

Würstchen의 추론 구조

 

Würstchen은 세 단계(Stage) 아키텍처를 통해 이러한 개선을 달성했습니다:

Stage A: VQGAN 기반 잠재 공간 생성

 

먼저 4:1 압축 비율로 작동하는 VQGAN을 학습했습니다. 이는 1024×1024 이미지를 256×256 discrete 토큰으로 인코딩합니다. 크기 8,192의 코드북을 사용하여 이미지의 압축된 표현을 만들었습니다.

수식적으로: $ X_q = f_Θ(X) $이며, 디코더는 $ f⁻¹_Θ(X_q) ≈ X $를 통해 원본 이미지를 재구성합니다.

Stage B: 조건부 잠재 확산 모델

 

Stage A의 비양자화된 잠재 공간에서 조건부 LDM을 학습했습니다. 핵심은 "Semantic Compressor"라는 인코더를 도입한 것입니다. 이 컴포넌트는 786×786 이미지를 16×24×24 크기의 매우 압축된 의미론적 표현으로 인코딩합니다(42:1 압축 비율).

Semantic Compressor는 ImageNet으로 사전 학습된 EfficientNetV2-Small을 초기화에 사용했으며, 학습 중에 가중치를 업데이트하여 고정밀 의미 정보를 담은 잠재 공간을 구축했습니다. 이 압축된 표현 Csc는 Stage B 디코더의 조건으로 제공되어 디코딩 과정을 안내합니다.

수식적으로: $ X̄_0 = f_ϑ(X̃_t, C_{sc}, C_{text}, t) $

여기서 $ X̃_t $는 노이즈가 추가된 표현, $ C_{sc} $는 Semantic Compressor 임베딩, $ C_{text} $는 텍스트 조건, $ t $는 타임스텝입니다.

Stage C: 텍스트 조건부 잠재 이미지 생성

 

Stage B에서 사용된 Semantic Compressor의 강력하게 압축된 잠재 표현에서 텍스트 조건부 LDM을 학습했습니다. 16개의 ConvNeXt 블록으로 구성된 아키텍처를 사용했으며, 다운샘플링 없이 작동합니다. U-Net 대신 이러한 구조를 선택한 이유는 이미지가 이미 42배로 압축되어 있어 추가 압축이 모델 품질에 해롭다는 것을 발견했기 때문입니다.

확산 과정은 다음과 같습니다: $ X_{sc,t} = √ᾱ_t · X_{sc} + √(1-ᾱ_t) · ϵ $

노이즈 예측을 위해 안정적인 학습을 위한 재공식화를 사용했습니다: $ ϵ̄ = \frac{(X_{sc,t} - A)}{(|1 - B| + 1e⁻⁵)} $ 여기서 $ A, B = f_θ(X_{sc,t}, C_{text}, t) $

$ p_2 $ 손실 가중치를 사용하여 높은 노이즈 수준이 손실에 더 많이 기여하도록 했습니다.

핵심 혁신 원리

이 접근법의 핵심은 의미론적 정보를 극도로 압축하여 확산 과정을 안내하는 것입니다. 언어의 잠재 표현에 비해 훨씬 더 상세한 가이던스를 제공하면서도 계산 요구사항을 크게 줄였습니다. 42:1이라는 높은 압축 비율의 잠재 공간에서 확산 모델을 학습함으로써, 매우 낮은 차원에서 작동하면서도 높은 품질의 이미지를 생성할 수 있었습니다.

4. 실험 결과 및 성능

4-1. 데이터셋 구조

학습 데이터: LAION-5B의 improved-aesthetic 하위 집합을 사용했습니다. 모든 단계는 이 데이터셋의 부분 집합에서 학습되었습니다.

 

평가 데이터셋:

• MS COCO-30K: 표준 zero-shot text-to-image 모델 평가 데이터셋입니다. COCO validation 셋에서 무작위로 선택한 30,000개의 프롬프트로 이미지를 생성했습니다. 다만 MS COCO의 프롬프트는 상당히 짧고 세부사항이 부족한 경향이 있습니다.
• Localized Narratives-COCO-5K: MS COCO의 짧은 프롬프트 문제를 보완하기 위해 Localized Narrative MS COCO 부분 집합에서 5,000개의 이미지를 생성했습니다. 이는 더 긴 설명적 프롬프트를 제공합니다.
• Parti-prompts: 1,633개의 매우 다양한 캡션으로 구성된 데이터셋입니다. 모델의 의도된 사용 시나리오를 밀접하게 반영하며, 다양한 스타일과 복잡한 요구사항을 포함합니다.

4-2. 성능 평가지표

자동화 지표

FID (Fréchet Inception Distance): 생성된 이미지와 실제 이미지 간의 분포 거리를 측정합니다. 모든 이미지를 256×256으로 다운샘플링하여 공정한 비교를 위해 평가했습니다. Würstchen은 COCO30K에서 23.6의 FID를 기록했습니다. SD 1.4(16.2), SD 2.1(15.1)보다는 높았지만, CogView2(24.0)와 Baseline LDM(43.5)을 크게 앞섰습니다.

 

IS (Inception Score): 생성된 이미지의 다양성과 품질을 측정합니다. Würstchen은 40.9로 모든 비교 모델 중 가장 높은 점수를 기록했습니다. SD 1.4(40.6), SD 2.1(40.1), LDM(30.3)을 앞질렀습니다.

 

PickScore: 인간 선호도를 모방하도록 설계된 주요 자동화 지표입니다. 동일한 프롬프트에서 생성된 이미지 세트 중 선택할 때 인간의 선호를 예측합니다. PickScore 결과는 일관되게 Würstchen의 우수성을 보여주었습니다:

• COCO-30K: Baseline LDM 대비 96.5%, DF-GAN 대비 99.8%, GALIP 대비 98.1%, SD 1.4 대비 78.1%, SD 2.1 대비 64.4% 선호
• Localized Narratives-COCO-5K: SD 2.1 대비 70.0% 선호
• Parti-prompts: SD 2.1 대비 74.6% 선호

인간 평가:

총 90명의 참가자가 3,343개(Parti-prompts)와 2,262개(COCO Captions)의 비교를 수행했습니다.

 

전체 선호도 (Overall Preference):

• MS COCO: Würstchen 41.3%, SD 2.1 40.6%, 동등 18.1%
• Parti-prompts: Würstchen 49.5%, SD 2.1 32.8%, 동등 17.7%

개인별 선호도 (Individual Preference): 50번째 백분위수 이상의 참가자(MS COCO 30회 이상, Parti-prompts 51회 이상 비교 수행)만 포함했을 때:

• MS COCO: Würstchen 선호 44.44%, SD 2.1 선호 27.78%
• Parti-prompts: Würstchen 선호 72.22%, SD 2.1 선호 5.56%

Parti-prompts에서 Würstchen이 명확하게 선호되었으며, 이는 모델의 의도된 사용 사례를 반영하므로 중요한 결과입니다.

 

효율성 지표:

• 학습 비용: Stage C는 24,602 GPU 시간만 필요했으며, 이는 SD 2.1(200,000 GPU 시간) 대비 약 8배 감소입니다. Stage B를 포함해도 총 35,602 GPU 시간으로 여전히 SD 모델들보다 훨씬 효율적입니다.
• 학습 샘플 수: Würstchen은 1.42B 샘플을 사용했으며, 이는 SD 1.4(4.8B), SD 2.1(3.9B)보다 적습니다.
• 탄소 배출: 약 2,276 kg CO2 eq.로 추정되며, SD 1.4(11,250 kg), SD 2.1(15,000 kg)에 비해 현저히 낮습니다.
• 추론 속도: A100 GPU에서 1024×1024 이미지 생성 시, Würstchen은 SD 2.1과 SD XL보다 2배 이상 빠른 속도를 보였습니다. torch.compile() 최적화를 적용하면 속도 차이가 더욱 증가했습니다.

성능 요약

Würstchen은 동일하거나 더 작은 크기의 모든 모델을 능가했으며, 훨씬 높은 계산 예산을 사용한 Stable Diffusion 모델들도 능가했습니다. SD XL만이 이미지 품질에서 여전히 우수했지만, SD XL은 더 높은 용량의 모델이며 데이터와 계산 예산이 알려지지 않아 공정한 비교가 아닙니다.

 

흥미롭게도 FID 점수는 다른 최첨단 모델들보다 상대적으로 높았는데, 연구진은 이를 고주파 특징의 차이로 설명했습니다. 시각적 검사 결과 Würstchen이 생성한 이미지가 다른 text-to-image 모델들보다 더 부드러운 경향이 있었으며, 이 차이는 COCO와 같은 실제 이미지에서 FID 메트릭을 계산할 때 가장 두드러졌습니다. 그러나 IS 점수가 가장 높았고, 인간 선호도 평가에서도 우수한 결과를 보여 실제 지각 품질은 높다는 것을 증명했습니다.

5. 향후 연구 방향성 및 해결하지 못한 과제

논문에서는 향후 연구 방향에 대해 다음과 같은 가능성을 제시했습니다:

• 더 큰 모델 파라미터로의 확장: 제안된 접근법이 효율적인 확장성을 제공하므로, 훨씬 더 큰 모델 파라미터로 확장할 가능성을 시사했습니다.
• 지속 가능한 생성 AI 연구: 이 연구가 더 지속 가능하고 계산적으로 효율적인 생성 AI 영역에 대한 추가 연구의 출발점이 되기를 희망했습니다.
• 소비자 하드웨어에서의 학습 및 미세조정: 소비자 하드웨어에서 대규모 모델을 학습하고 미세조정하며 배포하는 더 많은 가능성을 열어주기를 기대했습니다.

해결하지 못한 과제:

• 고주파 디테일: FID 점수가 상대적으로 높은 것으로 나타났듯이, 생성된 이미지가 다른 모델들에 비해 약간 더 부드러운 경향이 있습니다. 실제 이미지의 고주파 특징을 더 잘 포착하는 것이 개선점으로 남아있습니다.
• 최대 해상도: 모델이 한 변에 최대 1538픽셀까지 가변 종횡비의 이미지를 생성하도록 설계되었지만, 더 높은 해상도로의 확장 가능성은 명시적으로 다루지 않았습니다.
• Semantic Compressor 설계: ImageNet으로 사전 학습된 EfficientNetV2를 사용했지만, 대규모 텍스트-이미지 데이터셋의 광범위한 분포를 포착하기에는 초기 가중치가 이상적이지 않았습니다. 더 적합한 사전 학습 전략이 성능을 더욱 향상시킬 수 있습니다.
• 다양한 작업으로의 일반화: 논문은 주로 텍스트-이미지 생성에 초점을 맞췄으며, 이미지 편집, 인페인팅, 아웃페인팅 등 다른 조건부 생성 작업으로의 확장 가능성은 탐구하지 않았습니다.
• 연구진은 전체 소스 코드, 학습 및 추론 스크립트, 학습된 모델을 GitHub에 공개하여 재현성을 보장하고 커뮤니티의 추가 연구를 촉진하고자 했습니다.