1D序列扩散模型入门指南：从理论到PyTorch实战

1次阅读

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

传统序列生成模型如 RNN 和 Transformer 虽然强大，但在实际应用中仍存在一些局限性：

长程依赖问题 ：RNN 难以捕捉远距离的序列依赖关系
多样性不足 ：Transformer 倾向于生成保守、平均化的输出
模式坍塌 ：生成结果缺乏变化，容易陷入局部最优

扩散模型通过渐进式的加噪和去噪过程，很好地解决了这些问题，特别适合需要丰富多样性的生成任务。

前向过程逐步向数据添加高斯噪声，定义如下：

q(x_t|x_{t-1}) = N(x_t; √(1-β_t)x_{t-1}, β_tI)

其中 β_t 是噪声调度参数，控制每步添加的噪声量。

逆向过程学习如何逐步去除噪声：

p_θ(x_{t-1}|x_t) = N(x_{t-1}; μ_θ(x_t,t), Σ_θ(x_t,t))

神经网络 θ 需要预测噪声或均值，将带噪数据逐步恢复到干净数据。

实现一个灵活的噪声调度器，支持多种调度策略：

class NoiseScheduler:
    def __init__(self, num_timesteps=1000, schedule='cosine'):
        self.num_timesteps = num_timesteps

        if schedule == 'linear':
            self.betas = torch.linspace(1e-4, 0.02, num_timesteps)
        elif schedule == 'cosine':
            # Cosine schedule (improved DDPM)
            steps = torch.arange(num_timesteps + 1)
            alpha_bar = torch.cos((steps / num_timesteps + 0.008) / 1.008 * math.pi / 2) ** 2
            alpha_bar = alpha_bar / alpha_bar[0]
            betas = 1 - (alpha_bar[1:] / alpha_bar[:-1])
            self.betas = torch.clip(betas, 0, 0.999)

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

    def add_noise(self, x_0, t, noise=None):
        """
        为输入 x_0 在时间步 t 添加噪声
        参数:
            x_0: 原始输入 [batch_size, seq_len, dim]
            t: 时间步 [batch_size]
        返回:
            x_t: 加噪后的输入
            noise: 实际添加的噪声
        """
        if noise is None:
            noise = torch.randn_like(x_0)

        sqrt_alpha_bar = torch.sqrt(self.alpha_bars.to(x_0.device)[t])[:, None, None]
        sqrt_one_minus_alpha_bar = torch.sqrt(1. - self.alpha_bars.to(x_0.device)[t])[:, None, None]

        x_t = sqrt_alpha_bar * x_0 + sqrt_one_minus_alpha_bar * noise
        return x_t, noise

实现一个适合 1D 序列数据的 UNet 结构：

class ResBlock1D(nn.Module):
    def __init__(self, dim, dim_out, groups=8):
        super().__init__()
        self.proj = nn.Conv1d(dim, dim_out, 3, padding=1)
        self.norm = nn.GroupNorm(groups, dim_out)
        self.act = nn.SiLU()
        self.res_conv = nn.Conv1d(dim, dim_out, 1) if dim != dim_out else nn.Identity()

    def forward(self, x, time_emb=None):
        """
        参数:
            x: [batch_size, channels, seq_len]
            time_emb: [batch_size, dim]
        """
        h = self.proj(x)
        h = self.norm(h)

        if time_emb is not None:
            time_emb = time_emb.unsqueeze(-1)  # [batch_size, dim, 1]
            h = h + time_emb

        h = self.act(h)

        return h + self.res_conv(x)

class UNet1D(nn.Module):
    def __init__(self, dim=64, dim_mults=(1, 2, 4, 8), channels=1):
        super().__init__()

        # 时间嵌入
        time_dim = dim * 4
        self.time_mlp = nn.Sequential(SinusoidalPositionEmbeddings(dim),
            nn.Linear(dim, time_dim),
            nn.GELU(),
            nn.Linear(time_dim, time_dim)
        )

        # 下采样
        self.downs = nn.ModuleList([])
        self.ups = nn.ModuleList([])

        dims = [channels] + [dim * m for m in dim_mults]
        in_out = list(zip(dims[:-1], dims[1:]))

        for ind, (dim_in, dim_out) in enumerate(in_out):
            is_last = ind >= (len(in_out) - 1)

            self.downs.append(nn.ModuleList([ResBlock1D(dim_in, dim_out, time_emb_dim=time_dim),
                ResBlock1D(dim_out, dim_out, time_emb_dim=time_dim),
                Downsample1D(dim_out) if not is_last else nn.Identity()]))

        # 上采样
        for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):
            is_last = ind >= (len(in_out) - 1)

            self.ups.append(nn.ModuleList([ResBlock1D(dim_out * 2, dim_in, time_emb_dim=time_dim),
                ResBlock1D(dim_in, dim_in, time_emb_dim=time_dim),
                Upsample1D(dim_in) if not is_last else nn.Identity()]))

        self.final_conv = nn.Sequential(ResBlock1D(dim, dim),
            nn.Conv1d(dim, channels, 1)
        )

    def forward(self, x, time):
        """
        参数:
            x: [batch_size, channels, seq_len]
            time: [batch_size]
        """
        t = self.time_mlp(time)
        h = []

        # 下采样
        for block1, block2, downsample in self.downs:
            x = block1(x, t)
            x = block2(x, t)
            h.append(x)
            x = downsample(x)

        # 上采样
        for block1, block2, upsample in self.ups:
            x = torch.cat([x, h.pop()], dim=1)
            x = block1(x, t)
            x = block2(x, t)
            x = upsample(x)

        return self.final_conv(x)

完整的训练流程实现：

def train_loop(model, scheduler, dataloader, optimizer, device, epochs=1000):
    model.train()

    for epoch in range(epochs):
        for batch in dataloader:
            # 获取数据
            x_0 = batch.to(device)  # [batch_size, seq_len, dim]
            batch_size = x_0.shape[0]

            # 随机采样时间步
            t = torch.randint(0, scheduler.num_timesteps, (batch_size,), device=device)

            # 添加噪声
            noise = torch.randn_like(x_0)
            x_t, noise = scheduler.add_noise(x_0, t, noise)

            # 预测噪声
            noise_pred = model(x_t.transpose(1, 2), t).transpose(1, 2)

            # 计算损失
            loss = F.mse_loss(noise_pred, noise)

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

        if epoch % 100 == 0:
            print(f"Epoch {epoch} | Loss: {loss.item():.4f}")

小步数 (50-200)：适合快速原型验证，生成质量一般
中等步数 (200-500)：平衡质量和速度的常见选择
大步数 (500-1000)：高质量生成，但训练和推理速度慢

对于变长序列，使用掩码避免 padding 区域影响：

# 计算掩码损失
def masked_loss(pred, target, mask):
    """
    参数:
        pred: 预测值 [batch_size, seq_len, dim]
        target: 目标值 [batch_size, seq_len, dim]
        mask: 掩码 [batch_size, seq_len] (1 表示有效, 0 表示 padding)
    """
    loss = (pred - target).pow(2)
    loss = loss.mean(dim=-1)  # [batch_size, seq_len]
    loss = (loss * mask).sum() / mask.sum()
    return loss

使用 AMP(Automatic Mixed Precision) 加速训练时注意：

对损失值手动缩放
检查梯度是否出现 NaN
避免在时间嵌入中使用大数值

调度策略	优点	缺点
Linear	简单直接，易于实现	高步数时噪声变化剧烈
Cosine	平滑过渡，生成质量高	实现稍复杂，训练初期收敛慢
Sigmoid	灵活控制噪声节奏	需要调参较多