PRX 파트 3: 24시간 안에 텍스트-이미지 모델 훈련하기
픽셀 공간 훈련 및 X-예측
모델은 x-예측 방식을 사용하여 픽셀 공간에서 직접 훈련할 수 있게 하며, 변분 오토인코더(VAE)의 필요성을 없앱니다. 연산적으로 시퀀스 길이를 관리하기 위해 팀은 패치 크기를 32로, 초기 토큰 투영 레이어에서 256차원 병목을 사용했습니다.
512px 해상도에서는 시퀀스 길이가 256 토큰이며, 1024px에서는 1024 토큰으로 증가합니다. 훈련 일정은 512px에서 시작해 1024px에서 미세 조정으로 마무리되었습니다.
시각적 품질을 위한 지각 손실
픽셀을 직접 예측함으로써 고전적인 컴퓨터 비전 지각 손실을 활용하여 수렴 속도와 시각적 품질을 향상시킬 수 있습니다. 팀은 표준 흐름 매칭 목표 위에 두 가지 보조 손실을 구현했습니다:
- LPIPS: 가중치 0.1로 사용되어 저수준 지각 유사성을 포착합니다.
- DINO 기반 지각 손실: DINOv2를 사용하고 가중치 0.01로 더 강한 의미 신호를 제공합니다.
이 손실들은 패치별 특징이 아니라 모든 노이즈 레벨에 걸친 전체 이미지 풀링에 적용되어 보다 나은 실험 결과를 얻었습니다.
TREAD를 활용한 효율적인 토큰 라우팅
단계당 계산 비용을 줄이기 위해 팀은 TREAD(Token Routing for Efficient Architecture-agnostic Diffusion Training)를 사용했습니다. 이 방법은 토큰의 일부를 무작위로 선택해 연속적인 트랜스포머 블록을 건너뛰고 다시 삽입합니다.
구체적으로, 토큰의 50%가 트랜스포머의 두 번째 블록에서 마지막 전 블록으로 라우팅되었습니다. 기본 Classifier-Free Guidance(CFG) 하에서 품질 저하를 완화하기 위해, 밀집 조건 예측과 라우팅된 조건 예측을 비교하는 자체 가이드 스킴이 구현되었습니다.
표현 정렬 및 옵티마이저
표현 정렬은 REPA를 사용하고 DINOv3를 교사 모델로 하여 달성되었습니다. 정렬 손실은 8번째 트랜스포머 블록에서 한 번 적용되었으며 손실 가중치는 0.5였습니다. TREAD 라우팅과 일관성을 유지하기 위해 정렬 손실은 라우팅되지 않은 토큰에만 계산되었습니다.
옵티마이징을 위해 팀은 2D 파라미터(행렬)에는 Muon 옵티마이저를, 2D가 아닌 파라미터(바이어스, 정규화, 임베딩)에는 Adam을 사용했습니다.
| 옵티마이저 그룹 | 대상 파라미터 | 주요 하이퍼파라미터 |
|---|---|---|
| Muon | 2D 파라미터 | lr=1e-4, momentum=0.95, nesterov=true, ns_steps=5 |
| Adam | Non-2D 파라미터 | lr=1e-4, betas=(0.9, 0.95), eps=1e-8 |
훈련 데이터 및 일정
훈련은 세 개의 합성 데이터셋에서 수행되었습니다:
- Flux generated (1.7M 이미지)
- FLUX-Reason-6M (6M 이미지)
- midjourney-v6-llava (1M 이미지), 일관성을 위해 Gemini 1.5로 재캡션되었습니다.
훈련 일정은 512px에서 배치 크기 1024로 100k 스텝을 수행한 뒤, 1024px에서 배치 크기 512로 20k 스텝을 진행했습니다 (REPA 제외).
결과 및 시사점
최종 모델은 일부 텍스처 결함과 가끔 발생하는 해부학적 오류가 있지만, 팀은 이를 "명확히 사용 가능"하며 강력한 프롬프트 따름과 일관된 미학을 가진다고 설명합니다. 저자들은 남은 문제들이 레시피의 구조적 결함보다 훈련 부족 및 데이터 다양성 제한 때문일 가능성이 높다고 결론짓습니다.
Photoroom은 PRX GitHub repository를 통해 훈련 코드와 실험 프레임워크를 오픈소스화했습니다.