【week16学习周记】视频生成相关阅读，跟风调研并思考

本文链接：https://blog.csdn.net/alokag/article/details/136535642

DiT：Scalable Diffusion Models with Transformers（ICCV 23’）

这两年的生成式大模型进入了Transformer的时代，Sora的生成器也采用的是Transformer的结构。但在之前的基于diffusion model的生成式网络中，backbone的选择其实是UNet的结构。最初的扩散模型并不能适应Transformer，所以有一些早期的工作也将“扩散模型”和“自回归方法”完全划分开了（包括我前期的一些文章中也进行了切割）

为什么要用Transformer？
我们首先注意一下这篇文章的第一个词“Scalable”，Transformer相比于Unet的第一个优势就是scalable。简单来说就是对于大规模数据有更强的适应性。对于DDPM，小型任务的训练的效果是很好的，利用UNet进行像素级别的生成。然而UNet的结构在更加复杂，更加大规模的数据集上没办法很好的整体适应（参数量和结构的限制）
对于DDPM，由于整体模型是一个从噪声到噪声的等尺寸输入输出的网络架构，所以才采用了UNet。
为什么Transformer更加scalable？

本文主要贡献就是采用了Transformer，使得模型能够更好的适应大规模数据集。这篇文章首先说明了，传统扩散模型中UNet带来的归纳偏置对于生成模型的效果并不明显，所以可以被Transformer取代。基于统一的Transformer架构，可以让diffusion model有进一步的提升。

此前的一些工作

扩散模型

本文基本的扩散模型推理结构采用的是DDPM，但是DDPM对于分布的方差学习是采用的预设的方式。在IDDPM（这里我本人仅做了理论部分的学习，具体公式推导可以参考这篇文章：IDDPM）中，为了更好的学习方差信息，基于原本的损失函数有增加了额外了优化方差的项： $L_{vlb}$ ，从而使得网络构建的时候对方差也进行约束。

基本模型结构

直接来看DiT的基本框架。首先进入DiT的输入层，Transformer的输入必须是tokenized的串，对于视觉模态，是将标签的潜在空间噪声特征进行patch操作。
在这里插入图片描述
DiT block中给了三种基本结构，总共设计了四种方案来实现两个embedding的嵌入：一个是timestep embedding即具体的推理步数另一个是label embedding为具体的引导标签信息，当然这个部分的设计可以跟着基本需求的变化而改变。DiT的模型框架图给了三种主要的DiT Block的设计：In-context conditioning, CA block(15% Gflop的损失), Adaptive layer norm block(adaLN)和adaLN_Zero block(零初始化的adaLN)

上述四种做法中，效果最好的做法是基于adaLN-zero的（而且零初始化的效果比adaLN的效果好了很多！）但是这种好是在label引导的图像生成中的好，如果需要嵌入的信息更多时，cross attention的方法更加优越一些。
基于adaLN-zero的DiT block实现如下

class DiTBlock(nn.Module):
    """
    A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
    """
    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs):
        super().__init__()
        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, **block_kwargs)
        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        mlp_hidden_dim = int(hidden_size * mlp_ratio)
        approx_gelu = lambda: nn.GELU(approximate="tanh")
        self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
        )

        # zero init
        nn.init.constant_(adaLN_modulation[-1].weight, 0)
        nn.init.constant_(adaLN_modulation[-1].bias, 0)
	#
    def forward(self, x, c):
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)		# 通过silu+linear生成了六组hidden维度的基于condition的参数
        x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa))
        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
        return x

shift_msa 是注意力部分的偏移信息。
scale_msa 是注意力部分的缩放信息。
gate_msa 是注意力部分的门控信息，用于控制输入对注意力模块的影响程度。
shift_mlp 是MLP部分的偏移信息。
scale_mlp 是MLP部分的缩放信息。
gate_mlp 是MLP部分的门控信息，用于控制输入对MLP模块的影响程度。
整体对引入条件进行调控，即上图中的 $\alpha,\beta,\gamma参数$

此外，在DiT最后一层，为了使得最后输出和输入对齐，再添加了一个仅回归 $\gamma$ 和 $\beta$ 的最后一层：

class FinalLayer(nn.Module):
    """
    The final layer of DiT.
    """
    def __init__(self, hidden_size, patch_size, out_channels):
        super().__init__()
        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(hidden_size, 2 * hidden_size, bias=True)
        )

    def forward(self, x, c):
        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
        x = modulate(self.norm_final(x), shift, scale)
        x = self.linear(x)
        return x

为了放入类似ViT结构中，DiT同样需要做patch化。patch化的基本做法和ViT的操作方法相似（SPE方法）。首先，由LDM的VQGAN/VQVAE图像感知压缩后会生成潜在空间的特征图（以及后续扩散推理中的噪声latent），DiT的设计中有p=2,4,8的各种设计。（感觉后面没什么特别创新的了，基本框架和LDM是一样的）

DiT在大规模数据上有较好的效果，其中最主要的贡献就是基于adaptive LN将label/text信息进行嵌入的做法，这为transformer提供了除cross attention之外的另一种可能，并且这种做法可以不大量引入额外的计算量。

Video Diffusion Models（NIPS 22’）：基于扩散模型的视频生成方法

这篇工作是最早将扩散模型用在视频生成的任务上的。这篇工作有两个突出贡献，主要也将围绕这两个部分讲，首先是该工作将UNet结构拓展成时序的UNet结构，即3DUnet，基本结构如下
3dunet
这个UNet的输入是一个噪声视频（由噪声图像拓展到四维，即帧，HWC），对于上图的示意图，整体输入中 $\boldsymbol z_t,\boldsymbol c, \lambda_t$ 表示的是第t步推理的噪声视频，条件（label或其他）嵌入以及UNet输入的logSNR

之前DDPM并没有介绍这个logSNR，主要还是笔者自己的调研不够充分，后面看了一些有关扩散模型的实现方法的文章和源码，发现有很多是在学习过程中引入“log_snr”然后输入到UNet之中的做法
首先看一下SNR（信噪比）的定义：在DDPM中，信噪比其实很直观，因为整个过程是一个噪声叠加的过程，所以信噪比可以直接表示为： $\text{SNR}(t)=\frac {\overline \alpha_t}{1-\overline \alpha_t}$ ，其中overline表示的是累乘的积。那么损失函数有一种写法是如下：

有一种扩散模型的实现是基于log_snr来喂给Unet来做


	def q_sample(self, x_start, times, noise = None):
        noise = default(noise, lambda: torch.randn_like(x_start))

        log_snr = self.log_snr(times)
        # 统一img与时间T输入的维度
        log_snr_padded = right_pad_dims_to(x_start, log_snr)
        #
        alpha, sigma = sqrt(log_snr_padded.sigmoid()), sqrt((-log_snr_padded).sigmoid())
        x_noised = x_start * alpha + noise * sigma

        return x_noised, log_snr
        
	def p_losses(self, x_start, times, noise = None):
        noise = default(noise, lambda: torch.randn_like(x_start))

        x, log_snr = self.q_sample(x_start = x_start, times = times, noise = noise)
        model_out = self.model(x, log_snr)

        losses = self.loss_fn(model_out, noise, reduction = 'none')
        # Global Average Pooling
        losses = reduce(losses, 'b ... -> b', 'mean')

        if self.p2_loss_weight_gamma >= 0:
            # following eq 8. in https://arxiv.org/abs/2204.00227
            loss_weight = (self.p2_loss_weight_k + log_snr.exp()) ** -self.p2_loss_weight_gamma
            losses = losses * loss_weight

        return losses.mean()
    
    def forward(self, img, *args, **kwargs):
        b, c, h, w, device, img_size, = *img.shape, img.device, self.image_size
        assert h == img_size and w == img_size, f'height and width of image must be {img_size}'

        times = self.random_times(b)
        img = normalize_to_neg_one_to_one(img)
        return self.p_losses(img, times, *args, **kwargs)

对于学习过程，首先生成batch size的随机timestep，然后基于time step可以生成采样的对应的log_snr，UNet的输入也可以是log_snr。回到3D-UNet，除了将原本内容拓展到三维，还有一点是在每一块中除了卷积外还在卷积操作后加入时空注意力机制。
对于时空注意力机制，笔者是第一次接触视频理解与视频生成的具体相关项目，所以在此也进行一下认真的学习

视频理解：Is Space-Time Attention All You Need for Video Understanding?

时空注意力机制

以下内容主要借鉴自：https://zhuanlan.zhihu.com/p/372712811

首先：视频的本质和自然语言其实有相似之处，都是序列化的，且上下文信息是相关的，连贯的，所以在序列上学习效果很好的transformer以及其中的self- attention结构，在视频理解上也应该有较好的效果。但是视觉信息在之前的工作主要是像ViT一样进行patchify，这是在图像的空间维度上进行序列化，和基于时间维度的序列化的方法是不一样的。CNN的方法存在归纳偏置且仅仅能捕捉short-range关系，所以在视觉上来说，高分辨率和长视频条件下的CNN学习会效率很差。（这部份和DiT说的其实差不多，不过因为这个工作是在视频上的所以多了视频长度这个维度，此外由于重点其实在注意力机制上所以这部份可以略过）

直接看学习的过程，首先同样会进行每个视频帧的patchify变成一组patch： $\textbf{x}_{(p,t)}\in\mathbb R^{3P^2}$ ，然后每个patch前加一个positional embedding。然后就是对embedding的每个位进行一个self_attention操作，但如果直接这么做会带来巨大的开销，如果是N个patch和F个帧，就需要进行(NF+1)^2次的匹配，所以考虑仅仅在时间或者空间上做attention来降低复杂度的做法。
公式化的，首先输入token有
在这里插入图片描述
则可以通过设置不同层的不同head的参数，有多头自注意力机制的qkv为：

计算出的SA

如果变成spatial-only或者temporal-only的话，则可以表示为下（以spatial为例）：

最后再乘以v，在算上MLP和双重残差结构，计算的过程如下：

在这里插入图片描述

几种学习时序和空序的方法

刚刚讲到的，是联合了时间和空间的joint attention，但是有时间效率的影响所以不好进行大规模训练。
在这里插入图片描述
上图是几种不同的设计的基本思想，主要的差别就是MLP后的attention块的操作。（优于比较显然，一下内容为摘抄）

joint space-time attention(ST)：时空的所有patches两两做attention。
space-time self-attention(S)：只在同一帧内的patches间两两做attention。
divided space-time attention(T+S): 先考虑时间维度上，同一空间位置patches间做attention，然后再进行同一帧内空间上的attention。该策略在实验中证明最优。
sparse local global attention(L+G): 先算local attention，local的定义是时空上一定范围内的所有patches，dense地取，然后再在global上以一定的步长，间隔地取选取patches来算global attention。
axial attention(T+W+H): 分别在T, W, H三个维度上，挑出轴上的patches来做attention。

经过实验发现，divided的方法虽然看起来简单，但是效果却很不错：（实验是在若干视频数据集下的video classification工作，发现对于Divided方法，其效果是最好的，时间和参数量的增加也不是很大）
在这里插入图片描述

回到VDM…

为什么说时空分离注意力可以进行联合训练？
经过调研之后发现其实这个“联合训练”主打的是一个多功能，并不能对模型性能进行特别大的提升。基本的做法是将分离的部分通过掩码来隔离，举个例子是在视频后加一部分想学习的图片一起组合到一起，然后所有数据都会经过空序的attention block，然后图像不会经过时序的block（这也能够一定程度上提升空序attention的效果）

Reconstruction-guided sampling

对于VDM的条件生成情况，提出了重构指导的做法。视频生成的目标是生成较长的视频，一般来说至少也是三位数的帧数，即使是短视频也是。然而一般的训练的过程中都是在固定帧数下进行训练（如16帧），在采样的过程中，可以通过扩展采样结果，来通过一个视频来拓展生成另一个视频：如首先生成一个16帧的基础视频 $\textbf x^a\sim p_\theta(\textbf x)$ 则可以推理有 $\textbf x^a \sim p_\theta(\textbf x^b|\textbf x^a)$ 。这个推理过程需要训练一个额外的条件生成模型，或者通过插值方式来进行生成（听起来似乎有点像AI补帧不是吗？？）

之前的方法（如score based generative model）的做法是通过将视频进行部分替换，其中一部分进行完整的学习。这段内容我自己个人是没有太看明白的，但是知乎上又一个解释比较好：
在这里插入图片描述
这里对于“一致性”有一个可视化，即基本的替换法和加入了 $x_a$ 重建梯度更新的reconstruction-guided sampling的方法之间的生成效果差异：

一致性差异

Latte: Transformer为基础的视频生成扩散模型

之前的AIGC调研主要都集中在静态的图像生成上，但基于扩散模型的视频生成方法并没有那么多，主要是基本的效果跟不上。Latte的基本思想和Sora是较为一致的，并且这个工作是有开源代码的。
在这里插入图片描述
由于这些工作依旧都是以LDM为基本结构的，所以LDM中的 $\mathcal E$ 和 $\mathcal D$ 的部分均省略了，latte的主要贡献主要就集中在视频信息的处理和时间和空间的联合学习方法。

基本框架

latte提出了四种基本的变种，即如上图所示的四个variant：1. 将时序和空序的学习交错进行学习；2. 现进行空序的学习，然后再进行时序的学习；3. 通过两个MHA（多头注意力）来进行串行的学习，分别进行时序和空序学习；4. 通过拆分MHA的注意力头，来分别对时序和空序进行学习，最后通过一个融合模块进行整合。

数学化的表示应该如下来进行描述：

对于变种1和变种2
首先，推理过程中的视频信息为如下的格式：原始数据为 $V_L\in\mathbb R^{F\times H\times W\times C}$ ，经过tokenize后即喂给transformer的数据为 $\hat z\in \mathbb R^{n_f\times n_h\times n_w\times d}$ ，即被转化为 $n_f\times n_h\times n_w$ 个 $d$ 维的token，对于ViT，H和W会被扁平化，同样的，可以表示为 $\hat z\in \mathbb R^{n_f\times t\times d}$ 。同时会有一个时间-空间结合的positional embedding $\boldsymbol p$ 会被加入到原本的token中( $z=\hat z+\boldsymbol p$ )，最终形成了两种embedding：对于spatial Transformer（空间transformer），放入transformer的embedding形状为 $z_s\in \mathbb R^{n_f\times t\times d}$

变种1和变种2是对于多个transformer块的整体划分，而变种3和变种4是一个transformer块内的划分。具体的，变种3与4的实现是将block中的多头自注意力机制中的head进行分解，head可以理解为transformer种的不同通道，将其划分成一部分进行空间学习一部分进行时序学习，然后再整合，一种是串行一种是并行，两种方法分别对应了两个做法。

对于模型的具体构建，首先输入进行了两种不同方式的处理：1. 像ViT一样转化为patch，这样会保留帧的维度（直接取F帧）；2. 在帧的维度上进行一个固定步长的下采样（可以理解为在一个步长下做pooling？），这种方法得到的token可以结合时空信息。但是在进行输出的时候，需要对latent video进行一个上采样，然后才能去做最后的decoding。
在这里插入图片描述
对于Conditioning的做法，同样也是有两种做法，第一种和之前大部分做法类似，采用了直接视为token的策略；第二种则是将DiT中的ADALN方法引入（应该和adaln-zero的做法是完全一样的）对于时序的编码，可以通过绝对位置编码和相对位置编码来进行实现

绝对位置编码和相对位置编码法（留白）

训练方法

采用预训练模型

由于Latte的方法是transformer-based的，而主流的方法更多还是以UNet为骨干网络，或者就是某些大工程并没有开源预训练模型，所以采用的还是DiT在ImageNet上的预训练模型为初始参数。但DiT是在图像上的工作，所以有两个问题：

位置编码，对于ViT方法，每一个token仅仅有一个positional embedding，但是在latte中需要有一组positional embedding（ $n_f$ 个），所以可以进行的操作是将positional embedding直接复制拓展
标签引导下的标签编码：这个主要是ImageNet和视频生成的冲突，因为标签不一致，所以直接删除了label embedding改为零填充了。