一段階視覚生成のための表現分布マッチング

要旨

本稿では、我々が「表現分布マッチング（RDM）」と名付けたパラダイムの設計空間を解明する。このパラダイムは、凍結済みの事前学習済みエンコーダ下で生成された特徴分布と参照特徴分布をマッチングすることにより、一段階画像生成器を訓練するものである。我々は二つの設計軸、すなわち分布の比較方法と比較に用いる表現を特定し、それらに沿った統制実験により三つの知見を得た。第一に、十年前には説得力のある生成器を訓練できなかった古典的なMMD（最大平均不一致）は、適切に推定されれば強力でスケーラブルな目的関数となる。第二に、生成バッチがその際の操作変数となり、慣例的なバッチサイズをはるかに超える2048以上の最適値が存在する。第三に、任意の単一の表現は操作可能であり、画像が明らかに偽物であるにもかかわらず実スコアを下回るように駆動され得る。そのため、我々はバランスの取れたエンコーダ群に対してマッチングを行い、訓練損失とは独立で操作に耐性のある、14個のエンコーダに対するスライス化ワッサースタイン距離SW_r14を用いて評価する。これらの好ましい選択を組み合わせることで、改良型RDM（iRDM）が得られる。これはImageNet上でSW_r14スコア1.30という一段階生成の最先端を達成し、我々の目的関数が決して最適化しない人間の嗜好代理指標であるPickScoreによっても裏付けられており、マッチングされたサンプルの71.2%で従来の最良の一段階生成器よりも好まれる。同じ手法を適用して、四段階のFLUX.2 [klein]を一段階生成器にポスト訓練したところ、GenEvalスコアで0.826対0.794、PickScoreで22.76対22.58と、四段階版を上回り、所要時間はH200 GPUで90時間であった。プロジェクトページ: 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/.