텍스트-이미지 모델 훈련 설계: Ablation 연구를 통한 교훈

본 글은 텍스트-이미지(Text-to-Image) 파운데이션 모델을 처음부터 구축하고 효율적으로 훈련시키는 과정에 대한 심층적인 연구 보고서입니다. 이전 글에서 아키텍처 설계에 집중했다면, 이번 포스팅에서는 실제로 모델의 성능과 학습 안정성에 결정적 영향을 미치는 '훈련 방법론(Training Design)' 자체를 다룹니다.

연구진은 방대한 최신 훈련 트릭들 중 핵심적인 요소들을 선별하여 실험실 기록부(experimental logbook) 형태로 정리했습니다. 이는 단순히 기술을 나열하는 것이 아니라, 일관된 환경에서 여러 아이디어를 재현하고 그 효과를 분석하는 데 목적이 있습니다.

1. 엄격한 기준점 (Baseline) 설정:
모든 실험의 신뢰성을 확보하기 위해 매우 단순하고 투명한 Baseline을 구축했습니다. 이 기준 모델은 보조 목표(auxiliary objectives)나 아키텍처적 편법에 의존하지 않고, 순수한 Flow Matching (Lipman et al., 2022) 훈련 설정을 따릅니다.

• Baseline 환경: PRX-1.2B 모델을 사용하며, Flux VAE 잠재 공간에서 10만 스텝 동안 학습합니다.
• 주요 하이퍼파라미터: Optimizer는 AdamW를 사용하고, Learning Rate (lr)은 $1e-4$로 설정되었습니다.

이 기준점은 모든 개선 사항이나 성능 저하가 특정 개입(intervention)에 기인하는지 명확하게 추적할 수 있게 합니다.

2. 다각적인 평가 지표:
모델의 품질을 측정하기 위해 단일 지표에 의존하지 않고 여러 보조 지표를 사용합니다. 이들은 이미지 생성 분포와 실제 이미지 분포 간의 거리를 측정하는 데 중점을 둡니다:

• FID (Fréchet Inception Distance): 생성된 이미지와 실제 이미지 분포의 유사성을 측정하며, 값이 낮을수록 충실도가 높음을 의미합니다.
• CMMD (CLIP Maximum Mean Discrepancy): CLIP 임베딩을 사용하여 분포 간 거리를 측정하며, FID보다 지각적 품질(perceptual quality)을 더 잘 추적하는 경향이 있습니다.
• DINO-MMD: DINOv2 임베딩을 사용한 MMD로, 자기 지도 학습 비전 백본(self-supervised vision backbone) 관점에서의 분포 변화를 제공합니다.
• 네트워크 처리량 (Network throughput): 초당 평균 샘플 수(samples/s)로, 훈련의 엔드투엔드 효율성을 측정하는 핵심 지표입니다.

3. 주요 연구 방법론:
연구진은 탐색한 훈련 기법들을 네 가지 범주로 분류하여 분석합니다:

• Representation Alignment (표현 정렬): 기존 노이즈 제거 목적 외에, 강력하고 고정된(frozen) 비전 인코더를 사용하여 중간 특징을 직접 감독하는 보조 손실 함수를 추가합니다. 이 방법은 초기 학습 속도를 높이고 모델의 특징을 현대적인 자기 지도 인코더와 유사하게 만들어, 동일한 품질 달성에 필요한 컴퓨팅 자원을 줄이는 효과가 있습니다.
• Training Objectives (훈련 목표): 다양한 종류의 목적 함수(loss function)를 추가하여 모델이 더 풍부하고 안정적인 표현을 학습하도록 유도합니다.
• Token Routing and Sparsification: 계산 효율성을 높이기 위해 토큰 처리 방식을 최적화하는 기법들입니다.
• Data (데이터): 데이터셋 구성 및 전처리 단계에서 성능 향상을 도모하는 방법들입니다.

이러한 체계적인 실험 설계를 통해, 연구진은 어떤 훈련 설계가 모델의 수렴 속도와 최종 품질에 가장 큰 기여를 하는지 과학적으로 입증하고 있습니다. 다음 포스팅에서는 이러한 모든 실험 결과를 포함한 완전한 '훈련 레시피' 코드를 공개할 예정입니다.