AI扩散模型加噪原理详解：从数学基础到PyTorch实现

1次阅读

共计 2509 个字符，预计需要花费 7 分钟才能阅读完成。

最近 AI 生成领域最火的技术莫过于扩散模型（Diffusion Models）了。这种模型通过逐步加噪和去噪的过程，能够生成高质量的图像、音频等内容。作为一名刚开始接触这个领域的新手，我最困惑的就是其中的加噪原理——为什么要加噪？怎么加噪？今天我就结合自己的学习过程，用尽量通俗的语言和大家分享这方面的知识。

扩散模型的核心思想其实来源于物理学中的扩散现象。想象一下，把一滴墨水滴入清水中，墨水分子会逐渐扩散，直到完全均匀分布在水里。这个过程是不可逆的，而扩散模型就是模拟这个过程的逆过程。

从数学上看，这个过程可以用马尔可夫链来描述。马尔可夫链是指下一个状态只依赖于当前状态，而与之前所有状态无关的随机过程。在扩散模型中：

前向过程（加噪）：逐步向数据添加高斯噪声
反向过程（去噪）：学习如何逐步去除这些噪声

扩散模型可以看作是一种特殊的变分自编码器（VAE）。变分推断是一种近似复杂分布的方法，在这里用于估计真实数据分布。通过最小化变分下界（ELBO），模型可以学习到有效的去噪过程。

不同的噪声调度策略会影响模型的学习效果。常见的策略有：

线性调度：噪声强度随时间线性增加
余弦调度：噪声强度随时间遵循余弦曲线变化
平方调度：噪声强度随时间平方增加

数学表达式上，线性调度可以表示为：

βₜ = β_min + (β_max – β_min) * (t/T)

而余弦调度则表示为：

βₜ = cos(t/T * π/2)

import torch

def gaussian_noise(x, t, beta):
    """
    生成高斯噪声并加到输入数据上
    :param x: 输入数据 (batch_size, ...)
    :param t: 时间步 (batch_size,)
    :param beta: 噪声强度调度
    :return: 加噪后的数据
    """
    noise = torch.randn_like(x)
    alpha = 1 - beta[t]
    alpha_bar = torch.prod(alpha)
    return torch.sqrt(alpha_bar) * x + torch.sqrt(1 - alpha_bar) * noise

class NoiseScheduler:
    def __init__(self, num_steps=1000, schedule_type='linear', beta_start=1e-4, beta_end=0.02):
        self.num_steps = num_steps
        self.schedule_type = schedule_type

        if schedule_type == 'linear':
            self.betas = torch.linspace(beta_start, beta_end, num_steps)
        elif schedule_type == 'cosine':
            steps = torch.arange(num_steps) / num_steps
            self.betas = torch.cos(steps * math.pi / 2)
        else:
            raise ValueError(f"Unknown schedule type: {schedule_type}")

        self.alphas = 1. - self.betas
        self.alpha_bars = torch.cumprod(self.alphas, dim=0)

    def get_noise(self, x, t):
        return gaussian_noise(x, t, self.betas)

def train_diffusion(model, dataloader, scheduler, epochs=100):
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

    for epoch in range(epochs):
        for batch in dataloader:
            # 1. 随机采样时间步
            t = torch.randint(0, scheduler.num_steps, (batch.size(0),))

            # 2. 加噪
            noisy_data = scheduler.get_noise(batch, t)

            # 3. 预测噪声
            pred_noise = model(noisy_data, t)

            # 4. 计算损失
            loss = F.mse_loss(pred_noise, noisy_data - batch)

            # 5. 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        print(f"Epoch {epoch}, Loss: {loss.item()}")

通过实验发现，噪声尺度的选择对训练稳定性有很大影响：