Diffusion Immiscible : Accélération de l'entraînement par diffusion grâce à l'assignation de bruit

Immiscible Diffusion: Accelerating Diffusion Training with Noise Assignment

June 18, 2024

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

cs.AI

Résumé

Dans cet article, nous soulignons qu'une cartographie sous-optimale entre le bruit et les données entraîne un ralentissement de l'entraînement des modèles de diffusion. Lors de l'entraînement par diffusion, les méthodes actuelles diffusent chaque image dans l'ensemble de l'espace de bruit, ce qui aboutit à un mélange de toutes les images à chaque point de la couche de bruit. Nous insistons sur le fait que ce mélange aléatoire de la cartographie bruit-données complique l'optimisation de la fonction de débruîtage dans les modèles de diffusion. En nous inspirant du phénomène d'immiscibilité en physique, nous proposons la Diffusion Immiscible, une méthode simple et efficace pour améliorer ce mélange aléatoire. En physique, la miscibilité peut varier selon différentes forces intermoléculaires. Ainsi, l'immiscibilité signifie que le mélange des sources moléculaires reste distinct. Inspirés par cela, nous proposons une stratégie d'entraînement basée sur l'affectation puis la diffusion. Concrètement, avant de diffuser les données d'image en bruit, nous assignons un bruit cible de diffusion aux données d'image en minimisant la distance totale entre les paires image-bruit dans un mini-lot. Cette affectation fonctionne de manière analogue à des forces externes pour séparer les zones diffusables des images, atténuant ainsi les difficultés inhérentes à l'entraînement par diffusion. Notre approche est remarquablement simple, nécessitant seulement une ligne de code pour restreindre la zone diffusable de chaque image tout en préservant la distribution gaussienne du bruit. Cela garantit que chaque image est projetée uniquement vers un bruit proche. Pour réduire la complexité élevée de l'algorithme d'affectation, nous utilisons une méthode d'affectation quantifiée afin de ramener la surcharge computationnelle à un niveau négligeable. Les expériences montrent que notre méthode permet un entraînement jusqu'à 3 fois plus rapide pour les modèles de cohérence et DDIM sur le jeu de données CIFAR, et jusqu'à 1,3 fois plus rapide sur les jeux de données CelebA pour les modèles de cohérence. Par ailleurs, nous menons une analyse approfondie de la Diffusion Immiscible, qui éclaire la manière dont elle améliore la vitesse d'entraînement par diffusion tout en augmentant la fidélité.

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.