Diffusion Models Adapt to Low-Dimensional Structure Under Flexible Coefficient Choices

기술 요약: 유연한 계수 선택 하에서의 저차원 구조에 대한 확산 모델의 적응성

문제 정의
확산 모델은 학습된 스코어 함수(score function)에 의해 유도되는 반복적인 디노이징 업데이트를 통해 가우시안 노이즈를 점진적으로 데이터로 변환함으로써 고차원 데이터 생성에서 놀라운 경험적 성공을 입증해 왔다. 최근의 이론적 발전은 확산 모델이 알려지지 않은 저차원 데이터 구조에 적응하여, 샘플링 반복 복잡도가 주변 차원 $d$가 아닌 고유 차원 $k$에 의존하도록 할 수 있음을 입증했지만, 이러한 보장은 매우 제한적으로 규정된 업데이트 계수에 국한되어 왔다. 기존의 수렴 이론은 주로 역방향 시간 확률 미분 방정식(reverse-time SDE)의 계수와 정확히 일치하는 특정 이산화 방식(예: 드리프트 $\eta_t$ 및 분산 $\sigma_t^2$의 특정 선택)에 의존한다. 이는 근본적인 질문을 제기한다: 저차원 구조로의 적응은 확산 모델의 강건한(robust) 속성인가, 아니면 업데이트 계수의 정밀한 선택에 민가한(sensitive) 것인가? 서로 다른 계수 선택을 가진 샘플러들이 유사한 성능을 보이는 등 경험적 관찰은 강건함을 시사하지만, 이에 대한 통합된 이론적 정당화는 부족한 실정이었다.

방법론 및 프레임워크
저자들은 다음과 같은 일반적인 업데이트 규칙으로 정의되는 광범위한 클래스의 확산 샘플러를 분석하기 위한 통합된 이론적 프레임워크를 제안한다:
$$Y_{t-1} = \frac{1}{\sqrt{\alpha_t}} \left( Y_t + \eta_t s_t(Y_t) + \sigma_t Z_t \right)$$
여기서 $(\alpha_t)$는 노이즈 스케줄이며, $(\eta_t, \sigma_t^2)$는 업데이트 계수이다.

계수 선택의 유연성을 다루기 위해, 저자들은 다음과 같이 정의된 보조 순방향 과정(auxiliary forward process) $(\hat{X}_t)_{t=1}^T$를 도입한다:
$$\hat{X}_t = \sqrt{\alpha_t} X_0 + \sqrt{v_t} Z_t$$
여기서 $v_t$는 캐노니컬(canonical) 순방향 분산 $(1-\alpha_t)$가 아니라 업데이트 계수 $(\eta_t, \sigma_t^2)$에 의해 재귀적으로 결정되는 "유효 노이즈 분산"을 나타낸다. 유효 분산과 캐노니컬 분산 사이의 불일치는 노이즈 편차 시퀀스 $\delta_t = \frac{v_t}{1-\alpha_t} - 1$로 정량화된다.

분석의 핵심은 샘플러 출력 $Y_1$과 보조 주변 분포(auxiliary marginal) $\hat{X}_1$ 사이의 총 변동(Total Variation, TV) 거리를 바운딩하는 것이다. 저자들은 오차를 다음과 같이 분해한다:

1. 이산화 오차(Discretization Error): 참 스코어 함수를 추적하는 보조 시퀀스를 이용한 결합(coupling) 논증을 통해 분석된, 역방향 역학의 유한 단계 근사에서 발생하는 오차.
2. 스코어 추정 오차(Score Estimation Error): 참 스코어 $s_t^\star$ 대신 추정된 스코어 $s_t$를 사용함으로써 발생하는 오차.

이 분석은 데이터 지지 집합(support)의 커버링 수(covering numbers)를 통해 정의되는 고유 차원 $k$의 개념을 활용하며, 데이터가 저차원 매니폴드 상에 있거나 그 근처에 있다고 가정한다.

주요 기여 및 결과

1. 저차원 적응의 강건성:
   본 논문은 광범한 클래스의 업데이트 계수에 대해, $\varepsilon$-정확한 샘플을 생성하는 데 필요한 반복 횟수가 주변 차원 $d$와 무관하게 $\tilde{O}(k/\varepsilon)$임을 증명한다. 이는 $\delta_t$가 충분히 작은 상수 $c$에 대해 $|\delta_t| \leq c \frac{\log^2 T}{T}$를 만족하는 경우에 성립한다. 이 결과는 저차원 구조로의 적응이 특정 이산화 방식에 국한되지 않는 강건한 속성임을 보여준다.

2. 기존 보장의 일반화:
   이 프레임워크는 다음과 같이 널리 사용되는 여러 샘플러에 적용된다:

• 원형 DDPM (Ho et al., 2020)
• Improved DDPM (Nichol and Dhariwal, 2021)
• Analytic-DPM (Bao et al., 2022)
• $\eta_t \approx 1-\alpha_t$이고 $\sigma_t^2$가 합리적인 범위 내에서 변하는 다양한 기타 계수 선택들.
  따리 1(Corollary 1)은 이러한 방법들이 $\tilde{O}(k/T)$ 수렴율을 충족하는 조건을 만족함을 확인하며, 이를 통해 증명 가능한 저차원 보장을 갖는 샘플러의 범위를 크게 확장한다.

3. 결정론적 샘플러 처리:
   정확한 스코어 함수를 가정할 때(예: DDIM과 같이 $\sigma_t = 0$인 경우), 정리 2(Theorem 2)는 결과를 결정론적 샘플러로 확장한다. 이는 가우시안 노이즈가 주입되지 않더라도 이산화 오차가 고유 차원 $k$에 적응함을 보여준다.

4. 캐노니컬 주변 분포와의 화해:
   논문은 계수로 인한 분산 불일치 $\delta_t$ 때문에 보조 주변 분포 $\hat{X}_1$이 캐노컬 주변 분포 $X_1$과 다를 수 있다는 점을 다룬다. 따리 2(Corollary 2)는 모든 유효한 계수 선택에 대해, 수정된 노이즈 스케줄 $(\alpha'_t)$가 존재하여 샘플러의 정규화된 출력이 동일한 $\tilde{O}(k/T)$ 비율로 타겟 분포 $X_1$에 수렴함을 입증한다. 이는 실용적인 샘플러들이 최적의 비율을 달성하기 위해 특수한 계수 선택으로 교체될 필요가 없으며, 오히려 샘플러의 출력이 재스케일링된 타겟에 대해 이론적 보장이 성립함을 의미한다.

5. 하한(Lower Bounds)과의 비교:
   저자들은 계수가 완벽하게 일치하지 않으면 TV 거리가 $d$에 따라 스케일링될 수 있음을 시사하는 기존의 1단계 이산화 오차 하한과 자신들의 발견을 대조한다. 본 논문은 (계수로 인한 분산 불일치를 흡수하는) 보조 주변 분포를 추적함으로써, 훨씬 더 광범위한 클래스의 샘플러에 대해 고유 차원 적응형 보장을 설정할 수 있다고 주장한다.

의의 및 주장
본 논문은 서로 다른 계수 선택을 가진 확산 샘플러들이 구조화된 고차원 데이터에서 잘 작동한다는 경험적 관찰에 대한 이론적 정당성을 제공한다고 주장한다. 저차원 적응이 계수 선택에 대해 강건하다는 것을 증명함으로써, 이 연구는 다음과 같은 성과를 거둔다:

• 확산 모델 수렴에 대한 이론적 이해를 실질적으로 확장한다.
• ("완벽하게" 조정된 계수를 가진 모델을 넘어) 다양하고 흔히 사용되는 샘플러의 사용을 정당화한다.
• 계수의 유연한 이산화 하에서도 노이즈 수준이 캐노니컬 스케줄과 가깝다면, 고유 차원 $k$에 대한 의존성이 확산 메커니즘의 근본적인 속성임을 명확히 한다.

저자들은 본 결과가 정확하게 사전 학습된 스코어 함수의 가용성에 의존하며, 스코어 추정 오차가 보조 과정에 대해 평가된다는 점을 언급한다. 또한 결정론적 샘플러에 대해 계수 강건성과 스코어 추정 오차를 동시에 포착하는 통합된 이론의 개발을 향후 과제로 남겨둔다.