Immiscibele diffusie: Versnelling van diffusietraining met ruis Toewijzing

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

Samenvatting

In dit artikel wijzen we erop dat een suboptimale ruis-data mapping leidt tot trage training van diffusiemodellen. Tijdens de diffusietraining verspreiden huidige methoden elke afbeelding over de gehele ruisruimte, wat resulteert in een mengsel van alle afbeeldingen op elk punt in de ruislaag. We benadrukken dat deze willekeurige menging van ruis-data mapping de optimalisatie van de ruisverwijderingsfunctie in diffusiemodellen bemoeilijkt. Geïnspireerd door het onmengbare fenomeen in de natuurkunde, stellen we Immiscible Diffusion voor, een eenvoudige en effectieve methode om de willekeurige menging van ruis-data mapping te verbeteren. In de natuurkunde kan mengbaarheid variëren afhankelijk van verschillende intermoleculaire krachten. Onmengbaarheid betekent dus dat de menging van de moleculaire bronnen onderscheidbaar is. Hierdoor geïnspireerd, stellen we een toewijzing-dan-diffusie trainingsstrategie voor. Specifiek, voordat we de afbeeldingsdata in ruis diffunderen, wijzen we diffusiedoelruis toe aan de afbeeldingsdata door de totale afbeelding-ruispaarafstand in een mini-batch te minimaliseren. De toewijzing functioneert analoog aan externe krachten om de diffuseerbare gebieden van afbeeldingen te scheiden, waardoor de inherente moeilijkheden in diffusietraining worden verminderd. Onze aanpak is opmerkelijk eenvoudig, waarbij slechts één regel code nodig is om het diffuseerbare gebied voor elke afbeelding te beperken, terwijl de Gaussische verdeling van de ruis behouden blijft. Dit zorgt ervoor dat elke afbeelding alleen naar nabije ruis wordt geprojecteerd. Om de hoge complexiteit van het toewijzingsalgoritme aan te pakken, gebruiken we een gekwantiseerde-toewijzingsmethode om de rekenkosten tot een verwaarloosbaar niveau te reduceren. Experimenten tonen aan dat onze methode tot 3x snellere training bereikt voor consistentiemodellen en DDIM op de CIFAR-dataset, en tot 1.3x sneller op CelebA-datasets voor consistentiemodellen. Daarnaast voeren we een grondige analyse uit van Immiscible Diffusion, wat inzicht geeft in hoe het de snelheid van diffusietraining verbetert terwijl het de getrouwheid verbetert.

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.