SWIFT: Self-guided Weak-to-strong dIFfusion models via Target extrapolation

Training high-quality diffusion models is computationally expensive, and inference-time guidance often adds further sampling cost. In this work, we introduce Swift, a simple yet effective self-guided training method that uses a frozen weak denoiser to improve a stronger model without adding inference overhead. For each noisy input, Swift extrapolates the clean regression target away from the weak denoiser's prediction, amplifying components that remain underfit. The weak guide can be an early training checkpoint, requires no backpropagation during main-model training, and is discarded afterward. The resulting model retains the original sampler and supports both conditional and unconditional generation without label dropout or an unconditional branch. We explain this mechanism through a spectral analysis of finite-time population gradient flow in an RKHS. With an early-stopped weak guide, we prove that a nonempty range of extrapolation strengths yields strictly lower posterior-mean error than matched standard training, despite shifting the regression target away from the posterior mean. Experiments on CIFAR-10 and ImageNet-256 demonstrate improvements across model scales and representations. In one ImageNet-256 configuration, Swift surpasses the baseline's best FID using roughly half the total training compute, including guide preparation and training overhead. On RAEv2, it achieves FID 1.15 without inference-time guidance and 0.99 with guidance. These results establish target extrapolation as an effective approach to weak-to-strong diffusion training with efficient single-model inference.

한 줄 요약

SWIFT는 추론 과부하 없이 약한 디노이저를 활용해 강한 디퓨전 모델을 학습하는 자기 가이드 방식으로, ImageNet-256에서 FID 0.99를 달성한다.

핵심 기여도

**SWIFT**는 추론 시 가이드 없이, 학습 중에 고정된 약한 디노이저를 사용해 강한 모델을 훈련한다.
**타겟 엑스트라폴레이션**(target extrapolation)을 통해 약한 모델의 예측 방향을 따라 정답을 더 멀리 이동시켜 학습 효과를 높인다.
**ImageNet-256**에서 기존 기준 대비 약 50% 적은 컴퓨트로 FID 개선을 달성하며, RAEv2에서 가이드 없이 FID 1.15, 가이드 시 0.99를 달성.
**α = 0.4**의 엑스트라폴레이션 강도에서 후향 평균 오차가 기존 학습보다 낮음을 이론적으로 증명.

핵심 아이디어

SWIFT는 디퓨전 모델 학습에서 **추론 과부하 없이 약한 모델을 활용**하는 새로운 접근법이다. 기존 방법은 추론 시에 약한 모델을 가이드로 사용하지만, SWIFT는 학습 단계에서 약한 디노이저를 고정하고, 그 예측 방향을 따라 정답을 더 멀리 이동시켜 학습 타겟을 재설정한다. 이는 **후향 평균**(posterior mean)에서 벗어난 타겟을 사용하면서도, 학습 결과가 더 정확해질 수 있다는 통찰에 기반한다.

이론적으로는 **유한 시간 내 학습**(finite-time training) 상황에서, 약한 모델의 예측 방향을 따라 타겟을 이동시키면, **후향 평균 오차**(posterior-mean error)가 기존 학습보다 낮아질 수 있음을 증명(정리 4.1). 이는 **RKHS**(재현 커널 힐베르트 공간) 내 유한 시간 인구 그라디언트 흐름의 스펙트럼 분석을 통해 설명된다.

기술적 접근법

**SWIFT**는 **고정된 약한 디노이저**(weak denoiser)를 사용하며, 이는 훈련 중 **백프로파게이션 없이 고정**된다.
각 노이즈 입력에 대해, **정답 타겟을 약한 모델의 예측 방향으로 이동**(extrapolate)시켜 학습.
**α = 0.4**의 엑스트라폴레이션 강도를 사용하며, 이는 실험적으로 최적의 성능을 보인 값.
**CIFAR-10**, **ImageNet-256**, **RAEv2** 데이터셋에서 실험.
**DiffusionBench** 기준으로, **VAE**, **RAE**, **end-to-end**, **pixel-space** 등 다양한 디퓨전 표현 방식에서 적용 가능.
**가이드 모델**(guide model)은 주 모델의 훈련 시간의 절반만 학습한 초기 체크포인트를 사용.

주요 결과

**ImageNet-256**에서 SWIFT는 기존 기준 대비 약 50% 적은 컴퓨트로 **FID 개선**을 달성.
**RAEv2**에서 **가이드 없이 FID 1.15**, **가이드 시 FID 0.99**를 달성.
**CIFAR-10**과 **ImageNet-256**에서 **모든 DiffusionBench 설정**에서 **FID 개선**이 지속됨.
**α = 0.4**의 엑스트라폴레이션 강도에서 **후향 평균 오차 감소**가 이론적으로 증명됨.
**가이드 스케일**(guidance scale)은 기준 모델에서 1.78에서 SWIFT 후 1.20으로 감소.

의의 및 한계

SWIFT는 **단일 모델 추론**(single-model inference)을 유지하면서도, **약한 모델을 활용한 학습 개선**을 가능하게 하므로, **디퓨전 모델의 학습-추론 효율성 균형**에 기여한다. 특히, **가이드 없이도 높은 성능**을 보이는 점은 실용적 활용성을 높인다.

하지만, **타겟 엑스트라폴레이션**은 **정확한 방향성**(directional reference)에 의존하므로, **약한 모델의 예측 방향이 부정확하면 성능 저하**가 발생할 수 있다. 또한, **모든 데이터셋에서 동일한 α 값이 최적**이라는 보장은 없으며, **데이터셋별 최적화가 필요**할 수 있다.

실용적 활용

SWIFT는 **이미지 생성**(ImageNet-256, CIFAR-10) 및 **비디오, 오디오 생성**에 적용 가능한 **단일 모델 기반 디퓨전 학습 기법**이다. **가이드 없이도 높은 FID 성능**을 유지하므로, **실시간 생성**(real-time generation)이나 **자원 제한 환경**(low-resource setting)에서 유용하다. **모델 훈련 시간 절감**(약 50%)과 **추론 과부하 제거**는 산업적 배포에 적합한 기술로 평가된다.