[Gen AI] Diffusion 모델 샘플링 & 학습 트릭 정리

1. 샘플링(Sampling) 방법

1.1 DDIM

DDIM(Denoising Diffusion Implicit Models)은 DDPM의 stochastic sampling을 deterministic 방식으로 바꿔 적은 step으로도 고품질 샘플을 생성할 수 있게 한다.

• DDPM은 noise를 거슬러 올라갈 때 매 step에 randomness가 들어가지만, DDIM은 deterministic trajectory를 따라간다.
• 즉, DDIM은 한 스텝에 더 크게 노이즈를 벗겨내도 trajectory가 깨지지 않아서 더 적은 스텝으로도 좋은 품질의 이미지를 만들 수 있는 것이다.
• DDPM이 보통 1000 step을 필요로 하는데 DDIM은 50~100 step 정도로도 좋은 품질을 낼 수 있다.
• 예시: Stable Diffusion에서 num_inference_steps=50로 sampling 속도를 크게 줄이는 데 활용.

 

1.2 PNDM / LMSD / DPM-Solver

DDPM/DDIM에서 sampling을 할 때는 일반적으로 Euler step 같은 단순한 방식으로 noisy한 trajectory를 거슬러 올라간다. 하지만 이를 더 빠르고 고품질로 만들기 위해 multi-step ODE solver를 적용하는 것이 PNDM, LMSD, DPM-Solver이다.

• PNDM (Pseudo Numerical Methods for Diffusion Models)
  • DDIM의 single-step 방식 대신, Adams-Bashforth 계열의 multi-step solver를 써서 sampling path를 더 정확히 추정한다.
  • 2~4 step 이전의 noise prediction(gradient) 정보를 같이 사용해 현재 샘플링에 반영함으로써 trajectory를 보정한다.
• LMSD (Linear Multi-Step DPM)
  • 비슷하게 linear multi-step ODE solver를 적용해 안정성과 속도를 모두 잡는다.
  • 특히 timestep 수를 줄여도 blur나 artifact 없이 깨끗한 샘플을 뽑는 데 유리하다.
• DPM-Solver
  • adaptive step-size 방식을 도입해 timestep 수를 동적으로 조정하며 더 빠르고 안정적인 샘플링을 수행한다.
  • 이 때문에 DDIM을 50 step으로 돌릴 때보다 훨씬 적은 step (20~30 step)으로 비슷하거나 더 좋은 품질을 낼 수 있다.

PNDM, LMSD, DPM-Solver는 DDIM에서 더 나아간 multi-step ODE solver 방식으로, 같은 조건에서 더 적은 step으로 빠르고 깨끗하게 샘플을 생성할 수 있어 지금은 거의 표준처럼 사용된다. 대신 timestep schedule과 solver 하이퍼파라미터를 적절히 조절해야 최적 품질을 낼 수 있다는 trade-off가 있다.

 

1.3 Classifier-free Guidance (CFG)

조건 생성(class-conditional, text-conditional)에서 classifier를 별도로 학습시키지 않고도 샘플을 조건에 맞게 유도(guidance)하는 방식이다.

• 원래는 조건 생성에서 ∇_x log p(class|x) 형태의 classifier gradient를 사용해 샘플링을 유도했지만, 이 방식은 별도의 classifier를 따로 학습해야 하는 부담이 있었다.
• Classifier-free Guidance(CFG)는 이를 대체하기 위해, 같은 모델에서 조건이 있는 prediction과 조건이 없는 prediction을 둘 다 수행해 간단히 결합하는 방식이다.
• 이렇게 조건이 있는 prediction과 조건이 없는 prediction을 같이 학습하는 이유는, 나중에 샘플링(inference) 단계에서 조건이 있는 ε(x_t|c)와 없는 ε(x_t)를 동시에 사용해 guidance를 적용하기 위해서이다.
• 즉, 조건이 없는 branch도 함께 학습함으로써
  • 하나의 모델로 조건을 더 강하게 넣거나, 혹은 약하게 넣는 다양한 샘플을 유연하게 만들 수 있다.
  • 별도의 classifier gradient를 쓰지 않으면서도 원하는 정도로 condition을 강조하는 조절 가능한 생성이 가능해진다.

구체적으론 두 noise prediction을 계산한다.

• ε(x_t|c) : 조건(e.g class, text)이 주어진 noise prediction
• ε(x_t) : 조건을 주지 않은 unconditional noise prediction
• 이 둘을 다음과 같은 식으로 결합해 guidance scale w를 조정한다.
• 즉, 무조건적인 예측을 빼면서 조건에 의존하는 정도를 w배로 강조하는 것이다.

예를 들어 text-to-image에서 guidance scale=7.5는 text condition을 7.5배 더 강조한 샘플을 만드는 것과 같다.

• w=1 정도는 조건을 적당히 반영 (약한 guidance)
• w=7.5는 조건을 매우 강하게 반영 (보통 text2image에서 많이 씀)
• w=15는 오히려 불안정하거나 부자연스러운 결과를 낼 수 있다. (over guidance)

이러한 CFG를 사용하면 별도의 classifier를 따로 훈련할 필요가 없으므로 모델 개발이 훨씬 단순해진다. 또한 w를 쉽게 바꿔가며 guidance 강도를 실험적으로 조정할 수 있어, 연구나 서비스에서 prompt tuning을 빠르게 시도해볼 수 있다.

 

2. 학습 트릭과 안정화 기법

2.1 Noise Schedule & Beta Schedule

DDPM에서 forward process란, 깨끗한 이미지 x_0에 점점 noise를 추가해 x_t로 변형시켜 가는 과정이다.
이때 각 timestep에서 얼마나 noise를 추가할지를 결정하는 것이 바로 Noise Schedule (또는 Beta Schedule) 이다.

• Linear, Cosine, Exponential 등 다양한 beta schedule이 제안되었다.
  • Linear: 일정한 속도로 β_t를 증가시킨다. 구현은 간단하지만, 초기 step에서 noise가 너무 빠르게 증가해 이미지 정보를 잃을 수 있다.
  • Cosine (DDPM Improved, Nichol & Dhariwal): β_t를 천천히 증가시켜, 초기 구간에서 이미지 정보를 오래 보존하고 학습 안정성을 높인다.
  • Exponential: 앞부분에서 느리고 뒷부분에서 빠르게 증가한다. 일부 실험에서 성능 향상이 있지만 일반적으로 잘 쓰이지는 않는다.

EDM에서의 SNR 기반 설계

EDM(Elucidating the Design Space of Diffusion Models)은 β schedule 대신 SNR(Signal-to-Noise Ratio) 를 직접 설계하는 방식을 도입했다.

SNR(t) = (signal power) / (noise power) = (α(t)^2) / (σ(t)^2)

• SNR은 위와 같이 정의되며, timestep이 진행될수록 점진적으로 감소한다.
• EDM에서는 noise scale σ(t)를 power-law decay로 설계해 SNR이 일정한 기울기로 줄어들도록 설계했다.
• 이를 통해 timestep별 loss variance를 일정하게 유지해 학습을 안정화할 수 있다.

즉, EDM은 "시간에 따라 노이즈를 얼마나 넣을까?" 대신 "시그널과 노이즈의 비율(SNR)이 timestep에서 어떻게 줄어들어야 하는지"를 직접 설계해 학습 안정성과 효율을 극대화한 것

 

Noise schedule은 diffusion 모델에서 가장 핵심적인 하이퍼파라미터 중 하나이며, 학습의 안정성과 성능, 샘플링 품질에 직결되는 요소다. 특히 Cosine 스케줄과 SNR 기반 설계는 최근 거의 표준처럼 사용되고 있다.

 

2.2 V-Prediction 

기존 DDPM은 forward process에서 noisy image x_t를 주고, 모델이 직접 added noise ε(epsilon) 을 예측하도록 학습했다.

하지만, velocity prediction (v-prediction)은 x_t를 입력받아 velocity v를 예측하도록 학습하는 것이다.

v = (ε * σ) + (x₀ * α)
v = (x_t - α * x₀) / σ

• 여기서 velocity란 위와 같은 weighted sum이다.
• 즉 v는 noise와 data를 같이 섞어 노이즈와 원본 이미지 정보를 동시에 가진 latent velocity 같은 개념이다.
• 이런 방식을 사용하면,
  • sigma parameterization이 훨씬 안정적으로 작동한다.
  • timestep (noise scale)이 극단적으로 작거나 크더라도 loss variance가 일정해져 학습이 균일해진다.
  • 특히 고해상도 이미지나 다양한 β 스케줄을 쓸 때도 robust하게 동작한다.

 

 

기존에는 noisy image에서 noise만 직접 예측했지만, EDM 이후 velocity v를 예측함으로써 noise scale(sigma)에 강건하고 loss variance를 일정하게 유지할 수 있게 되었다. 

 

2.3 SNR weighting 

SNR weighting은 EDM에서 제안된 학습 안정화 기법으로, timestep 별로 loss를 조절해, 특정 노이즈 레벨에만 과도하게 치우치지 않고 전 구간에서 균형있게 학습이 되도록 만드는 기법이다.

 

✅ 기존 Diffusion 모델의 문제점

• diffusion 모델은 timestep이 작을 땐 이미지가 거의 원본이라 loss가 작게 나오고, timestep이 클 땐 완전 noise라 의미 있는 gradient가 잘 안 나온다.
• 그래서 특정 구간에만 loss가 집중되거나, timestep별로 gradient variance가 크게 달라져 학습이 불안정해진다.

✅ SNR weighting 개선

• timestep t에서 signal-to-noise ratio (SNR) 를 계산해, SNR에 기반한 weight를 loss에 곱해준다.
• SNR이 큰 깨끗한 구간에서는 loss를 작게 만들어서 과도하게 학습되지 않게 하고,
• SNR이 작은 noisy한 구간에서는 loss를 더 크게 주어 충분히 학습되도록 한다.
• 깨끗한 구간은 모델이 이미 잘 reconstruct하기 쉽기 때문에 굳이 loss를 크게 줄 필요가 없고,
• noisy한 구간에서도 잘 작동하도록 loss를 보강해줘야 diffusion process 전체를 안정적으로 학습할 수 있다.

*SNR은 t가 작을때 크고, t가 커질수록 작아진다.

 

2.4 Gradient Clipping & EMA

✅ Gradient Clipping

• diffusion 모델은 timestep마다 loss scale이 달라 gradient exploding이 발생하기 쉽다.
• 특히 고해상도나 large batch로 학습할 때 gradient norm이 순간적으로 너무 커져 optimizer step에서 파라미터가 이상치로 튈 위험이 있다.
• 이를 방지하기 위해 gradient가 특정 threshold를 넘으면 gradient를 비율적으로 축소해 threshold를 넘지 않도록 하는 기법이다.
• 이렇게 하면 optimizer가 파라미터를 너무 크게 업데이트하지 않아 학습이 안정적이다.

 

✅ Exponential Moving Average (EMA)

θ_ema ← 0.999 * θ_ema + 0.001 * θ

 

EMA는 학습 중 모델 파라미터를 부드럽게 평균내는 방식이다. 기본 파라미터 θ를 학습하면서, 동시에 아래처럼 smoothed 파라미터 θ_ema를 업데이트한다.

• 이렇게 하면 noisy한 gradient update의 단기적 요동을 완화해 더 부드러운 파라미터를 얻을 수 있다.
• 실제로 샘플링(inference)에서는 EMA 파라미터를 사용하는 것이 일반적이다.
  • 예를 들어 DDPM, LDM, EDM 코드에서도 마지막 sampling 단계는 θ_ema를 불러와 이미지 생성에 쓴다.