연속 시간 일관성 모델의 단순화, 안정화 및 확장

Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models

October 14, 2024

저자: Cheng Lu, Yang Song

cs.AI

초록

일관성 모델(CM)은 빠른 샘플링을 위해 최적화된 확산 기반 생성 모델의 강력한 클래스입니다. 대부분의 기존 CM은 이산화된 시간 단계를 사용하여 훈련되며, 이는 추가 하이퍼파라미터를 도입하고 이산화 오류에 취약합니다. 연속 시간 공식은 이러한 문제를 완화할 수 있지만, 훈련 불안정성으로 인해 성공이 제한되었습니다. 이에 대응하여, 우리는 확산 모델과 CM의 이전 매개변수화를 통합하고 불안정성의 근본 원인을 식별하는 간소화된 이론적 프레임워크를 제안합니다. 이 분석을 기반으로, 확산 과정 매개변수화, 네트워크 아키텍처 및 훈련 목표에서 주요 개선 사항을 소개합니다. 이러한 변경으로 연속 시간 CM을 전례없이 대규모로 훈련할 수 있게 되었으며, ImageNet 512x512에서 15억 개의 매개변수에 도달했습니다. 우리가 제안하는 훈련 알고리즘은 단 두 개의 샘플링 단계만 사용하여 CIFAR-10에서 2.06, ImageNet 64x64에서 1.48, ImageNet 512x512에서 1.88의 FID 점수를 달성하여, 최고의 기존 확산 모델과 FID 점수 간의 격차를 10% 이내로 좁히고 있습니다.

English

Consistency models (CMs) are a powerful class of diffusion-based generative models optimized for fast sampling. Most existing CMs are trained using discretized timesteps, which introduce additional hyperparameters and are prone to discretization errors. While continuous-time formulations can mitigate these issues, their success has been limited by training instability. To address this, we propose a simplified theoretical framework that unifies previous parameterizations of diffusion models and CMs, identifying the root causes of instability. Based on this analysis, we introduce key improvements in diffusion process parameterization, network architecture, and training objectives. These changes enable us to train continuous-time CMs at an unprecedented scale, reaching 1.5B parameters on ImageNet 512x512. Our proposed training algorithm, using only two sampling steps, achieves FID scores of 2.06 on CIFAR-10, 1.48 on ImageNet 64x64, and 1.88 on ImageNet 512x512, narrowing the gap in FID scores with the best existing diffusion models to within 10%.