Nicht mischbare Diffusion: Beschleunigung des Diffusionstrainings mit Rauschen Zuweisung

Immiscible Diffusion: Accelerating Diffusion Training with Noise Assignment

June 18, 2024

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

cs.AI

Zusammenfassung

In diesem Paper weisen wir darauf hin, dass eine suboptimale Rauschdaten-Zuordnung zu einem langsamen Training von Diffusionsmodellen führt. Während des Diffusionstrainings verbreiten aktuelle Methoden jedes Bild über den gesamten Rauschraum, was zu einer Mischung aller Bilder an jedem Punkt in der Rauschschicht führt. Wir betonen, dass diese zufällige Mischung von Rauschdaten-Zuordnungen die Optimierung der Entrauschungsfunktion in Diffusionsmodellen erschwert. Inspiriert von dem unvermischbaren Phänomen in der Physik schlagen wir die Immiscible Diffusion vor, eine einfache und effektive Methode, um die zufällige Mischung von Rauschdaten-Zuordnungen zu verbessern. In der Physik kann die Mischbarkeit je nach verschiedenen intermolekularen Kräften variieren. Somit bedeutet Unvermischbarkeit, dass die Vermischung der molekularen Quellen unterscheidbar ist. Inspiriert davon schlagen wir eine Zuweisungs-und-Diffusion-Trainingsstrategie vor. Spezifisch weisen wir vor der Diffusion der Bilddaten in den Rausch den Bilddaten das Zielrauschen zu, indem wir die Gesamtdistanz der Bild-Rausch-Paare in einem Mini-Batch minimieren. Die Zuweisungsfunktionen wirken analog zu externen Kräften, um die diffundierbaren Bereiche der Bilder zu trennen und somit die inhärenten Schwierigkeiten im Diffusionstraining zu mildern. Unser Ansatz ist bemerkenswert einfach und erfordert nur eine Codezeile, um den diffundierbaren Bereich für jedes Bild zu begrenzen, während die Gaußsche Verteilung des Rauschens erhalten bleibt. Dies stellt sicher, dass jedes Bild nur in die nahegelegenen Rauschbereiche projiziert wird. Um die hohe Komplexität des Zuweisungsalgorithmus zu bewältigen, verwenden wir eine quantisierte Zuweisungsmethode, um den Rechenaufwand auf ein vernachlässigbares Niveau zu reduzieren. Experimente zeigen, dass unsere Methode eine bis zu 3-fach schnellere Schulung für Konsistenzmodelle und DDIM auf dem CIFAR-Datensatz erreicht und bis zu 1,3-fach schneller auf CelebA-Datensätzen für Konsistenzmodelle. Darüber hinaus führen wir eine gründliche Analyse zur Immiscible Diffusion durch, die aufzeigt, wie sie die Geschwindigkeit des Diffusionstrainings verbessert und gleichzeitig die Treue erhöht.

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.