PRX 팀이 Flow Matching 기반 텍스트-이미지 모델의 학습 파이프라인을 체계적으로 검증해 REPA 정렬, 토큰 라우팅, BF16 정밀도 관리 등으로 수렴 속도 및 이미지 품질 개선

Context

텍스트-이미지 생성 모델의 학습은 아키텍처 선택 이후 학습 효율성, 수렴 안정성, 표현 품질을 동시에 확보해야 한다. 업계의 "학습 트릭" 목록이 빠르게 늘어나고 있으나 어떤 기법이 실제로 효과를 내는지, 조합했을 때 어떤 상호작용을 보이는지 체계적으로 검증된 결과가 부족했다.

Technical Solution

• 기준선 설정: Flow Matching 기반 순수 학습 환경 구축 (보조 목적함수 제외, 100k 스텝, 256×256 해상도, 배치 크기 256)
• 표현 정렬(REPA): 초기 학습 단계에 표현 대조 목적함수 적용 후 비활성화하는 번인 전략으로 조기 수렴 개선
• 토큰 라우팅(TREAD/SPRINT): 고해상도(1024²) 학습에서 불필요한 토큰을 선택적으로 처리해 계산량 감소
• x-예측 목적함수 도입: 직접 예측 방식으로 변경해 1024² 해상도에서 안정적인 학습 가능
• 정밀도 관리: BF16 자동 연산 사용, 가중치와 옵티마이저 상태는 FP32 유지로 수치 안정성 확보
• 데이터 최적화: 긴 캡션 사용, 합성 이미지 vs 실제 이미지의 텍스처-구조 트레이드오프 검증, 지도 미세조정(SFT) 적용
• 평가 메트릭 다원화: FID, CLIP-MMD, DINOv2-MMD로 다각도 성능 검증, 처리량(samples/s)으로 효율성 측정

Impact

FID 기준선 18.20에서 REPA-E 적용 시 개선, FLUX2-AE 토크나이저 도입으로 이미지 품질 점프, BF16 가중치 저장 버그로 인한 FID 악화(18.20 → 21.87), CMMD(0.41 → 0.61), DINO-MMD(0.39 → 0.57) 수치적 손상 확인.

Key Takeaway

대규모 생성 모델 학습에서 개별 최적화 기법만큼 중요한 것은 깨끗한 기준선 환경 구축 후 체계적 검증과 정밀도 같은 "조용한" 세부사항 관리이다. 특히 혼합 정밀도 연산 시 가중치 저장 형식 오류는 여러 날의 학습 시간을 낭비하게 하므로 초기에 엄격하게 확인해야 한다.