Difusión Inmiscible: Aceleración del Entrenamiento de Difusión con Asignación de Ruido

Immiscible Diffusion: Accelerating Diffusion Training with Noise Assignment

June 18, 2024

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

cs.AI

Resumen

En este artículo, señalamos que el mapeo subóptimo entre ruido y datos conduce a un entrenamiento lento de los modelos de difusión. Durante el entrenamiento de difusión, los métodos actuales difunden cada imagen a través de todo el espacio de ruido, lo que resulta en una mezcla de todas las imágenes en cada punto de la capa de ruido. Destacamos que esta mezcla aleatoria del mapeo ruido-datos complica la optimización de la función de eliminación de ruido en los modelos de difusión. Inspirándonos en el fenómeno de la inmisibilidad en física, proponemos Difusión Inmiscible, un método simple y efectivo para mejorar la mezcla aleatoria del mapeo ruido-datos. En física, la miscibilidad puede variar según diversas fuerzas intermoleculares. Por lo tanto, la inmisibilidad significa que la mezcla de las fuentes moleculares es distinguible. Inspirados por esto, proponemos una estrategia de entrenamiento de asignación-then-difusión. Específicamente, antes de difundir los datos de la imagen en ruido, asignamos un ruido objetivo de difusión para los datos de la imagen minimizando la distancia total entre pares imagen-ruido en un mini-lote. La asignación funciona de manera análoga a fuerzas externas para separar las áreas difusibles de las imágenes, mitigando así las dificultades inherentes en el entrenamiento de difusión. Nuestro enfoque es notablemente simple, requiriendo solo una línea de código para restringir el área difusible de cada imagen mientras se preserva la distribución gaussiana del ruido. Esto asegura que cada imagen se proyecte solo a ruido cercano. Para abordar la alta complejidad del algoritmo de asignación, empleamos un método de asignación cuantizada para reducir la sobrecarga computacional a un nivel insignificante. Los experimentos demuestran que nuestro método logra un entrenamiento hasta 3 veces más rápido para modelos de consistencia y DDIM en el conjunto de datos CIFAR, y hasta 1.3 veces más rápido en los conjuntos de datos CelebA para modelos de consistencia. Además, realizamos un análisis exhaustivo sobre la Difusión Inmiscible, que arroja luz sobre cómo mejora la velocidad de entrenamiento de difusión mientras aumenta la fidelidad.

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.