DiffusionBench: ImageNet과 Text-to-Image를 연결하는 통합 훈련 프레임워크 NanoGen을 통한 DiT(디퓨전 트랜스포머) 전방위 평가

Diffusion transformer(DiT) 연구는 ImageNet FID 중심으로 수렴했으나, 이 점수가 텍스트-투-이미지(T2I) 성능을 예측하는지 불확실하다. NanoGen으로 동일한 코드·레시피로 ImageNet과 T2I를 모두 학습·평가하면 ImageNet 개선이 T2I 개선으로 일반화되지 않음을 확인하고, 두 축을 합친 DiffusionBench를 제안할 수 있다.

출발점과 기존 한계: Diffusion transformer(DiT)는 Transformer로 패치 기반 이미지를 처리해 생성 성능을 높였으나, 커뮤니티는 주로 class-conditional ImageNet 평가(FID)에 의존했다. FID 중심 개발은 특정 벤치마크에 과적합된 개선을 유도할 수 있으며, T2I와 같은 실제 응용 축에서의 성능 변화를 반영하지 못할 위험이 있다.

해결 원리: NanoGen은 동일한 DiT 백본과 최적화 레시피를 유지한 채로 데이터 파이프라인과 conditioning 토큰만 교체하는 접근을 채택한다. 구체적으로 encoder–decoder로 분리된 DDT 구조, in-context conditioning(조건을 토큰으로 인코더에 prepend), 및 동일한 AdamW/학습률 스케줄/EMA/timestep 샘플링을 공유한다. 이로 인해 ImageNet→T2I 전환은 class-embedder ↔ frozen text encoder(예: Qwen3-0.6B) 교체와 데이터 로더 변경으로 가능해져 엔지니어링 비용을 크게 줄인다.

달라지는 점과 영향: 동일한 모델·레시피로 ImageNet과 T2I를 비교하면, ImageNet FID가 낮아진다고 해서 T2I 지표가 항상 향상되지는 않는다는 것이 실험적으로 확인된다(21개 모델, Pearson r ∈ [-0.377, -0.580]). 즉, ImageNet 기반 단일 축 개선은 범용적 모델 향상을 보장하지 않으며, DiT 연구는 두 작업 축에서의 평가를 병행해야 한다.

전체 접근과 핵심 아이디어: NanoGen은 하나의 DiT 백본(DDT 기반), 하나의 training loop, 하나의 config 포맷으로 ImageNet과 T2I를 모두 지원한다. 작업 전환은 (i) 데이터 파이프라인(클래스 레이블 이미지 ↔ 캡션 이미지)과 (ii) conditioning 모듈(4 timestep tokens + 8 class tokens ↔ 4 timestep tokens + 256 text tokens)만 변경하면 된다.

핵심 메커니즘 상세: 백본은 encoder와 decoder로 분리된 DDT를 사용한다. encoder는 noisy input과 conditioning tokens를 받아 semantic representation을 생성하고, decoder는 얕지만 넓은 구조로 그 표현과 noisy entity를 입력받아 예측 대상을 출력한다. AdaLN은 decoder에만 적용되고 encoder에서는 제거된다. 모든 conditioning(timestep 포함)은 인코더 토큰으로 전처리되어 in-context conditioning을 실현한다.

학습·샘플링·하이퍼세부: Optimiser는 AdamW(β1=0.9, β2=0.95), learning rate는 linear warmup→decay로 2e-4에서 2e-5로 스케줄한다. gradient clipping=1.0, EMA decay=0.9995, timestep 샘플링은 logit-normal(mean=0, std=1). 기본 예측 목표는 v-prediction이며, 원래 논문이 x-prediction을 요구하는 경우에는 보존한다. 샘플러는 기본 Euler(50 NFEs)를 사용한다. MeanFlow는 flow-matching과 MeanFlow loss 혼합(75%/25%)으로 학습한다.

평가 프로토콜: ImageNet에서는 FID, IS, FDr, MIND를 지원하고 T2I에서는 GenEval, DPG-Bench, GenAIBench를 사용한다. ImageNet은 50,000 이미지(클래스당 50 샘플)로 평가하고, T2I는 CFG scale=6.0 등 세부 설정을 고정해 비교한다.

메인 벤치마크 결과: NanoGen으로 재현한 ImageNet-256 실험에서 latent-space 방법(특히 FLUX.2-VAE, REPA-E 계열)이 FID 1~2 영역에서 최고 성능을 보였고, RAE 계열은 약간 높은 FID를 기록했다(Table 2). T2I 측정에서는 GenEval, DPG-Bench, GenAIBench 간 상이한 순위 차이가 관찰되었고(Table 3), ImageNet FID와 T2I 지표 간의 Pearson 상관계수는 -0.377에서 -0.580으로 강한 양의 예측을 보이지 않았다(Fig. 1).

Ablation·수렴·효율성: 동일 백본·레시피 하에서 RAE 기반 모델은 VAE 기반 모델보다 수렴이 빠른 경향을 보였고, pixel-space 방법은 80 epoch에서 여전히 latent-space보다 FID가 높았다. MeanFlow의 one-/two-step 설정은 FID가 크게 낮아지지는 않았으나(예: NFE=1에서 FID≈6.60) 빠른 샘플링 용도로 가능성을 보였다. 학습 시간은 T2I가 ImageNet과 비교해 동일한 규모의 계산으로 학습 가능했으며(예: 100K steps 기준, 32 H200 사용 시 약 10시간 예시), pixel-space 방법은 latents 계산 비용이 없어 ImageNet에서 더 저렴한 경향을 보였다.

추가 관찰: 공개된 대형 T2I 모델(예: SD-3.5-Large, FLUX-1/2 등)의 GenEval 등 점수는 NanoGen으로 사전학습만 수행한 모델보다 일반적으로 높았고, BLIP-3o-60K 같은 데이터셋에서의 감독적 파인튜닝은 GenEval을 크게 끌어올릴 수 있어 현행 T2I 지표의 취약성(metric hacking) 문제가 존재한다.

전체 아키텍처 구조: 백본은 Decoupled Diffusion Transformer(DDT)로 encoder와 decoder로 분리된다. encoder는 noisy input과 conditioning tokens를 하나의 토큰 시퀀스로 받아 semantic representation을 출력하고, decoder는 얕은(depth 2) 대신 넓은(width 증가) 구조로 그 표현과 noisy tokens를 입력받아 예측을 수행한다. 논문에서 사용한 대표 설정은 encoder 28층, width 1152, decoder 2층 width 2048로 약 615M 파라미터 규모이다.

수학적/알고리즘적 기반: 학습은 v-prediction을 기본으로 하고 timestep 샘플링은 logit-normal 분포를 따른다. 추가적으로 차원 의존적 timestep shifting t_m = (α t_n) / (1+(α−1)t_n) (α = sqrt(m/n), n=4096, m=effective input dim)을 적용한다. MeanFlow 모델은 flow-matching loss와 MeanFlow loss를 혼합(75%/25%)해 최적화한다. 샘플링은 Euler 기반 ODE 통합을 사용하며 기본 NFEs는 50이다.

Prior work 대비 차별점: NanoGen 자체는 새로운 모델이 아니라 통합된 훈련·평가 인프라다. 기존 DiT 연구(DDT, RAE, REPA, VAE 기반 방법 등)의 다양한 구성요소를 하나의 레시피로 통합해 비교 가능한 실험을 제공한다. 이로 인해 동일 하이퍼·샘플러·백본 조건에서 방법 간 상대적 성능을 공정하게 평가할 수 있다.

구현 및 학습 세부사항: Optimiser는 AdamW(β1=0.9, β2=0.95), learning-rate는 linear warmup→decay(2e-4→2e-5), gradient clipping=1.0, EMA decay=0.9995를 사용한다. ImageNet 학습은 보통 80 epochs, batch size 1024를 사용했고 T2I pre-training은 100K iterations, batch size 1024, 10% conditioning dropout, CFG scale=6.0로 설정했다. 평가 스위트는 HuggingFace Transformers 의존으로 통일되어 온라인 모니터링을 지원한다.

논문에 명시된 한계는 다음과 같다. (1) ImageNet↔T2I 상관관계는 본 연구에서 사용한 모델 규모와 계산 예산(예: 100K iterations, batch size 1024, 32 H200 등)에서 측정되었으며 다른 규모에서는 결과가 달라질 수 있다. (2) 벤치마크 실험 대부분은 100K iterations로 제한되어 있어 장기 학습에서의 수렴 차이는 반영되지 않을 수 있다. (3) 현행 T2I 지표는 작업 특화 fine-tuning에 의해 높은 점수를 얻을 수 있어 metric hacking 위험이 존재한다.