PRX Part 3 - Text-to-Image 모델을 24시간 만에 훈련하기

Photoroom이 H200 GPU 32장으로 text-to-image 모델을 24시간 안에 훈련시킨 스피드런 결과를 공개함. 총 컴퓨팅 비용은 약 $1,500(GPU당 시간당 $2). Part 1, Part 2에서 검증한 트릭들을 전부 합쳐서 실전 레시피를 만든 것이 핵심임.

핵심 아키텍처: 픽셀 스페이스 직접 훈련

• X-prediction 방식을 사용해서 VAE 없이 픽셀 공간에서 직접 훈련함. 디코더가 필요 없으니 파이프라인이 훨씬 깔끔해짐
• 패치 사이즈 32, 초기 토큰 프로젝션에 256차원 보틀넥을 둬서 시퀀스 길이를 관리함. 512px에서 시퀀스 길이 256, 1024px에서 1024
• 보통 256px -> 512px -> 1024px 순서로 올리는데, 여기서는 512px부터 바로 시작해서 1024px로 파인튜닝하는 전략을 씀

Perceptual Loss와 토큰 라우팅

• LPIPS(가중치 0.1) + DINOv2 기반 perceptual loss(가중치 0.01)를 flow matching 목적함수 위에 추가함. 풀 이미지를 풀링해서 모든 노이즈 레벨에 적용
• TREAD 토큰 라우팅으로 50%의 토큰이 2번째~끝에서 두 번째 블록을 건너뛰게 해서 연산량을 절감함
• 라우팅된 모델은 vanilla CFG에서 품질이 떨어질 수 있어서, dense vs. routed conditional prediction을 사용하는 self-guidance 방식을 적용함

학습 세부 설정

• REPA로 representation alignment 적용. DINOv3를 teacher로, 8번째 transformer 블록에서 alignment loss 가중치 0.5로 설정
• 옵티마이저는 Muon(FSDP 구현)을 2D 파라미터에, 나머지(bias, norm, embedding 등)는 Adam 사용
• 훈련 데이터는 합성 데이터셋 3개 합쳐서 약 870만 장: Flux generated(170만), FLUX-Reason-6M(600만), midjourney-v6-llava(100만, Gemini 1.5로 재캡셔닝)
• 스케줄: 512px에서 100k 스텝(배치 1024) -> 1024px에서 20k 스텝(배치 512, REPA 없이). EMA smoothing 0.999, 업데이트 간격 10

결과

• 24시간 훈련치고 상당히 쓸 만한 수준이 나옴. 프롬프트 팔로잉이 강하고 전반적인 미학적 일관성도 있음
• 텍스처 글리치, 간헐적인 해부학 오류 등 아직 결함이 있지만, 구조적 문제가 아니라 언더트레이닝 아티팩트로 판단됨. 데이터와 컴퓨팅을 더 넣으면 예측 가능하게 개선될 것으로 봄
• 코드 전체 오픈소스로 공개: github.com/Photoroom/PRX

불과 몇 년 전만 해도 경쟁력 있는 diffusion 모델 훈련에 수백만 달러가 들었는데, 이제 $1,500이면 쓸 만한 모델이 나오는 시대가 됨. 엔지니어링 최적화가 얼마나 중요한지 보여주는 사례임.