이미지 로딩 중...
AI Generated
2025. 11. 8. · 3 Views
AI 이미지 생성 2편 - Diffusion Model 원리
Stable Diffusion, DALL-E의 핵심인 Diffusion Model의 작동 원리를 단계별로 학습합니다. Forward Process, Reverse Process, U-Net 구조, Noise Scheduling까지 실제 코드로 구현하며 이해합니다.
목차
- Diffusion Model 개요
- Forward Process
- Reverse Process
- Noise Scheduling
- U-Net Architecture
- Training Process
- Sampling Methods
- Conditional Generation
1. Diffusion Model 개요
시작하며
여러분이 AI 이미지 생성 도구를 사용해 본 적이 있나요? Stable Diffusion이나 DALL-E 같은 서비스로 텍스트 프롬프트만 입력하면 놀라운 품질의 이미지가 생성되는 것을 보셨을 겁니다.
하지만 내부에서 어떻게 "아무것도 없는 노이즈"에서 "완성된 이미지"를 만들어내는지 궁금하지 않으셨나요? Diffusion Model은 바로 이 마법 같은 일을 가능하게 하는 핵심 기술입니다.
2020년 이후 GAN을 제치고 이미지 생성 분야의 최강자로 떠올랐죠. 이 개념을 이해하면 여러분도 Stable Diffusion을 커스터마이징하거나, 자신만의 이미지 생성 모델을 훈련시킬 수 있습니다.
단순히 도구를 사용하는 것을 넘어 원리를 이해하고 제어하는 수준으로 도약할 수 있습니다.
개요
간단히 말해서, Diffusion Model은 "이미지에 점진적으로 노이즈를 추가하는 과정을 학습한 뒤, 그 역과정을 통해 노이즈에서 이미지를 생성하는 생성 모델"입니다. 왜 이런 복잡한 방식을 사용할까요?
GAN처럼 한 번에 이미지를 생성하는 것보다 훨씬 안정적이고 고품질의 결과를 얻을 수 있기 때문입니다. 예를 들어, 대규모 데이터셋으로 학습할 때 mode collapse(특정 패턴만 반복 생성)나 학습 불안정성 문제가 거의 없습니다.
기존 GAN에서는 Generator와 Discriminator가 경쟁하며 학습했다면, Diffusion Model은 노이즈 제거 과정만 학습하면 됩니다. 훨씬 단순하고 안정적인 학습 목표입니다.
Diffusion Model의 핵심 특징은 크게 세 가지입니다: (1) 점진적 변화(Gradual Transformation) - 수백~수천 단계에 걸쳐 천천히 변환, (2) 확률적 접근(Stochastic Process) - 매 단계가 확률 분포로 정의됨, (3) 역과정 학습(Reverse Process Learning) - 노이즈 추가의 역과정을 신경망으로 근사. 이러한 특징들이 높은 품질과 다양성을 동시에 달성할 수 있게 합니다.
코드 예제
import torch
import torch.nn as nn
class SimpleDiffusionModel:
def __init__(self, num_timesteps=1000):
# 총 디퓨전 스텝 수
self.num_timesteps = num_timesteps
# 베타 스케줄: 노이즈 추가 강도를 점진적으로 증가
self.betas = torch.linspace(0.0001, 0.02, num_timesteps)
self.alphas = 1 - self.betas
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
def add_noise(self, x0, t, noise=None):
# Forward process: 원본 이미지에 노이즈 추가
if noise is None:
noise = torch.randn_like(x0)
sqrt_alphas_cumprod_t = self.alphas_cumprod[t] ** 0.5
sqrt_one_minus_alphas_cumprod_t = (1 - self.alphas_cumprod[t]) ** 0.5
# 공식: x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * noise
return sqrt_alphas_cumprod_t * x0 + sqrt_one_minus_alphas_cumprod_t * noise
설명
이것이 하는 일: Diffusion Model은 크게 두 과정으로 나뉩니다. Forward process에서는 깨끗한 이미지를 점점 노이즈로 만들고, Reverse process에서는 노이즈를 다시 이미지로 복원합니다.
첫 번째로, 초기화 단계에서는 총 몇 단계에 걸쳐 노이즈를 추가할지 결정합니다(보통 1000 스텝). betas는 각 타임스텝마다 얼마나 노이즈를 추가할지 정의하는데, 처음에는 아주 조금(0.0001)씩, 나중에는 더 많이(0.02) 추가하도록 선형으로 증가시킵니다.
왜 이렇게 할까요? 초반에는 이미지 구조를 유지하면서 천천히, 후반에는 빠르게 완전한 노이즈로 만들기 위해서입니다.
그 다음으로, alphas_cumprod를 계산합니다. 이것은 누적곱으로, t번째 스텝까지 노이즈를 추가했을 때 원본 이미지가 얼마나 남아있는지를 나타냅니다.
예를 들어 t=500일 때 0.3이라면, 원본 이미지의 30%만 남고 70%는 노이즈가 된 상태입니다. 마지막으로, add_noise 함수가 실제 노이즈 추가를 수행합니다.
수학적으로 증명된 공식을 사용해서, 중간 스텝을 일일이 거치지 않고도 "t번째 스텝의 상태"를 한 번에 계산할 수 있습니다. 이는 학습 시 효율성을 크게 높여줍니다.
여러분이 이 코드를 사용하면 임의의 타임스텝 t에서의 노이즈가 섞인 이미지를 즉시 생성할 수 있습니다. 학습 시에는 랜덤한 t를 선택해 노이즈를 추가한 뒤, 모델이 그 노이즈를 예측하도록 훈련시킵니다.
추론 시에는 완전한 노이즈(t=1000)에서 시작해 역방향으로 단계별로 노이즈를 제거하며 이미지를 생성합니다.
실전 팁
💡 betas 스케줄은 성능에 큰 영향을 줍니다. Linear 외에도 Cosine, Quadratic 등 다양한 스케줄을 실험해보세요. 최근 연구에서는 Cosine 스케줄이 더 나은 결과를 보이는 경우가 많습니다.
💡 타임스텝 수를 줄이면(예: 50 스텝) 생성 속도가 빨라지지만 품질이 떨어질 수 있습니다. DDIM 같은 가속 샘플링 방법을 사용하면 스텝 수를 10~50으로 줄여도 품질을 유지할 수 있습니다.
💡 alphas_cumprod를 미리 계산해두면 매번 누적곱을 계산할 필요가 없어 학습 속도가 향상됩니다. 실무에서는 이를 buffer로 등록해 GPU 메모리에 상주시킵니다.
💡 노이즈는 표준정규분포에서 샘플링합니다. torch.randn_like를 사용하면 입력과 같은 shape의 노이즈를 쉽게 생성할 수 있습니다. 재현성을 위해 시드를 고정하는 것도 잊지 마세요.
2. Forward Process
시작하며
여러분이 머신러닝 모델을 학습시킬 때 "학습 데이터를 어떻게 만들지"가 가장 중요한 고민이죠? Diffusion Model에서는 깨끗한 이미지를 학습 데이터로 사용하는데, 여기에 특별한 변환을 가합니다.
Forward Process(순방향 과정)는 이미지를 점진적으로 파괴하는 과정입니다. 마치 사진을 햇빛에 오래 방치하면 색이 바래듯이, 이미지에 노이즈를 조금씩 추가해 결국 완전한 잡음으로 만듭니다.
이 과정이 중요한 이유는 나중에 신경망이 이 역과정을 배울 때 "어느 정도 노이즈가 추가된 상태"를 다양하게 경험할 수 있기 때문입니다. 1000단계로 나눠서 천천히 파괴하면, 각 단계별로 어떻게 복원해야 하는지 세밀하게 학습할 수 있습니다.
개요
간단히 말해서, Forward Process는 깨끗한 이미지 x₀에 가우시안 노이즈를 T번 반복해서 추가하여 완전한 노이즈 x_T로 만드는 마르코프 체인입니다. 왜 마르코프 체인이냐고요?
각 단계 x_t는 오직 이전 단계 x_{t-1}에만 의존하기 때문입니다. 예를 들어, t=500 단계의 상태는 t=499의 상태에만 영향을 받고, 그 이전 단계들은 직접적으로 고려하지 않습니다.
이런 구조 덕분에 수학적으로 다루기 쉽고 효율적인 계산이 가능합니다. 기존 데이터 증강에서는 회전, 크롭, 색상 변경 등을 했다면, Diffusion의 Forward Process는 오직 노이즈 추가만 합니다.
하지만 이 단순한 방법이 놀랍게도 강력한 생성 모델을 만드는 기반이 됩니다. Forward Process의 핵심 특징: (1) 결정론적 스케줄 - 노이즈 추가량이 미리 정해진 β_t에 따라 결정됨, (2) 가역성 - 수학적으로는 역과정이 존재 (하지만 실제로는 신경망으로 근사해야 함), (3) 간단한 구현 - 단 하나의 공식으로 임의의 타임스텝 계산 가능.
이러한 특징이 Diffusion Model을 실용적으로 만듭니다.
코드 예제
import torch
import matplotlib.pyplot as plt
def forward_diffusion_demo(image, num_steps=1000):
"""
이미지에 Forward Diffusion을 적용하는 데모
"""
# 베타 스케줄 정의
betas = torch.linspace(0.0001, 0.02, num_steps)
alphas = 1.0 - betas
alphas_cumprod = torch.cumprod(alphas, dim=0)
# 시각화할 타임스텝 선택
timesteps_to_show = [0, 100, 300, 500, 700, 999]
noisy_images = []
for t in timesteps_to_show:
# 각 타임스텝에서의 노이즈 추가
noise = torch.randn_like(image)
sqrt_alpha_bar = alphas_cumprod[t] ** 0.5
sqrt_one_minus_alpha_bar = (1 - alphas_cumprod[t]) ** 0.5
# x_t 계산
noisy_image = sqrt_alpha_bar * image + sqrt_one_minus_alpha_bar * noise
noisy_images.append(noisy_image)
return noisy_images, timesteps_to_show
# 사용 예시
original_image = torch.randn(3, 256, 256) # 예시 이미지
noisy_versions, steps = forward_diffusion_demo(original_image)
print(f"생성된 노이즈 이미지: {len(noisy_versions)}개")
설명
이것이 하는 일: 이 코드는 하나의 이미지가 1000단계에 걸쳐 어떻게 노이즈로 변해가는지 시각적으로 보여줍니다. 마치 타임랩스 사진처럼 각 단계의 스냅샷을 저장합니다.
첫 번째로, 베타 스케줄을 정의합니다. 0.0001에서 0.02까지 선형적으로 증가하는 값들인데, 이는 초반에는 이미지를 거의 유지하면서 미세한 노이즈만 추가하고, 후반으로 갈수록 급격하게 노이즈를 증가시킨다는 의미입니다.
실무에서는 이 값들을 조정해서 생성 품질을 최적화할 수 있습니다. 그 다음으로, alphas_cumprod를 계산합니다.
이것은 수학의 누적곱 개념인데, 예를 들어 t=3일 때는 (1-β₀) × (1-β₁) × (1-β₂)입니다. 이 값이 0.9라면 "원본 이미지의 90%가 남아있다"는 의미입니다.
핵심은 이 누적곱 덕분에 중간 단계를 일일이 거치지 않고도 한 번에 "t단계 후의 상태"를 계산할 수 있다는 점입니다. 세 번째로, 선택된 타임스텝마다 노이즈를 추가합니다.
torch.randn_like로 원본과 같은 크기의 랜덤 노이즈를 생성하고, 수학적으로 증명된 공식 x_t = √(ᾱ_t) × x₀ + √(1-ᾱ_t) × ε을 적용합니다. 이 공식은 "원본의 일부"와 "노이즈의 일부"를 적절한 비율로 섞는 것인데, 시간이 지날수록 노이즈 비율이 커집니다.
마지막으로, 여러 타임스텝의 결과를 리스트로 반환합니다. t=0에서는 거의 원본 그대로, t=100에서는 약간 흐릿하게, t=500에서는 이미지가 많이 손상되고, t=999에서는 완전한 노이즈만 남는 것을 볼 수 있습니다.
여러분이 이 코드를 실행하면 Diffusion의 핵심 아이디어를 직관적으로 이해할 수 있습니다. 실제 학습 시에는 배치 내 각 이미지마다 랜덤한 t를 선택해서 다양한 노이즈 레벨을 경험하게 합니다.
추론 시에는 이 과정의 역방향(t=999부터 t=0까지)을 진행하며 이미지를 생성하게 됩니다.
실전 팁
💡 타임스텝을 시각화하면 모델이 학습하는 내용을 이해하기 쉽습니다. t=0200은 디테일, t=200700은 구조, t=700~1000은 전반적인 형태를 학습한다고 볼 수 있습니다.
💡 alphas_cumprod가 0에 가까워질수록 원본 정보가 거의 사라집니다. 실무에서는 마지막 타임스텝에서 이 값이 0.0001 이하가 되도록 β 스케줄을 조정합니다.
💡 배치 단위로 처리할 때는 각 샘플마다 다른 t를 할당하세요. t = torch.randint(0, num_steps, (batch_size,))로 쉽게 구현할 수 있습니다.
💡 동일한 노이즈를 재사용하려면 시드를 고정하세요. torch.manual_seed(42)를 사용하면 재현 가능한 결과를 얻을 수 있어 디버깅이 쉬워집니다.
💡 Forward Process는 학습 시에만 사용되고 추론 시에는 필요 없습니다. 하지만 원리를 이해하는 것이 Reverse Process를 설계하는 데 필수적입니다.
3. Reverse Process
시작하며
여러분이 노이즈에서 이미지를 만들어낸다고 하면 믿기 어렵죠? 하지만 Forward Process가 이미지를 노이즈로 만드는 과정이라면, Reverse Process는 그 정반대로 노이즈에서 이미지를 복원하는 과정입니다.
이것이 바로 Diffusion Model의 핵심이자 가장 어려운 부분입니다. Forward는 단순히 노이즈를 추가하면 되지만, Reverse는 "어떤 노이즈를 제거해야 할지"를 정확히 알아야 하기 때문입니다.
수학적으로는 역방향 조건부 확률 분포 p(x_{t-1} | x_t)를 알아야 하는데, 이를 직접 계산하는 것은 불가능합니다. 그래서 신경망을 사용해 이 분포를 근사하는 것이 Diffusion Model의 핵심 아이디어입니다.
개요
간단히 말해서, Reverse Process는 완전한 노이즈 x_T에서 시작해 점진적으로 노이즈를 제거하며 깨끗한 이미지 x₀로 복원하는 과정이고, 이를 신경망 ε_θ가 학습합니다. 왜 신경망이 필요할까요?
Forward Process는 정해진 수식이 있지만, Reverse는 "현재 상태에서 어떤 노이즈가 추가되었는지 추정"해야 하는 복잡한 작업이기 때문입니다. 예를 들어, 고양이 이미지에 노이즈를 추가한 결과와 강아지 이미지에 노이즈를 추가한 결과는 비슷해 보일 수 있지만, 각각 다른 방식으로 복원되어야 합니다.
기존 VAE나 GAN에서는 잠재 공간에서 한 번에 이미지를 생성했다면, Diffusion은 1000번의 작은 스텝으로 나누어 진행합니다. 이렇게 하면 각 스텝이 매우 간단한 작업(조금의 노이즈 제거)만 수행하면 되므로, 전체적으로 더 안정적이고 고품질의 결과를 얻을 수 있습니다.
Reverse Process의 핵심 특징: (1) 신경망 기반 - U-Net 등의 구조로 노이즈 예측, (2) 조건부 생성 가능 - 텍스트나 레이블로 생성 제어 가능, (3) 점진적 개선 - 각 스텝마다 조금씩 품질 향상. 이러한 특징들이 Stable Diffusion 같은 서비스를 가능하게 합니다.
코드 예제
import torch
import torch.nn as nn
class ReverseProcess:
def __init__(self, model, betas):
"""
model: 노이즈를 예측하는 신경망 (U-Net 등)
betas: Forward process에서 사용한 베타 스케줄
"""
self.model = model
self.betas = betas
self.alphas = 1.0 - betas
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
@torch.no_grad()
def denoise_step(self, xt, t):
"""
한 타임스텝만큼 노이즈 제거
"""
# 신경망으로 추가된 노이즈 예측
predicted_noise = self.model(xt, t)
# 알파 값들 가져오기
alpha_t = self.alphas[t]
alpha_bar_t = self.alphas_cumprod[t]
# x_{t-1} 계산을 위한 계수들
coef1 = 1 / torch.sqrt(alpha_t)
coef2 = (1 - alpha_t) / torch.sqrt(1 - alpha_bar_t)
# 평균 계산: μ_θ(x_t, t)
mean = coef1 * (xt - coef2 * predicted_noise)
# t > 0일 때만 노이즈 추가 (stochastic)
if t > 0:
noise = torch.randn_like(xt)
variance = self.betas[t]
return mean + torch.sqrt(variance) * noise
else:
return mean # t=0일 때는 노이즈 추가 안 함
설명
이것이 하는 일: 이 코드는 노이즈가 섞인 이미지 x_t에서 한 단계 이전의 덜 노이즈한 이미지 x_{t-1}을 생성합니다. 1000번 반복하면 완전한 노이즈에서 깨끗한 이미지가 만들어집니다.
첫 번째로, 초기화 단계에서 학습된 신경망 모델과 베타 스케줄을 저장합니다. 이 신경망은 수백만~수억 장의 이미지로 학습되어 "현재 상태 x_t와 타임스텝 t가 주어졌을 때 어떤 노이즈가 추가되었는지"를 예측할 수 있습니다.
Forward와 동일한 베타 스케줄을 사용하는 것이 중요한데, 학습 시와 추론 시의 노이즈 분포가 일치해야 하기 때문입니다. 그 다음으로, denoise_step 함수가 핵심 작업을 수행합니다.
먼저 신경망에 현재 상태 x_t와 타임스텝 t를 입력하면, 신경망이 "이 이미지에 추가된 노이즈"를 예측합니다. 이것이 정말 놀라운 점인데, 모델이 수백만 장의 이미지를 보며 학습한 덕분에 매우 정확하게 노이즈를 분리해낼 수 있습니다.
세 번째로, 예측된 노이즈를 사용해 이전 상태의 평균(mean)을 계산합니다. 수학적으로 증명된 공식 μ_θ = (1/√α_t) × (x_t - ((1-α_t)/√(1-ᾱ_t)) × ε_θ)을 사용하는데, 이는 "현재 상태에서 예측된 노이즈를 제거한 값"을 의미합니다.
계수들이 복잡해 보이지만, 본질은 "노이즈 부분을 빼라"는 것입니다. 마지막으로, 약간의 랜덤 노이즈를 추가합니다.
왜 노이즈를 제거하면서 다시 노이즈를 추가할까요? 이것은 확률적 샘플링을 위한 것으로, 매번 조금씩 다른 결과를 생성해 다양성을 높입니다.
단, 마지막 스텝(t=0)에서는 노이즈를 추가하지 않아 최종 이미지가 깨끗하게 나옵니다. 여러분이 이 코드를 사용하면 실제 이미지 생성 과정을 경험할 수 있습니다.
실무에서는 t=999부터 t=0까지 루프를 돌며 이 함수를 반복 호출합니다. 초반에는 큰 변화가 없어 보이지만, 중반부터 이미지의 형태가 드러나고, 후반에는 디테일이 추가됩니다.
텍스트 조건을 추가하면 원하는 내용의 이미지를 생성할 수 있습니다.
실전 팁
💡 @torch.no_grad() 데코레이터는 필수입니다. 추론 시에는 그래디언트를 계산할 필요가 없으므로 메모리를 크게 절약할 수 있습니다.
💡 배치 처리 시 모든 샘플이 동일한 t에서 시작해야 합니다. t=999에서 시작해 동시에 t를 감소시키며 진행하세요.
💡 마지막 스텝에서 노이즈를 추가하지 않는 것을 잊지 마세요. if t > 0 조건이 빠지면 최종 이미지에도 노이즈가 남아 품질이 떨어집니다.
💡 샘플링 속도를 높이려면 DDIM(Denoising Diffusion Implicit Models)을 사용하세요. 1000 스텝 대신 50 스텝만으로도 비슷한 품질을 얻을 수 있습니다.
💡 실무에서는 CFG(Classifier-Free Guidance)를 사용해 텍스트 조건의 영향력을 조절합니다. guidance_scale을 7~15 정도로 설정하면 프롬프트에 더 충실한 이미지가 생성됩니다.
4. Noise Scheduling
시작하며
여러분이 Diffusion Model을 직접 학습시켜봤는데 결과가 흐릿하거나 디테일이 부족한 경험이 있나요? 모델 구조나 데이터셋 문제라고 생각하기 쉽지만, 의외로 노이즈 스케줄이 원인인 경우가 많습니다.
Noise Scheduling은 "각 타임스텝에서 얼마나 노이즈를 추가할지" 결정하는 전략입니다. 너무 급격하게 노이즈를 추가하면 모델이 학습하기 어렵고, 너무 천천히 추가하면 불필요하게 많은 스텝이 필요합니다.
이것은 마치 요리에서 불 조절과 같습니다. 처음부터 센 불로 조리하면 겉은 타고 속은 익지 않듯이, 초반부터 강한 노이즈를 주면 이미지 구조가 너무 빨리 파괴됩니다.
적절한 스케줄이 고품질 생성의 핵심입니다.
개요
간단히 말해서, Noise Scheduling은 타임스텝 t에 따라 노이즈 추가 강도 β_t를 정의하는 함수이며, Linear, Cosine, Quadratic 등 다양한 방식이 있습니다. 왜 여러 스케줄이 필요할까요?
데이터셋의 특성에 따라 최적의 스케줄이 다르기 때문입니다. 예를 들어, 얼굴 이미지처럼 구조가 명확한 경우는 Cosine 스케줄이, 자연 이미지처럼 복잡한 경우는 Linear 스케줄이 더 나을 수 있습니다.
작은 이미지(32×32)와 큰 이미지(512×512)도 다른 스케줄이 필요합니다. 기존 Linear 스케줄에서는 β가 균일하게 증가했다면, Cosine 스케줄은 초반에 더 천천히 증가합니다.
이렇게 하면 모델이 세밀한 디테일을 더 오래 유지하며 학습할 수 있어 품질이 향상됩니다. Noise Scheduling의 핵심 특징: (1) 하이퍼파라미터 - 모델 성능에 직접적 영향, (2) 데이터 의존적 - 이미지 해상도와 복잡도에 따라 조정 필요, (3) 수학적 제약 - ᾱ_T가 거의 0에 가까워야 함.
이러한 특징들을 고려해 스케줄을 선택하고 조정해야 합니다.
코드 예제
import torch
import numpy as np
import matplotlib.pyplot as plt
def linear_schedule(num_timesteps, beta_start=0.0001, beta_end=0.02):
"""
선형 노이즈 스케줄: DDPM 논문의 기본 방식
"""
return torch.linspace(beta_start, beta_end, num_timesteps)
def cosine_schedule(num_timesteps, s=0.008):
"""
코사인 노이즈 스케줄: 초반 디테일 보존에 유리
Improved DDPM 논문에서 제안
"""
steps = num_timesteps + 1
x = torch.linspace(0, num_timesteps, steps)
# 알파바 누적곱을 코사인 함수로 정의
alphas_cumprod = torch.cos(((x / num_timesteps) + s) / (1 + s) * np.pi * 0.5) ** 2
alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
# 베타 계산: β_t = 1 - (ᾱ_t / ᾱ_{t-1})
betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
# 베타를 [0.0001, 0.9999] 범위로 클리핑
return torch.clip(betas, 0.0001, 0.9999)
def compare_schedules(num_timesteps=1000):
"""
여러 스케줄 비교 시각화
"""
linear_betas = linear_schedule(num_timesteps)
cosine_betas = cosine_schedule(num_timesteps)
# 누적곱 계산
linear_alphas_cumprod = torch.cumprod(1 - linear_betas, dim=0)
cosine_alphas_cumprod = torch.cumprod(1 - cosine_betas, dim=0)
return {
'linear': {'betas': linear_betas, 'alphas_cumprod': linear_alphas_cumprod},
'cosine': {'betas': cosine_betas, 'alphas_cumprod': cosine_alphas_cumprod}
}
# 사용 예시
schedules = compare_schedules()
print(f"Linear 스케줄 마지막 ᾱ: {schedules['linear']['alphas_cumprod'][-1]:.6f}")
print(f"Cosine 스케줄 마지막 ᾱ: {schedules['cosine']['alphas_cumprod'][-1]:.6f}")
설명
이것이 하는 일: 이 코드는 두 가지 주요 노이즈 스케줄을 구현하고 비교합니다. 각 스케줄은 모델이 이미지를 어떻게 학습하는지에 큰 영향을 미칩니다.
첫 번째로, Linear 스케줄은 가장 단순한 방식입니다. β를 0.0001부터 0.02까지 균일하게 증가시키는데, 이는 DDPM 원본 논문에서 사용된 방법입니다.
구현이 매우 간단하고 많은 경우에 잘 작동하지만, 초반에 노이즈가 너무 빨리 증가해 디테일을 잃을 수 있다는 단점이 있습니다. 특히 고해상도 이미지(512×512 이상)에서는 이 문제가 두드러집니다.
그 다음으로, Cosine 스케줄은 더 정교한 접근입니다. β를 직접 정의하지 않고, 대신 ᾱ_t(누적 알파)를 코사인 함수로 정의한 뒤 역산해서 β를 구합니다.
왜 이렇게 복잡하게 할까요? 코사인 함수의 특성상 초반에는 천천히, 후반에는 빠르게 감소하기 때문에 이미지 구조를 더 오래 유지할 수 있습니다.
실험 결과 ImageNet 같은 복잡한 데이터셋에서 품질이 크게 향상되었습니다. 세 번째로, 클리핑 과정이 중요합니다.
계산된 β가 0보다 작거나 1보다 크면 확률 분포가 무너지므로, 안전한 범위인 [0.0001, 0.9999]로 제한합니다. 특히 Cosine 스케줄에서는 초반에 β가 음수가 될 수 있어 이 처리가 필수적입니다.
마지막으로, compare_schedules 함수는 두 스케줄의 ᾱ_t 값을 비교합니다. Linear는 t=1000에서 ᾱ ≈ 0.0002, Cosine은 ᾱ ≈ 0.00001로, 둘 다 거의 0에 가깝지만 도달하는 경로가 다릅니다.
Cosine이 초반 200 스텝에서 훨씬 높은 ᾱ 값을 유지해 디테일 보존에 유리합니다. 여러분이 이 코드로 스케줄을 시각화하면 어느 것이 자신의 데이터에 맞는지 판단할 수 있습니다.
실무에서는 먼저 Cosine으로 시작해보고, 학습이 불안정하면 Linear로 변경하는 것을 추천합니다. 작은 이미지(64×64)는 Linear, 큰 이미지(512×512)는 Cosine이 일반적으로 더 나은 결과를 보입니다.
실전 팁
💡 s 파라미터는 Cosine 스케줄의 오프셋입니다. 기본값 0.008이 대부분 잘 작동하지만, 매우 작은 이미지(32×32)에서는 0.01~0.02로 증가시키면 더 나을 수 있습니다.
💡 ᾱ_T (마지막 타임스텝의 누적 알파)가 너무 크면(> 0.01) 노이즈가 충분하지 않아 학습이 어려워집니다. 반드시 0.0001 이하로 떨어지는지 확인하세요.
💡 스케줄을 바꾸면 모델을 처음부터 다시 학습해야 합니다. 중간에 변경하면 학습이 불안정해지므로, 초기에 신중히 선택하세요.
💡 Quadratic 스케줄(β_t = t²에 비례)도 실험해볼 가치가 있습니다. 일부 데이터셋에서는 Cosine보다 나은 결과를 보입니다.
💡 PyTorch로 스케줄을 구현할 때는 미리 계산해서 buffer로 등록하세요. self.register_buffer('betas', betas)를 사용하면 모델 저장/로딩 시 자동으로 포함됩니다.
5. U-Net Architecture
시작하며
여러분이 Diffusion Model의 원리는 이해했지만, "실제로 노이즈를 예측하는 신경망은 어떻게 생겼을까?"라는 의문이 들 수 있습니다. 모든 Diffusion Model의 심장부에는 U-Net이라는 특별한 구조가 있습니다.
U-Net은 원래 의료 이미지 세그멘테이션을 위해 개발되었지만, 이미지의 공간 정보를 보존하며 처리하는 능력 덕분에 Diffusion에 완벽하게 맞아떨어졌습니다. 픽셀 단위로 노이즈를 예측해야 하는 Diffusion의 특성과 정확히 맞아떨어지는 것이죠.
Stable Diffusion, DALL-E 2, Imagen 등 모든 주요 Diffusion 모델이 U-Net 기반입니다. 이 구조를 이해하면 모델을 커스터마이징하거나 자신만의 개선 방법을 시도할 수 있습니다.
개요
간단히 말해서, U-Net은 입력을 점진적으로 다운샘플링했다가 다시 업샘플링하는 인코더-디코더 구조에 스킵 연결(skip connection)을 추가한 신경망입니다. 왜 U자 형태일까요?
왼쪽 절반(인코더)에서는 해상도를 줄이며 고수준 특징을 추출하고, 오른쪽 절반(디코더)에서는 해상도를 복원하며 세밀한 출력을 생성합니다. 예를 들어, 256×256 이미지가 128→64→32로 줄어들었다가 다시 64→128→256으로 복원됩니다.
각 단계에서 "무엇을 그릴지"(고수준)와 "어디에 그릴지"(저수준) 정보를 모두 활용합니다. 기존 단순 CNN에서는 정보가 일방향으로만 흐른다면, U-Net은 스킵 연결로 초기 레이어의 디테일을 후기 레이어에 직접 전달합니다.
이렇게 하면 다운샘플링 과정에서 잃어버린 공간 정보를 복구할 수 있어 정확한 픽셀 단위 예측이 가능합니다. U-Net의 핵심 특징: (1) 대칭 구조 - 인코더와 디코더가 거울상, (2) 스킵 연결 - 동일 해상도 레이어 간 직접 연결, (3) 타임스텝 임베딩 - 각 레이어에 t 정보를 주입해 시간 인식.
이러한 특징들이 노이즈 예측에 최적화된 구조를 만듭니다.
코드 예제
import torch
import torch.nn as nn
import math
class TimestepEmbedding(nn.Module):
"""
타임스텝을 고차원 벡터로 임베딩 (Transformer의 위치 인코딩과 유사)
"""
def __init__(self, dim):
super().__init__()
self.dim = dim
def forward(self, t):
device = t.device
half_dim = self.dim // 2
# 주파수 계산
emb = math.log(10000) / (half_dim - 1)
emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
# sin/cos 임베딩
emb = t[:, None] * emb[None, :]
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
return emb
class SimpleUNet(nn.Module):
"""
Diffusion을 위한 간소화된 U-Net 구조
"""
def __init__(self, in_channels=3, out_channels=3, time_dim=256):
super().__init__()
# 타임스텝 임베딩
self.time_embed = nn.Sequential(
TimestepEmbedding(time_dim),
nn.Linear(time_dim, time_dim),
nn.ReLU()
)
# 인코더 (다운샘플링)
self.enc1 = nn.Conv2d(in_channels, 64, 3, padding=1)
self.enc2 = nn.Conv2d(64, 128, 3, padding=1, stride=2)
self.enc3 = nn.Conv2d(128, 256, 3, padding=1, stride=2)
# 병목 (가장 낮은 해상도)
self.bottleneck = nn.Conv2d(256, 256, 3, padding=1)
# 디코더 (업샘플링)
self.dec3 = nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1)
self.dec2 = nn.ConvTranspose2d(256, 64, 4, stride=2, padding=1) # 256 = 128 + 128 (스킵)
self.dec1 = nn.Conv2d(128, out_channels, 3, padding=1) # 128 = 64 + 64 (스킵)
def forward(self, x, t):
# 타임스텝 임베딩
t_emb = self.time_embed(t)
# 인코더
e1 = torch.relu(self.enc1(x))
e2 = torch.relu(self.enc2(e1))
e3 = torch.relu(self.enc3(e2))
# 병목
b = torch.relu(self.bottleneck(e3))
# 디코더 (스킵 연결 포함)
d3 = torch.relu(self.dec3(b))
d3 = torch.cat([d3, e2], dim=1) # 스킵 연결
d2 = torch.relu(self.dec2(d3))
d2 = torch.cat([d2, e1], dim=1) # 스킵 연결
output = self.dec1(d2)
return output
# 사용 예시
model = SimpleUNet()
x = torch.randn(4, 3, 64, 64) # 배치 4, RGB, 64x64
t = torch.randint(0, 1000, (4,)) # 타임스텝
noise_pred = model(x, t)
print(f"입력 shape: {x.shape}, 출력 shape: {noise_pred.shape}")
설명
이것이 하는 일: 이 코드는 Diffusion Model의 핵심 신경망인 U-Net을 간소화해 구현합니다. 실제 Stable Diffusion은 훨씬 복잡하지만, 기본 원리는 동일합니다.
첫 번째로, 타임스텝 임베딩이 핵심입니다. 단순히 t=500 같은 숫자를 넣으면 신경망이 학습하기 어려우므로, Transformer의 위치 인코딩처럼 sin/cos 함수로 고차원 벡터(256차원)로 변환합니다.
이렇게 하면 "비슷한 타임스텝은 비슷한 임베딩"을 가지게 되어 일반화 성능이 향상됩니다. 예를 들어, t=499와 t=500의 임베딩은 매우 가까워서 모델이 부드러운 예측을 학습할 수 있습니다.
그 다음으로, 인코더 부분이 이미지를 분석합니다. 첫 번째 Conv는 RGB 3채널을 64채널로 확장해 기본 특징을 추출합니다.
두 번째와 세 번째 Conv는 stride=2로 해상도를 절반씩 줄이며 128채널, 256채널로 늘립니다. 64×64 입력이 32×32, 16×16으로 줄어드는데, 해상도는 줄지만 채널 수는 늘어나 "무엇이 있는지"에 대한 고수준 정보를 추출합니다.
세 번째로, 병목(bottleneck) 레이어는 가장 압축된 표현을 만듭니다. 16×16×256 크기인데, 원본 64×64×3보다 공간은 작지만 훨씬 추상적인 정보를 담고 있습니다.
여기서 "대략 어떤 이미지인지"를 파악합니다. 네 번째로, 디코더가 해상도를 복원합니다.
ConvTranspose2d로 업샘플링하며, 스킵 연결이 핵심입니다. torch.cat([d3, e2], dim=1)는 디코더의 출력과 인코더의 동일 해상도 특징을 합치는데, 이렇게 하면 다운샘플링 때 잃어버린 공간 정보를 복구할 수 있습니다.
예를 들어, "고양이 귀가 어디 있었는지"를 기억해 정확한 위치에 복원할 수 있죠. 마지막으로, 최종 출력은 입력과 동일한 shape(3×64×64)를 가집니다.
이것이 "예측된 노이즈"이며, 실제 추가된 노이즈와의 차이를 최소화하도록 학습됩니다. 여러분이 이 구조를 이해하면 Stable Diffusion의 U-Net 변형(Attention 추가, ResNet 블록 등)도 쉽게 이해할 수 있습니다.
실무에서는 이 기본 구조에 Self-Attention 레이어(장거리 의존성), ResNet 블록(깊은 네트워크), Group Normalization(안정적 학습) 등을 추가해 성능을 극대화합니다.
실전 팁
💡 스킵 연결은 반드시 동일한 해상도 레이어 간에만 연결하세요. 크기가 다르면 torch.cat에서 에러가 발생합니다.
💡 실제 구현에서는 타임스텝 임베딩을 각 레이어에 더하거나 곱합니다. 위 코드는 단순화했지만, 실무에서는 AdaGN(Adaptive Group Normalization) 같은 기법을 사용합니다.
💡 메모리가 부족하면 채널 수를 줄이세요. 64→128→256 대신 32→64→128로 하면 메모리 사용량이 1/4로 줄어듭니다.
💡 Attention 레이어는 낮은 해상도(16×16 이하)에만 추가하세요. 높은 해상도에 추가하면 메모리와 연산량이 폭발적으로 증가합니다.
💡 체크포인팅(gradient checkpointing)을 사용하면 메모리를 절반으로 줄일 수 있습니다. torch.utils.checkpoint.checkpoint로 쉽게 적용 가능합니다.
6. Training Process
시작하며
여러분이 지금까지 Diffusion의 구조를 이해했다면, 이제 "어떻게 학습시킬까?"라는 실전 단계입니다. 아무리 좋은 아키텍처라도 제대로 학습하지 못하면 무용지물이죠.
Diffusion Model의 학습은 의외로 단순합니다. GAN처럼 두 네트워크를 동시에 학습할 필요도 없고, VAE처럼 복잡한 손실 함수를 설계할 필요도 없습니다.
단지 "추가된 노이즈를 정확히 예측하도록" 학습하면 됩니다. 이 단순함이 Diffusion이 GAN을 제치고 주류가 된 이유입니다.
학습이 안정적이고, 모드 붕괴 같은 문제도 거의 없으며, 하이퍼파라미터에 덜 민감합니다. 실무에서 바로 적용할 수 있는 수준입니다.
개요
간단히 말해서, Diffusion 학습은 "랜덤 타임스텝에서 이미지에 노이즈를 추가하고, 모델이 그 노이즈를 예측하도록 하며, MSE 손실로 최적화"하는 과정입니다. 왜 이렇게 간단할까요?
Forward Process가 수학적으로 명확히 정의되어 있어서, "정답"인 노이즈를 알고 있기 때문입니다. 예를 들어, 고양이 이미지에 ε ~ N(0,1) 노이즈를 추가했다면, 모델이 정확히 그 ε을 예측하도록 학습하면 됩니다.
GAN처럼 "진짜 같은지 가짜 같은지" 애매한 목표가 아닙니다. 기존 이미지 분류에서는 레이블을 예측했다면, Diffusion은 노이즈를 예측합니다.
둘 다 지도 학습이지만, Diffusion은 각 픽셀마다 정답이 있어 훨씬 풍부한 학습 신호를 제공합니다. Diffusion 학습의 핵심 특징: (1) 단순한 손실 함수 - MSE 또는 L1만 사용, (2) 안정적인 학습 - GAN처럼 붕괴 위험 없음, (3) 확장 가능 - 데이터가 많을수록 계속 개선됨.
이러한 특징들이 대규모 학습을 가능하게 합니다.
코드 예제
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
class DiffusionTrainer:
def __init__(self, model, dataloader, num_timesteps=1000, device='cuda'):
self.model = model.to(device)
self.dataloader = dataloader
self.num_timesteps = num_timesteps
self.device = device
# 노이즈 스케줄 설정
self.betas = torch.linspace(0.0001, 0.02, num_timesteps).to(device)
self.alphas = 1.0 - self.betas
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
# 옵티마이저 (AdamW 추천)
self.optimizer = optim.AdamW(model.parameters(), lr=1e-4)
def train_step(self, x0):
"""
한 배치에 대한 학습 스텝
"""
batch_size = x0.shape[0]
x0 = x0.to(self.device)
# 랜덤 타임스텝 샘플링
t = torch.randint(0, self.num_timesteps, (batch_size,), device=self.device)
# 가우시안 노이즈 생성
noise = torch.randn_like(x0)
# Forward process: 노이즈 추가
sqrt_alphas_cumprod_t = self.alphas_cumprod[t][:, None, None, None] ** 0.5
sqrt_one_minus_alphas_cumprod_t = (1 - self.alphas_cumprod[t])[:, None, None, None] ** 0.5
xt = sqrt_alphas_cumprod_t * x0 + sqrt_one_minus_alphas_cumprod_t * noise
# 모델로 노이즈 예측
predicted_noise = self.model(xt, t)
# MSE 손실 계산
loss = nn.functional.mse_loss(predicted_noise, noise)
# 역전파 및 최적화
self.optimizer.zero_grad()
loss.backward()
# 그래디언트 클리핑 (학습 안정화)
torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
self.optimizer.step()
return loss.item()
def train_epoch(self):
"""
한 에폭 학습
"""
self.model.train()
total_loss = 0
for batch_idx, (images, _) in enumerate(self.dataloader):
loss = self.train_step(images)
total_loss += loss
if batch_idx % 100 == 0:
print(f"Batch {batch_idx}, Loss: {loss:.4f}")
avg_loss = total_loss / len(self.dataloader)
return avg_loss
# 사용 예시 (가상의 데이터로더)
# model = SimpleUNet()
# dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
# trainer = DiffusionTrainer(model, dataloader)
#
# for epoch in range(100):
# loss = trainer.train_epoch()
# print(f"Epoch {epoch}, Avg Loss: {loss:.4f}")
설명
이것이 하는 일: 이 코드는 Diffusion Model의 전체 학습 파이프라인을 구현합니다. 실제 ImageNet이나 LAION 데이터셋 학습도 이 기본 구조를 따릅니다.
첫 번째로, 초기화 단계에서 모델, 데이터로더, 노이즈 스케줄을 준비합니다. AdamW 옵티마이저를 사용하는데, 이는 Adam에 weight decay를 추가한 것으로 Diffusion에서 가장 좋은 성능을 보입니다.
학습률 1e-4는 대부분의 경우에 잘 작동하는 안전한 값입니다. 너무 크면 학습이 불안정하고, 너무 작으면 수렴이 느려집니다.
그 다음으로, train_step이 핵심 학습 로직입니다. 먼저 배치의 각 이미지마다 0부터 999 사이의 랜덤 타임스텝을 샘플링합니다.
왜 랜덤일까요? 모든 타임스텝을 골고루 학습하기 위해서입니다.
만약 항상 t=500만 학습하면 다른 타임스텝에서는 성능이 떨어집니다. 랜덤 샘플링으로 1000개 스텝 모두를 균등하게 학습할 수 있습니다.
세 번째로, Forward Process를 적용해 노이즈를 추가합니다. [:, None, None, None]은 브로드캐스팅을 위한 차원 확장인데, (batch_size,) 크기의 알파 값을 (batch_size, channels, height, width)에 맞추기 위함입니다.
이렇게 하면 배치의 각 이미지가 서로 다른 타임스텝에서 노이즈가 추가됩니다. 네 번째로, 모델이 노이즈를 예측하고 실제 노이즈와 비교합니다.
MSE 손실을 사용하는데, 이는 픽셀별 제곱 오차의 평균입니다. 간단하지만 매우 효과적이며, 일부 연구에서는 L1 손실(절대값 오차)도 시도하지만 MSE가 일반적으로 더 안정적입니다.
마지막으로, 역전파와 그래디언트 클리핑을 수행합니다. 클리핑은 그래디언트 폭발을 방지하는 중요한 기법인데, 특히 깊은 U-Net에서는 필수적입니다.
최대 노름을 1.0으로 제한해 안정적인 학습을 보장합니다. 여러분이 이 코드를 실행하면 수만~수십만 스텝 후에 모델이 의미 있는 이미지를 생성하기 시작합니다.
초반에는 흐릿한 형태만 나오지만, 학습이 진행될수록 디테일이 추가되고 다양성이 증가합니다. 실무에서는 FID(Fréchet Inception Distance) 같은 메트릭으로 품질을 정량적으로 평가합니다.
실전 팁
💡 EMA(Exponential Moving Average)를 사용하면 생성 품질이 크게 향상됩니다. ema_model.update(model) 방식으로 가중치의 이동 평균을 유지하세요.
💡 그래디언트 축적(gradient accumulation)으로 큰 배치 크기 효과를 낼 수 있습니다. GPU 메모리가 부족할 때 유용합니다.
💡 학습 중 주기적으로 샘플을 생성해 시각적으로 확인하세요. 손실이 줄어도 품질이 나빠질 수 있어 눈으로 확인이 중요합니다.
💡 타임스텝 샘플링을 균등 분포 대신 로그 스케일로 하면 중요한 스텝(초반과 후반)에 더 집중할 수 있습니다.
💡 멀티 GPU 학습 시 DistributedDataParallel을 사용하세요. DataParallel보다 훨씬 빠르고 효율적입니다.
7. Sampling Methods
시작하며
여러분이 모델 학습을 완료했다면, 이제 가장 흥미진진한 단계입니다. 바로 이미지 생성이죠.
하지만 여기서 문제가 하나 있습니다. 1000 스텝을 모두 거치려면 시간이 너무 오래 걸립니다.
Stable Diffusion으로 이미지를 생성해본 분들은 알겠지만, 한 장 만드는데 몇 초에서 수십 초가 걸립니다. 만약 실시간 애플리케이션이나 대량 생성이 필요하다면 어떻게 할까요?
Sampling Methods는 바로 이 문제를 해결합니다. 품질을 유지하면서 생성 속도를 10~100배 가속시키는 다양한 기법들이 있습니다.
DDIM, DPM-Solver, PNDM 등이 대표적입니다.
개요
간단히 말해서, Sampling Methods는 Reverse Process를 더 빠르게 수행하는 기법들로, 스텝 수를 줄이거나 더 효율적인 경로를 찾아 생성 속도를 향상시킵니다. 왜 가속이 가능할까요?
원래 DDPM은 1000 스텝이 필요하지만, 실제로는 많은 중간 단계가 중복되거나 불필요합니다. 예를 들어, t=500과 t=501은 거의 차이가 없으므로, 큰 스텝(t=500 → t=400 → t=300...)으로 건너뛰어도 품질이 크게 떨어지지 않습니다.
DDIM 같은 방법은 이를 수학적으로 정당화합니다. 기존 DDPM이 확률적(stochastic) 샘플링이었다면, DDIM은 결정론적(deterministic) 샘플링입니다.
동일한 시작 노이즈에서 항상 같은 결과를 만들어 재현성이 보장됩니다. Sampling의 핵심 특징: (1) 속도-품질 트레이드오프 - 스텝이 적을수록 빠르지만 품질 저하 가능, (2) 다양한 알고리즘 - DDIM, DPM++, Euler 등 선택 가능, (3) 조건부 제어 - CFG로 프롬프트 영향력 조절.
이러한 특징들이 실용적인 이미지 생성을 가능하게 합니다.
코드 예제
import torch
import numpy as np
class DDIMSampler:
"""
DDIM: Denoising Diffusion Implicit Models
DDPM보다 10~100배 빠른 샘플링
"""
def __init__(self, model, betas, num_train_timesteps=1000):
self.model = model
self.num_train_timesteps = num_train_timesteps
# 알파 계산
alphas = 1.0 - betas
self.alphas_cumprod = torch.cumprod(alphas, dim=0)
@torch.no_grad()
def sample(self, shape, num_inference_steps=50, eta=0.0, device='cuda'):
"""
DDIM 샘플링
Args:
shape: 생성할 이미지 크기 (batch, channels, height, width)
num_inference_steps: 추론 스텝 수 (기본 50)
eta: 확률성 조절 (0=결정론적, 1=DDPM과 동일)
"""
# 타임스텝 서브샘플링 (1000 -> 50)
step_ratio = self.num_train_timesteps // num_inference_steps
timesteps = np.arange(0, self.num_train_timesteps, step_ratio)[::-1].copy()
timesteps = torch.from_numpy(timesteps).long().to(device)
# 초기 노이즈 (순수 가우시안)
xt = torch.randn(shape, device=device)
# 역방향 진행
for i, t in enumerate(timesteps):
# 현재와 이전 타임스텝
t_current = t
t_prev = timesteps[i+1] if i < len(timesteps) - 1 else torch.tensor(0)
# 알파 값
alpha_bar_t = self.alphas_cumprod[t_current]
alpha_bar_t_prev = self.alphas_cumprod[t_prev] if t_prev > 0 else torch.tensor(1.0)
# 노이즈 예측
t_batch = t_current.unsqueeze(0).repeat(shape[0])
predicted_noise = self.model(xt, t_batch)
# x0 예측 (원본 이미지 추정)
x0_pred = (xt - (1 - alpha_bar_t) ** 0.5 * predicted_noise) / alpha_bar_t ** 0.5
# 방향 벡터 계산
direction = (1 - alpha_bar_t_prev) ** 0.5 * predicted_noise
# DDIM 업데이트 (결정론적)
xt = alpha_bar_t_prev ** 0.5 * x0_pred + direction
# eta > 0이면 약간의 확률성 추가 (선택적)
if eta > 0 and i < len(timesteps) - 1:
noise = torch.randn_like(xt)
variance = eta * ((1 - alpha_bar_t_prev) / (1 - alpha_bar_t)) * (1 - alpha_bar_t / alpha_bar_t_prev)
xt = xt + variance ** 0.5 * noise
return xt
# 사용 예시
# model = SimpleUNet()
# betas = torch.linspace(0.0001, 0.02, 1000)
# sampler = DDIMSampler(model, betas)
#
# # 4장의 64x64 RGB 이미지 생성, 50 스텝만 사용
# images = sampler.sample(shape=(4, 3, 64, 64), num_inference_steps=50)
# print(f"생성된 이미지: {images.shape}")
설명
이것이 하는 일: 이 코드는 DDIM 알고리즘을 구현해 빠른 이미지 생성을 가능하게 합니다. 원래 1000 스텝이 필요한 DDPM을 50 스텝으로 단축시킵니다.
첫 번째로, 타임스텝 서브샘플링이 핵심 아이디어입니다. 1000개의 타임스텝 중에서 균등하게 50개만 선택합니다 (예: 0, 20, 40, ..., 980).
왜 균등 간격일까요? 모든 구간(초반, 중반, 후반)을 골고루 커버하기 위해서입니다.
일부 방법(DPM++)은 비균등 샘플링을 사용하기도 하는데, 중요한 구간(후반부)에 더 많은 스텝을 할당합니다. 그 다음으로, 초기 상태를 순수 가우시안 노이즈로 설정합니다.
torch.randn으로 생성하는데, 여기서 시드를 고정하면 동일한 노이즈에서 시작해 재현 가능한 생성이 됩니다. Stable Diffusion의 "Seed" 파라미터가 바로 이것입니다.
세 번째로, 각 스텝에서 x0(원본 이미지)를 예측합니다. 이것이 DDIM의 핵심 차별점인데, DDPM은 이전 스텝 x_{t-1}만 예측하지만, DDIM은 최종 목표인 x0를 직접 추정합니다.
공식 x0 = (x_t - √(1-ᾱ_t) × ε) / √(ᾱ_t)는 "현재 상태에서 노이즈를 제거한 것이 원본"이라는 직관적인 의미입니다. 네 번째로, 방향 벡터를 계산해 다음 상태로 이동합니다.
DDIM은 x0 예측값을 향해 직선으로 이동하는 것이 아니라, 적절한 노이즈를 남겨두며 점진적으로 접근합니다. 이렇게 하는 이유는 중간 스텝들을 건너뛰어도 최종 분포가 유지되도록 수학적으로 설계되었기 때문입니다.
마지막으로, eta 파라미터로 확률성을 조절할 수 있습니다. eta=0이면 완전 결정론적(동일 시드 → 동일 결과), eta=1이면 DDPM과 동일한 확률적 샘플링입니다.
실무에서는 보통 0~0.3 사이를 사용해 약간의 다양성을 유지하면서도 재현성을 보장합니다. 여러분이 이 코드를 사용하면 생성 시간을 대폭 단축할 수 있습니다.
50 스텝이면 대부분 충분하지만, 더 빠르게 하려면 2030 스텝으로 줄일 수도 있습니다(품질 약간 저하). 반대로 최고 품질을 원하면 100200 스텝을 사용할 수 있습니다.
Stable Diffusion에서 "Steps" 슬라이더가 바로 이 파라미터입니다.
실전 팁
💡 num_inference_steps는 20~100 사이에서 선택하세요. 20 이하는 품질 저하가 심하고, 100 이상은 개선이 미미합니다.
💡 DPM-Solver++는 DDIM보다 더 빠릅니다. 동일 품질을 10~15 스텝만으로 달성할 수 있어, 실시간 애플리케이션에 적합합니다.
💡 타임스텝 스케줄을 로그 스케일로 하면 후반부에 더 집중해 디테일이 향상됩니다. np.linspace(0, 1, N)**2 * 1000 같은 방식을 시도해보세요.
💡 Classifier-Free Guidance와 함께 사용하면 텍스트 조건을 더 잘 따릅니다. 조건부/무조건부 예측을 섞어 제어력을 높입니다.
💡 배치 생성 시 동일한 노이즈를 재사용하지 마세요. 각 이미지마다 독립적인 노이즈를 사용해야 다양성이 보장됩니다.
8. Conditional Generation
시작하며
여러분이 지금까지 학습한 Diffusion Model은 랜덤한 이미지를 생성합니다. 하지만 실제로 원하는 것은 "귀여운 고양이", "미래적인 도시" 같은 특정 내용의 이미지죠.
어떻게 제어할 수 있을까요? Conditional Generation은 텍스트, 레이블, 이미지 등의 조건으로 생성을 제어하는 기술입니다.
Stable Diffusion이 텍스트 프롬프트로 원하는 이미지를 만드는 것이 바로 이 기술 덕분입니다. CLIP, T5 같은 텍스트 인코더와 Cross-Attention 메커니즘을 결합하면, 모델이 "프롬프트의 의미"를 이해하고 그에 맞는 이미지를 생성할 수 있습니다.
단순한 이미지 생성을 넘어 창작 도구로 발전하는 결정적 요소입니다.
개요
간단히 말해서, Conditional Generation은 외부 정보(텍스트, 레이블 등)를 모델에 주입해 생성 과정을 제어하며, Classifier-Free Guidance로 조건의 영향력을 조절합니다. 왜 조건이 필요할까요?
무조건부 생성은 다양하지만 제어가 불가능합니다. 예를 들어, 100장을 생성해도 원하는 포즈의 고양이가 안 나올 수 있습니다.
조건부 생성은 "검은 고양이, 앉은 자세, 파란 배경"처럼 구체적으로 지정할 수 있어 실용성이 극대화됩니다. 기존 클래스 조건(ImageNet 레이블)에서는 "고양이" 같은 단순한 범주만 가능했다면, 텍스트 조건은 "눈 내리는 밤에 창가에 앉아있는 흰 고양이"처럼 복잡한 설명이 가능합니다.
자연어의 표현력이 생성 품질을 한 단계 끌어올렸습니다. Conditional Generation의 핵심 특징: (1) Cross-Attention - 텍스트와 이미지 특징을 연결, (2) CFG - guidance_scale로 조건 강도 조절, (3) 다양한 조건 타입 - 텍스트, 이미지, 세그멘테이션 등.
이러한 특징들이 Stable Diffusion, DALL-E를 가능하게 합니다.
코드 예제
import torch
import torch.nn as nn
class CrossAttention(nn.Module):
"""
텍스트 조건을 이미지 특징에 주입하는 Cross-Attention
"""
def __init__(self, dim, context_dim):
super().__init__()
self.to_q = nn.Linear(dim, dim) # Query: 이미지 특징
self.to_k = nn.Linear(context_dim, dim) # Key: 텍스트 특징
self.to_v = nn.Linear(context_dim, dim) # Value: 텍스트 특징
self.scale = dim ** -0.5
def forward(self, x, context):
"""
x: 이미지 특징 (B, N_img, dim)
context: 텍스트 특징 (B, N_text, context_dim)
"""
q = self.to_q(x)
k = self.to_k(context)
v = self.to_v(context)
# Attention 계산: softmax(Q·K^T / √d) · V
attn = torch.einsum('bnd,bmd->bnm', q, k) * self.scale
attn = torch.softmax(attn, dim=-1)
out = torch.einsum('bnm,bmd->bnd', attn, v)
return out
class ConditionalUNet(nn.Module):
"""
텍스트 조건을 받는 U-Net (간소화 버전)
"""
def __init__(self, img_channels=3, text_dim=768):
super().__init__()
# 이미지 인코더
self.img_conv = nn.Conv2d(img_channels, 256, 3, padding=1)
# Cross-Attention 레이어
self.cross_attn = CrossAttention(dim=256, context_dim=text_dim)
# 출력 레이어
self.out_conv = nn.Conv2d(256, img_channels, 3, padding=1)
def forward(self, x, t, text_embedding):
"""
x: 노이즈 이미지 (B, C, H, W)
t: 타임스텝 (B,)
text_embedding: 텍스트 임베딩 (B, N_tokens, text_dim)
"""
# 이미지 특징 추출
h = self.img_conv(x) # (B, 256, H, W)
# Reshape for attention: (B, 256, H, W) -> (B, H*W, 256)
B, C, H, W = h.shape
h_flat = h.view(B, C, H*W).permute(0, 2, 1)
# Cross-Attention으로 텍스트 정보 주입
h_attended = self.cross_attn(h_flat, text_embedding)
# Reshape back: (B, H*W, 256) -> (B, 256, H, W)
h = h_attended.permute(0, 2, 1).view(B, C, H, W)
# 노이즈 예측
noise_pred = self.out_conv(h)
return noise_pred
@torch.no_grad()
def classifier_free_guidance_sample(model, text_emb, null_emb, guidance_scale=7.5):
"""
Classifier-Free Guidance를 사용한 샘플링
Args:
model: 조건부 U-Net
text_emb: 텍스트 임베딩 (조건부)
null_emb: 빈 텍스트 임베딩 (무조건부)
guidance_scale: CFG 강도 (높을수록 프롬프트를 더 따름)
"""
# 조건부 예측
noise_pred_cond = model(x, t, text_emb)
# 무조건부 예측
noise_pred_uncond = model(x, t, null_emb)
# CFG 공식: ε = ε_uncond + s * (ε_cond - ε_uncond)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
return noise_pred
# 사용 예시
# model = ConditionalUNet()
# text_emb = clip_model.encode_text("a cute cat") # CLIP으로 텍스트 인코딩
# null_emb = clip_model.encode_text("") # 빈 프롬프트
#
# # CFG를 사용한 이미지 생성 (guidance_scale=7.5)
# noise = classifier_free_guidance_sample(model, text_emb, null_emb, 7.5)
설명
이것이 하는 일: 이 코드는 텍스트 조건으로 이미지 생성을 제어하는 핵심 메커니즘인 Cross-Attention과 CFG를 구현합니다. Stable Diffusion의 작동 원리와 동일합니다.
첫 번째로, Cross-Attention이 텍스트와 이미지를 연결합니다. Query는 이미지 특징에서, Key와 Value는 텍스트 특징에서 생성됩니다.
"이미지의 각 위치가 텍스트의 어느 부분에 주목해야 하는지" 계산하는 것이죠. 예를 들어, "빨간 사과"라는 프롬프트가 있으면, 이미지의 사과 영역이 "빨간"과 "사과" 토큰에 높은 attention을 갖게 됩니다.
그 다음으로, torch.einsum으로 효율적인 행렬 연산을 수행합니다. 'bnd,bmd->bnm'는 배치(b)별로 N개 이미지 토큰과 M개 텍스트 토큰 간의 유사도를 계산하는데, 일반적인 행렬곱보다 가독성이 좋고 GPU 최적화가 잘 됩니다.
Softmax로 정규화하면 각 이미지 토큰마다 텍스트 토큰들에 대한 확률 분포가 됩니다. 세 번째로, Attention 가중치로 Value를 가중합해 최종 출력을 만듭니다.
이것이 "텍스트 정보가 주입된 이미지 특징"이 되어, 모델이 프롬프트 내용을 이해하고 반영할 수 있게 합니다. 실제 Stable Diffusion은 U-Net의 여러 레이어마다 Cross-Attention을 추가해 다양한 추상도에서 텍스트를 반영합니다.
네 번째로, Classifier-Free Guidance(CFG)가 핵심 기법입니다. 조건부와 무조건부 예측을 모두 수행한 뒤, 그 차이를 증폭시켜 조건을 더 강하게 반영합니다.
수학적으로는 ε = ε_u + s × (ε_c - ε_u)인데, s(guidance_scale)가 클수록 조건부 예측 방향으로 더 크게 이동합니다. s=1이면 일반 조건부, s=7.5면 조건을 7.5배 강조한 것입니다.
마지막으로, guidance_scale 조절이 매우 중요합니다. 낮은 값(13)은 자유롭고 창의적이지만 프롬프트를 잘 안 따르고, 높은 값(1520)은 프롬프트에 충실하지만 과포화되거나 부자연스러울 수 있습니다.
실무에서는 7~10이 가장 균형잡힌 결과를 보입니다. 여러분이 이 메커니즘을 이해하면 Stable Diffusion의 다양한 파라미터를 효과적으로 조절할 수 있습니다.
"Prompt strength", "CFG scale" 같은 설정들이 내부적으로 어떻게 작동하는지 알게 되어, 원하는 결과를 더 쉽게 얻을 수 있습니다. ControlNet, IP-Adapter 같은 고급 조건부 기법들도 모두 이 기본 원리를 확장한 것입니다.
실전 팁
💡 텍스트 인코더는 CLIP이 가장 많이 쓰입니다. transformers 라이브러리의 CLIPTextModel을 사용하면 쉽게 구현할 수 있습니다.
💡 무조건부 임베딩은 빈 문자열 ""로 생성합니다. 학습 시 10~15% 확률로 조건을 드롭해야 CFG가 제대로 작동합니다.
💡 Cross-Attention은 메모리를 많이 사용합니다. Flash Attention이나 xFormers를 사용하면 메모리를 절반으로 줄일 수 있습니다.
💡 guidance_scale이 너무 높으면 색상이 과포화되고 디테일이 이상해집니다. 10을 넘지 않는 것을 추천합니다.
💡 멀티 조건(텍스트 + 이미지, 텍스트 + 포즈)은 여러 Cross-Attention을 병렬로 사용합니다. ControlNet이 대표적인 예시입니다.