불혼합 확산: 노이즈 할당을 통한 확산 학습 가속화

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

초록

본 논문에서는 잡음-데이터 매핑의 비최적화가 확산 모델의 학습 속도를 저하시킨다는 점을 지적합니다. 현재의 확산 학습 방법은 각 이미지를 전체 잡음 공간에 걸쳐 확산시켜, 잡음 레이어의 모든 지점에서 모든 이미지가 혼합되도록 합니다. 우리는 이러한 잡음-데이터 매핑의 무작위 혼합이 확산 모델의 잡음 제거 함수 최적화를 복잡하게 만든다고 강조합니다. 물리학에서의 비혼화 현상에서 영감을 얻어, 우리는 잡음-데이터 매핑의 무작위 혼합을 개선하기 위한 간단하면서도 효과적인 방법인 비혼화 확산(Immiscible Diffusion)을 제안합니다. 물리학에서 혼화성은 다양한 분자 간 힘에 따라 달라질 수 있습니다. 따라서 비혼화성은 분자 소스의 혼합이 구별 가능함을 의미합니다. 이를 바탕으로, 우리는 할당-후-확산 학습 전략을 제안합니다. 구체적으로, 이미지 데이터를 잡음으로 확산시키기 전에, 미니 배치 내에서 이미지-잡음 쌍의 총 거리를 최소화하여 이미지 데이터에 대한 확산 목표 잡음을 할당합니다. 이 할당은 이미지의 확산 가능 영역을 분리하기 위한 외부 힘과 유사하게 작용하여, 확산 학습의 고유한 어려움을 완화합니다. 우리의 접근 방식은 매우 간단하며, 각 이미지에 대해 확산 가능 영역을 제한하면서도 잡음의 가우시안 분포를 유지하기 위해 단 한 줄의 코드만 필요합니다. 이를 통해 각 이미지는 인접한 잡음에만 투영됩니다. 할당 알고리즘의 높은 복잡성을 해결하기 위해, 우리는 계산 오버헤드를 무시할 수 있는 수준으로 줄이기 위해 양자화 할당 방법을 사용합니다. 실험 결과, 우리의 방법은 CIFAR 데이터셋에서 일관성 모델과 DDIM에 대해 최대 3배 빠른 학습 속도를 달성했으며, CelebA 데이터셋에서 일관성 모델에 대해 최대 1.3배 빠른 학습 속도를 보였습니다. 또한, 우리는 비혼화 확산에 대한 철저한 분석을 수행하여, 이 방법이 어떻게 확산 학습 속도를 향상시키면서도 충실도를 개선하는지에 대한 통찰을 제공합니다.

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.