기술 요약: 메트로폴리스 조정 확산 모델 (MADM)

문제 제기

점수 기반 확산 모델로부터 표본을 추출하는 과정은 본질적으로 두 가지 원인으로 인한 편향을 수반합니다. 하나는 학습된 점수 함수의 근사 오차이고, 다른 하나는 수치 적분기에 의해 도입된 시간 이산화 오차입니다. 표준 예측자 - 교정자 (Predictor-Corrector) 프레임워크에서는 예측자 단계 (종종 확률 흐름 ODE) 가 노이즈 수준 간에 표본을 이동시킨 후, 현재 노이즈 수준에서 표본을 정제하기 위해 교정자 단계 (일반적으로 수정되지 않은 랑주뱅 알고리즘, ULA) 를 수행합니다.

ULA 는 점수 함수를 사용하여 이동을 제안하지만, 연속적인 랑주뱅 확산 과정을 이산화하기 때문에 그 자체로 편향된 표본 추출기입니다. 결과적으로 예측자 - 교정자 방식의 "교정자"는 자체적인 이산화 편향을 도입하여 목표 주변 분포를 불변으로 유지하지 못합니다. 마르코프 연쇄 몬테 카를로 (MCMC) 에서 이러한 편향을 보정하는 표준 메트로폴리스 - 헤이스팅스 (MH) 또는 바커 수용 - 거부 메커니즘은 확산 모델에 직접 적용할 수 없습니다. 이는 MH 가 목표 밀도 비율 $p_t(\tilde{x})/p_t(x)$를 평가해야 하는데, 확산 모델에서는 점수 함수 $s_\phi(x, t) \approx \nabla_x \log p_t(x)$만 이용 가능하여 밀도 비율을 계산하는 것이 불가능하기 때문입니다.

방법론

저자들은 편향된 ULA 교정자를 수용 - 거부 보정 랑주뱅 단계로 대체하는 메트로폴리스 조정 확산 모델 (MADM) 프레임워크를 제안합니다. 핵심 혁신은 밀도 비율이 필요하지 않고 점수 함수만을 사용하여 필요한 수용 확률을 계산하는 방법을 개발한 것입니다.

1. 밀도 비율 항등식

저자들은 목표 밀도 비율의 로그가 현재 상태 $x$와 제안된 상태 $\tilde{x}$ 사이의 경로 따라 점수 함수의 선적분으로 표현될 수 있음을 입증했습니다:
$$\log \left( \frac{p_t(\tilde{x})}{p_t(x)} \right) = \int_0^1 \langle s(x + u(\tilde{x} - x), t), \tilde{x} - x \rangle \, du$$
이 항등식을 통해 밀도 비율을 점수 평가만을 사용하여 추정할 수 있습니다.

2. 베르누이 공장을 통한 정확한 보정

밀도를 평가하지 않고 정확한 보정을 달성하기 위해, 저자들은 두 개의 동전 베르누이 공장 (two-coin Bernoulli factory) (Gonçalves et al., 2017) 에 기반한 알고리즘을 제안합니다.

• 가정: 선적분 피적분 함수의 절대값이 $C$로 제한되는 상한 $C(x, \tilde{x})$가 존재합니다.
• 메커니즘: 알고리즘은 포아송 절단 멱급수를 사용하여 적분 (상수로 스케일링됨) 의 지수 함수에 대한 불변 추정치 $W$를 구성합니다.
• 절차: 알고리즘은 고정된 확률 $\alpha'$로 먼저 거부하는 "두 개의 동전" 절차를 사용하여 바커 수용 - 거부 결정을 시뮬레이션합니다. 거부되지 않은 경우, 확률 $W$로 동전을 던집니다. 동전이 성공하면 제안이 수용되고, 그렇지 않으면 라운드가 다시 시작됩니다.
• 결과: 이 절차는 점수가 정확하고 상한 가정이 성립하는 경우, 목표 분포 $p_t$를 불변으로 유지하는 정확한 바커 조정 랑주뱅 교정자를 생성합니다.

3. 심슨의 법칙을 통한 효율적인 근사

정확한 베르누이 공장 방법이 무작위적이고 잠재적으로 많은 수의 점수 평가를 필요로 한다는 점을 인식한 저자들은 확장 가능한 근사법을 제안합니다.

• 방법: 선적분을 심슨 1/3 법칙 (뉴턴 - 코츠 구적법) 을 사용하여 결정론적으로 근사합니다. 이를 위해서는 현재 상태, 제안된 상태, 그리고 중점에서 점수를 평가해야 합니다.
• 수용: 생성된 추정치를 표준 메트로폴리스 - 헤이스팅스 수용 확률에 대입합니다.
• 특성: 이는 구적법 절단으로 인해 작은 편향을 도입하지만, 근사는 단계 크기 $h$에 대해 $O(h^{5/2})$ 차수의 정확도를 달성합니다. 계산 비용이 저렴하여 제안당 중점 점수 평가 한 번만 추가로 필요합니다.

주요 기여

1. MADM 프레임워크: 확산 모델을 위한 메트로폴리스 조정 교정자 프레임워크를 도입하여 편향된 ULA 단계를 체계적으로 대체합니다.
2. 정확한 방법: 두 개의 동전 베르누이 공장을 사용하여 확산 모델에서 랑주뱅 보정을 조정하는 첫 번째 정확한 방법을 제시하며, 점수 함수와 선적분의 상한만 필요합니다.
3. 확장 가능한 근사: 고차 정확도 ($O(h^{5/2})$) 를 달성하면서도 거의 추가 비용이 들지 않는 (점수 평가 한 번 추가) 실용적인 심슨 법칙 근사법입니다.
4. 실증적 검증: 합성 데이터셋과 이미지 데이터셋 모두에서 향상된 표본 품질을 입증했습니다.

결과

저자들은 카라스 (Karras) 등 (2022) 의 사전 훈련된 모델을 사용하여 합성 2D 데이터셋과 표준 이미지 벤치마크 (CIFAR-10, FFHQ, AFHQv2, ImageNet-64) 에서 MADM 을 평가했습니다.

• 합성 데이터셋: 2D 데이터셋 (나선, 깔때기, 시에르핀스키 삼각형, 핀휠) 에서 표준 조상 표본 추출기 (ULA 교정자가 포함된 예측자와 동일) 는 실제 데이터 분포의 지지 영역 밖에서 상당한 이상치를 가진 표본을 생성했습니다. MADM 교정자 단계는 이러한 표본을 성공적으로 데이터 지지 영역으로 되돌려 이상치를 제거했습니다.
• 이미지 데이터셋:
  • 무조건부 모델 (CIFAR-10, FFHQ, AFHQv2) 의 경우, 예측자 단계당 하나의 MADM 교정자 단계를 적용하면 결정론적 기준선 (Heun 및 오일러 ODE 솔버) 에 비해 프레체트 인셉션 거리 (FID) 에서 일관되지만 modest 한 개선을 보였습니다.
  • 클래스 조건부 ImageNet-64 모델의 경우, 개선 효과가 더 뚜렷했습니다. 오일러 예측자와 함께 MADM 을 사용한 것이 보정 없이 고차 Heun 예측자를 사용한 것보다 더 큰 FID 개선을 보였습니다. 이는 단계 수가 많을 때 ODE 솔버의 차수를 단순히 높이는 것보다 MH 조정이 이산화 오차를 완화하는 데 더 효과적임을 시사합니다.
• 계산 비용: 정확한 방법 (알고리즘 2) 은 필요한 점수 평가 횟수가 무작위적이기 때문에 계산 비용이 매우 높은 것으로 나타났습니다. 따라서 실험은 주로 일관된 개선을 제공하면서도 관리 가능한 오버헤드를 가진 심슨 법칙 근사 (알고리즘 3) 를 활용했습니다.

중요성과 주장

본 논문은 ULA 교정자의 본질적인 편향을 해결하기 위해 확산 모델에서 랑주뱅 보정을 조정하는 첫 번째 정확한 방법을 제공한다고 주장합니다. 점수 함수를 활용하여 선적분을 통해 밀도 비율을 추정함으로써, MADM 은 접근 불가능한 데이터 밀도에 의존하지 않고도 생성 모델링에서 엄격한 MCMC 수용 기준을 사용할 수 있게 합니다.

저자들은 MADM 을 표본 품질을 향상시키는 예측자 - 교정자 프레임워크의 자연스러운 확장으로 위치시킵니다. 그들은 정확한 방법이 이론적으로 중요하지만, 심슨 법칙 근사가 이미지 벤치마크에서 일관되게 FID 를 개선하는 실용적이고 확장 가능한 해결책을 제공한다고 지적합니다. 이 작업은 완벽한 점수 추정을 가정하며, 이는 확산 모델의 고유한 한계이지 특정 방법의 결함이 아님을 인정합니다. 향후 연구는 정확한 방법이나 그 서브루틴의 효율적인 변형을 개발하는 데 초점을 맞출 것으로 제안됩니다.