[論文レビュー] Improved Techniques for Training Consistency Models

Yang Song, Prafulla Dhariwal|arXiv (Cornell University)|Oct 22, 2023

Generative Adversarial Networks and Image Synthesis被引用数 12

ひとこと要約

本論文は、教師から EMA を除去し、LPIPS の代わりに Pseudo-Huber 損失を用い、対数正規分布ノイズスケジュールと離散化カリキュラムを導入することで、1ステップおよび2ステップのサンプリングにおいて蒸留なしで FID の状態的改善を達成し、一貫性トレーニングを向上させる。

ABSTRACT

Consistency models are a nascent family of generative models that can sample high quality data in one step without the need for adversarial training. Current consistency models achieve optimal sample quality by distilling from pre-trained diffusion models and employing learned metrics such as LPIPS. However, distillation limits the quality of consistency models to that of the pre-trained diffusion model, and LPIPS causes undesirable bias in evaluation. To tackle these challenges, we present improved techniques for consistency training, where consistency models learn directly from data without distillation. We delve into the theory behind consistency training and identify a previously overlooked flaw, which we address by eliminating Exponential Moving Average from the teacher consistency model. To replace learned metrics like LPIPS, we adopt Pseudo-Huber losses from robust statistics. Additionally, we introduce a lognormal noise schedule for the consistency training objective, and propose to double total discretization steps every set number of training iterations. Combined with better hyperparameter tuning, these modifications enable consistency models to achieve FID scores of 2.51 and 3.25 on CIFAR-10 and ImageNet $64\times 64$ respectively in a single sampling step. These scores mark a 3.5$\times$ and 4$\times$ improvement compared to prior consistency training approaches. Through two-step sampling, we further reduce FID scores to 2.24 and 2.77 on these two datasets, surpassing those obtained via distillation in both one-step and two-step settings, while narrowing the gap between consistency models and other state-of-the-art generative models.

研究の動機と目的

• 事前学習済み拡散モデルに依存せず、CT を Consistency Distillation (CD) に匹敵または上回る改善へ動機づける。
• 学習済み指標 LPIPS などの訓練および評価におけるバイアスと依存性を排除する。
• 訓練カリキュラムとノイズスケジュールを開発し、スケーラビリティを保ちながら1ステップのサンプリング品質を向上させる。
• CT のより深い理論的理解を提供し、従来の分析で見落とされた欠点に対処する。

提案手法

• 以前の CT 分析の欠陥を修正するため、教師ネットワークから指数移動平均 (EMA) を排除する。
• 訓練のためのロバストで指標に依存しない目的関数として、LPIPS を Pseudo-Huber 損失に置換する。
• CT 目的関数のノイズスケジュールとして対数正規分布ベースを導入する。
• 訓練中の区分手順を区間ごとに総離散化ステップを倍増させるカリキュラムを採用し、データ効率を改善する。
• これらの変更を加えた CT が CIFAR-10 および ImageNet 64x64 で1ステップおよび2ステップのサンプリングにおいて CD を上回ることを実証する。

実験結果

リサーチクエスチョン

• RQ1 拡散モデルの事前訓練に依存せず、Consistency Training (CT) は Consistency Distillation (CD) より高いサンプル品質を達成できるか？
• RQ2 教師ネットワークから EMA を除去することで CT 分析の理論的欠陥を解消し、実際の性能が向上するか？
• RQ3 Pseudo-Huber のようなロバスト損失と対数正規ノイズスケジュールは、LPIPS のような学習済み指標なしで CT を改善するか？
• RQ4 離散化ステップスケジュールと訓練カリキュラムが1ステップおよび多段サンプリング品質に与える影響は？

主な発見

• 提案された変更を加えた CT は、一回のサンプリングステップで CIFAR-10 の FID 2.51、ImageNet 64x64 の FID 3.25 を達成。
• 二段階サンプリングで CIFAR-10 が 2.24、ImageNet 64x64 が 2.77 に改善し、1ステップおよび2ステップの設定で CD を上回る。
• LPIPS を Pseudo-Huber 損失に置換することで、学習済み指標への依存を排除し、評価バイアスを低減。
• 対数正規分布ノイズスケジュールと離散化ステップのカリキュラムは、サンプル品質と訓練効率を大きく向上させる。
• これらの改善を加えた CT 手法は、従来の CT を上回り、蒸留なしでトップの拡散モデルや GAN に匹敵する。