扩散模型完全指南：从数学原理到工程实践

# 代码：前向加噪过程（可直接采样任意时刻）
def forward_diffusion(x0, t, sigma_q):
    """
    Args:
        x0: 原始数据 [B, d]
        t: 时间步 [B] 或标量，范围 [0, 1]
        sigma_q: 终端噪声标准差
    Returns:
        x_t: 加噪后的数据 [B, d]
        noise: 添加的噪声 [B, d]
    """
    sigma_t = sigma_q * torch.sqrt(t)  # σ_t = σ_q * √t
    noise = torch.randn_like(x0)
    return x0 + sigma_t * noise, noise

# 代码：扩散训练损失（基础版）
def diffusion_loss(model, x0, schedule, batch_size):
    """
    Args:
        model: 神经网络 f_θ(x_t, t) → μ_t(x_t)
        schedule: 噪声调度器，提供 σ_t
        x0: 原始数据批次 [B, d]
    """
    # 1. 随机采样时间和噪声
    t = torch.rand(batch_size, device=x0.device)  # t ~ Uniform[0,1]
    sigma_t = schedule(t)  # σ_t
    eps = torch.randn_like(x0)  # ε ~ N(0, I)
    
    # 2. 前向加噪：利用可直接采样性质 (2)
    x_t = x0 + sigma_t.unsqueeze(1) * eps  # [B, d]
    
    # 3. 预测去噪后的值（即条件期望）
    mu_pred = model(x_t, t)  # f_θ(x_t, t) ≈ E[x_{t-Δt}|x_t]
    
    # 4. 回归损失（L2）
    # 注意：实际训练中通常预测的是 x_{t+Δt} 给定 x_t
    loss = torch.mean((mu_pred - x0)**2)
    return loss

# 代码：DDPM采样（随机版本）
@torch.no_grad()
def ddpm_sample(model, schedule, steps=1000, batch_size=1):
    """
    DDPM随机采样器实现
    """
    device = next(model.parameters()).device
    sigma_q = schedule.sigma_max
    
    # 1. 从纯噪声开始: x₁ ~ N(0, σ_q²)
    x = torch.randn(batch_size, *model.input_shape, device=device) * sigma_q
    
    # 2. 时间步调度
    sigmas = schedule.get_sigmas(steps)  # [steps+1]
    
    # 3. 反向迭代
    for i in range(steps):
        t = (steps - i) / steps  # 当前时间
        sigma_t = sigmas[i]
        sigma_next = sigmas[i+1]
        
        # 预测条件期望: μ_{t-Δt}(x_t)
        mu_pred = model(x, torch.full((batch_size,), t, device=device))
        
        # 添加随机噪声（关键！保证采样多样性）
        noise = torch.randn_like(x) * torch.sqrt(sigma_t**2 - sigma_next**2)
        x = mu_pred + noise
    
    return x

# 代码：DDIM采样（确定性版本）
@torch.no_grad()
def ddim_sample(model, schedule, steps=50, batch_size=1):
    """
    DDIM确定性采样器实现
    """
    device = next(model.parameters()).device
    sigma_q = schedule.sigma_max
    
    # 1. 从纯噪声开始
    x = torch.randn(batch_size, *model.input_shape, device=device) * sigma_q
    sigmas = schedule.get_sigmas(steps)
    
    for i in range(steps):
        t = (steps - i) / steps
        sigma_t = sigmas[i]
        sigma_next = sigmas[i+1]
        
        # 预测条件期望
        mu_pred = model(x, torch.full((batch_size,), t, device=device))
        
        # DDIM更新：确定性传输
        lambda_factor = sigma_t / (sigma_t - sigma_next + 1e-8)
        x = x + lambda_factor * (mu_pred - x)
    
    return x

# 代码：统一参数化框架
class DiffusionModel(nn.Module):
    def __init__(self, prediction_type='epsilon'):
        """
        prediction_type: 'x0', 'epsilon', 'v', or 'x_prev'
        """
        super().__init__()
        self.prediction_type = prediction_type
        self.net = UNet(...)  # 实际网络架构
    
    def compute_loss(self, x0, t, schedule):
        """计算训练损失"""
        sigma_t = schedule(t)
        eps = torch.randn_like(x0)
        
        # 前向加噪
        x_t = x0 + sigma_t.unsqueeze(1) * eps
        
        # 根据参数化类型确定目标
        if self.prediction_type == 'x0':
            target = x0
        elif self.prediction_type == 'epsilon':
            target = eps
        elif self.prediction_type == 'v':
            # v-prediction: v = α_t·ε - σ_t·x₀
            alpha_t = torch.sqrt(1 - sigma_t**2)
            target = alpha_t.unsqueeze(1) * eps - sigma_t.unsqueeze(1) * x0
        elif self.prediction_type == 'x_prev':
            # 预测 x_{t-Δt}
            sigma_prev = schedule(t - 1/1000)  # 简化
            target = (sigma_prev/sigma_t) * x0 + ...  # 略
        
        # 网络预测
        pred = self.net(x_t, t)
        return torch.mean((pred - target)**2)

# 代码：Flow Matching训练（线性流）
def flow_matching_loss(model, x0, x1, t):
    """
    Args:
        x0: 目标数据 ~ p
        x1: 源噪声 ~ q (如高斯)
        t: 时间 ~ Uniform[0,1]
    """
    # 线性插值轨迹
    x_t = t * x1 + (1 - t) * x0
    
    # 线性流的速度场（常数）
    v_target = x0 - x1
    
    # 回归损失
    v_pred = model(x_t, t)
    return torch.mean((v_pred - v_target)**2)

# 代码：ODE求解器采样（使用torchdiffeq）
from torchdiffeq import odeint

class ProbabilityFlowODE:
    def __init__(self, score_model):
        self.score_model = score_model  # 学习 ∇log p_t(x)
    
    def reverse_ode(self, t, x):
        """ODE的右端函数: dx/dt = f(x,t)"""
        score = self.score_model(x, t)  # ≈ ∇log p_t(x)
        return -0.5 * score  # 简化版本（σ_q=1）
    
    def sample(self, x1, t_span):
        """从t=1积分到t=0"""
        return odeint(self.reverse_ode, x1, t_span, method='dopri5')

# 代码：多种噪声调度实现
class NoiseSchedule:
    """噪声调度基类"""
    def __call__(self, t):
        raise NotImplementedError

class ScheduleLogLinear(NoiseSchedule):
    """对数线性调度：σ ∈ [σ_min, σ_max] 对数均匀 [Karras et al., 2022]"""
    def __init__(self, N, sigma_min=0.02, sigma_max=10.0):
        self.sigmas = torch.logspace(
            math.log10(sigma_min), 
            math.log10(sigma_max), 
            N
        )
    
    def __call__(self, t):
        idx = (t * (len(self.sigmas)-1)).long().clamp(0, len(self.sigmas)-2)
        return self.sigmas[idx]

class ScheduleDDPM(NoiseSchedule):
    """DDPM方差保持调度：Var(x_t) ≈ 常数 [Ho et al., 2020]"""
    def __init__(self, N, beta_min=0.0001, beta_max=0.02):
        betas = torch.linspace(beta_min, beta_max, N)
        alphas = 1 - betas
        self.alphas_cumprod = torch.cumprod(alphas, dim=0)
    
    def __call__(self, t):
        idx = (t * (len(self.alphas_cumprod)-1)).long()
        return torch.sqrt(1 - self.alphas_cumprod[idx])

class ScheduleKarras(NoiseSchedule):
    """Karras调度：σ(t) = t, 配合整体缩放 [Karras et al., 2022]"""
    def __init__(self, N, sigma_min=0.002, sigma_max=80.0):
        self.sigmas = ((sigma_max**(1/3) + 
                       torch.linspace(0, 1, N) * (sigma_min**(1/3) - sigma_max**(1/3)))**3)
    
    def __call__(self, t):
        idx = (t * (len(self.sigmas)-1)).long()
        return self.sigmas[idx]

# 概念代码：潜空间扩散训练流程
class LatentDiffusion:
    def __init__(self, autoencoder, unet, schedule):
        self.autoencoder = autoencoder  # 预训练，冻结
        self.unet = unet                 # 可训练
        self.schedule = schedule
    
    def train_step(self, x0, text_condition=None):
        # 1. 编码到潜空间（冻结）
        with torch.no_grad():
            z0 = self.autoencoder.encode(x0)
        
        # 2. 潜空间扩散训练
        t = torch.rand(z0.shape[0], device=z0.device)
        sigma_t = self.schedule(t)
        eps = torch.randn_like(z0)
        z_t = z0 + sigma_t.unsqueeze(1) * eps
        
        # 3. UNet预测（可加入text condition）
        eps_pred = self.unet(z_t, t, context=text_condition)
        
        return torch.mean((eps_pred - eps)**2)
    
    @torch.no_grad()
    def sample(self, condition=None, steps=50):
        # 1. 潜空间采样
        z_1 = torch.randn(1, *self.latent_shape)
        z_0 = ddim_sample(self.unet, self.schedule, steps)(z_1, condition)
        
        # 2. 解码回像素空间
        x_0 = self.autoencoder.decode(z_0)
        return x_0

# 完整可运行示例：2D螺旋数据扩散
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

# 1. 螺旋数据生成
def swiss_roll(n_samples=1000):
    """生成2D螺旋数据集"""
    n = np.sqrt(np.random.rand(n_samples)) * 5 * np.pi
    d = np.random.rand(n_samples) * 2 - 1
    x = n * np.cos(n) + 0.5 * d * np.sin(n)
    y = n * np.sin(n) + 0.5 * d * np.cos(n)
    return torch.stack([torch.tensor(x), torch.tensor(y)], dim=1).float()

# 2. 时间嵌入 + MLP去噪器
class TimeEmbedding(nn.Module):
    def __call__(self, sigma):
        """简单的2维时间嵌入"""
        sigma = sigma.unsqueeze(1)
        return torch.cat([
            torch.sin(torch.log(sigma) / 2),
            torch.cos(torch.log(sigma) / 2)
        ], dim=1)

class DenoiserMLP(nn.Module):
    def __init__(self, dim=2, hidden_dims=[16, 128, 128, 16]):
        super().__init__()
        layers = []
        dims = [dim + 2] + hidden_dims  # +2 for time embedding
        for i in range(len(dims)-1):
            layers.extend([
                nn.Linear(dims[i], dims[i+1]),
                nn.GELU()
            ])
        layers.append(nn.Linear(hidden_dims[-1], dim))
        self.net = nn.Sequential(*layers)
        self.time_emb = TimeEmbedding()
    
    def forward(self, x, sigma):
        t_emb = self.time_emb(sigma)  # [B, 2]
        inp = torch.cat([x, t_emb], dim=1)  # [B, dim+2]
        return self.net(inp)

# 3. 训练循环
def train_diffusion(model, data_loader, schedule, epochs=5000, lr=1e-3):
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    
    for epoch in range(epochs):
        for x0 in data_loader:
            optimizer.zero_grad()
            
            # 采样训练样本
            t = torch.rand(x0.shape[0])
            sigma_t = schedule(t)
            eps = torch.randn_like(x0)
            x_t = x0 + sigma_t.unsqueeze(1) * eps
            
            # 预测噪声（这里用epsilon参数化）
            eps_pred = model(x_t, sigma_t)
            loss = torch.mean((eps_pred - eps)**2)
            
            loss.backward()
            optimizer.step()
        
        if epoch % 500 == 0:
            print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

# 4. 采样与可视化
@torch.no_grad()
def sample_and_plot(model, schedule, n_samples=200, steps=20):
    model.eval()
    sigmas = schedule.get_sigmas(steps)
    
    # 从噪声开始
    x = torch.randn(n_samples, 2) * sigmas[0]
    
    # DDIM采样轨迹记录
    trajectory = [x.clone()]
    for i in range(steps):
        t = (steps - i) / steps
        sigma_t = sigmas[i]
        sigma_next = sigmas[i+1]
        
        # 预测x0（epsilon参数化转x0）
        eps_pred = model(x, torch.full((n_samples,), t))
        x0_pred = (x - sigma_t * eps_pred) / torch.sqrt(1 - sigma_t**2 + 1e-8)
        
        # DDIM更新
        direction = (x - sigma_t * x0_pred) / torch.sqrt(1 - sigma_t**2 + 1e-8)
        x = sigma_next * x0_pred + torch.sqrt(1 - sigma_next**2 + 1e-8) * direction
        trajectory.append(x.clone())
    
    # 可视化
    fig, axes = plt.subplots(1, 4, figsize=(16, 4))
    for idx, t_idx in enumerate([0, 7, 14, 20]):
        ax = axes[idx]
        ax.scatter(trajectory[t_idx][:,0], trajectory[t_idx][:,1], 
                  s=2, alpha=0.5, c='blue')
        ax.set_title(f'Step {t_idx}/{steps} (t={1-t_idx/steps:.2f})')
        ax.set_xlim(-10, 10)
        ax.set_ylim(-10, 10)
        ax.grid(alpha=0.3)
    plt.tight_layout()
    plt.show()

# 技巧1：梯度裁剪（防止梯度爆炸）
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# 技巧2：学习率调度（余弦退火）
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=epochs, eta_min=1e-5
)

# 技巧3：损失加权（不同噪声水平不同权重）
def weighted_loss(eps_pred, eps, sigma_t, weighting='uniform'):
    if weighting == 'uniform':
        weights = torch.ones_like(sigma_t)
    elif weighting == 'snr':  # Signal-to-Noise Ratio weighting
        weights = 1 / (sigma_t**2 + 1e-6)
    elif weighting == 'min_snr':  # Min-SNR weighting [Hang et al.]
        snr = 1 / sigma_t**2
        weights = torch.minimum(snr, torch.ones_like(snr) * 5)
    return torch.mean(weights * (eps_pred - eps)**2)

# 技巧1：增加采样步数（但计算成本↑）
# 技巧2：使用更高级的ODE求解器（DPM-Solver, Heun, etc.）
# 技巧3：添加引导（classifier-free guidance）

def classifier_free_guidance(model, x_t, t, cond, uncond, guidance_scale=7.5):
    """CFG: 增强条件生成的质量"""
    # 条件预测
    eps_cond = model(x_t, t, context=cond)
    # 无条件预测（训练时随机drop condition）
    eps_uncond = model(x_t, t, context=uncond)
    # 引导组合
    return eps_uncond + guidance_scale * (eps_cond - eps_uncond)

# 技巧4：动态阈值（防止采样发散）
def dynamic_thresholding(x, percentile=99.5):
    """防止采样值过大导致图像异常"""
    s = torch.quantile(x.abs().reshape(x.shape[0], -1), percentile/100, dim=1)
    s = torch.clamp(s, min=1.0)
    return x.clamp(-s.unsqueeze(1), s.unsqueeze(1))

# 概念验证：记忆化检测
def check_memorization(generated_samples, train_samples, threshold=0.1):
    """检测生成样本是否记忆训练数据"""
    from scipy.spatial.distance import cdist
    
    # 计算生成样本与训练样本的最小L2距离
    distances = cdist(
        generated_samples.reshape(len(generated_samples), -1),
        train_samples.reshape(len(train_samples), -1),
        metric='euclidean'
    )
    min_distances = distances.min(axis=1)
    
    # 统计"太接近"训练样本的比例
    memorized_ratio = (min_distances < threshold).mean()
    return memorized_ratio, min_distances

# 实践建议：
# • 记忆化比例 < 1% 通常可接受
# • 高记忆化 → 增加数据/正则化/早停