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

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

왜 중요한가

핵심 기여

NanoGen: ImageNet과 T2I를 동일 코드베이스에서 학습 가능한 통합 프레임워크

하나의 DiT 백본, 하나의 optimiser, 하나의 training loop, 하나의 config로 ImageNet(class-conditional)과 text-to-image(T2I)를 모두 지원한다. 데이터 파이프라인과 conditioning 모듈 교체(약 12줄의 config 변경)만으로 두 작업 전환이 가능하다.

교차-태스크 비교: ImageNet FID와 T2I 지표 간 상관관계 부족

NanoGen으로 21개의 latent diffusion 모델을 ImageNet과 T2I로 동일 조건 하에 학습·평가한 결과, ImageNet FID와 T2I 지표(GenEval, DPG-Bench, GenAIBench) 사이의 Pearson 상관계수는 -0.377에서 -0.580으로 강한 양의 예측력을 보이지 않았다.

DiffusionBench: ImageNet과 T2I를 통합한 벤치마크 제안

ImageNet 단일 축이 아닌 ImageNet+T2I 결합 지표 모음인 DiffusionBench를 제시해 DiT 방법의 범용성을 평가하도록 권장한다. 공개 코드와 모델(https://github.com/End2End-Diffusion/diffusion-bench, Hugging Face 허브 링크)을 함께 제공한다.

다양한 DiT 방법의 재현성과 비교

RAE, VAE, pixel-space, REG, MeanFlow 등 여러 기존 방법을 NanoGen으로 재구현해 ImageNet 및 T2I에서 성능을 비교하고, 동일 백본·레시피 하에서의 상대적 성능과 수렴 속도 차이를 보고했다.

핵심 아이디어 이해하기

출발점과 기존 한계: 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 위험이 존재한다.

실무 활용

NanoGen은 동일한 코드·레시피로 ImageNet과 T2I 모델을 학습·평가할 수 있어 DiT 연구자가 추가적인 엔지니어링 없이 두 작업에서의 성능을 비교할 수 있게 한다. GitHub 저장소와 Hugging Face 모델 허브를 통해 코드와 체크포인트가 공개되어 있다.

연구자: 새로운 DiT 아키텍처나 학습 기법을 ImageNet과 T2I에서 동일 환경으로 재현·비교
개발자: 텍스트-투-이미지 모델의 초기 프로토타입을 ImageNet 설정에서 빠르게 전환해 학습
벤치마크 운영자: DiffusionBench 기반으로 다축 평가 파이프라인 구축

코드 공개 여부: 공개

코드 저장소 보기

키워드

diffusion transformer(DiT)NanoGen(통합 훈련 프레임워크)DiffusionBench(통합 벤치마크)text-to-image generation(텍스트-투-이미지)ImageNet FID(ImageNet FID)latent diffusion(잠재 확산)

용어 해설

분류기-프리 가이던스(Classifier-Free Guidance (CFG)): — 조건 입력과 무조건 입력을 무작위로 섞어 샘플링 시 가이던스 스케일을 적용해 생성 품질과 다양성의 절충을 조절하는 기술이다. 모델 출력에 조건과 비조건 예측을 선형 결합해 최종 샘플을 얻으며, 가이던스 스케일이 크면 조건 일치성이 증가하고 다양성은 감소한다.
잠재 공간 확산(Latent Diffusion): — 사전학습된 VAE나 인코더로 이미지를 저차원 잠재공간으로 변환한 뒤 그 공간에서 확산 과정을 학습하는 방법이다. 픽셀 직접 처리 대비 연산·메모리 비용을 줄여 고해상도 학습을 실현한다.
표현 오토인코더 (RAE)(Representation Autoencoder (RAE)): — 사전학습된 비동결(또는 고정) 비전 인코더의 출력 표현을 토크나이저로 사용해 디퓨전 모델이 그 표현을 재생성하도록 학습하는 방식이다. 고수준 의미 구조를 보존하는 잠재가 생성 수렴 속도를 높인다.
표현 정렬 (REPA)(REPA (representation alignment)): — DiT 내부 표현을 미리 학습된 비전 인코더의 내부 특성과 정렬시키는 기법이다. 이 정렬 신호를 활용해 VAE와 DiT를 end-to-end로 함께 튜닝하면 학습 안정성과 품질 향상이 관찰된다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 피드

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 피드