Three Creates All: You Only Sample 3 Steps

TL;DR

This paper introduces MTEO, a method that distills layer-wise time embeddings to enable fast, high-quality diffusion sampling with only three steps, without increasing inference time.

Contribution

MTEO is a novel, plug-and-play approach that distills small, layer-specific time embeddings, significantly improving few-step diffusion sampling performance.

Findings

Achieves state-of-the-art results in few-step sampling.

Narrowed the gap between distillation and lightweight methods.

No additional inference overhead introduced.

Abstract

Diffusion models deliver high-fidelity generation but remain slow at inference time due to many sequential network evaluations. We find that standard timestep conditioning becomes a key bottleneck for few-step sampling. Motivated by layer-dependent denoising dynamics, we propose Multi-layer Time Embedding Optimization (MTEO), which freeze the pretrained diffusion backbone and distill a small set of step-wise, layer-wise time embeddings from reference trajectories. MTEO is plug-and-play with existing ODE solvers, adds no inference-time overhead, and trains only a tiny fraction of parameters. Extensive experiments across diverse datasets and backbones show state-of-the-art performance in the few-step sampling and substantially narrow the gap between distillation-based and lightweight methods. Code will be available.