分布バックトラッキングは、ワンステップ拡散蒸留の収束速度を向上させるより速い軌道を構築します。

Distribution Backtracking Builds A Faster Convergence Trajectory for One-step Diffusion Distillation

August 28, 2024

著者: Shengyuan Zhang, Ling Yang, Zejian Li, An Zhao, Chenye Meng, Changyuan Yang, Guang Yang, Zhiyuan Yang, Lingyun Sun

cs.AI

要旨

拡散モデルのサンプリング速度を加速することは依然として重要な課題です。最近のスコア蒸留手法は、重い教師モデルを1ステップの生徒生成器に蒸留し、その生徒モデルによって生成されたサンプル上で2つのスコア関数の差を計算することで最適化されます。しかし、蒸留プロセスの初期段階にはスコアの不一致問題があります。なぜなら、既存の手法は主に事前学習された拡散モデルのエンドポイントを教師モデルとして使用することに焦点を当てており、生徒生成器と教師モデルの収束軌跡の重要性を見落としているからです。この問題に対処するために、私たちは教師モデルの完全な収束軌跡を導入し、生徒生成器を蒸留するための「分布バックトラッキング蒸留（DisBack）」を提案します。DisBackはDegradation RecordingとDistribution Backtrackingの2つの段階で構成されています。Degradation Recordingは、教師モデルの収束軌跡を取得するために設計されており、訓練された教師モデルから未訓練の初期生徒生成器までの劣化経路を記録します。劣化経路は、教師モデルの中間分布を暗示的に表します。そして、Distribution Backtrackingは、生徒生成器を訓練して、教師モデルの収束軌跡を近似するために中間分布をバックトラックします。幅広い実験により、DisBackが既存の蒸留手法よりも速くかつより良い収束を達成し、同等の生成パフォーマンスを達成することが示されました。特筆すべきは、DisBackは実装が容易であり、パフォーマンスを向上させるために既存の蒸留手法に一般化できる点です。当社のコードはhttps://github.com/SYZhang0805/DisBackで公開されています。

English

Accelerating the sampling speed of diffusion models remains a significant challenge. Recent score distillation methods distill a heavy teacher model into an one-step student generator, which is optimized by calculating the difference between the two score functions on the samples generated by the student model. However, there is a score mismatch issue in the early stage of the distillation process, because existing methods mainly focus on using the endpoint of pre-trained diffusion models as teacher models, overlooking the importance of the convergence trajectory between the student generator and the teacher model. To address this issue, we extend the score distillation process by introducing the entire convergence trajectory of teacher models and propose Distribution Backtracking Distillation (DisBack) for distilling student generators. DisBask is composed of two stages: Degradation Recording and Distribution Backtracking. Degradation Recording is designed to obtain the convergence trajectory of teacher models, which records the degradation path from the trained teacher model to the untrained initial student generator. The degradation path implicitly represents the intermediate distributions of teacher models. Then Distribution Backtracking trains a student generator to backtrack the intermediate distributions for approximating the convergence trajectory of teacher models. Extensive experiments show that DisBack achieves faster and better convergence than the existing distillation method and accomplishes comparable generation performance. Notably, DisBack is easy to implement and can be generalized to existing distillation methods to boost performance. Our code is publicly available on https://github.com/SYZhang0805/DisBack.