Abra: 디퓨전 이미지 학습의 스케일링

초록

컴퓨팅 최적 스케일링 법칙(compute-optimal scaling laws)은 최첨단 언어 모델의 훈련을 안내하지만, 시각 생성 분야에서는 여전히 대부분 탐구되지 않았다. 본 연구는 Abra를 사용한 텍스트-이미지 확산 모델에 대한 체계적인 스케일링 법칙 연구를 제시한다. Abra는 통제된 플로우 매칭 트랜스포머(flow-matching transformer) 계열로, 세 자릿수에 달하는 계산량(10^19~10^22 FLOPs)에 걸쳐 훈련되었으며, 이는 이전 연구들보다 훨씬 큰 컴퓨팅 예산에 해당한다. 우리는 확산 모델이 언어 모델만큼 예측 가능하게 스케일링되지만, 최적 훈련을 위해 훨씬 더 많은 데이터를 필요로 함을 입증한다. 컴퓨팅 최적성은 파라미터당 약 200개의 이미지 토큰에서 달성되며, 이는 LLM에 대한 Chinchilla 컴퓨팅 최적 처방의 10배에 해당한다. 또한 언어 모델과 달리 확산 모델은 과도한 훈련(overtraining)에 강건하며, 실무자들은 더 큰 모델보다 더 많은 데이터를 선택하는 쪽으로 오차를 두어야 함을 보여준다. 마지막으로, 이러한 예측 가능성이 훈련 손실을 넘어 생성 품질 지표, 최적 CFG 설정, 표현 품질, 심지어 훈련 곡선의 형태에까지 확장되며, 훈련 곡선은 보편적인 형태로 수렴함을 입증한다.

English

Compute-optimal scaling laws guide the training of frontier language models yet remain largely unexplored for visual generation. We present a systematic scaling law study for text-to-image diffusion models using Abra, a controlled family of flow-matching transformers trained across three orders of magnitude worth of compute (10^{19} to 10^{22} FLOPs), reaching significantly larger compute budgets than previous works. We demonstrate that diffusion models scale just as predictably as language models but require far more data to train optimally: compute optimality occurs at approximately 200 image tokens per parameter, ten times the Chinchilla compute-optimal prescription for LLMs. We show that unlike language models, diffusion models are robust to overtraining and that practitioners should err on the side of more data rather than a larger model. Finally, we show that this predictability extends beyond training loss to generative quality metrics, optimal CFG settings, representation quality, and even the shape of the training curves, which collapse onto a universal form.