扩散模型原理¶

q(xₜ|xₜ₋₁) = N(xₜ; √(1-βₜ)xₜ₋₁, βₜI)

pθ(xₜ₋₁|xₜ) = N(xₜ₋₁; μθ(xₜ,t), Σθ(xₜ,t))

L = E[||ε - εθ(xₜ, t)||²]

import torch
import torch.nn as nn
import numpy as np

# ===== 1. 噪声调度（线性 β 调度）=====
def linear_beta_schedule(timesteps=1000, beta_start=1e-4, beta_end=0.02):
    """
    定义每步加噪的幅度 β
    timesteps: 总扩散步数（通常 1000）
    """
    return torch.linspace(beta_start, beta_end, timesteps)  # 从小到大线性增加

# 预计算扩散过程的参数
timesteps = 1000
betas = linear_beta_schedule(timesteps)      # β: 每步噪声幅度 [T]
alphas = 1.0 - betas                          # α = 1 - β
alphas_cumprod = torch.cumprod(alphas, dim=0) # ᾱ = ∏α，累积乘积

# ===== 2. 前向过程：给数据加噪 =====
def q_sample(x_0, t, noise=None):
    """
    一步到位地给 x_0 加 t 步的噪声（利用数学推导的解析式）
    x_0: 原始数据 [batch, dim]
    t: 时间步索引 [batch]
    返回: x_t，即加了 t 步噪声后的数据
    """
    if noise is None:
        noise = torch.randn_like(x_0)  # 从标准正态分布采样噪声

    # 根据时间步 t 提取对应的 ᾱ
    alpha_t = alphas_cumprod[t].unsqueeze(-1)  # [batch, 1]

    # 解析式：x_t = √ᾱ * x_0 + √(1-ᾱ) * ε
    x_t = torch.sqrt(alpha_t) * x_0 + torch.sqrt(1 - alpha_t) * noise
    return x_t, noise

# ===== 3. 噪声预测网络（简化版 MLP，实际用 U-Net）=====
class SimpleNoisePredictor(nn.Module):
    """
    预测给定 x_t 和时间步 t 时加入的噪声
    实际应用中替换为 U-Net 或 Transformer
    """
    def __init__(self, data_dim, time_dim=64):
        super().__init__()
        # 时间步嵌入：把整数 t 转化为向量
        self.time_embed = nn.Embedding(1000, time_dim)

        self.net = nn.Sequential(
            nn.Linear(data_dim + time_dim, 256),  # 数据 + 时间步
            nn.SiLU(),                              # SiLU 激活（扩散模型常用）
            nn.Linear(256, 256),
            nn.SiLU(),
            nn.Linear(256, data_dim)               # 输出预测的噪声
        )

    def forward(self, x_t, t):
        t_emb = self.time_embed(t)              # 时间嵌入 [batch, time_dim]
        x_in = torch.cat([x_t, t_emb], dim=-1) # 拼接数据和时间信息
        return self.net(x_in)                   # 预测噪声

# ===== 4. 训练循环 =====
data_dim = 512    # 假设是蛋白质特征维度
model = SimpleNoisePredictor(data_dim)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# 模拟训练数据（实际替换为蛋白质坐标等）
fake_proteins = torch.randn(1000, data_dim)

for step in range(200):
    # 随机采样批次和时间步
    idx = torch.randint(0, 1000, (32,))
    x_0 = fake_proteins[idx]                           # 原始数据
    t = torch.randint(0, timesteps, (32,))             # 随机时间步

    # 前向：添加噪声
    x_t, true_noise = q_sample(x_0, t)

    # 预测噪声
    pred_noise = model(x_t, t)

    # 损失：预测噪声 vs 真实噪声（L2 损失）
    loss = nn.MSELoss()(pred_noise, true_noise)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if step % 50 == 0:
        print(f'Step {step}, Loss: {loss.item():.4f}')

# ===== 5. 反向过程（DDPM 采样）=====
@torch.no_grad()
def ddpm_sample(model, shape, device='cpu'):
    """从纯噪声逐步去噪生成数据"""
    x = torch.randn(shape).to(device)  # 从标准正态开始（纯噪声）

    for t in reversed(range(timesteps)):  # 从 T=999 逐步到 t=0
        t_batch = torch.full((shape[0],), t, dtype=torch.long)

        pred_noise = model(x, t_batch)  # 预测当前步的噪声

        # 计算去噪后的 x_{t-1}（简化版，实际包含方差项）
        alpha = alphas[t]
        alpha_bar = alphas_cumprod[t]
        x = (x - (1 - alpha) / torch.sqrt(1 - alpha_bar) * pred_noise) / torch.sqrt(alpha)

        if t > 0:  # 最后一步不加噪声
            x += torch.sqrt(betas[t]) * torch.randn_like(x)

    return x  # 生成的数据

generated = ddpm_sample(model, shape=(5, data_dim))
print(f'生成的蛋白质特征形状: {generated.shape}')