약한 확산 모델을 훈련 목표 외삽에 활용해 추론 비용 없이 성능을 높인다
SWIFT: Self-guided Weak-to-strong dIFfusion models via Target extrapolation
무엇인가
이 논문은 확산 모델 학습이 계산적으로 비싸고, 추론 시 가이던스를 쓰면 샘플링 비용이 더 늘어나는 문제를 다룬다. 기존 autoguidance는 약한 디노이저를 샘플링 중에 함께 평가해 강한 모델의 예측을 외삽하는데, SWIFT는 같은 약-to-강 외삽을 추론이 아니라 훈련 목표에 옮겨 최종 디노이저 자체가 보정을 학습하게 하고 약한 모델은 버리는 것을 목표로 한다.
어떻게 동작하나
방법은 단순히 회귀 목표를 바꾸는 것이다. 표준 학습이 잡음 입력 x_t에서 깨끗한 x0를 직접 예측하도록 훈련하는 반면, SWIFT는 y_alpha = x0 + alpha(x0 - D_weak(x_t, sigma_t))를 목표로 쓴다. alpha=0이면 기존 학습과 같고, alpha>0이면 약한 디노이저의 예측에서 멀어지는 방향으로 정답을 넘어서 외삽한다. D_weak는 같은 과제로 미리 학습한 뒤 고정한 약한 디노이저이며, 기본 설정은 주 모델의 약 4분의 1 파라미터와 절반 학습 예산을 쓴다. 주 모델 학습 중에는 역전파 없이 한 번의 순전파만 추가되고, 학습 후에는 버려져 원래 샘플러와 단일 디노이저로 추론한다. 조건부·무조건부 생성 모두 라벨 드롭아웃이나 무조건부 분기를 추가하지 않아도 된다.
무엇과 다른가
이론적으로는 고정된 잡음 수준에서 RKHS의 유한 시간 population gradient flow를 분석한다. 조기 종료한 약한 가이드를 쓰면 alpha가 0보다 큰 비어 있지 않은 구간에서 SWIFT의 posterior-mean 오차가 동일 시간만큼 학습한 표준 모델보다 엄격히 낮아질 수 있음을 증명한다. 목표 자체는 posterior mean에서 의도적으로 편향되지만, 유한 시간 학습에서 남아 있는 미학습 성분을 상쇄한다는 설명이다. 다만 이 정리는 가이드 준비 비용을 제외한 동일 주 모델 최적화 시간의 population denoising MSE에 대한 것이며 FID 개선을 직접 보장하지 않고, 모든 약한 디노이저에 대해 성립하는 보장도 아니다.
어떻게 쓰나
실험은 CIFAR-10과 ImageNet-256에서 진행됐고, ImageNet-256은 DiffusionBench를 참고한 19개 구성에서 평가됐다. RAE/DINOv2-B 특징에서는 80에폭 후 무가이드 FID가 2.09에서 1.75로, 5개 특징 평균 FDr이 4.13에서 3.69로 낮아졌다. RAEv2/DINOv3-L에서는 160에폭 예산 안에서 무가이드 FID가 1.56에서 1.15로 개선됐고, SWIFT 값은 140에폭에서 선택됐다. 잠재 표현과 픽셀 표현을 아우르는 18개 DiffusionBench 구성에서 하나의 고정 레시피가 모든 경우의 무가이드 FID를 16.3%에서 49.4%까지 낮췄고 중앙값은 31.3%였다. 가장 낮은 FID를 가진 세 베이스라인도 16.3%에서 17.4% 개선됐다.
전제와 한계
추론 시 가이던스와도 경쟁하지 않고 보완된다. 사용 가능한 최적 가이던스 설정에서 SWIFT는 18개 DiffusionBench 구성 모두의 guided FID를 낮췄다. RAEv2에서는 내부 가이던스 FID가 1.04에서 0.99로, 무가이드 FID가 1.56에서 1.15로 좋아졌다. autoguidance와 결합한 5개 구성에서는 SWIFT가 무가이드 베이스라인을 개선하고, 추론 가이던스가 그 위를 다시 개선하며, SWIFT에 맞춰 가이던스 세기를 재조정했을 때 가장 좋은 FID가 나왔다. RAEv2의 선호 내부 가이던스 세기는 1.78에서 1.20으로 낮아졌는데, 저자들은 훈련 중 보정과 추론 시 보정이 일부 겹치지만 완전히 대체되지는 않는다고 해석한다.
계산량 측면에서 SWIFT는 고정 가이드 때문에 에폭당 약 9% FLOPs를 더 쓰고, 가이드 사전 학습에 대략 베이스라인 10에폭에 해당하는 비용이 든다. 이 비용을 모두 포함해도 DINOv2-B에서 SWIFT 120에폭이 320에폭까지 관찰한 베이스라인 최고치 1.53을 넘었고, 그 베이스라인 체크포인트에 도달하는 데 필요한 FLOPs의 약 0.45배만 썼다. RAEv2에서는 50에폭 SWIFT가 160에폭 베이스라인의 1.56보다 좋은 1.34를 기록했고, 가이드 준비를 포함해 0.44배 FLOPs가 들었다. DINOv2-B를 240에폭까지 늘린 비교에서는 20에폭부터 240에폭까지 36번의 평가 모두에서 SWIFT가 앞섰고, 격차는 0.12에서 0.26 FID로 커졌다. 240에폭에서 훈련 중 평가 FID는 1.63에서 1.38로 개선됐고, 320에폭까지 간 별도 베이스라인은 최고 1.53으로 SWIFT가 120에폭에 도달한 1.50보다 나빴다.
하이퍼파라미터 선택도 보고한다. alpha에 대한 성능은 대체로 U자형이다. CIFAR-10에서는 0.7에서 0.9 근처의 양수 값이 넓은 유리 구간을 형성했고, ImageNet DINOv2-B에서는 0.5에서 1.0이 베이스라인보다 크게 좋았다. 다만 최적값은 모델에 따라 달라서 폭넓은 실험에서는 튜닝 없이 alpha=0.4를 고정 기본값으로 썼다. 가이드는 학생보다 약하지만 유능해야 한다. CIFAR-10에서 alpha=0.8일 때 채널 기본값을 학생의 128에서 64로 줄이면 FID가 2.58에서 1.60으로 좋아졌지만, 훨씬 작은 가이드는 빠르게 나빠졌다. 반폭 가이드가 테스트한 alpha 전반에서 가장 좋았고, 가이드 학습 길이는 일정 역량을 넘으면 비교적 둔감했으며 ImageNet에서는 더 일찍 고정하는 것이 도움이 되기도 했다. 공동 학습 가이드는 테스트한 비교에서 뚜렷한 최종 이점이 없었고, 일부 용량 전환이나 시간·잡음 수준별 계수는 CIFAR-10에서 개선을 줬지만 다른 스케줄은 이득을 잃었다.
개발자 관점에서 SWIFT는 추론 시 별도 가이드 모델이나 추가 샘플링 평가 없이 단일 확산 모델의 품질을 올리고 싶을 때 시도할 만한 훈련 목표 변경이다. 기존 학습 루프에서 목표만 바꾸고 약한 체크포인트를 고정해 한 번 더 순전파하면 되므로 구현 부담은 작지만, alpha와 가이드 용량을 데이터셋·모델별로 확인해야 하고 가이드 사전 학습 비용까지 포함해 time-to-quality를 봐야 한다. 저자들이 밝힌 한계는 이론이 population denoising MSE에 대한 것이고 FID 개선을 직접 보장하지 않는다는 점, 모든 약한 디노이저에서 성립하지 않는다는 점, 가이드가 너무 작으면 성능이 급격히 나빠진다는 점, alpha 최적값이 모델 의존적이라는 점, 음수 alpha는 초기 FID 하락을 앞당길 수 있어도 일관된 최종 이득이 없다는 점이다. 또한 공동 학습 가이드의 이점은 불명확하고 일부 CIFAR-10 세부 스케줄 차이는 반복 실험으로도 해소되지 않았다. 구현은 게재 후 공개 예정이라고만 밝혔고 원문에 공개 URL은 제시되지 않았다.