擴散模型(Diffusion Model)是一類以逐步加噪與去噪過程為核心的生成模型,其中最具代表性的版本為 Ho 等人於 2020 年提出的去噪擴散機率模型(Denoising Diffusion Probabilistic Model,DDPM)。相較於 VAE 透過重建損失引導 latent space 的學習,以及 GAN 透過對抗訓練來驅動生成器改善輸出,Diffusion Model 的目標則是讓模型能看著一張被加了數次噪聲的影像,猜出當初被加入了怎樣的噪聲,進而能夠從純高斯噪聲出發,逐步還原出真實的資料分布。這種設計使訓練過程相對穩定,且生成樣本的多樣性通常優於 GAN。本篇章將以 Fashion-MNIST 為例,介紹 DDPM 的基本原理與實作流程。

DDPM 的運作分為兩個方向相反的過程。前向過程(forward process)是固定的,對原始資料 x0 在每個時間步 t 逐步疊加高斯噪聲,得到一系列逐漸模糊的樣本 x1, x2, ..., xT;當 T 夠大時,xT 的分布將趨近於標準常態分布。每一步的條件分布定義為:

q(xt | xt−1) = N(xt; (1−βt)1/2 xt−1, βtI)

其中 N 為常態分佈;βt 為事先設計好的噪聲排程(noise schedule),是一組從小到大的固定常數;I 是單位矩陣(identity matrix),表示各像素的噪聲是獨立的,rows 和 columns 的數目皆為影像像素數量的平方。前向過程有一個重要的特性:可以不經過逐步計算,直接從 x0 一次跳到任意時間步 t 的結果。令 αt = 1 − βt,ᾱt = ∏s=1t αs,則:

q(xt | x0) = N(xt; ᾱt1/2 x0, (1−ᾱt)I)

亦即給定 x0 與時間步 t,可以直接對噪聲 ε ~ N(0, I) 取樣並計算出 xt

xt = ᾱt1/2 x0 + (1−ᾱt)1/2 ε

此性質使得訓練時能夠高效率地對任意 t 取樣,而不必逐步執行前向過程。反向過程(reverse process)則是由神經網路驅動的,其目標是學習從 xt 估計出 xt−1 的條件分布。DDPM 將此分布參數化為:

pθ(xt−1 | xt) = N(xt−1; μθ(xt, t), σt2I)

其中,pθ 代表模型學出的機率分布,μθ 代表該分布的平均值,σt2 則是該分布的變異數;下標的 θ 代表分布參數是由模型所學出,σ 在原始論文中直接使用 βt 或其他固定值,後續亦有人改用模型來學。而實作上,網路並非直接預測均值 μθ,而是預測加入 xt 中的噪聲 εθ(xt, t),並透過以下關係式將兩者互換:

μθ(xt, t) = (1 / αt1/2) (xt − (βt / (1−ᾱt)1/2) εθ(xt, t))

因此訓練目標簡化為預測噪聲的均方誤差:

L = Et, x0, ε [‖ε − εθ(xt, t)‖2]

其中 t 在每次訓練迭代中從 {1, ..., T} 均勻取樣,xt 則由上述的 closed form 公式計算而來。εθ 的網路架構需要同時接收帶噪影像 xt 與時間步 t 作為輸入,其中 t 的注入方式通常是透過時間步嵌入(timestep embedding),其做法類似 Transformer 的位置編碼,以正弦與餘弦函數將純量 t 轉換為向量表示,再注入至網路的各個區塊中。

以下是根據上述過程,以 Fashion-MNIST 為例的完整 DDPM 訓練與生成範例:

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

T = 300
beta = torch.linspace(0.0001, 0.02, T).to(device)
alpha = 1.0 - beta
alpha_bar = torch.cumprod(alpha, dim=0)


def get_timestep_embedding(t, t_dim):
	half = t_dim // 2
	freq = torch.exp(-torch.arange(half, device=t.device) * (torch.log(torch.tensor(10000.0)) / (half - 1)))
	args = t[:, np.newaxis].float() * freq[np.newaxis, :]
	return torch.cat([torch.sin(args), torch.cos(args)], dim=-1)


def q_sample(x0, t):
	ab = alpha_bar[t][:, np.newaxis, np.newaxis, np.newaxis]
	eps = torch.randn_like(x0)
	return ab ** 0.5 * x0 + (1 - ab) ** 0.5 * eps, eps


class ResBlock(nn.Module):
	def __init__(self, in_channels, out_channels, t_dim):
		super().__init__()
		self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
		self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
		self.norm1 = nn.GroupNorm(4, out_channels)
		self.norm2 = nn.GroupNorm(4, out_channels)
		self.t_proj = nn.Linear(t_dim, out_channels)
		self.relu = nn.ReLU()
		if in_channels != out_channels:
			self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1)
		else:
			self.shortcut = nn.Identity()

	def forward(self, x, t_emb):
		h = self.relu(self.norm1(self.conv1(x)))
		h = h + self.t_proj(t_emb)[:, :, np.newaxis, np.newaxis]
		h = self.relu(self.norm2(self.conv2(h)))
		return h + self.shortcut(x)


class UNet(nn.Module):
	def __init__(self, t_dim=128):
		super().__init__()
		self.t_dim = t_dim
		self.t_mlp = nn.Sequential(
			nn.Linear(t_dim, t_dim * 2),
			nn.ReLU(),
			nn.Linear(t_dim * 2, t_dim)
		)
		self.enc1 = ResBlock(1, 32, t_dim)
		self.enc2 = ResBlock(32, 64, t_dim)
		self.down = nn.MaxPool2d(2)
		self.mid = ResBlock(64, 64, t_dim)
		self.up = nn.Upsample(scale_factor=2, mode='nearest')
		self.dec1 = ResBlock(64 + 32, 32, t_dim)
		self.dec2 = ResBlock(32, 32, t_dim)
		self.out = nn.Conv2d(32, 1, kernel_size=1)

	def forward(self, x, t):
		t_emb = get_timestep_embedding(t, self.t_dim)
		t_emb = self.t_mlp(t_emb)
		h1 = self.enc1(x, t_emb)
		h2 = self.enc2(self.down(h1), t_emb)
		h = self.mid(h2, t_emb)
		h = self.dec1(torch.cat([self.up(h), h1], dim=1), t_emb)
		h = self.dec2(h, t_emb)
		return self.out(h)


train_set = datasets.FashionMNIST(root='./data', train=True, download=True, transform=transforms.ToTensor())
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)

model = UNet().to(device)
optim = torch.optim.Adam(model.parameters(), lr=1e-3)

model.train()
for i in range(10):
	print(f'Epoch {i+1}/10')
	for x0, _ in train_loader:
		x0 = x0.to(device)
		t = torch.randint(0, T, (x0.shape[0],), device=device)
		xt, eps = q_sample(x0, t)
		loss = F.mse_loss(model(xt, t), eps)
		optim.zero_grad()
		loss.backward()
		optim.step()

n_show = 8
show_steps = np.linspace(T, 1, 7).astype(int)

model.eval()
with torch.no_grad():
	xt = torch.randn(n_show, 1, 28, 28).to(device)
	frames = {}
	for step in range(T, 0, -1):
		t_batch = torch.full((n_show,), step - 1, device=device, dtype=torch.long)
		eps_pred = model(xt, t_batch)
		a = alpha[step - 1]
		ab = alpha_bar[step - 1]
		mu = (1 / a ** 0.5) * (xt - (beta[step - 1] / (1 - ab) ** 0.5) * eps_pred)
		if step > 1:
			xt = mu + beta[step - 1] ** 0.5 * torch.randn_like(xt)
		else:
			xt = mu
		if step in show_steps:
			frames[step] = xt.cpu()

col_steps = list(show_steps) + [0]
cols = [frames[s] if s != 0 else xt.cpu() for s in col_steps]
grid_rows = []
for row in range(n_show):
	row_imgs = [c[row, 0].numpy() for c in cols]
	grid_rows.append(np.concatenate(row_imgs, axis=1))
grid = np.concatenate(grid_rows, axis=0)

tick_positions = [28 * i + 14 for i in range(len(col_steps))]
labels = [f't={s}' for s in col_steps]

plt.imshow(grid, cmap='gray')
plt.xticks(tick_positions, labels)
plt.yticks([])
plt.show()

在上述範例中:

Diffusion model 也跟 AE 和 GAN 一樣,可以加上 condition。由於改動部分不多,因此此處只示範 model class 的更動,如下:

class UNet(nn.Module):
	def __init__(self, t_dim=128, num_classes=10):
		super().__init__()
		self.t_dim = t_dim
		self.label_emb = nn.Embedding(num_classes, t_dim)
		self.t_mlp = nn.Sequential(
			nn.Linear(t_dim, t_dim * 2),
			nn.ReLU(),
			nn.Linear(t_dim * 2, t_dim)
		)
		self.enc1 = ResBlock(1, 32, t_dim)
		self.enc2 = ResBlock(32, 64, t_dim)
		self.down = nn.MaxPool2d(2)
		self.mid = ResBlock(64, 64, t_dim)
		self.up = nn.Upsample(scale_factor=2, mode='nearest')
		self.dec1 = ResBlock(64 + 32, 32, t_dim)
		self.dec2 = ResBlock(32, 32, t_dim)
		self.out = nn.Conv2d(32, 1, kernel_size=1)

	def forward(self, x, t, label):
		t_emb = get_timestep_embedding(t, self.t_dim)
		t_emb = self.t_mlp(t_emb) + self.label_emb(label)
		h1 = self.enc1(x, t_emb)
		h2 = self.enc2(self.down(h1), t_emb)
		h = self.mid(h2, t_emb)
		h = self.dec1(torch.cat([self.up(h), h1], dim=1), t_emb)
		h = self.dec2(h, t_emb)
		return self.out(h)

在上述範例中,加入 condition 的方式跟 timestep embedding 一樣;而維度也設定為相符,則主要是為了實作上的方便,你也可以試試看使用不同的維度,但會需要另外撰寫一些轉換用的模組。

另外一個值得一提的變體是生成新圖案的方式。我們先前的做法,大致上是逐步的預測噪聲並去除,但也可以改為在預測噪聲後,先估計出 x0,再利用 x0 的估計值與預測的噪聲,直接內插出噪聲程度更低的某個 xt'(t' < t),來讓生成過程可以不用乖乖的一步一步一步跑。這個作法只需要修改預測的部分,如下:

step_seq = np.linspace(T, 1, 50).astype(int)
step_pairs = list(zip(step_seq[:-1], step_seq[1:]))

model.eval()
with torch.no_grad():
	xt = torch.randn(n_show, 1, 28, 28).to(device)
	frames = {}
	for step, step_prev in step_pairs:
		t_batch = torch.full((n_show,), step - 1, device=device, dtype=torch.long)
		eps_pred = model(xt, t_batch)
		ab = alpha_bar[step - 1]
		ab_prev = alpha_bar[step_prev - 1]
		x0_pred = (xt - (1 - ab) ** 0.5 * eps_pred) / ab ** 0.5
		xt = ab_prev ** 0.5 * x0_pred + (1 - ab_prev) ** 0.5 * eps_pred
		if step in show_steps:
			frames[step] = xt.cpu()
	frames[0] = xt.cpu()

Diffusion model 還有其他不少的著名使用案例與延伸變體,如: