来源:计算机书童论文信息
题目:Diffusion Vision Transformers: Generating Images with Transformers
扩散视觉Transformer:基于Transformer生成图像
作者:Hossein Talebi, Peyman Milanfar
论文创新点
时间依赖的多头自注意力(TMSA)机制:作者提出了一种新颖的时间依赖的多头自注意力(TMSA)机制,该机制能够动态适应去噪过程的不同阶段,同时捕捉空间和时间依赖性及其相互作用。TMSA通过在自注意力中集成时间成分,使得键、查询和值权重在去噪过程中按时间步长调整,从而提高了去噪网络的灵活性和性能。
DiffiT模型架构:作者引入了一种新的基于视觉Transformer(ViT)的扩散模型,称为DiffiT,该模型统一了去噪网络的设计模式,并可用于潜在空间和图像空间的各种图像生成任务。DiffiT在多个数据集上展示了最先进的性能,特别是在ImageNet-256数据集上达到了新的最先进FID分数。
细粒度控制去噪过程:作者提出了一种细粒度控制去噪过程的方法,通过TMSA机制实现对去噪过程的动态调整。这种方法显著提高了生成图像的质量和参数效率,使得DiffiT在生成高保真图像方面表现出色。
潜在空间和图像空间模型:作者提出了潜在空间和图像空间的DiffiT模型,并在不同分辨率的类条件和无条件合成任务中展示了最先进的性能。潜在空间DiffiT模型在ImageNet-256数据集上达到了新的最先进FID分数,且参数数量显著减少。
摘要
扩散模型以其强大的表达能力和高质量的样本生成,在生成领域取得了最先进的性能。开创性的视觉Transformer(ViT)也在识别任务中展示了强大的建模能力和可扩展性。本文中,作者研究了ViT在基于扩散的生成学习中的有效性,并提出了一种新的模型,称为扩散视觉Transformer(DiffiT)。具体来说,作者提出了一种细粒度控制去噪过程的方法,并引入了时间依赖的多头自注意力(TMSA)机制。DiffiT在生成高保真图像方面表现出色,且参数效率显著提高。作者还提出了潜在空间和图像空间的DiffiT模型,并在不同分辨率的类条件和无条件合成任务中展示了最先进的性能。潜在空间DiffiT模型在ImageNet-256数据集上达到了新的最先进FID分数1.73,其参数比其他基于Transformer的扩散模型如MDT和DiT分别减少了19.85%和16.88%。
关键字
扩散模型,视觉Transformer,生成学习,高保真图像生成,时间依赖的多头自注意力
III. 方法
训练
尽管从不同角度推导,扩散模型通常可以写成学习以下去噪自编码器目标[76]
直观地说,给定来自的噪声样本(通过生成),神经网络被训练来预测添加的噪声量(即)。等效地,神经网络也可以被训练来预测[64, 28]。上述目标也称为去噪分数匹配[76],其目标是尝试用神经网络拟合数据分数(即),也称为分数网络。分数网络可以通过关系与相关。采样
扩散模型的样本可以通过以下解决从到的随机微分方程族来模拟[19, 26, 34, 86]:其中是反向标准维纳过程,是一个描述采样过程中随机噪声量的函数。如果对所有,则过程变为概率常微分方程[3](ODE),可以通过ODE积分器如去噪扩散隐式模型(DDIM[69])求解。否则,可以使用随机微分方程(SDE)求解器,包括原始去噪扩散概率模型(DDPM[28])的求解器。通常,ODE求解器可以在较少步数内收敛到高质量样本,而SDE求解器对不准确的分数模型更为鲁棒[34]。DiffiT模型
时间依赖的自注意力
在每一层,作者的Transformer块接收,一组在2D网格上排列的空间标记作为输入。它还接收,一个表示时间步的时间标记。与[28]类似,作者通过将位置时间嵌入馈送到带有swish激活[21]的小MLP中获得时间标记。这个时间标记被传递到去噪网络的所有层。作者引入了时间依赖的多头自注意力,通过在共享空间中投影特征和时间标记嵌入来捕捉长程空间和时间依赖性。具体来说,共享空间中的时间依赖查询、键和值通过空间和时间嵌入和的线性投影计算:其中、、、、、分别表示查询、键和值的相应空间和时间线性投影权重。作者注意到,公式3到5中的操作等价于将每个空间标记与时间标记连接后的线性投影。因此,键、查询和值都是时间和空间标记的线性函数,它们可以自适应地修改不同时间步的注意力行为。作者定义、和,它们是查询、键和值的堆叠形式。自注意力计算如下:其中是键的缩放因子,对应于相对位置偏置[66]。在计算注意力时,相对位置偏置允许在每个注意力头中编码信息。请注意,尽管相对位置偏置隐含地受到输入时间嵌入的影响,但直接将其与此组件集成可能会导致次优性能,因为它需要捕捉空间和时间信息。请参见第5.4节了解更多分析。DiffiT Transformer块
Transformer块是提出的DiffiT架构的核心构建块,定义为:其中TMSA表示时间依赖的多头自注意力,如上所述,是时间嵌入标记,是空间标记,LN和MLP分别表示层归一化和MLP。潜在空间
最近,潜在扩散模型已被证明在生成高质量的大分辨率图像方面有效[59, 74]。在图3中,作者展示了潜在空间DiffiT模型的架构。作者首先使用预训练的变分自编码器网络[59]对图像进行编码。然后将特征图转换为非重叠的补丁并投影到新的嵌入空间。与DiT模型[56]类似,作者使用没有上采样或下采样层的视觉Transformer作为潜在空间中的去噪网络。此外,作者还利用三通道的无分类器指导来提高生成样本的质量。最后阶段是一个线性层来解码输出。图像空间
DiffiT架构
如图4所示,DiffiT在图像空间中使用对称的U形编码器-解码器架构,其中收缩路径和扩展路径通过每个分辨率的跳跃连接相互连接。具体来说,编码器或解码器路径的每个分辨率由个连续的DiffiT块组成,包含作者提出的时间依赖的自注意力模块。在每个路径的开始,对于编码器和解码器,使用卷积层来匹配特征图的数量。在每个分辨率之间转换时,还使用卷积上采样或下采样层。作者推测,这些卷积层的使用嵌入了归纳图像偏置,可以进一步提高性能。在本文的其余部分,作者讨论了DiffiT Transformer块和作者提出的时间依赖的自注意力机制。作者在构建U形去噪架构时使用提出的Transformer块作为残差单元。局部注意力
当空间标记数量较大时,注意力的二次成本扩展性较差,特别是在大特征图的情况下。在不失一般性的前提下,上述Transformer块可以应用于局部区域,其中自注意力在非重叠的分区窗口内计算。尽管这些分区窗口不允许信息在不同区域之间传播,但带有瓶颈层的U-Net结构允许不同区域之间的信息共享。DiffiT ResBlock
作者通过将提出的DiffiT Transformer块与额外的卷积层结合来定义最终的残差单元:其中GN表示组归一化操作[78],DiffiT-Transformer在公式7和8中定义。作者的图像空间扩散模型的残差单元是结合了卷积层和Transformer块的混合单元。IV. 结果