- 디퓨전 모델은 이미지 생성을 넘어 오디오·비디오·3D·단백질 설계·로봇 경로 계획처럼 다봉 분포 샘플링이 필요한 문제에 쓰이며, 이 튜토리얼은 최적화 관점으로 학습과 샘플링을 연결함
- 학습 과정은 데이터에 노이즈를 섞은 (x_\sigma=x_0+\sigma\epsilon)를 만들고, 신경망 (\epsilon_\theta(x,\sigma))가 노이즈 방향을 예측하도록 평균제곱오차를 최소화함
- 학습된 denoiser는 데이터 집합 (\mathcal{K})로의 근사 투영으로 해석되며, 이상적 denoiser는 (\sigma)-스무딩된 제곱 거리 함수의 그래디언트와 연결됨
- DDIM 샘플링은 (f(x)=\frac{1}{2}\mathrm{dist}_{\mathcal{K}}(x)^2)에 대한 근사 그래디언트 하강으로 볼 수 있고, (\sigma_t) 스케줄이 반복 횟수와 denoiser 평가 비용을 결정함
- 그래디언트 추정 업데이트와 노이즈 추가를 결합하면 DDIM, DDPM, 저자들의 개선 샘플러를
gam·mu파라미터로 함께 다룰 수 있으며, toy 모델과 latent diffusion 예제로 이어짐
최적화 관점의 디퓨전 모델
- 디퓨전 모델은 다봉 분포에서 샘플을 생성하는 데 강점을 보이며, Stable Diffusion 같은 텍스트-이미지 생성 도구뿐 아니라 오디오, 비디오, 3D 생성, 단백질 설계, 로봇 경로 계획에도 적용됨
- 튜토리얼의 이론적 기반은 ICML 2024 논문과 관련 논문의 최적화 해석임
- 구현은
smalldiffusion을 주로 참조하며, 본문 코드는 원 라이브러리보다 교육용으로 단순화됨
학습: 노이즈 방향 예측
- 디퓨전 모델은 학습 예시에서 데이터 집합 (\mathcal{K})를 배우고, 그 집합에서 샘플을 생성하는 것을 목표로 함
- 이미지라면 (\mathcal{K} \subset \mathbb{R}^{c\times h \times w})는 현실적인 이미지에 해당하는 픽셀값 집합임
- 같은 틀은 오디오, 비디오, 로봇 궤적, 텍스트 같은 이산 영역에도 적용됨
- 학습 절차는 세 단계로 볼 수 있음
- (x_0 \sim \mathcal{K}), (\sigma), (\epsilon \sim N(0,I))를 샘플링함
- (x_\sigma=x_0+\sigma\epsilon)로 노이즈가 섞인 데이터를 만듦
- (\epsilon_\theta(x_\sigma,\sigma))가 (\epsilon)를 예측하도록 제곱 손실을 최소화함
- 코드에서는
training_loop가 배치x0마다generate_train_sample로sigma와eps를 만들고,model(x0 + sigma * eps, sigma)의 출력과eps사이의 MSE를 최적화함 - (\sigma)는 연속 구간에서 균일 샘플링하기보다 (N)개의 값으로 이산화한 (\sigma) 스케줄에서 뽑음
Schedule클래스는 가능한sigmas목록을 감싸고 학습 중 배치별로 값을 샘플링함- 본문 예제는
ScheduleLogLinear(N, sigma_min=0.02, sigma_max=10)을 사용함 ScheduleDDPM은 픽셀 공간 디퓨전 모델,ScheduleLDM은 Stable Diffusion 같은 latent diffusion 모델용 스케줄임
Swissroll toy 예제
- toy 데이터셋은 초기 디퓨전 논문 중 하나인 Sohl-Dickstein et al. 2015에서 사용한 나선형 점 집합이며, (\mathcal{K}\subset\mathbb{R}^2)임
- 단순 데이터셋에서는 denoiser를 MLP로 구현함
- 입력은 (x\in\mathbb{R}^2)와 (\sigma)의 2차원 임베딩을 이어붙인 값임
- 출력은 노이즈 (\epsilon\in\mathbb{R}^2) 예측값임
- 많은 디퓨전 모델이 (\sigma)에 sinusoidal positional embedding을 쓰지만, 이 예제에서는 단순한 2차원 임베딩도 잘 동작함
- 예제 학습 설정은
ScheduleLogLinear(N=200, sigma_min=0.005, sigma_max=10)과epochs=15000을 사용함 - 학습된 denoiser는 (x-\sigma\epsilon_\theta(x,\sigma))를 그려 벡터장으로 시각화할 수 있음
- (\sigma)가 클 때 denoiser는 데이터 평균을 예측하는 경향이 있음
- (\sigma)가 낮고 입력 (x)가 데이터에 가까우면 실제 데이터 포인트를 예측함
Denoising을 투영으로 해석하기
- 데이터 집합 (\mathcal{K})에 대한 거리 함수는 (\mathrm{dist}_{\mathcal{K}}(x)=\min{|x-x_0|:x_0\in\mathcal{K}})로 정의됨
- (x)의 투영 (\mathrm{proj}_{\mathcal{K}}(x))는 이 거리를 달성하는 (\mathcal{K}) 안의 점 집합임
- (\mathcal{K})가 닫힌 집합이고 (x\notin\mathcal{K})이며 투영이 유일하면, 제곱 거리 함수의 그래디언트는 (x-\mathrm{proj}_{\mathcal{K}}(x))가 됨
- 거리 함수 (\mathrm{dist}_{\mathcal{K}})는 모든 곳에서 미분 가능하지 않기 때문에,
min대신 softmin을 사용해 (\sigma)로 스무딩한 제곱 거리 함수를 도입함 - 스무딩된 거리 함수의 그래디언트는 (x)가 정하는 가중치에 따라 (\mathcal{K})의 점들의 가중 평균 쪽을 향함
이상적 denoiser와 상대 오차 모델
- 이상적 denoiser (\epsilon^*)는 특정 (\sigma)에서 학습 손실을 정확히 최소화하는 denoiser임
- 데이터가 유한 집합 (\mathcal{K}) 위의 이산 균일 분포라면, 이상적 denoiser는 닫힌형 식으로 표현됨
- 각 데이터 포인트의 가중치는 (x_\sigma)와 해당 점 사이의 거리에 따라 정해짐
- 작은 데이터셋에서는
IdealDenoiser로 직접 계산 가능함
- toy 데이터에서 이상적 denoiser는 (\sigma)가 클 때 데이터 평균을 향하고, (\sigma)가 작을 때 가장 가까운 데이터 포인트를 향함
- 핵심 정리는 모든 (\sigma>0), (x\in\mathbb{R}^n)에 대해 (\frac{1}{2}\nabla_x \mathrm{dist}^2_{\mathcal{K}}(x,\sigma)=\sigma\epsilon^*(x,\sigma))라는 관계를 둠
- 상대 오차 모델은 (x-\sigma\epsilon_\theta(x,\sigma))가 (\mathrm{proj}_{\mathcal{K}}(x))를 잘 근사하는 조건을 사용함
- (\sqrt{n}\sigma)가 (\mathrm{dist}_{\mathcal{K}}(x))를 상수 배 안에서 잘 추정할 때 적용됨
- 오차는 (\eta\mathrm{dist}_{\mathcal{K}}(x)) 이하로 제한된다고 가정함
- 낮은 노이즈에서는 manifold hypothesis 아래 추가 노이즈 대부분이 데이터 다양체에 직교하므로 denoising이 투영을 근사함
- 높은 노이즈에서는 (\sigma)가 (\mathcal{K})의 지름보다 크면 데이터의 가중 평균을 예측하는 denoiser도 작은 상대 오차를 가짐
- CIFAR-10은 이상적 denoiser 계산이 가능한 크기이며, 실험에서 샘플링 궤적 위의 정확한 투영과 이상적 denoiser 출력 사이 상대 오차가 작게 나타남
샘플링: 반복 denoising과 DDIM
- 학습된 denoiser가 있으면, 노이즈가 섞인 (x_t)와 노이즈 수준 (\sigma_t)에서 (\hat{x}0^t=x_t-\sigma_t\epsilon\theta(x_t,\sigma_t))로 (x_0)를 예측함
- 시작점은 (\sigma_T)를 (\mathcal{K})의 지름에 비해 크게 잡고, (x_T)를 (N(0,\sigma_T))에서 독립 샘플링해 (\mathcal{K})에서 멀리 떨어지게 함
- 높은 노이즈에서 한 번의 denoiser 호출은 상대 오차가 작아도 절대 오차가 클 수 있으며, 이상적 denoiser의 예측은 데이터 평균에 가까움
- 그래서 샘플링은 (\sigma_t) 스케줄을 따라 denoiser를 반복 호출해 (x_T,\ldots,x_0) 시퀀스를 만듦
- 업데이트 (x_{t-1}=x_t-(\sigma_t-\sigma_{t-1})\epsilon_\theta(x_t,\sigma_t))는 좌표 변환을 거친 결정론적 DDIM 샘플링 알고리듬과 같음
- DDIM과의 동등성 증명은 논문 Appendix A에 있음
거리 최소화로 본 DDIM
- DDIM은 (f(x)=\frac{1}{2}\mathrm{dist}_{\mathcal{K}}(x)^2)에 대한 근사 그래디언트 하강으로 해석됨
- 스텝 크기는 (1-\sigma_{t-1}/\sigma_t)임
- (\nabla f(x_t))는 (\epsilon_\theta(x_t,\sigma_t))로 추정됨
- (\sigma_t) 스케줄은 샘플링 중 그래디언트 스텝 수와 크기를 결정함
- 스텝이 너무 적으면 (\mathrm{dist}_{\mathcal{K}}(x_t))가 줄지 않아 수렴하지 않을 수 있음
- 작은 스텝을 많이 쓰면 denoiser 평가 횟수가 늘어나 계산 비용이 커짐
- admissible schedule은 각 반복에서 (\sqrt{n}\sigma_t)가 (\mathrm{dist}_{\mathcal{K}}(x_t))와 상수 배 안에서 맞도록 하는 스케줄임
- 기하급수적으로 감소하는 log-linear (\sigma_t) 시퀀스는 admissible schedule임
- 정리에 따르면 DDIM으로 생성된 (x_t)에서 (\nabla\mathrm{dist}{\mathcal{K}}(x))가 존재하고 (\mathrm{dist}{\mathcal{K}}(x_T)=\sqrt{n}\sigma_T)이면, (x_t)가 제곱 거리 함수의 그래디언트 하강으로 생성되고 (\mathrm{dist}_{\mathcal{K}}(x_t)/\sqrt{n}\approx\sigma_t)가 유지됨
- toy 예제에서는 원래 log-linear 스케줄에서 부분 샘플링해 20개 스텝의 DDIM 샘플러를 구현하며, 대부분의 샘플이 원 데이터에 가깝지만 개선 여지가 남아 있음
그래디언트 추정 기반 개선 샘플러
- (\nabla\mathrm{dist}{\mathcal{K}}(x))가 (x)와 (\mathrm{proj}{\mathcal{K}}(x)) 사이에서 불변이라는 점을 이용해, 현재 추정과 이전 추정을 섞는 업데이트를 사용함
- 업데이트 (\bar{\epsilon}t=\gamma\epsilon\theta(x_t,\sigma_t)+(1-\gamma)\epsilon_\theta(x_{t+1},\sigma_{t+1}))는 이전 스텝의 오차를 현재 추정으로 보정하는 방식임
- toy 모델 샘플에서는 이 방식이 DDIM보다 빠르게 수렴하고, 샘플이 원 데이터에 더 가까워짐
- DDIM과 비교하면 이 샘플러는 모멘텀을 추가한 해석이 가능하며, 궤적이 overshoot할 수 있지만 더 빠르게 수렴할 수 있음
- 생성 과정에서 노이즈를 추가하면 샘플링 품질이 경험적으로 개선됨
- 원래 (\sigma_t) 스케줄을 유지하려면 더 작은 (\sigma_{t'})까지 denoise한 뒤 (w_t\sim N(0,I)) 노이즈를 다시 추가함
- (\mu=\frac{1}{2})일 때 DDPM sampler를 정확히 복원함
- 전체 업데이트 (x_{t-1}=x_t-(\sigma_t-\sigma_{t'})\bar{\epsilon}_t+\eta w_t)는 세 샘플러를 일반화함
- DDIM:
gam=1, mu=0 - DDPM:
gam=1, mu=0.5 - 그래디언트 추정 샘플러:
gam=2, mu=0
- DDIM:
더 큰 모델과 참고 자료
- 앞의 학습 코드는 toy 데이터뿐 아니라 이미지 디퓨전 모델을 처음부터 학습하는 데도 사용 가능함
- FashionMNIST 예제는 FashionMNIST 데이터셋에서 학습해 Papers with Code 리더보드의 FID 기준 2위 점수를 얻는 예제로 제공됨
- 샘플링 코드는 수정 없이 사전학습된 latent diffusion 모델에도 사용 가능함
- 예시는
ScheduleLDM(1000)과ModelLatentDiffusion('stabilityai/stable-diffusion-2-1-base')를 사용함 - 텍스트 조건은
An astronaut riding a horse로 설정하고, 50개 (\sigma) 스텝으로 샘플링한 뒤 latent를 디코딩함
- 예시는
- (\gamma) 모멘텀 항의 효과는 고해상도 텍스트-이미지 생성에서 비교 시각화됨
- 추가로 볼 만한 자료
- What are diffusion models: Markov process를 되돌리는 이산 시간 관점의 디퓨전 모델 소개
- Generative modeling by estimating gradients of the data distribution: 확률 미분 방정식을 되돌리는 연속 시간 관점의 디퓨전 모델 소개
- The annotated diffusion model: PyTorch 디퓨전 모델 구현 상세 설명