[논문 리뷰] Three Creates All: You Only Sample 3 Steps

Yuren Cai, Guangyi Wang외 4명|arXiv (Cornell University)|2026. 03. 23.

Generative Adversarial Networks and Image Synthesis인용 수 0

한 줄 요약

MTEO는 확산 모델에서 단일 전역 시간 조건화를 대체하기 위해 소형 층별 시간 임베딩을 학습시키며, 추가 추론 비용 없이 최소한의 학습 가능 매개변수로 고품질의 소수 단계 샘플링(3–6 NFE)을 가능하게 한다.

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.

연구 동기 및 목표

• 추론 시 단계를 줄여 품질을 해치지 않으면서 확산 샘플링을 더 빠르게 만들도록 동기를 부여한다.
• 소수 단계 영역에서의 기존 단일 시간 조건화의 한계를 식별한다.
• 경량화된 훈련 가능한 다층 시간 임베딩 프레임워크(MTEO)를 제안하여 궤적을 증류한다.
• 다양한 백본 및 데이터세트에서 3–6 NFE에 걸쳐 최첨단 성능을 입증한다.

제안 방법

• 확산 백본 전반에서 시간 조건화와 FiLM 모듈링이 어떻게 사용되는지 분석한다.
• 샘플링 단계마다 층별 시간 임베딩을 할당하는 다층 시간 임베딩(MTEO)을 도입한다.
• 사전 학습된 확산 백본을 고정하고 소수의 층별 시간 임베딩만 학습한다.
• 고충실도 교사 궤적에 대한 궤적 증류를 통해 임베딩을 학습한다.
• 학습 오버헤드를 최소화하기 위해 초기 종료를 포함한 효율적이고 단계별 학습 스케줄을 사용한다.

실험 결과

연구 질문

• RQ1 층별 시간 조건화가 공유된 전역 시간 임베딩보다 소수 단계 확산 샘플링을 개선할 수 있는가?
• RQ2 층마다 시간 조건화가 해제될 때 FiLM 모듈의 용량을 얼마나 활용할 수 있는가?
• RQ3 MTEO를 사용할 때 학습 오버헤드와 샘플링 품질 간의 트레이드오프는 무엇인가?
• RQ4 MTEO가 3–6 NFE의 최첨단 성능을 달성하면서 추론 시 무오버헤드를 유지하는가?
• RQ5 다른 스텝 수와 백본에 걸쳐 MTEO 임베딩의 전달 가능성은 어느 정도인가?

주요 결과

• MTEO는 다수의 백본 및 데이터세트에서 3–6 NFE 구간에서 최첨단 성능을 달성한다.
• 층별 시간 임베딩은 층 간 조건화를 분리하여 추론 비용을 증가시키지 않으면서 더 풍부한 FiLM 모듈화를 가능하게 한다.
• FiLM은 적절한 층별 모듈화가 주어지면 중간 표현을 보정할 충분한 용량을 가진다.
• 매개변수의 극소수만 학습하고(<0.2%), 궤적 증류를 사용하면 증류 기반 방법에 비해 상당한 학습 효율 이득을 얻는다.
• MTEO는 다른 스텝 수에 대해 여전히 전달 가능하며 더 무거운 증류 방법과의 격차를 크게 좁힐 수 있다.