Несмешиваемая диффузия: ускорение обучения диффузии с помощью шума Назначение

Immiscible Diffusion: Accelerating Diffusion Training with Noise Assignment

June 18, 2024

Авторы: Yiheng Li, Heyang Jiang, Akio Kodaira, Masayoshi Tomizuka, Kurt Keutzer, Chenfeng Xu

cs.AI

Аннотация

В данной статье мы указываем, что неоптимальное отображение шума на данные приводит к медленному обучению моделей диффузии. В процессе обучения диффузии текущие методы диффундируют каждое изображение по всему пространству шума, что приводит к смешиванию всех изображений в каждой точке слоя шума. Мы подчеркиваем, что случайное смешивание отображения шума на данные усложняет оптимизацию функции денойзинга в моделях диффузии. Вдохновляясь немешаемым явлением в физике, мы предлагаем метод "Немешаемая диффузия" - простой и эффективный способ улучшить случайное смешивание отображения шума на данные. В физике совместимость может меняться в зависимости от различных межмолекулярных сил. Таким образом, немешаемость означает, что смешивание молекулярных источников различимо. Вдохновленные этим, мы предлагаем стратегию обучения с присвоением, а затем диффузией. Конкретно, перед диффузией данных изображения в шум мы присваиваем целевой шум диффузии для данных изображения, минимизируя общее расстояние пары изображение-шум в мини-пакете. Присвоение функционирует аналогично внешним силам для разделения областей, подлежащих диффузии на изображениях, тем самым смягчая врожденные трудности в обучении диффузии. Наш подход чрезвычайно прост, требуя всего одну строку кода для ограничения области диффузии для каждого изображения, сохраняя при этом гауссовское распределение шума. Это гарантирует, что каждое изображение проецируется только на близкий шум. Для решения высокой сложности алгоритма присвоения мы используем метод квантованного присвоения для снижения вычислительной нагрузки до незначительного уровня. Эксперименты показывают, что наш метод достигает ускорения обучения до 3 раз для моделей согласованности и DDIM на наборе данных CIFAR, и до 1,3 раза быстрее на наборах данных CelebA для моделей согласованности. Кроме того, мы проводим тщательный анализ метода "Немешаемая диффузия", который проливает свет на то, как он улучшает скорость обучения диффузии, сохраняя достоверность.

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.