One Step Diffusion via Shortcut Models

Diffusion models and flow-matching models have enabled generating diverse and realistic images by learning to transfer noise to data. However, sampling from these models involves iterative denoising over many neural network passes, making generation slow and expensive. Previous approaches for speeding up sampling require complex training regimes, such as multiple training phases, multiple networks, or fragile scheduling. We introduce shortcut models, a family of generative models that use a single network and training phase to produce high-quality samples in a single or multiple sampling steps. Shortcut models condition the network not only on the current noise level but also on the desired step size, allowing the model to skip ahead in the generation process. Across a wide range of sampling step budgets, shortcut models consistently produce higher quality samples than previous approaches, such as consistency models and reflow. Compared to distillation, shortcut models reduce complexity to a single network and training phase and additionally allow varying step budgets at inference time.

한 줄 요약

Shortcut 모델은 단일 네트워크와 학습 단계로 1단계 생성이 가능한 확산 기반 생성 모델이다.

핵심 기여도

Shortcut 모델은 단일 네트워크와 단일 학습 단계로 1단계 또는 다단계 생성이 가능하다.
기존의 2단계 디스틸레이션 방법 대비 16% 이하의 추가 연산량만으로 동일 또는 우수한 생성 성능을 달성한다.
CelebA-HQ와 ImageNet-256 데이터셋에서 1단계 생성 성능이 기존 2단계 방법과 유사하거나 더 우수하다.
로봇 제어 분야에도 적용 가능하며, 10배 이하의 추론 비용으로 유사한 성능을 유지한다.

핵심 아이디어

Shortcut 모델은 기존 확산 모델이 반복적인 노이즈 제거를 통해 생성하는 방식과 달리, 생성 과정에서 특정 단계를 건너뛸 수 있도록 설계되었다. 이는 네트워크가 현재 노이즈 수준뿐 아니라 **원하는 스텝 크기(step size)**에 대해서도 조건을 받도록 함으로써 가능하다. 이 조건을 통해 모델은 생성 과정에서 "단축"을 수행할 수 있으며, 이는 **self-bootstrapping** 방식으로 학습된다. 즉, 작은 스텝 크기에서 데이터에 정착된 모델이, 더 큰 스텝 크기에서 자기 자신을 기반으로 학습하게 된다. 이는 별도의 디스틸레이션 단계 없이 단일 네트워크와 단일 학습 과정으로 가능하다는 점에서 혁신적이다.

기술적 접근법

**Shortcut 모델**은 단일 네트워크와 단일 학습 단계로 구성된다.
학습 시, **step-size conditioned model**로, 작은 스텝 크기에서 데이터에 정착한 후, 더 큰 스텝 크기에서 self-consistency loss를 통해 학습한다.
**L2 regression**을 사용하며 별도의 학습 스케줄이 필요 없다.
기존 확산 모델 대비 약 16%의 추가 연산량만으로 학습 가능.
**CelebA-HQ**, **ImageNet-256** 데이터셋에서 실험 수행.
**로봇 제어** 분야에도 적용 가능하며, 확산 정책 대신 shortcut 정책을 사용하여 추론 비용을 10배 이상 줄임.

주요 결과

**CelebA-HQ**와 **ImageNet-256** 데이터셋에서 1단계 생성 성능이 기존 2단계 디스틸레이션 방법과 유사하거나 더 우수함.
다단계 생성 시, 기존 확산 모델과 동일한 수준의 생성 품질 유지.
로봇 제어 분야에서 10배 이하의 추론 비용으로 유사한 성능 유지.
**16% 이하의 추가 연산량**으로 단일 네트워크 학습 가능.

의의 및 한계

Shortcut 모델은 기존 확산 모델의 반복적 추론 문제를 해결하면서도, 복잡한 학습 절차 없이 단일 네트워크로 생성 가능하다는 점에서 학술적·실용적 가치가 있다. 특히, 기존 디스틸레이션 방법 대비 학습 절차가 단순하며, 추론 시 스텝 수를 유연하게 조절할 수 있다. 그러나, 이 모델은 노이즈-데이터 매핑이 데이터셋의 기대값에만 의존하므로, GAN 또는 VAE와 같은 다른 생성 모델에 비해 학습 문제를 단순화하는 방식이 제한적이다. 또한, 1단계 생성 품질과 다단계 생성 품질 간의 격차는 여전히 존재한다.

실용적 활용

Shortcut 모델은 이미지 생성뿐 아니라 **로봇 제어**와 같은 실시간 추론이 필요한 분야에 적용 가능하다. 특히, 추론 비용이 낮은 환경에서 기존 확산 모델의 성능을 유지하면서도 추론 속도를 크게 향상시킬 수 있다. 산업 현장에서의 빠른 이미지 생성, 실시간 제어 정책 생성 등에 활용 가능하다.