Consistency Models Made Easy

Consistency models (CMs) offer faster sampling than traditional diffusion models, but their training is resource-intensive. For example, as of 2024, training a state-of-the-art CM on CIFAR-10 takes one week on 8 GPUs. In this work, we propose an effective scheme for training CMs that largely improves the efficiency of building such models. Specifically, by expressing CM trajectories via a particular differential equation, we argue that diffusion models can be viewed as a special case of CMs. We can thus fine-tune a consistency model starting from a pretrained diffusion model and progressively approximate the full consistency condition to stronger degrees over the training process. Our resulting method, which we term Easy Consistency Tuning (ECT), achieves vastly reduced training times while improving upon the quality of previous methods: for example, ECT achieves a 2-step FID of 2.73 on CIFAR10 within 1 hour on a single A100 GPU, matching Consistency Distillation trained for hundreds of GPU hours. Owing to this computational efficiency, we investigate the scaling laws of CMs under ECT, showing that they obey the classic power law scaling, hinting at their ability to improve efficiency and performance at larger scales. Our code (https://github.com/locuslab/ect) is publicly available, making CMs more accessible to the broader community.

한 줄 요약

ECT는 기존 CM 학습 시간을 1/4~1/3로 줄이며, 2-step FID 2.73을 1시간 내 A100 단일 GPU에서 달성하는 CM 학습 효율화 기법이다.

핵심 기여도

ECT(Easy Consistency Tuning)를 제안하여, CM 학습 시간을 기존 iCT 대비 33%로 줄임.
CIFAR-10에서 2-step FID 2.73을 1시간 내 A100 단일 GPU에서 달성.
CM의 스케일링 법칙을 최초로 분석, 파워 로우 규칙을 확인.
ECT를 통해 학습된 ECM 모델은 iCT 대비 50%+6.25%의 학습 비용으로 더 높은 FD_DINOv2 성능을 보임.

핵심 아이디어

기존 CM은 일관성 조건을 강하게 만족시키기 위해 복잡한 학습 과정이 필요했으나, ECT는 CM의 일관성 조건을 미분 방정식으로 표현함으로써, 기존의 DM(확산 모델)이 CM의 특별한 경우임을 밝혔다. 이는 기존에 학습된 DM을 기반으로 CM을 fine-tuning하면서 점진적으로 일관성 조건을 강화하는 방식으로 학습을 진행할 수 있음을 의미한다. 이는 기존 CM 학습에서 요구되던 복잡한 하이퍼파라미터 조정과 오랜 학습 시간을 대폭 줄이는 핵심 아이디어이다. ECT는 pretraining과 tuning의 2단계 구조를 가지며, tuning 단계는 전체 학습 비용의 10% 이하로 구성된다.

기술적 접근법

ECT는 미분 방정식을 기반으로 CM의 일관성 조건을 표현.
기존 EDM2 모델을 pretraining 단계로 사용.
tuning 단계에서는 점진적으로 일관성 조건을 강화하며 학습.
c_skip(t) = σ_data² / (t² + σ_data²), c_out(t) = tσ_data / √(t² + σ_data²) 공식 사용.
모델 크기: ECM-S/M/L/XL (280M ~ 1.1B 파라미터), FLOPs: 102G ~ 406G.
tuning 단계는 12.8M 이미지, batch size 128×100k로 학습.

주요 결과

CIFAR-10에서 ECT는 2-step FID 2.73을 1시간 내 A100 단일 GPU에서 달성.
iCT 대비 33%의 학습 비용으로 2-step 성능 개선.
ImageNet 64×64에서 ECM은 FD_DINOv2 152.21을 기록, iCT(205.11)보다 우수.
ECT는 전체 학습 비용의 50%+6.25%로 iCT와 동일한 조건에서 더 높은 성능 달성.

의의 및 한계

ECT는 CM 학습의 계산 효율성을 획기적으로 개선하여, 기존 CM의 학습 시간(8 GPU, 1주일)을 1 GPU, 1시간으로 줄이는 데 성공하였다. 이는 CM의 대규모 실험 및 연구 확산을 촉진할 수 있는 중요한 기여이다. 또한, ECT를 통해 CM의 스케일링 법칙을 최초로 분석하여, 파워 로우 규칙이 적용됨을 밝혀, 모델 크기와 계산량의 관계를 명확히 하였다. 다만, ECT는 기존 DM을 기반으로 학습하기 때문에, DM이 아닌 다른 모델 구조와의 호환성은 아직 검증되지 않았으며, tuning 단계에서의 하이퍼파라미터 선택은 여전히 중요하다는 한계가 있다.

실용적 활용

ECT는 이미지 생성 분야에서 빠른 샘플링이 필요한 산업(예: 콘텐츠 생성, 디자인 자동화)에 적용 가능하며, 연구자들이 CM을 보다 쉽게 실험하고 확장할 수 있도록 도와주는 실용적 도구이다. 특히, ECT의 계산 효율성은 대규모 모델 학습 및 다양한 애플리케이션 개발에 유리한 조건을 제공한다.