Representation Distribution Matching for One-Step Visual Generation

We elucidate the design space of Representation Distribution Matching (RDM), our name for the paradigm that trains a one-step image generator by matching generated and reference feature distributions under frozen pretrained encoders. We identify two design axes, how the distributions are compared and the representations they are compared in, and controlled studies along them yield three findings. First, the classical MMD, which could not train convincing generators a decade ago, becomes a strong and scalable objective once estimated right. Second, the generated batch is then the operative variable, with an optimum above 2048, far beyond customary batch sizes. Third, any single representation can be gamed, driven below the real score while images stay visibly fake, so we match against a balanced battery of encoders and evaluate with SW_r14, a Sliced-Wasserstein distance over 14 encoders that is independent of the training loss and resists gaming. Combining the preferred choices yields improved RDM (iRDM): it sets the one-step state of the art on ImageNet at SW_r14 1.30, corroborated by PickScore, a human-preference proxy our objective never optimizes, which prefers it over the prior best one-step generator on 71.2% of matched samples. The same recipe post-trains the four-step FLUX.2 [klein] into a one-step generator, surpassing the four-step version on GenEval, 0.826 to 0.794, and on PickScore, 22.76 to 22.58, in 90 H200 GPU-hours. Project page: https://alan-lanfeng.github.io/rdm/.

한 줄 요약

RDM을 기반으로 한 iRDM이 ImageNet에서 SW_r14 1.30 성능을 달성하며, 4단계 FLUX.2를 1단계로 변환해 GenEval 0.826을 기록한다.

핵심 기여도

RDM의 두 설계 축(분포 비교 방식, 표현 공간)을 체계적으로 분석하고, iRDM이라는 최적화된 레시피를 제안.
MMD가 올바르게 추정되면 강력한 목표 함수가 되는 것을 밝힘: Nyström 기반 reference와 대규모 batch(2048 이상) 사용.
단일 인코더 대신 14개 인코더 패널을 사용한 SW_r14 평가 지표를 제안, 게이밍 저항성 확보.
iRDM이 ImageNet에서 SW_r14 1.30, PickScore 71.2% 우승, FLUX.2를 1단계로 변환해 GenEval 0.826 달성.

핵심 아이디어

RDM은 이미지 생성기의 출력 분포가 참 분포와 일치하도록, 동결된 사전학습 인코더의 특징 분포를 일치시키는 방식이다. 이 연구는 RDM의 성능을 좌우하는 두 축을 분리해 분석한다: (1) 분포 비교 방법, (2) 비교에 사용되는 표현 공간.

첫 번째 축에서, 과거에는 MMD가 부족한 목표로 간주되었으나, 연구는 MMD가 올바르게 추정되면 강력한 목표가 될 수 있음을 밝힌다. 구체적으로, Nyström 방법을 사용해 1.28M 이미지로 구성된 reference를 4096개의 랜드마크로 압축하고, 생성된 배치는 2048 이상으로 확장해 정확도를 높인다.

두 번째 축에서는 단일 인코더에 과적합되는 문제를 해결하기 위해 14개 인코더 패널을 사용하고, 제약 최적화를 통해 인코더 간 균형을 유지한다. 이는 "가장 약한 막대" 원칙에 기반하며, 생성된 이미지가 모든 인코더에서 실현 가능한 수준에 도달해야 한다는 아이디어다.

기술적 접근법

**분포 비교**: MMD를 사용, Nyström 기반 reference와 대규모 생성 배치(2048 이상)로 추정.
**표현 공간**: 14개 인코더 패널을 사용, 제약 최적화를 통해 균형 유지.
**조건 생성**: 이미지-텍스트 결합 분포를 일치시켜, 텍스트 조건에 대한 정확도를 높임.
**평가 지표**: SW_r14, 14개 인코더 기반 Sliced-Wasserstein 거리, 학습 손실과 독립적.
**하이퍼파라미터**: 생성 배치 크기 2048 이상, Nyström 랜드마크 수 4096, 인코더 수 14.

주요 결과

**ImageNet 1단계 생성**: iRDM은 SW_r14 1.30 성능을 달성, 이전 최고 기록(2.05) 대비 +37.5% 개선.
**PickScore**: 71.2%의 샘플에서 이전 최고 1단계 생성기보다 우수함.
**FLUX.2 4단계 → 1단계 변환**: iRDM을 적용한 모델이 GenEval 0.826 (기존 0.794 대비 +3.7%) 달성.
**PickScore**: 22.76 (기존 22.58 대비 +0.18) 개선.
**GPU 시간**: 90 H200 GPU-hour 소요.

의의 및 한계

iRDM은 RDM 패러다임의 설계 공간을 체계적으로 탐색해, 1단계 생성기의 성능 한계를 높이는 데 기여한다. 특히, MMD의 재평가와 인코더 패널 기반 평가 지표는 기존 생성 모델 평가 방식의 한계를 보완한다. 또한, FLUX.2를 1단계로 변환한 사례는 RDM의 유연성과 확장성을 입증한다.

그러나, iRDM이 생성된 이미지가 참 분포와 완전히 일치하지는 않으며, SW_r14 1.30은 1.0(참 분포)에 비해 여전히 차이가 있다. 이는 향후 연구에서 다중 스케일 커널, 태스크별 인코더 패널, 조건 결합 방식 등을 통해 개선될 수 있다.

실용적 활용

iRDM은 이미지 생성뿐 아니라 텍스트-이미지 생성, 멀티모달 생성 등에 적용 가능하다. 특히, 사전학습 인코더가 존재하는 모든 도메인(예: 음성, 영상)에 RDM 패러다임을 확장할 수 있어, 생성 모델의 단계 수를 줄이며 효율성을 높이는 데 유용하다.