آموزش مدل‌های دیفیوژن (Diffusion Models): پیاده‌سازی گام‌به‌گام DDPM و DDIM

تولید تصویر در یادگیری عمیق سال‌ها تحت سلطه‌ی شبکه‌های زایای رقابتی (GAN) بود؛ مدل‌هایی که یادگیری آن‌ها به‌شدت ناپایدار بود و با کوچک‌ترین تغییر پارامترها دچار فروپاشی حالت (mode collapse) می‌شدند. مدل‌های دیفیوژن (Diffusion Models) با تبدیل فرآیند تولید به نویززدایی گام‌به‌گام، پایداری آموزش را حل کردند و مبنای سیستم‌هایی مثل Stable Diffusion و Midjourney شدند.

درک معماری این مدل‌ها دیگر یک موضوع نظری صرف نیست؛ برای سفارشی‌سازی پایپ‌لاین‌های تولید، پیاده‌سازی روش‌های هدایت (guidance) یا بهینه‌سازی سرعت اینفرانس، باید دقیقاً بدانید نویز چگونه تزریق و چگونه توسط شبکه پیش‌بینی می‌شود. در این درس، کل منطق ریاضی DDPM و جهش سریع‌تر DDIM را با کد تمیز و ساختاری سبک پیاده‌سازی می‌کنیم.

برای اجرای کدهای این درس و ساخت نسخه‌ی کوچک U-Net نیازی به کارت گرافیک‌های ابری گران‌قیمت یا پرداخت‌های ارزی ندارید؛ این پیاده‌سازی پایه‌ای طوری طراحی شده که به‌سادگی روی CPU یا تراشه‌های شخصی اجرا شود.

مدل دیفیوژن (diffusion model) نویززدایی را یاد می‌گیرد. کافی است به آن آموزش دهید تا مقدار کمی نویز را از یک تصویر نویزدار حذف کند؛ این کار را رو به عقب هزار بار تکرار کنید تا یک مدل تولید تصویر داشته باشید.

اهداف یادگیری

• استخراج فرآیند نویزدهی رو به جلو (forward noising process) به صورت x_0 -> x_1 -> ... -> x_T و توضیح این‌که چرا فرم بسته‌ی (closed-form) q(x_t | x_0) برای هر t برقرار است
• پیاده‌سازی هدف آموزشی به سبک DDPM که نویز اضافه‌شده در هر گام را رگرسیون می‌کند، همراه با یک نمونه‌گیر (sampler) که از نویز خالص رو به عقب حرکت کرده و به تصویر می‌رسد
• ساخت یک U-Net شرطی‌شده بر زمان (آن‌قدر کوچک که روی CPU آموزش ببیند) برای پیش‌بینی نویز در هر گام زمانی
• توضیح تفاوت بین نمونه‌گیری به روش DDPM و DDIM، و کاربرد مناسب هر کدام (درس 23 به مباحث تطبیق جریان [flow matching] و جریان تصحیح‌شده [rectified flow] به‌طور مفصل می‌پردازد)

صورت مسئله

شبکه‌های زایای رقابتی یا همان GANها خروجی را یک‌باره تولید می‌کنند: نویز وارد می‌شود، تصویر بیرون می‌آید؛ فقط در یک گذر رو به جلو (forward pass). آن‌ها سریع‌اند ولی آموزش دادنشان بسیار دشوار است. در مقابل، مدل‌های دیفیوژن تصویر را به‌صورت تکرارشونده تولید می‌کنند: از نویز خالص شروع می‌کنند، در گام‌های کوچک نویززدایی را پیش می‌برند و تصویر کم‌کم پدیدار می‌شود. این مدل‌ها کند هستند، اما آموزش آن‌ها بسیار آسان‌تر است. در پنج سال گذشته همین ویژگی دوم بازی را برده است: هر تیم کوچکی می‌تواند یک مدل دیفیوژن را آموزش دهد و خروجی‌های معقولی بگیرد؛ در حالی که آموزش GAN مهارتی است که بعد از سال‌ها آزمون و خطای ناموفق به دست می‌آید.

گذشته از پایداری آموزش، ساختار گام‌به‌گام دیفیوژن همان چیزی است که قفل تمام قابلیت‌های مدرن تولید تصویر را باز می‌کند: شرطی‌سازی بر اساس متن (text conditioning)، این‌پینتینگ (inpainting)، ویرایش تصویر، وضوح‌افزایی (super-resolution) و کنترل سبک تصویر. هر گام از حلقه‌ی نمونه‌گیری، نقطه‌ای است که می‌توانید قید جدیدی به مدل تزریق کنید. همین امکان اتصال (hook) دلیل اصلی این است که Stable Diffusion، Imagen، DALL-E 3، Midjourney و در واقع هر مدل تولید تصویر قابل کنترلی که با آن سر و کار دارید، همگی مبتنی بر دیفیوژن هستند.

در این درس، حداقل پیاده‌سازی ممکن از DDPM را می‌سازیم: نویزدهی رو به جلو، نویززدایی رو به عقب و حلقه‌ی آموزش. در درس بعدی (Stable Diffusion) آن را با یک VAE، یک رمزگذار متن (text encoder) و هدایت بدون طبقه‌بند (classifier-free guidance) ترکیب می‌کنیم تا یک سیستم آماده برای محیط عملیاتی بسازیم.

مفاهیم پایه

فرآیند رو به جلو

یک تصویر x_0 بردارید. مقدار کمی نویز گوسی به آن اضافه کنید تا x_1 به دست آید. دوباره مقدار کمی نویز بیشتر اضافه کنید تا به x_2 برسید. این کار را برای $T$ گام ادامه دهید تا جایی که x_T عملاً از نویز خالص گوسی قابل تشخیص نباشد.

$$q(x_t \mid x_{t-1}) = \mathcal{N}\left(x_t; \sqrt{1 - \beta_t} \, x_{t-1}, \; \beta_t I\right)$$

پارامتر beta_t یک برنامه‌ی زمانی واریانس (variance schedule) کوچک است که معمولاً به‌صورت خطی از 0.0001 تا 0.02 طی $T=1000$ گام تغییر می‌کند. هر گام اندکی از سیگنال می‌کاهد و نویز تازه‌ای تزریق می‌کند.

جهش با فرم بسته

افزودن نویز گام‌به‌گام در واقع یک زنجیره‌ی مارکوف (Markov chain) است، اما فرمول ریاضی جمع می‌شود: می‌توانید x_t را مستقیماً از روی x_0 و در یک گام نمونه‌گیری کنید.

$$\begin{aligned} &\text{Define } \alpha_t = 1 - \beta_t \\ &\text{Define } \bar{\alpha}_t = \prod_{s=1}^{t} \alpha_s \\[6pt] &\text{Then: } q(x_t \mid x_0) = \mathcal{N}\left(x_t; \sqrt{\bar{\alpha}_t} \, x_0, \; (1 - \bar{\alpha}_t) I\right) \\[6pt] &\text{Equivalently: } x_t = \sqrt{\bar{\alpha}_t} \, x_0 + \sqrt{1 - \bar{\alpha}_t} \, \epsilon \quad \text{where } \epsilon \sim \mathcal{N}(0, I) \end{aligned}$$

تمام کاربردی بودن مدل‌های دیفیوژن در عمل فقط مدیون همین یک رابطه است. در طول آموزش، یک t تصادفی برمی‌دارید، x_t را مستقیماً از روی x_0 می‌سازید و مدل را در یک مرحله آموزش می‌دهید؛ بدون این‌که نیازی به شبیه‌سازی کل زنجیره‌ی مارکوف باشد.

فرآیند معکوس

فرآیند رو به جلو ثابت است. این فرآیند معکوس p(x_{t-1} | x_t) است که شبکه عصبی آن را یاد می‌گیرد. مدل‌های دیفیوژن x_{t-1} را مستقیماً پیش‌بینی نمی‌کنند؛ بلکه نویز epsilon را که در گام $t$ اضافه شده بود پیش‌بینی می‌کنند و محاسبات ریاضی x_{t-1} را از روی آن به دست می‌آورد.

تابع زیان آموزش

برای هر گام از آموزش:

1. یک تصویر واقعی x_0 نمونه‌گیری کنید.
2. یک گام زمانی t را به‌صورت یکنواخت از $[1, T]$ بردارید.
3. نویز epsilon ~ N(0, I) را نمونه‌گیری کنید.
4. مقدار x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * epsilon را محاسبه کنید.
5. مقدار epsilon_theta(x_t, t) را با شبکه پیش‌بینی کنید.
6. تابع زیان (loss) را با رابطه‌ی || epsilon - epsilon_theta(x_t, t) ||^2 کمینه کنید.

کل ماجرا همین است. شبکه عصبی یاد می‌گیرد که نویز را در هر گام زمانی پیش‌بینی کند. تابع زیان همان خطای میانگین مربعات (MSE) است. نه خبری از بازی تخاصمی است، نه فروپاشی حالت (mode collapse) رخ می‌دهد و نه نوسان‌های عجیب می‌بینید.

نمونه‌گیر (DDPM)

برای تولید خروجی: از x_T ~ N(0, I) شروع کنید و گام‌به‌گام به عقب برگردید.

for t = T, T-1, ..., 1:
    eps = model(x_t, t)
    x_{t-1} = (1 / sqrt(alpha_t)) * (x_t - (beta_t / sqrt(1 - alpha_bar_t)) * eps) + sqrt(beta_t) * z
    where z ~ N(0, I) if t > 1, else 0
return x_0

نکته‌ی کلیدی اینجاست: با این‌که توزیع شرطی معکوس در حالت کلی فرم بسته ندارد، برای این فرآیند خاصِ رو به جلوی گوسی دارای فرم بسته است. این ضرایب به‌ظاهر بدترکیب، همان روابطی هستند که قاعده‌ی بیز (Bayes' rule) پیش رویتان می‌گذارد.

چرا 1000 گام؟

برنامه‌ی زمانی نویز رو به جلو طوری انتخاب می‌شود که در هر گام فقط به اندازه‌ای نویز اضافه شود که مرحله‌ی معکوس تقریباً گوسی بماند. اگر تعداد گام‌ها خیلی کم باشد، فرآیند معکوس از حالت گوسی فاصله می‌گیرد و شبکه نمی‌تواند آن را به‌خوبی مدل کند. اگر تعداد گام‌ها بیش از حد زیاد باشد، نمونه‌گیری بسیار پرهزینه می‌شود و بازدهی چندانی هم اضافه نمی‌کند. انتخاب $T=1000$ با یک برنامه‌ی خطی، حالت پیش‌فرض DDPM است.

نمونه‌گیری DDIM: 20 برابر سریع‌تر

آموزش تغییری نمی‌کند؛ فقط فرآیند نمونه‌گیری عوض می‌شود. روش DDIM (سانگ و همکاران، 2020) یک فرآیند معکوس معین (deterministic) تعریف می‌کند که بدون نیاز به آموزش مجدد، از روی گام‌های زمانی می‌پرد. نمونه‌گیری در 50 گام با DDIM کیفیتی نزدیک به 1000 گام DDPM ارائه می‌دهد. تمام سیستم‌های عملیاتی از DDIM یا گونه‌های حتی سریع‌تر از آن (مثل DPM-Solver یا Euler ancestral) استفاده می‌کنند.

شرطی‌سازی زمانی

شبکه‌ی epsilon_theta(x_t, t) باید بداند در حال نویززدایی از کدام گام زمانی است. مدل‌های مدرن دیفیوژن، t را از طریق امبدینگ‌های زمانی سینوسی (همان ایده‌ی کدگذاری مکانی [positional encoding] در ترنسفورمرها) وارد می‌کنند؛ به این صورت که این امبدینگ‌ها به نقشه‌های ویژگی (feature maps) در تمام سطوح U-Net اضافه می‌شوند.

t_embedding = sinusoidal(t)
feature_map += MLP(t_embedding)

بدون شرطی‌سازی زمانی، شبکه مجبور است سطح نویز را از روی خود تصویر حدس بزند؛ روشی که کار می‌کند ولی از نظر بازدهی داده و نمونه‌ها (sample-efficiency) بسیار ضعیف‌تر است.

پیاده‌سازی

گام 1: برنامه‌ی زمانی نویز

import torch

def linear_beta_schedule(T=1000, beta_start=1e-4, beta_end=2e-2):
    return torch.linspace(beta_start, beta_end, T)


def precompute_schedule(betas):
    alphas = 1.0 - betas
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    return {
        "betas": betas,
        "alphas": alphas,
        "alphas_cumprod": alphas_cumprod,
        "sqrt_alphas_cumprod": torch.sqrt(alphas_cumprod),
        "sqrt_one_minus_alphas_cumprod": torch.sqrt(1.0 - alphas_cumprod),
        "sqrt_recip_alphas": torch.sqrt(1.0 / alphas),
    }

schedule = precompute_schedule(linear_beta_schedule(T=1000))

یک بار از قبل محاسبه کنید و در زمان آموزش و نمونه‌گیری با ایندکس فراخوانی نمایید.

گام 2: دیفیوژن رو به جلو (q_sample)

def q_sample(x0, t, noise, schedule):
    sqrt_a = schedule["sqrt_alphas_cumprod"][t].view(-1, 1, 1, 1)
    sqrt_one_minus_a = schedule["sqrt_one_minus_alphas_cumprod"][t].view(-1, 1, 1, 1)
    return sqrt_a * x0 + sqrt_one_minus_a * noise

فرم بسته‌ی تک‌خطی. t دسته‌ای از گام‌های زمانی است؛ یکی برای هر تصویر درون بچ (batch).

گام 3: یک U-Net کوچکِ شرطی‌شده بر زمان

import torch.nn as nn
import torch.nn.functional as F
import math

def timestep_embedding(t, dim=64):
    half = dim // 2
    freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / half)
    args = t[:, None].float() * freqs[None]
    emb = torch.cat([args.sin(), args.cos()], dim=-1)
    return emb


class TinyUNet(nn.Module):
    def __init__(self, img_channels=3, base=32, t_dim=64):
        super().__init__()
        self.t_mlp = nn.Sequential(
            nn.Linear(t_dim, base * 4),
            nn.SiLU(),
            nn.Linear(base * 4, base * 4),
        )
        self.t_dim = t_dim
        self.enc1 = nn.Conv2d(img_channels, base, 3, padding=1)
        self.enc2 = nn.Conv2d(base, base * 2, 4, stride=2, padding=1)
        self.mid = nn.Conv2d(base * 2, base * 2, 3, padding=1)
        self.dec1 = nn.ConvTranspose2d(base * 2, base, 4, stride=2, padding=1)
        self.dec2 = nn.Conv2d(base * 2, img_channels, 3, padding=1)
        self.time_proj = nn.Linear(base * 4, base * 2)

    def forward(self, x, t):
        t_emb = timestep_embedding(t, self.t_dim)
        t_emb = self.t_mlp(t_emb)
        t_proj = self.time_proj(t_emb)[:, :, None, None]

        h1 = F.silu(self.enc1(x))
        h2 = F.silu(self.enc2(h1)) + t_proj
        h3 = F.silu(self.mid(h2))
        d1 = F.silu(self.dec1(h3))
        d2 = torch.cat([d1, h1], dim=1)
        return self.dec2(d2)

یک U-Net دو سطحی که شرطی‌سازی زمانی در گلوگاه (bottleneck) به آن تزریق شده است. برای تصاویر واقعی، عمق و عرض لایه‌ها را بالا ببرید.

گام 4: حلقه‌ی آموزش

def train_step(model, x0, schedule, optimizer, device, T=1000):
    model.train()
    x0 = x0.to(device)
    bs = x0.size(0)
    t = torch.randint(0, T, (bs,), device=device)
    noise = torch.randn_like(x0)
    x_t = q_sample(x0, t, noise, schedule)
    pred = model(x_t, t)
    loss = F.mse_loss(pred, noise)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()

این کل حلقه‌ی آموزش است. نه از بازی‌های GAN خبری هست و نه تابع زیان اختصاصی؛ فقط یک فراخوانی ساده‌ی MSE.

گام 5: نمونه‌گیر (DDPM)

@torch.no_grad()
def sample(model, schedule, shape, T=1000, device="cpu"):
    model.eval()
    x = torch.randn(shape, device=device)
    betas = schedule["betas"].to(device)
    sqrt_one_minus_a = schedule["sqrt_one_minus_alphas_cumprod"].to(device)
    sqrt_recip_alphas = schedule["sqrt_recip_alphas"].to(device)

    for t in reversed(range(T)):
        t_batch = torch.full((shape[0],), t, dtype=torch.long, device=device)
        eps = model(x, t_batch)
        coef = betas[t] / sqrt_one_minus_a[t]
        mean = sqrt_recip_alphas[t] * (x - coef * eps)
        if t > 0:
            x = mean + torch.sqrt(betas[t]) * torch.randn_like(x)
        else:
            x = mean
    return x

1000 گذر رو به جلو برای تولید یک بچ از نمونه‌ها. در کد واقعی، این بخش را با یک نمونه‌گیر 50 گامی DDIM جایگزین می‌کنید.

گام 6: نمونه‌گیر DDIM (معین، حدود 20 برابر سریع‌تر)

@torch.no_grad()
def sample_ddim(model, schedule, shape, steps=50, T=1000, device="cpu", eta=0.0):
    model.eval()
    x = torch.randn(shape, device=device)
    alphas_cumprod = schedule["alphas_cumprod"].to(device)

    ts = torch.linspace(T - 1, -1, steps + 1).round().long()
    for i in range(steps):
        t = ts[i]
        t_prev = ts[i + 1]
        t_batch = torch.full((shape[0],), t, dtype=torch.long, device=device)
        eps = model(x, t_batch)
        a_t = alphas_cumprod[t]
        a_prev = alphas_cumprod[t_prev] if t_prev >= 0 else torch.tensor(1.0, device=device)
        x0_pred = (x - torch.sqrt(1 - a_t) * eps) / torch.sqrt(a_t)
        sigma = eta * torch.sqrt((1 - a_prev) / (1 - a_t) * (1 - a_t / a_prev))
        dir_xt = torch.sqrt(1 - a_prev - sigma ** 2) * eps
        noise = sigma * torch.randn_like(x) if eta > 0 else 0
        x = torch.sqrt(a_prev) * x0_pred + dir_xt + noise
    return x

برنامه‌ی زمانی در t_prev = -1 تمام می‌شود؛ بنابراین در آخرین گام a_prev = 1.0 خواهد بود و تصویر تمیز x0_pred بدون نویز باقی‌مانده تحویل داده می‌شود. با eta=0 فرآیند کاملاً معین است (ورودی نویز یکسان همیشه دقیقاً همان خروجی را تولید می‌کند). مقدار eta=1 در صورتی که تمام گام‌های زمانی پیموده شوند (steps=T)، مجدداً همان DDPM را بازسازی می‌کند.

کاربرد در محیط عملیاتی

برای کارهای سطح پروداکشن، از کتابخانه‌ی diffusers استفاده کنید:

from diffusers import DDPMScheduler, UNet2DModel

unet = UNet2DModel(sample_size=32, in_channels=3, out_channels=3, layers_per_block=2)
scheduler = DDPMScheduler(num_train_timesteps=1000)

این کتابخانه زمان‌بندهای آماده (DDPM، DDIM، DPM-Solver، Euler و Heun)، ساختارهای U-Net قابل تنظیم، پایپ‌لاین‌های تبدیل متن به تصویر و تصویر به تصویر، و ابزارهای کمکی فاین‌تیونینگ با LoRA را همراه خود دارد.

برای کارهای پژوهشی، مخزن k-diffusion (اثر کاترین کراوسون) دقیق‌ترین پیاده‌سازی‌های مرجع و بهترین نسخه‌های نمونه‌گیری را در اختیارتان می‌گذارد.

خروجی نهایی

این درس موارد زیر را به عنوان خروجی تحویل می‌دهد:

• outputs/prompt-diffusion-sampler-picker.md — یک پرامپت که بر اساس کیفیت هدف، محدودیت تأخیر زمانی (latency) و نوع شرطی‌سازی، مناسب‌ترین گزینه را بین DDPM / DDIM / DPM-Solver / Euler انتخاب می‌کند.
• outputs/skill-noise-schedule-designer.md — ابزاری که با دریافت T و سطح تخریب مد نظر، برنامه‌ی بتای خطی، کسینوسی یا سیگموئید می‌سازد و نمودارهای نسبت سیگنال به نویز (SNR) را در طول زمان رسم می‌کند.

تمرین‌ها

1. (ساده) فرآیند رو به جلو را مصورسازی کنید: یک تصویر بردارید و x_t را در t in [0, 100, 250, 500, 750, 1000] رسم کنید. بررسی کنید که آیا x_1000 شبیه نویز خالص گوسی هست یا نه.
2. (متوسط) مدل TinyUNet را روی مجموعه‌داده‌ی دایره‌های ساختگی (synthetic-circles) به مدت 20 ایپاک آموزش دهید و 16 دایره نمونه‌گیری کنید. نمونه‌گیری DDPM (با 1000 گام) و DDIM (با 50 گام) را با هم مقایسه کنید؛ آیا از یک سید نویز (noise seed) یکسان تصاویر مشابهی تولید می‌کنند؟
3. (سخت) یک برنامه‌ی زمانی نویز کسینوسی را پیاده‌سازی کنید (نیکول و دهاریوال، 2021): alpha_bar_t = cos^2((t/T + s) / (1 + s) * pi / 2). همین مدل را یک بار با برنامه‌ی خطی و یک بار با کسینوسی آموزش دهید و نشان دهید که زمان‌بندی کسینوسی در تعداد گام‌های کم، نمونه‌های بهتری تولید می‌کند.

اصطلاحات کلیدی

برای مطالعه‌ی بیشتر

• Denoising Diffusion Probabilistic Models (Ho et al., 2020) — مقاله‌ای که دیفیوژن را کاربردی کرد و از نظر معیار FID مدل‌های GAN را شکست داد
• Improved DDPM (Nichol & Dhariwal, 2021) — معرفی برنامه‌ی زمانی کسینوسی و بازپارامتری‌سازی v (v-parameterisation)
• DDIM (Song, Meng, Ermon, 2020) — نمونه‌گیر معینی که اینفرانس بلادرنگ (real-time) را ممکن کرد
• Elucidating the Design Space of Diffusion (Karras et al., 2022) — نگاهی یکپارچه به تمام تصمیم‌های طراحی دیفیوژن؛ در حال حاضر بهترین مرجع در این حوزه