Diffusione Immiscibile: Accelerare l'Addestramento della Diffusione con l'Assegnazione del Rumore

Immiscible Diffusion: Accelerating Diffusion Training with Noise Assignment

June 18, 2024

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

cs.AI

Abstract

In questo articolo, evidenziamo come una mappatura subottimale tra rumore e dati porti a un addestramento lento dei modelli di diffusione. Durante l'addestramento di diffusione, i metodi attuali diffondono ogni immagine in tutto lo spazio del rumore, risultando in una miscela di tutte le immagini in ogni punto dello strato di rumore. Sottolineiamo che questa miscela casuale di mappatura rumore-dati complica l'ottimizzazione della funzione di denoising nei modelli di diffusione. Traendo ispirazione dal fenomeno dell'immiscibilità in fisica, proponiamo la Diffusione Immiscibile, un metodo semplice ed efficace per migliorare la miscela casuale della mappatura rumore-dati. In fisica, la miscibilità può variare in base a diverse forze intermolecolari. Pertanto, l'immiscibilità significa che la miscela delle sorgenti molecolari è distinguibile. Ispirati da ciò, proponiamo una strategia di addestramento basata su assegnazione e successiva diffusione. Nello specifico, prima di diffondere i dati dell'immagine nel rumore, assegniamo un rumore target di diffusione per i dati dell'immagine minimizzando la distanza totale tra coppie immagine-rumore in un mini-batch. L'assegnazione funziona analogamente a forze esterne per separare le aree diffusibili delle immagini, mitigando così le difficoltà intrinseche nell'addestramento di diffusione. Il nostro approccio è notevolmente semplice, richiedendo solo una riga di codice per limitare l'area diffusibile di ogni immagine preservando la distribuzione gaussiana del rumore. Ciò garantisce che ogni immagine sia proiettata solo su rumore vicino. Per affrontare l'elevata complessità dell'algoritmo di assegnazione, utilizziamo un metodo di assegnazione quantizzata per ridurre il sovraccarico computazionale a un livello trascurabile. Gli esperimenti dimostrano che il nostro metodo raggiunge un addestramento fino a 3 volte più veloce per i modelli di consistenza e DDIM sul dataset CIFAR, e fino a 1,3 volte più veloce sui dataset CelebA per i modelli di consistenza. Inoltre, conduciamo un'analisi approfondita sulla Diffusione Immiscibile, che getta luce su come essa migliori la velocità di addestramento di diffusione aumentando al contempo la fedeltà.

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.