扩散模型简介1：从物理原理到代码实现

扩散模型（Diffusion Models）近年来在生成式人工智能领域取得了显著进展，尤其在图像生成任务中表现优异。本文将从扩散模型的物理原理入手，深入解析其工作机制，并通过代码实现展示其实际应用。

一、扩散模型的基本概念

扩散模型是一类生成模型，其核心思想是通过逐步向数据添加噪声，直到数据变为纯噪声，然后学习如何反向去噪，从而生成与训练数据分布相似的新样本。这个过程可以分为两个阶段：

  1. 正向扩散过程（Forward Diffusion Process）：将数据逐步添加噪声，直至数据完全被噪声淹没。
  2. 反向去噪过程（Reverse Denoising Process）：学习如何从噪声中逐步恢复出原始数据。

二、物理原理解析

1. 正向扩散过程

正向扩散过程可以被看作是一个马尔可夫链，每一步向数据添加少量的高斯噪声。数学上，假设我们有一个数据分布 $ q(\mathbf{x}_0) $，正向过程定义为：

其中，$ \beta_t $ 是一个调控噪声程度的参数，通常随时间步 $ t $ 线性或非线性增长。随着 $ t $ 增加，数据逐步被噪声覆盖，最终在 $ t = T $ 时，数据近似于纯噪声。

2. 反向去噪过程

反向过程的目标是从噪声中逐步恢复出原始数据。这个过程同样可以通过一个马尔可夫链实现：

其中，$ \mu_\theta $ 和 $ \Sigma_\theta $ 是需要学习的参数，通过神经网络参数化。训练的目标是最小化正向过程和反向过程之间的差异，通常采用变分下界（Variational Lower Bound, VLB）作为损失函数。

3. 训练目标

扩散模型的训练目标是最大化数据在反向过程下的似然：

其中，$ D_{\mathrm{KL}} $ 是KL散度，用于衡量两个分布之间的差异。

三、代码实现

下面，我们将通过PyTorch框架实现一个简单的扩散模型，用于生成MNIST手写数字。

1. 环境准备

首先，确保安装了必要的库：

1
pip install torch torchvision matplotlib

2. 导入库

1
2
3
4
5
6
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np

3. 定义扩散过程参数

1
2
3
4
5
6
T = 1000  # 总时间步数
beta_start = 1e-4
beta_end = 0.02
beta = torch.linspace(beta_start, beta_end, T)
alpha = 1.0 - beta
alpha_bar = torch.cumprod(alpha, dim=0)

4. 定义神经网络

使用一个简单的U-Net结构来预测噪声：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
class SimpleUNet(nn.Module):
    def __init__(self):
        super(SimpleUNet, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(28*28 + 1, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, 28*28)
        )
    
    def forward(self, x, t):
        # x: [batch_size, 28*28]
        # t: [batch_size, 1]
        t = t.unsqueeze(1)
        x = torch.cat([x, t], dim=1)
        return self.net(x)

5. 定义正向过程和损失函数

1
2
3
4
5
6
7
8
9
10
def q_sample(x0, t, noise=None):
    if noise is None:
        noise = torch.randn_like(x0)
    return torch.sqrt(alpha_bar[t]) * x0 + torch.sqrt(1 - alpha_bar[t]) * noise

def loss_fn(model, x0, t):
    noise = torch.randn_like(x0)
    xt = q_sample(x0, t, noise)
    predicted_noise = model(xt, t.float())
    return nn.MSELoss()(predicted_noise, noise)

6. 训练模型

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
# 数据加载
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.view(-1))
])

train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True)

# 初始化模型和优化器
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = SimpleUNet().to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# 训练循环
for epoch in range(10):
    model.train()
    epoch_loss = 0
    for batch_idx, (x, _) in enumerate(train_loader):
        x = x.to(device)
        batch_size = x.size(0)
        t = torch.randint(0, T, (batch_size,), device=device).long()
        optimizer.zero_grad()
        loss = loss_fn(model, x, t)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
    avg_loss = epoch_loss / len(train_loader)
    print(f'Epoch {epoch+1}, Loss: {avg_loss:.4f}')

7. 生成样本

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
@torch.no_grad()
def sample(model, T, size):
    model.eval()
    x = torch.randn(size, 28*28).to(device)
    for t in reversed(range(T)):
        t_batch = torch.full((size,), t, device=device, dtype=torch.long)
        predicted_noise = model(x, t_batch.float())
        if t > 0:
            x = (x - torch.sqrt(1 - alpha_bar[t]) * predicted_noise) / torch.sqrt(alpha[t])
            noise = torch.randn_like(x) * torch.sqrt(beta[t])
            x += noise
        else:
            x = (x - torch.sqrt(1 - alpha_bar[t]) * predicted_noise) / torch.sqrt(alpha[t])
    return x

# 生成并显示样本
samples = sample(model, T, 16).cpu()
samples = samples.view(-1, 1, 28, 28)
grid = torchvision.utils.make_grid(samples, nrow=4)
plt.imshow(grid.permute(1, 2, 0).squeeze())
plt.show()

四、总结

扩散模型通过模拟物理中的扩散过程，将数据逐步加噪声直至纯噪声，然后学习如何逆向去噪生成新数据。其核心在于设计有效的正向和反向过程，以及选择合适的神经网络结构来预测噪声。本文通过理论解析和代码实现，展示了扩散模型的基本原理和应用方法。随着研究的深入，扩散模型在图像生成、语音合成等领域展现出广阔的应用前景。