深入解析CFG扩散模型：从原理到工程实践

1次阅读

没有评论

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

扩散模型近年来在生成任务中表现出色，但 CFG（Conditional Free Guidance）扩散模型的训练过程却常常让工程师们头疼。主要问题集中在两个方面：

模式崩溃 ：模型容易陷入局部最优，生成结果多样性急剧下降
梯度爆炸 ：在训练后期，梯度值可能突然变得极大，导致训练不稳定

这些问题在条件生成任务中尤为突出，因为模型需要同时学习数据分布和条件信息的映射关系。

CFG 的核心思想是通过条件概率来引导生成过程。其关键公式可以表示为：

$$\nabla_{x_t} \log p(x_t|c) = \nabla_{x_t} \log p(x_t) + \gamma \nabla_{x_t} \log p(c|x_t)$$

其中：

$x_t$ 表示第 t 步的噪声图像
$c$ 是条件信息
$\gamma$ 是 guidance scale，控制条件对生成过程的影响程度

这个公式告诉我们，CFG 实际上是在无条件生成的梯度基础上，加上了一个条件梯度的加权项。

class NoiseScheduler:
    def __init__(self, num_timesteps=1000, beta_start=1e-4, beta_end=0.02):
        self.num_timesteps = num_timesteps
        self.betas = torch.linspace(beta_start, beta_end, num_timesteps)
        self.alphas = 1. - self.betas
        self.alpha_bars = torch.cumprod(self.alphas, dim=0)

    def add_noise(self, x0, t, noise=None):
        """根据时间步 t 向原始图像 x0 添加噪声"""
        if noise is None:
            noise = torch.randn_like(x0)

        sqrt_alpha_bar = torch.sqrt(self.alpha_bars[t])
        sqrt_one_minus_alpha_bar = torch.sqrt(1. - self.alpha_bars[t])

        return sqrt_alpha_bar * x0 + sqrt_one_minus_alpha_bar * noise

class ConditionEmbedder(nn.Module):
    def __init__(self, cond_dim=128, hidden_dim=256):
        super().__init__()
        self.mlp = nn.Sequential(nn.Linear(cond_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim)
        )

    def forward(self, c):
        """将条件信息 c 映射到隐空间"""
        return self.mlp(c)

def train_step(model, x0, c, noise_scheduler, optimizer, max_grad_norm=1.0):
    """包含梯度裁剪的训练步骤"""
    optimizer.zero_grad()

    # 随机采样时间步
    t = torch.randint(0, noise_scheduler.num_timesteps, (x0.size(0),))

    # 添加噪声
    noise = torch.randn_like(x0)
    xt = noise_scheduler.add_noise(x0, t, noise)

    # 前向传播
    pred_noise = model(xt, t, c)

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

    # 反向传播
    loss.backward()

    # 梯度裁剪
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)

    optimizer.step()

    return loss.item()

我们使用 FID（Frechet Inception Distance）指标评估不同 guidance scale 下的生成质量：