단일 단계 시각 생성을 위한 표현 분포 매칭

초록

우리는 고정된 사전 학습 인코더 하에서 생성된 특징 분포와 참조 특징 분포를 매칭하여 단일 단계 이미지 생성기를 훈련하는 패러다임인 표현 분포 매칭(Representation Distribution Matching, RDM)의 설계 공간을 규명한다. 우리는 두 가지 설계 축, 즉 분포를 비교하는 방식과 비교 대상이 되는 표현을 식별하고, 이들에 대한 통제된 연구를 통해 세 가지 발견을 도출한다. 첫째, 10년 전만 해도 설득력 있는 생성기를 훈련할 수 없었던 고전적 MMD가, 올바르게 추정되면 강력하고 확장 가능한 목적 함수가 된다. 둘째, 생성된 배치가 그때 작동 변수가 되며, 관례적인 배치 크기를 훨씬 넘어서는 2048 이상에서 최적점을 가진다. 셋째, 단일 표현은 조작될 수 있으며, 실제 점수 이하로 낮아져도 이미지는 여전히 가짜로 보이므로, 우리는 균형 잡힌 인코더 배터리에 대해 매칭하고, 훈련 손실과 독립적이며 조작에 저항하는 14개 인코더에 대한 슬라이스드-바서스타인 거리(Sliced-Wasserstein distance)인 SW_r14로 평가한다. 선호되는 선택들을 결합하면 개선된 RDM(iRDM)이 도출된다: ImageNet에서 SW_r14 1.30으로 단일 단계 최고 성능(state of the art)을 달성하며, 이는 우리의 목적 함수가 절대 최적화하지 않는 인간 선호도 대리 지표인 PickScore에 의해 뒷받침되며, 이는 일치된 샘플의 71.2%에서 이전 최고 단일 단계 생성기보다 선호된다. 동일한 방법으로 4단계 FLUX.2 [klein]을 사후 훈련하여 단일 단계 생성기로 전환하며, GenEval에서 0.826 대 0.794, PickScore에서 22.76 대 22.58로 4단계 버전을 능가하며, 90시간의 H200 GPU 시간이 소요된다. 프로젝트 페이지: https://alan-lanfeng.github.io/rdm/.

English

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/.