Difusão Imiscível: Acelerando o Treinamento de Difusão com Atribuição de Ruído

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

Resumo

Neste artigo, destacamos que o mapeamento subótimo entre ruído e dados leva a um treinamento lento de modelos de difusão. Durante o treinamento de difusão, os métodos atuais difundem cada imagem por todo o espaço de ruído, resultando em uma mistura de todas as imagens em cada ponto da camada de ruído. Enfatizamos que essa mistura aleatória do mapeamento ruído-dados complica a otimização da função de remoção de ruído em modelos de difusão. Inspirados pelo fenômeno da imiscibilidade na física, propomos a Difusão Imiscível, um método simples e eficaz para melhorar a mistura aleatória do mapeamento ruído-dados. Na física, a miscibilidade pode variar de acordo com várias forças intermoleculares. Assim, a imiscibilidade significa que a mistura das fontes moleculares é distinguível. Inspirados por isso, propomos uma estratégia de treinamento de atribuição seguida de difusão. Especificamente, antes de difundir os dados da imagem em ruído, atribuímos um ruído alvo de difusão para os dados da imagem, minimizando a distância total entre pares imagem-ruído em um mini-lote. A atribuição funciona de forma análoga a forças externas para separar as áreas difusíveis das imagens, mitigando assim as dificuldades inerentes ao treinamento de difusão. Nossa abordagem é notavelmente simples, exigindo apenas uma linha de código para restringir a área difusível de cada imagem, enquanto preserva a distribuição gaussiana do ruído. Isso garante que cada imagem seja projetada apenas para ruídos próximos. Para lidar com a alta complexidade do algoritmo de atribuição, empregamos um método de atribuição quantizada para reduzir a sobrecarga computacional a um nível insignificante. Experimentos demonstram que nosso método alcança um treinamento até 3x mais rápido para modelos de consistência e DDIM no conjunto de dados CIFAR, e até 1,3x mais rápido no conjunto de dados CelebA para modelos de consistência. Além disso, realizamos uma análise detalhada sobre a Difusão Imiscível, que esclarece como ela melhora a velocidade de treinamento de difusão enquanto aumenta a fidelidade.

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.