기술 요약: 확산 모델의 노이즈 스케줄 설계: 최적 제어 관점

문제 제기

확산 모델은 생성 과정을 학습된 스코어 함수에 의해 주도되는 시간 역전 확률 미분 방정식 (SDE) 의 시뮬레이션으로 형식화하여, 근사 샘플링을 위한 선도적인 프레임워크로 부상했습니다. 훈련 목적 (스코어 매칭) 은 잘 이해되고 있지만, 전방 및 역방향 과정을 지배하는 함수 $f(t)$와 $g(t)$인 노이즈 스케줄의 설계는 여전히 대부분 휴리스틱에 의존하고 있습니다.

기존 이론적 연구는 특정, 종종 일정한 노이즈 스케줄에 대해 샘플링 오차 상한 (클룰백 - 라이블러 발산으로 측정) 을 확립했습니다. 그러나 이러한 상한은 일반적으로 데이터 차원 $d$와 문제 상수에 대해 비최적으로 스케일링되며, 종종 $\tilde{O}(d^2/n)$ 또는 $\tilde{O}(\max\{J^\star, d\}/n)$로 나타납니다. 여기서 $J^\star$는 데이터 분포의 피셔 정보이고 $n$은 이산화 단계 수입니다. $J^\star$가 $d$보다 훨씬 클 수 있으므로, 이러한 보장은 심각한 문제 의존적 페널티를 겪습니다. 반면, $\tilde{O}(d/n)$의 최첨단 이론적 상한이 존재하지만, 이는 실제에서 거의 배포되지 않는 특정 노이즈 스케줄에 의존하여 확산 모델 이론과 경험적 구현 사이의 간극을 만듭니다.

방법론

저자들은 노이즈 스케줄 설계를 최적 제어 문제로 재구성하는 원칙적인 프레임워크를 제안합니다. 핵심 혁신은 확률 밀도 진화 (포커 - 플랑크 편미분 방정식에 의해 지배됨) 의 처리 불가능한 최적화를 우회하고 피셔 정보 궤적에 초점을 맞추는 데 있습니다.

1. 최적 제어 형식화:

   • 상태: 보흐너 공식을 통해 유도된 1 차 상미분 방정식 (ODE) 에 따라 진화하는 확산 과정의 피셔 정보 $J_t = J(p_t)$:
     $$\dot{J}_t = 2f(t)J_t - g(t)\|\nabla s_t\|_{L^2}^2$$
   • 제어: 노이즈 스케줄 함수 $(f, g)$.
   • 목적: KL 샘플링 오차의 상한이 되는 피셔 정보를 포함하는 함수를 최소화합니다. 저자들은 기르사노프 정리, 보흐너 공식, 엔트로피 파워의 오목성, 그리고 블라치만 - 스탐 부등식을 포함한 확률 및 정보 이론 도구를 사용하여 이 상한을 유도합니다.

2. 폐쇄형 스케줄 유도:

   • 목표 데이터 분포가 **강한 로그 볼록 (SLC)**이거나 **가우시안 혼합 모델 (GMM)**이라는 추가 가정 하에, 저자들은 최적 노이즈 스케줄에 대한 폐쇄형 특성을 유도합니다.
   • 그들은 $f(t)$와 $g(t)$가 아핀 관계를 만족하는 스케줄 클래스인 **아핀 결합 스케줄 (Affine-Coupled Schedules, ACS)**을 도입합니다: $f(t) = \theta g(t) + \omega$.
   • 결과적으로 얻어진 ODE 를 풀면 $(\theta, \omega, g_0, \lambda^*)$로 매개변수화된 $f(t)$와 $g(t)$에 대한 명시적 표현이 도출됩니다.

3. 적응형 매개변수화:

   • 이 프레임워크는 유도된 매개변수를 계산 예산 (신경망 함수 평가, NFE, $n$으로 표기) 에 따라 조정할 수 있는 하이퍼파라미터로 취급합니다.
   • 최적 감쇠율 $\lambda^*$는 람베르트 W 함수를 통해 $n$에 따라 동적으로 스케일링되는 것으로 나타나며, 이를 통해 스케줄이 초기화 오차와 이산화 오차 사이의 균형을 자동으로 조정할 수 있습니다.

주요 기여

1. 최적 제어 프레임워크: 본 논문은 상태가 피셔 정보 궤적인 최적 제어 문제로 노이즈 스케줄 설계를 형식화합니다. 이는 처리 불가능한 편미분 방정식 (PDE) 제약을 처리 가능한 상미분 방정식 (ODE) 으로 대체하여 최적 오차 상한에 대한 충분 조건을 유도할 수 있게 합니다.
2. $\tilde{O}(d/n)$ 오차 상한: 저자들은 최적 제어 문제의 해가 KL 발산에서 $\tilde{O}(d/n)$의 샘플링 오차 상한을 달성함을 증명합니다. 이는 실제에서 사용되는 스케줄을 포함한 광범위한 노이즈 스케줄 클래스로 최첨단 이론적 보장을 확장합니다.
3. 아핀 결합 스케줄 (ACS): 그들은 표준 경험적 스케줄을 일반화하는 새로운 매개변수화 노이즈 스케줄 클래스 (ACS) 를 유도합니다:
   • VE-Exponential: 매개변수 $(\theta, \omega) = (0, 0)$으로 설정하여 복원됩니다.
   • VP-Sigmoid: $(\theta, \omega) = (1/2, 0)$으로 설정하여 복원됩니다.
   • ACS 프레임워크는 이러한 고정된 휴리스틱을 넘어 체계적인 최적화를 가능하게 하는 추가 조정 가능 매개변수 ($g_0, c$) 를 도입합니다.
4. 경험적 검증: CIFAR-10 및 ImageNet(64x64) 에 대한 실험은 ACS 하이퍼파라미터를 체계적으로 조정하면 다양한 NFE 예산에 걸쳐 표준 베이스라인 (VP-Linear, VP-Cosine, VP-Sigmoid) 보다 우수한 프레이트 인셉션 거리 (FID) 점수를 산출함을 보여줍니다.

결과

• 이론적: 본 논문은 SLC 및 GMM 분포에 대해 제안된 ACS 프레임워크가 최적의 $\tilde{O}(d/n)$ 수렴 속도를 달성함을 확립합니다. 또한 데이터 분포가 매우 이방성 ($J^\star \gg d$) 일 때 특히 비최적의 절충을 겪을 수 있는 상수 또는 선형 스케줄 (실제에서 일반적) 이 왜 발생하는지 분석합니다.
• 경험적: 무조건부 CIFAR-10 과 클래스 조건부 ImageNet 에서 적응형 ACS 프레임워크는 일관되게 표준 베이스라인을 능가합니다. 예를 들어, ImageNet 에서 NFE 가 50 일 때 ACS 는 5.30의 FID 를 달성하여 VP-Linear(7.46), VP-Cosine(5.67), VP-Sigmoid(6.77) 보다 우수한 성능을 보였습니다. 이러한 결과는 ACS 가 계산 예산에 효과적으로 적응하여 곡률과 분산 궤적을 조정하는 반면, 고정된 스케줄은 단계 크기가 변함에 따라 품질이 저하됨을 나타냅니다.

중요성

본 논문은 노이즈 스케줄 설계를 위한 원칙적인 방법을 제공함으로써 확산 모델 이론과 실무 간의 간극을 해소한다고 주장합니다. 경험적 스윕이나 고정된 휴리스틱에 의존하는 대신, 저자들은 특정 스케줄이 작동하는 이유에 대한 이론적 정당성과 새로운 최적화된 스케줄을 생성하는 메커니즘을 제공합니다.

그 중요성은 다음과 같습니다:

1. 일반화: 표준 경험적 스케줄 (지수, 시그모이드) 이 더 넓고 이론적으로 최적인 클래스 (ACS) 의 특수한 경우임을 보여줍니다.
2. 적응성: 정적 스케줄을 사용하는 대신 사용 가능한 계산 예산 ($n$) 과 고유한 데이터 특성 (피셔 정보) 에 기반하여 스케줄을 동적으로 조정하는 방법을 제공합니다.
3. 성능: 이론적으로 근거한 이 접근 방식이 표준 벤치마크에서 생성 품질 (FID 점수) 의 실질적인 향상으로 이어짐을 입증하여, 생성 모델링에서 최적 제어 관점의 유용성을 검증합니다.

저자들은 폐쇄형 유도가 SLC 또는 GMM 분포를 가정하지만, 결과적으로 얻어진 스케줄은 복잡한 실제 데이터에서도 잘 작동한다고 지적하며, 이는 이론적 통찰력이 유도를 위해 가정한 특정 가정을 넘어 일반화됨을 시사합니다. 향후 작업은 이러한 특성을 더 넓은 범주의 복잡한 데이터 분포로 확장하는 것으로 식별됩니다.