国科网

2025-06-04 21:25:10  星期三
立足国科融媒,服务先进科技
使用PyTorch实现去噪扩散模型

点赞

0
发布时间:2024年01月08日 浏览量:145次 所属栏目:人工智能 发布者:田佳恬

在深入研究去噪扩散概率模型(DDPM)如何工作的细节之前,让我们先看看生成式人工智能的一些发展,也就是DDPM的一些基础研究。

VAE

VAE 采用了编码器、概率潜在空间和解码器。在训练过程中,编码器预测每个图像的均值和方差。然后从高斯分布中对这些值进行采样,并将其传递到解码器中,其中输入的图像预计与输出的图像相似。这个过程包括使用KL Divergence来计算损失。VAEs的一个显著优势在于它们能够生成各种各样的图像。在采样阶段简单地从高斯分布中采样,解码器创建一个新的图像。

GAN

在变分自编码器(VAEs)的短短一年之后,一个开创性的生成家族模型出现了——生成对抗网络(GANs),标志着一类新的生成模型的开始,其特征是两个神经网络的协作:一个生成器和一个鉴别器,涉及对抗性训练过程。生成器的目标是从随机噪声中生成真实的数据,例如图像,而鉴别器则努力区分真实数据和生成数据。在整个训练阶段,生成器和鉴别器通过竞争性学习过程不断完善自己的能力。生成器生成越来越有说服力的数据,从而比鉴别器更聪明,而鉴别器又提高了辨别真实样本和生成样本的能力。这种对抗性的相互作用在生成器生成高质量、逼真的数据时达到顶峰。在采样阶段,经过GAN训练后,生成器通过输入随机噪声产生新的样本。它将这些噪声转换为通常反映真实示例的数据。

为什么我们需要另一个模型架构

两种模型都有不同的问题,虽然GANs擅长于生成与训练集中的图像非常相似的逼真图像,但VAEs擅长于创建各种各样的图像,尽管有产生模糊图像的倾向。但是现有的模型还没有成功地将这两种功能结合起来——创造出既高度逼真又多样化的图像。这一挑战给研究人员带来了一个需要解决的重大障碍。

在第一篇GAN论文发表六年后,在VAE论文发表七年后,一个开创性的模型出现了:去噪扩散概率模型(DDPM)。DDPM结合了两个世界的优势,擅长于创造多样化和逼真的图像。

在本文中,我们将深入研究DDPM的复杂性,涵盖其训练过程,包括正向和逆向过程,并探索如何执行采样。在整个探索过程中,我们将使用PyTorch从头开始构建DDPM,并完成其完整的训练。

这里假设你已经熟悉深度学习的基础知识,并且在深度计算机视觉方面有坚实的基础。我们不会介绍这些基本概念,我们的目标是生成人类确信其真实性的图像。

DDPM

去噪扩散概率模型(DDPM)是生成模型领域的一种前沿方法。与依赖显式似然函数的传统模型不同,DDPM通过对扩散过程进行迭代去噪来运行。这包括逐渐向图像中添加噪声并试图去除该噪声。其基本理论是基于这样一种思想:通过一系列扩散步骤转换一个简单的分布,例如高斯分布,可以得到一个复杂而富有表现力的图像数据分布。或者说通过将样本从原始图像分布转移到高斯分布,我们可以创建一个模型来逆转这个过程。这使我们能够从全高斯分布开始,以图像分布结束,有效地生成新图像。

DDPM的训练包括两个基本步骤:产生噪声图像这是固定和不可学习的正向过程,以及随后的逆向过程。逆向过程的主要目标是使用专门的机器学习模型对图像进行去噪。

正向扩散过程

正向过程是一个固定且不可学习的步骤,但是它需要一些预定义的设置。在深入研究这些设置之前,让我们先了解一下它是如何工作的。

这个过程的核心概念是从一个清晰的图像开始。在用“T”表示的特定步长上,少量噪声按照高斯分布逐渐引入。

从图像中可以看出,噪声是在每一步递增的,我们深入研究这种噪音的数学表示。

噪声是从高斯分布中采样的。为了在每一步引入少量的噪声,我们使用马尔可夫链。要生成当前时间戳的图像,我们只需要上次时间戳的图像。马尔可夫链的概念在这里是关键的,并对随后的数学细节至关重要。

马尔可夫链是一个随机过程,其中过渡到任何特定状态的概率仅取决于当前状态和经过的时间,而不是之前的事件序列。这一特性简化了噪声添加过程的建模,使其更易于数学分析。

用beta表示的方差参数被有意地设置为一个非常小的值,目的是在每个步骤中只引入最少量的噪声。

步长参数“T”决定了生成全噪声图像所需的步长。在本文中,该参数被设置为1000,这可能显得很大。我们真的需要为数据集中的每个原始图像创建1000个噪声图像吗?马尔可夫链方面被证明有助于解决这个问题。由于我们只需要上一步的图像来预测下一步,并且每一步添加的噪声保持不变,因此我们可以通过生成特定时间戳的噪声图像来简化计算。采用对的再参数化技巧使我们能够进一步简化方程。

将式(3)中引入的新参数纳入式(2)中,对式(2)进行了发展,得到了结果。

逆向扩散过程

我们已经为图像引入了噪声下一步就是执行逆操作了。除非我们知道初始条件,即t = 0时的未去噪图像,否则无法从数学上实现对图像进行逆向处理去噪。我们的目标是直接从噪声中采样以创建新图像,这里缺乏关于结果的信息。所以我需要设计一种在不知道结果的情况下逐步去噪图像的方法。所以就出现了使用深度学习模型来近似这个复杂的数学函数的解决方案。

有了一点数学背景,模型将近似于方程(5)。一个值得注意的细节是,我们将坚持DDPM原始论文并保持固定的方差,尽管也有可能使模型学习它。

该模型的任务是预测当前时间戳和前一个时间戳之间添加的噪声的平均值。这样做可以有效地去除噪音,达到预期的效果。但是如果我们的目标是让模型预测从“原始图像”到最后一个时间戳添加的噪声呢?

除非我们知道没有噪声的初始图像,否则在数学上执行逆向过程是具有挑战性的,让我们从定义后方差开始。

模型的任务是预测从初始图像添加到时间戳t的图像的噪声。正向过程使我们能够执行这个操作,从一个清晰的图像开始,并在时间戳t处进展到一个有噪声的图像。

训练算法

我们假设用于进行预测的模型体系结构将是一个U-Net。训练阶段的目标是:对于数据集中的每个图像,在[0,T]范围内随机选择一个时间戳,并计算正向扩散过程。这产生了一个清晰的,有点噪声的图像,以及实际使用的噪声。然后利用我们对逆向过程的理解,使用该模型来预测添加到图像中的噪声。有了真实的和预测的噪声,我们似乎已经进入了一个有监督的机器学习问题。

最主要的问题来了,我们应该用哪个损失函数来训练我们的模型呢?由于处理的是概率潜在空间,Kullback-Leibler (KL)散度是一个合适的选择。

KL散度衡量两个概率分布之间的差异,在我们的例子中,是模型预测的分布和期望分布。在损失函数中加入KL散度不仅可以指导模型产生准确的预测,还可以确保潜在空间表示符合期望的概率结构。

KL散度可以近似为L2损失函数,所以可以得到以下损失函数:

最终我们得到了论文中提出的训练算法。

采样

逆向流程已经解释完成了,下面就是如何使用了。从时刻T的一个完全随机的图像开始,并使用逆向过程T次,最终到达时刻0。这构成了本文中概述的第二种算法

参数

我们有很多不同的参数beta,beta_tildes,alpha, alpha_hat 等等。目前都不知道如何选择这些参数。但是此时已知的唯一参数是T,它被设置为1000。

对于所有列出的参数,它们的选择取决于beta。从某种意义上说,Beta决定了我们要在每一步中添加的噪声量。因此,为了确保算法的成功,仔细选择beta是至关重要的。其他的参数因为太多,请参考论文。

在原始论文的实验阶段探索了各种抽样方法。最初的线性采样方法图像要么接收到的噪声不足,要么变得过于嘈杂。为了解决这个问题,采用了另一种更常用的方法,即余弦采样。余弦采样提供了更平滑和更一致的噪声添加。

模型Pytorch实现

我们将利用U-Net架构进行噪声预测,之所以选择U-Net,是因为U-Net是图像处理、捕获空间和特征地图以及提供与输入相同的输出大小的理想架构。

考虑到任务的复杂性和对每一步使用相同模型的要求(其中模型需要能够以相同的权重去噪完全有噪声的图像和稍微有噪声的图像),调整模型是必不可少的。这包括合并更复杂的块,并通过正弦嵌入步骤引入对所用时间戳的感知。这些增强的目的是使模型成为去噪任务的专家。在继续构建完整的模型之前,我们将介绍每个块。

ConvNext块

为了满足提高模型复杂度的需要,卷积块起着至关重要的作用。这里不能仅仅依赖于u-net论文中的基本块,我们将结合ConvNext。

输入由代表图像的“x”和大小为“time_embedding_dim”的嵌入的时间戳可视化“t”组成。由于块的复杂性以及与输入和最后一层的残差连接,在整个过程中,块在学习空间和特征映射方面起着关键作用。

class ConvNextBlock(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        mult=2,
        time_embedding_dim=None,
        norm=True,
        group=8,
    ):
        super().__init__()
        self.mlp = (
            nn.Sequential(nn.GELU(), nn.Linear(time_embedding_dim, in_channels))
            if time_embedding_dim
            else None
        )
 
        self.in_conv = nn.Conv2d(
            in_channels, in_channels, 7, padding=3, groups=in_channels
        )
 
        self.block = nn.Sequential(
            nn.GroupNorm(1, in_channels) if norm else nn.Identity(),
            nn.Conv2d(in_channels, out_channels * mult, 3, padding=1),
            nn.GELU(),
            nn.GroupNorm(1, out_channels * mult),
            nn.Conv2d(out_channels * mult, out_channels, 3, padding=1),
        )
 
        self.residual_conv = (
            nn.Conv2d(in_channels, out_channels, 1)
            if in_channels != out_channels
            else nn.Identity()
        )
 
    def forward(self, x, time_embedding=None):
        h = self.in_conv(x)
        if self.mlp is not None and time_embedding is not None:
            assert self.mlp is not None, "MLP is None"
            h = h + rearrange(self.mlp(time_embedding), "b c -> b c 1 1")
        h = self.block(h)
        return h + self.residual_conv(x)

分享说明:转发分享请注明出处。

    相关图讯
    网站简介  |   联系我们  |   广告服务  |   监督电话
    本网站由国科网运营维护 国科网讯(北京)技术有限公司版权所有  咨询电话:010-88516927
    地址:北京市海淀区阜石路甲69号院1号楼1层一单元114
    ICP备案号:京ICP备15066964号-8   违法和不良信息举报电话:010-67196565
    12300电信用户申诉受理中心   网络违法犯罪举报网站   中国互联网举报中心   12321网络不良与垃圾信息举报中心   12318全国文化市场举报网站
    代理域名注册服务机构:阿里巴巴云计算(北京)有限公司