不混和拡散：ノイズ割り当てによる拡散学習の高速化

Immiscible Diffusion: Accelerating Diffusion Training with Noise Assignment

June 18, 2024

著者: Yiheng Li, Heyang Jiang, Akio Kodaira, Masayoshi Tomizuka, Kurt Keutzer, Chenfeng Xu

cs.AI

要旨

本論文では、拡散モデルの学習速度が遅い原因として、ノイズとデータのマッピングが最適でない点を指摘します。現在の拡散学習手法では、各画像をノイズ空間全体に拡散させるため、ノイズ層の各点で全ての画像が混合された状態になります。このランダムなノイズ-データマッピングの混合が、拡散モデルのノイズ除去関数の最適化を複雑にしていることを強調します。物理学における不混和現象に着想を得て、我々はノイズ-データマッピングのランダムな混合を改善するシンプルかつ効果的な手法「不混和拡散（Immiscible Diffusion）」を提案します。物理学において、混和性は様々な分子間力によって変化します。したがって、不混和性とは分子源の混合が識別可能であることを意味します。これに着想を得て、我々は「割り当て→拡散」という学習戦略を提案します。具体的には、画像データをノイズに拡散する前に、ミニバッチ内の画像-ノイズペアの総距離を最小化することで、画像データに対する拡散目標ノイズを割り当てます。この割り当ては、画像の拡散可能領域を分離する外部力として機能し、拡散学習に内在する困難を軽減します。我々のアプローチは非常にシンプルで、各画像の拡散可能領域を制限するためにたった1行のコードを追加するだけで、ノイズのガウス分布を維持します。これにより、各画像は近傍のノイズにのみ投影されます。割り当てアルゴリズムの高い計算複雑性に対処するため、量子化割り当て法を採用し、計算オーバーヘッドを無視できるレベルまで削減します。実験結果は、CIFARデータセットにおいてConsistency ModelとDDIMの学習速度を最大3倍、CelebAデータセットにおいてConsistency Modelの学習速度を最大1.3倍向上させることを示しています。さらに、不混和拡散に関する詳細な分析を行い、学習速度を向上させながら忠実度を高めるメカニズムを明らかにしました。

English

In this paper, we point out suboptimal noise-data mapping leads to slow training of diffusion models. During diffusion training, current methods diffuse each image across the entire noise space, resulting in a mixture of all images at every point in the noise layer. We emphasize that this random mixture of noise-data mapping complicates the optimization of the denoising function in diffusion models. Drawing inspiration from the immiscible phenomenon in physics, we propose Immiscible Diffusion, a simple and effective method to improve the random mixture of noise-data mapping. In physics, miscibility can vary according to various intermolecular forces. Thus, immiscibility means that the mixing of the molecular sources is distinguishable. Inspired by this, we propose an assignment-then-diffusion training strategy. Specifically, prior to diffusing the image data into noise, we assign diffusion target noise for the image data by minimizing the total image-noise pair distance in a mini-batch. The assignment functions analogously to external forces to separate the diffuse-able areas of images, thus mitigating the inherent difficulties in diffusion training. Our approach is remarkably simple, requiring only one line of code to restrict the diffuse-able area for each image while preserving the Gaussian distribution of noise. This ensures that each image is projected only to nearby noise. To address the high complexity of the assignment algorithm, we employ a quantized-assignment method to reduce the computational overhead to a negligible level. Experiments demonstrate that our method achieve up to 3x faster training for consistency models and DDIM on the CIFAR dataset, and up to 1.3x faster on CelebA datasets for consistency models. Besides, we conduct thorough analysis about the Immiscible Diffusion, which sheds lights on how it improves diffusion training speed while improving the fidelity.