Multi-student Diffusion Distillation for Better One-step Generators

要約

拡散モデルは、長時間にわたる複数ステップの推論手順を犠牲にして、高品質のサンプル生成を実現します。
これを克服するために、拡散蒸留技術は、ワンステップで教師に匹敵する、または教師を超えることができる生徒生成器を生成します。
ただし、学生モデルの推論速度は教師アーキテクチャのサイズによって制限されるため、計算量の多いアプリケーションのリアルタイム生成が妨げられます。
この作業では、条件付き教師拡散モデルを複数の単一ステップ ジェネレーターに蒸留するフレームワークである Multi-Student Distillation (MSD) を導入します。
各スチューデント ジェネレータはコンディショニング データのサブセットを担当するため、同じ容量でより高い生成品質が得られます。
MSD は、複数の抽出された生徒をトレーニングし、サイズを小さくできるため、推論が高速になります。
また、MSD は、同じアーキテクチャを使用した単一スチューデント蒸留よりも軽量で品質の向上を実現します。
私たちは、分布マッチングと敵対的蒸留技術を使用して、同じサイズ以下の複数の生徒を一段階蒸留で訓練することにより、MSD が効果的であることを実証しました。
MSD は、スチューデントの数が少ない場合、シングルステップ生成のより高速な推論により、競争力のある結果を獲得します。
MSD は、同じサイズの 4 人の生徒を使用して、ワンステップ画像生成の新しい最先端技術を確立しました。ImageNet-64×64 では FID 1.20、ゼロショット COCO2014 では 8.20 です。

要約(オリジナル)

Diffusion models achieve high-quality sample generation at the cost of a lengthy multistep inference procedure. To overcome this, diffusion distillation techniques produce student generators capable of matching or surpassing the teacher in a single step. However, the student model’s inference speed is limited by the size of the teacher architecture, preventing real-time generation for computationally heavy applications. In this work, we introduce Multi-Student Distillation (MSD), a framework to distill a conditional teacher diffusion model into multiple single-step generators. Each student generator is responsible for a subset of the conditioning data, thereby obtaining higher generation quality for the same capacity. MSD trains multiple distilled students, allowing smaller sizes and, therefore, faster inference. Also, MSD offers a lightweight quality boost over single-student distillation with the same architecture. We demonstrate MSD is effective by training multiple same-sized or smaller students on single-step distillation using distribution matching and adversarial distillation techniques. With smaller students, MSD gets competitive results with faster inference for single-step generation. Using 4 same-sized students, MSD sets a new state-of-the-art for one-step image generation: FID 1.20 on ImageNet-64×64 and 8.20 on zero-shot COCO2014.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google