Training Design for Text-to-Image Models: Lessons from Ablations

🖼️ 인포그래픽

Training Design for Text-to-Image Models: Lessons from Ablations 내용을 설명하는 본문 이미지

🖼️ 4컷 인포그래픽

💡 한 줄 요약

Photoroom 팀은 PRX-1.2B 텍스트-이미지 모델을 기준점으로 삼아, 표현 정렬과 보조 목적 함수가 훈련 수렴·품질·처리량에 어떤 영향을 주는지 절제된 ablation 방식으로 검증했다.

📌 핵심 요약

이 글은 텍스트-이미지 모델을 처음부터 효율적으로 학습시키는 시리즈의 두 번째 글로, 이전 글이 PRX 모델의 아키텍처 선택을 다뤘다면 이번 글은 실제 훈련 설계에서 효과가 있었던 기법들을 실험 기록 형식으로 정리한다.
기준 실험은 PRX-1.2B 모델을 Flux VAE latent space에서 순수 Flow Matching 방식으로 100K 스텝 학습시키는 설정이며, 보조 목적이나 아키텍처 기반 속도 절약 장치를 넣지 않아 이후 비교의 기준점으로 사용된다.
평가는 FID, CMMD, DINO-MMD, 네트워크 처리량을 중심으로 진행되며, 저자들은 이 지표들이 지각 품질의 완벽한 대리 지표는 아니지만 반복 실험에서 실용적인 점수판 역할을 한다고 설명한다.
표현 정렬, 특히 REPA는 동결된 비전 인코더의 패치 표현을 학생 모델의 중간 표현과 맞추는 보조 손실을 추가해 초기 표현 학습을 돕고, 실험에서는 DINOv2와 DINOv3 교사를 모두 사용했을 때 기준선보다 품질 지표가 개선됐다.
REPA는 품질 향상과 동시에 추가 teacher forward 및 패치 수준 유사도 손실 때문에 처리량 감소를 동반했으며, DINOv3는 가장 좋은 지표를 보였지만 더 느렸고 DINOv2는 상대적으로 효율적인 절충안으로 제시됐다.

🧩 주요 포인트

이 글은 텍스트-이미지 모델을 처음부터 효율적으로 학습시키는 시리즈의 두 번째 글로, 이전 글이 PRX 모델의 아키텍처 선택을 다뤘다면 이번 글은 실제 훈련 설계에서 효과가 있었던 기법들을 실험 기록 형식으로 정리한다.
기준 실험은 PRX-1.2B 모델을 Flux VAE latent space에서 순수 Flow Matching 방식으로 100K 스텝 학습시키는 설정이며, 보조 목적이나 아키텍처 기반 속도 절약 장치를 넣지 않아 이후 비교의 기준점으로 사용된다.
평가는 FID, CMMD, DINO-MMD, 네트워크 처리량을 중심으로 진행되며, 저자들은 이 지표들이 지각 품질의 완벽한 대리 지표는 아니지만 반복 실험에서 실용적인 점수판 역할을 한다고 설명한다.
표현 정렬, 특히 REPA는 동결된 비전 인코더의 패치 표현을 학생 모델의 중간 표현과 맞추는 보조 손실을 추가해 초기 표현 학습을 돕고, 실험에서는 DINOv2와 DINOv3 교사를 모두 사용했을 때 기준선보다 품질 지표가 개선됐다.
REPA는 품질 향상과 동시에 추가 teacher forward 및 패치 수준 유사도 손실 때문에 처리량 감소를 동반했으며, DINOv3는 가장 좋은 지표를 보였지만 더 느렸고 DINOv2는 상대적으로 효율적인 절충안으로 제시됐다.

🧠 상세 정리

1. 시리즈의 초점 전환: 아키텍처에서 훈련 설계로

이 글은 Photoroom 팀이 공개적으로 대규모 텍스트-이미지 기반 모델을 처음부터 학습시키는 과정을 다루는 시리즈의 두 번째 글이다. 첫 번째 글에서는 PRX라는 모델의 핵심 아키텍처 선택과 설계 동기를 설명했고, 1.2B 파라미터 규모의 초기 소형 버전도 공개했다. 이번 글에서는 모델 구조 자체보다 훈련 과정에서 실제로 수렴을 빠르게 하고, 학습을 안정화하며, 더 나은 표현을 만들게 한 요인들을 기록한다. 저자들은 최신 훈련 기법을 전부 조사하려는 것이 아니라, 최근 아이디어들을 일관된 설정에서 재현하거나 변형해 실험하고 최적화와 수렴에 미친 영향을 보고하는 실험 로그북 형식을 택했다.

2. 실험의 목적과 다음 단계

저자들의 목표는 개별 훈련 기법을 나열하는 데 그치지 않고, 어떤 방법이 실제 기준선 대비 수렴이나 훈련 효율을 개선하는지 확인하는 것이다. 또한 각 기법을 따로 평가한 뒤, 여러 기법을 조합했을 때에도 여전히 유용한지 살피겠다고 밝힌다. 글 말미의 후속 계획으로는 이번 글에서 다룬 실험을 포함한 전체 훈련 레시피를 코드로 공개하고, 가장 좋은 구성 요소들을 하나의 설정으로 묶어 공개적인 속도 실험을 진행하겠다고 예고한다. 이 과정은 제한된 조건에서 신중한 훈련 설계가 어디까지 효과를 낼 수 있는지 보여주는 스트레스 테스트로도 제시된다.

3. 기준선: 단순하고 투명한 Flow Matching 설정

모든 훈련 효율 기법을 논하기 전에 저자들은 깨끗한 기준 실험을 먼저 세운다. 기준선은 순수 Flow Matching 학습 설정이며, 추가 목적 함수나 컴퓨팅 자원을 아끼기 위한 아키텍처상의 단축 기법을 사용하지 않는다. 모델은 이전 글에서 소개한 PRX-1.2B로, 이미지 토큰과 텍스트 토큰에 대해 global attention을 사용하는 single stream 구조이며 Flux VAE latent space에서 학습된다. 이 기준선의 역할은 이후 관찰되는 개선이나 악화를 특정 개입의 효과로 돌릴 수 있게 하는 안정적인 비교 지점이다.

4. 기준 학습 구성과 고정된 비교 조건

기준 실험은 100K 스텝, MidJourneyV6로 생성된 공개 1M synthetic image 데이터셋, 256×256 해상도, global batch size 256으로 설정된다. 옵티마이저는 AdamW이며 학습률은 1e-4, weight decay는 0.0, eps는 1e-15, betas는 (0.9, 0.95)로 제시된다. 텍스트 인코더는 GemmaT5, positional encoding은 Rotary, attention mask는 padding mask를 사용하고 EMA는 비활성화된다. 저자들은 특별히 언급하지 않는 한 이 구성을 고정해, 하이퍼파라미터 변경이나 숨은 설정 변화가 결과 해석을 흐리지 않도록 했다.

5. 평가 지표: 품질과 효율을 함께 보는 점수판

저자들은 체크포인트를 추적하기 위해 FID, CMMD, DINO-MMD, 네트워크 처리량을 사용한다. FID는 Inception-v3 특징의 평균과 공분산을 이용해 생성 이미지와 실제 이미지 분포가 얼마나 가까운지 측정하며, 낮을수록 일반적으로 샘플 충실도가 높다고 본다. CMMD는 CLIP 이미지 임베딩과 Maximum Mean Discrepancy를 이용하는 거리 측정으로, Gaussian feature distribution을 가정하지 않고 표본 효율이 더 좋을 수 있다고 설명된다. DINO-MMD는 같은 MMD 기반 거리 측정을 DINOv2 이미지 임베딩 위에서 계산해 자기지도 비전 백본 관점의 분포 이동을 보완적으로 보여주며, 네트워크 처리량은 초당 처리 샘플 수로 전체 훈련 효율을 나타낸다.

6. 표현 정렬의 문제의식: denoising 전에 표현을 잘 만들기

확산 모델과 flow 모델은 보통 손상된 입력에서 noise-like target 또는 vector field를 예측하는 단일 목적 함수로 학습된다. 저자들은 학습 초기에 이 하나의 목적이 유용한 내부 표현을 만드는 일과 그 위에서 denoising을 배우는 일을 동시에 떠안는다고 설명한다. 표현 정렬은 이 문제를 명시화해, 기본 denoising 목적을 유지하면서 강력한 동결 비전 인코더가 만든 중간 특징을 직접 감독하는 보조 손실을 추가한다. 핵심 가정은 표현 학습이 병목이며, compute가 제한된 상황에서 diffusion이나 flow transformer가 판별적 특징을 배우더라도 foundation vision encoder보다 뒤처질 수 있다는 점이다.

7. REPA의 방식: teacher 패치 표현과 student 중간 표현 맞추기

REPA는 기본 flow-matching objective 위에 representation matching term을 더한다. 깨끗한 샘플 x0는 사전학습된 비전 인코더를 통과해 패치 임베딩을 만들고, 모델은 보간된 상태 xt를 처리해 각 패치에 대응하는 중간 hidden token을 만든다. 작은 projection head는 student hidden token을 teacher embedding space로 보내며, 보조 손실은 대응되는 teacher 토큰과 student 토큰 사이의 패치별 유사도, 보통 cosine similarity를 최대화하도록 구성된다. 최종 손실은 flow matching loss와 REPA loss의 합이며, λ가 두 손실의 trade-off를 조절한다.

8. REPA 실험 결과와 iREPA의 부분적 관찰

저자들은 기준 PRX 학습 위에 REPA를 적용하고 동결 교사로 DINOv2와 DINOv3를 사용했다. 기준선은 FID 18.2, CMMD 0.41, DINO-MMD 0.39, 3.95 batches/sec였고, REPA-DINOv3는 FID 14.64, CMMD 0.35, DINO-MMD 0.30, 3.46 batches/sec를 기록했다. REPA-DINOv2도 FID 16.6, CMMD 0.39, DINO-MMD 0.31, 3.66 batches/sec로 기준선보다 개선됐지만 처리량 저하가 있었다. 이어서 iREPA는 전역 의미보다 spatial structure를 정렬해야 한다는 문제의식에서 3×3 convolutional projection과 teacher patch token의 spatial normalization을 제안하며, 원문 발췌 범위에서는 DINOv2 위에서 수렴이 더 매끄럽고 지표가 더 꾸준히 개선됐지만 같은 변화가 다른 경우에 얼마나 잘 전이됐는지는 뒷부분이 제시되지 않는다.

🧾 핵심 주장 / 시사점

저자들이 강조하는 핵심은 훈련 효율이 단순히 처리량만의 문제가 아니라, 같은 품질에 도달하는 데 필요한 시간과 안정성까지 포함한다는 점이다.
REPA 결과는 더 강한 teacher가 더 좋은 품질 지표를 낼 수 있지만, 추가 계산 비용 때문에 실제 선택은 품질과 처리량의 절충으로 결정된다는 점을 보여준다.
iREPA의 문제의식은 텍스트-이미지 생성에서 단순한 전역 의미 표현보다 패치 단위의 공간 구조 보존이 학습 품질에 중요할 수 있음을 시사한다.

✅ 액션 아이템

PRX-1.2B 기준 실험을 복제해 Flux VAE 잠재공간의 순수 Flow Matching 100K 스텝 설정에서 REPA 미적용 기준선과 지표 분포를 재확인한다.
DINOv2·DINOv3 교사를 포함한 REPA 변형을 동일 조건에서 실행해 FID, CMMD, DINO-MMD, 처리량을 병행 수집한다.
보조 손실 적용 시 처리량 저하를 반영해 품질 지표 개선폭과 속도 손실을 동시에 비교해 DINOv2·DINOv3의 절충점을 정량화한다.