DASH: Dual-Branch Score Distillation for Guidance-Calibrated Compact Diffusion Models

기술 요약: DASH: 가이던스 보정된 컴팩트 확산 모델을 위한 이중 브랜치 스코어 증류(Dual-Branch Score Distillation)

문제 정의
Classifier-Free Guidance(CFG)를 사용하는 클래스 조건부 확산 모델은 조건부($\epsilon_y$) 및 무조건부($\epsilon_\emptyset$) 예측을 모두 생성하기 위해 단일 네트워크에 의존한다. 가이던스 메커니즘은 생성을 유도하기 위해 그 차이인 $\Delta = \epsilon_y - \epsilon_\emptyset$를 증폭시킨다. 그러나 파라미터 축소를 통해 이러한 모델을 압축할 때, 기존의 증류 목적 함수에서 결정적인 구조적 실패가 발생한다.

표준 증류 방식은 합성된 가이드 예측($\tilde{\epsilon}_S \approx \tilde{\epsilon}_T$) 또는 조건부 브랜치만을 감독한다. 이는 **부정-결정 문제(underdetermined optimization problem)**를 야기한다. 즉, 학생 모델의 브랜치들을 구성하는 여러 가지 분해 방식이 동일한 손실 함수를 만족할 수 있으며, 여기에는 가이던스 간극이 붕괴되는($\Delta_S \to 0$) 퇴보적인 해(degenerate solution)가 포함된다. 이러한 상태에서는 조건부 및 무조건부 예측이 수렴하게 되어, 추론 시 가이던스 가중치 $w$가 효과를 잃게 된다. 이는 합성 훈련 손실이 최소화되더라도 발생하는 현상이다. 또한, 기존 방법들은 교사 모델의 학습된 타임스텝별 훈련 커리큘럼(중요도 가중치)을 전이하는 데 실패하며, 이로 인해 학생 모델은 제한된 증류 예산 내에서 이 스케줄을 다시 학습해야 하며, 이는 종종 최적의 훈련 궤적을 손상시킨다.

방법론: DASH 프레임워크
본 논문은 부정-결정을 해결하고 가이던스 보정을 유지하기 위해 두 가지 주요 메커니즘을 갖춘 DASH (Dual-Branch Score Distillation) 프레임워크를 제안한다.

1. 이중 브랜치 독립 감독 (Dual-Branch Independent Supervision):
   합성된 출력을 매칭하는 대신, DASH는 별도의 손실 항을 사용하여 두 스코어 브랜치를 독립적으로 감독한다:

   • 모방 손실 ($L_{im}$): 조건부 브랜치를 매칭함 ($\epsilon^y_S \approx \epsilon^y_T$).
   • 무조건부 손실 ($L_{un}$): 무조건부 브랜치를 매칭함 ($\epsilon^\emptyset_S \approx \epsilon^\emptyset_T$).
   • 앵커 손실 ($L_{an}$): 초기 훈련 단계에서 드리프트를 방지하기 위해 조건부 브랜치를 실제 노이즈($\epsilon$)에 고정하는 경량화된 정규화 항.

   두 브랜치 모두에 제약을 가함으로써, 이 프레임워크는 각 샘플에 대한 목표 출력을 유일하게 규정하여 가이던스 간극 붕괴를 허용하는 가용 해 공간을 제거한다. 저자들은 보정 상태를 모니터링하기 위해 가이던스 간극 비율 $\rho = E[\|\Delta_S\|] / E[\|\Delta_T\|]$를 정의하며, $\rho \approx 1$은 성공적인 보정을 나타낸다.

2. TIRT 전이 (Timestep-Importance Rebalanced Training Transfer):
   교사 모델은 콘텐츠가 어려운 중간 궤적 영역에 훈련을 집중시키기 위해 타임스텝별 스케일링 인자($\lambda_t$)를 학습하는 Min-SNR 가중치의 확장형인 TIRT를 사용한다. DASH는 교사 모델이 수렴한 타임스텝별 중요도 커리큘럼($\hat{\omega}_T$)을 **동결된 사전 확률(frozen prior)**로서 학생 모델에 직접 복사하는 TIRT 전이를 도입한다. 이는 학생 모델이 예측 매칭의 용이성이 아닌 콘텐츠의 난이도에 따라 보상받는 증류 그래디언트에 의해 커리큘럼을 손상시키는 것을 방지한다.

주요 기여

• 구조적 부정-결정 해결: 본 논문은 파라미터 압축된 확산 모델에서 발생하는 가이던스 간극의 퇴보적 붕괴를 식별하고 해결한다. 저자들은 무조건부 감독이 단순히 유익한 수준을 넘어 필수적이며, 전체 증류 이득의 60% 이상을 차지함을 입증한다.
• 커리큘럼 전이: 학습된 타임스텝 가중치를 동결된 사전 확률로 사용하는 TIRT 전이를 도입하여, 학생 모델이 커리큘럼을 다시 학습할 필요 없이 교사의 학습된 데이터 적응형 난이도 가중치를 활용하도록 한다.
• 이중 브랜치 프레임워크: 조건부 및 무조건부 브랜치를 독립적으로 제약하는 통합된 증류 파이프라인을 제시하며, 이는 압축 시 가이던스 충실도를 유지하는 데 경험적으로 필수적임을 검증하였다.

실험 결과
실험은 교사 모델(35.8M 파라미터)과 압축된 학생 모델(6.1M 파라미터, 5.9배 압축)을 사용하여 CIFAR-10 및 CIFAR-100 데이터셋에서 수행되었다.

• 성능: DASH 학생 모델은 50-step DDIM 샘플링에서 CIFAR-10의 경우 8.87, CIFAR-100의 경우 10.47의 FID를 달성했다. 이는 교사 모델(FID 5.47 / 6.80)과 4 FID 포인트 차이 내에 있으며, 처음부터 학습하는 방식(FID 20.47 / 26.14) 및 단일 제약 베이스라인보다 크게 우수한 성능을 보인다.
• 가이던스 보정: DASH는 CIFAR-10에서 0.91, CIFAR-100에서 0.89의 가이던스 간극 비율 $\rho$를 유지하며, 방향 코사인 $\cos(\Delta)$는 0.94 근처를 기록했다. 반면, 합성 증류 베이스라인은 부분적 붕괴($\rho \approx 0.68$)를 겪으며, 조건부 전용 감독은 거의 완전한 붕괴($\rho \approx 0.09$)로 이어져 처음부터 학습하는 것보다 낮은 성능을 보였다.
• 절제 연구 (Ablation Studies):
  • $L_{un}$(무조건부 손실)을 제거하면 CIFAR-10에서 7.77 포인트의 FID 저하가 발생하고 $\rho$가 0.61로 붕괴되어, 이것이 지배적인 구성 요소임을 확인했다.
  • TIRT 전이를 제거(고정된 Min-SNR으로 대체)하면 2.60 포인트의 FID 저하가 발생하여, 학습된 커리큘럼 전이의 필요성을 입증했다.
  • 교차 데이터셋 성능 비율(DASH의 경우 1.18배 vs 베이스라인의 경우 1.51배)은 레이블 복잡도에 따른 견고한 일반화 성능을 확인시켜 준다.

의의 및 주장
본 논문은 DASH가 기존 증류 목적 함수의 구조적 결함을 해결함으로써 파라미터 효율적인 확산 모델의 새로운 표준을 세운다고 주장한다. 저자들의 주장은 다음과 같다:

1. 무조건부 감독은 타협 불가능하다: 이는 합성 또는 조건부 목적 함수로 대체될 수 없다. 무조건부 감독 없이는 압축 시 CFG가 효과를 잃게 된다.
2. 커리큘럼 전이는 특징 증류와 구별된다: 학습 스케줄(TIRT)을 동결된 사전 확률로 전달하는 것은 교사의 데이터 적응형 난이도 가중치를 손상시키지 않는 참신한 기여이다.
3. 타협 없는 효율성: 이 프레임워크는 상당한 파라미터 감소(5.9배)를 가능하게 하면서도 보정된 가이던스 메커니즘을 보존하여, 추론 단계를 늘리지 않고도 자원이 제한된 환경에서 고품질 생성을 가능하게 한다.

이 연구는 파라미터 압축과 브랜치 수준의 보정 사이의 접점에 위치하며, 단계 감소 방법론과 결합 가능한 직교적인 솔루션을 제공한다.