CogVideo与CogVideoX模型结构

PixelCNN

要追溯VQ-VAE的思想，就不得不谈到自回归模型。可以说，VQ-VAE做生成模型的思路，源于PixelRNN、PixelCNN之类的自回归模型，这类模型留意到我们要生成的图像，实际上是离散的而不是连续的。以cifar10的图像为例，它是32×32大小的3通道图像，换言之它是一个32×32×3的矩阵，矩阵的每个元素是0～255的任意一个整数，这样一来，我们可以将它看成是一个长度为32×32×3=3072的句子，而词表的大小是256，从而用语言模型的方法，来逐像素地、递归地生成一张图片（传入前面的所有像素，来预测下一个像素），这就是所谓的自回归方法：

其中每一个都是256分类问题，只不过所依赖的条件有所不同。

VQ-VAE

针对自回归模型的固有毛病，VQ-VAE提出的解决方案是：先降维，然后再对编码向量用PixelCNN建模。

降维离散化

因为PixelCNN生成的离散序列，你想用PixelCNN建模编码向量，那就意味着编码向量也是离散的才行。而我们常见的降维手段，比如自编码器，生成的编码向量都是连续性变量，无法直接生成离散变量。同时，生成离散型变量往往还意味着存在梯度消失的问题。

最邻近重构

在VQ-VAE中，一张n×n×3的图片x先被传入一个encoder中，得到连续的编码向量z：

这里的z是一个大小为d的向量。另外，VQ-VAE还维护一个Embedding层，我们也可以称为编码表，记为

这里每个都是一个大小为d的向量。接着，VQ-VAE通过最邻近搜索，将z映射为这K个向量之一：

我们可以将z对应的编码表向量记为，我们认为才是最后的编码结果。最后将传入一个decoder，希望重构原图。

整个流程是：

最邻近

这样一来，因为是编码表E中的向量之一，所以它实际上就等价于这K个整数之一，因此这整个流程相当于将整张图片编码为一个整数。

当然，上述过程是比较简化的，如果只编码为一个向量，重构时难免失真，而且泛化性难以得到保证。所以实际编码时直接用多层卷积将x编码为m×m个大小为d的向量：

也就是说，z的总大小为m×m×d，它依然保留着位置结构，然后每个向量都用前述方法映射为编码表中的一个，就得到一个同样大小的，然后再用它来重构。这样一来，也等价于一个m×m的整数矩阵，这就实现了离散型编码。

自行设计梯度

我们知道，如果是普通的自编码器，直接用下述loss进行训练即可：

但是，在VQ-VAE中，我们用来重构的是而不是z，那么似乎应该用这个loss才对：

但问题是的构建过程包含了argmin，这个操作是没梯度的，所以如果用第二个loss的话，我们没法更新encoder。

换言之，我们的目标其实是最小，但是却不好优化，而容易优化，但却不是我们的优化目标。

VQ-VAE使用了一个很精巧也很直接的方法，称为Straight-Through Estimator，你也可以称之为"直通估计"。Straight-Through的思想很简单，就是前向传播的时候可以用想要的变量（哪怕不可导），而反向传播的时候，用你自己为它所设计的梯度。根据这个思想，我们设计的目标函数是：

其中sg是stop gradient的意思，就是不要它的梯度。这样一来，前向传播计算（求loss）的时候，就直接等价于，然后反向传播（求梯度）的时候，由于不提供梯度，所以它也等价于，这个就允许我们对encoder进行优化了。

维护编码表

要注意，根据VQ-VAE的最邻近搜索的设计，我们应该期望和z是很接近的（事实上编码表E的每个向量类似各个z的聚类中心出现），但事实上未必如此，即使和都很小，也不意味着和z差别很小（即不意味着）。

所以，为了让和z更接近，我们可以直接地将加入到loss中：

除此之外，还可以做得更仔细一些。由于编码表（）相对是比较自由的，而z要尽力保证重构效果，所以我们应当尽量"让去靠近z"而不是"让z去靠近"，而因为的梯度等于对的梯度加上对z的梯度，所以我们将它等价地分解为

第一项相等于固定z，让靠近z，第二项则反过来固定，让z靠近。注意这个"等价"是对于反向传播（求梯度）来说的，对于前向传播（求loss）它是原来的两倍。根据我们刚才的讨论，我们希望"让去靠近z"多于"让z去靠近"，所以可以调一下最终的loss比例：

其中，在原论文中使用的是。

拟合编码分布

经过上述一大通设计之后，我们终于将图片编码为m×m的整数矩阵了，由于这个m×m的矩阵一定程度上也保留了原来输入图片的位置信息，所以我们可以用自回归模型比如PixelCNN，来对编码矩阵进行拟合（即建模先验分布）。通过PixelCNN得到编码分布后，就可以随机生成一个新的编码矩阵，然后通过编码表E映射为3维的实数矩阵（行×列×编码维度），最后经过decoder得到一张图片。

CogVideo

CogVideo 是基于大规模预训练 Transformer 进行视频生成的工作，也是近期推出的 CogVideoX 的前身。相比于文生图任务，文生视频的主要难点在于两个方面：首先是数据更加稀缺，视频-文本配对数据比较少；其次是视频多了时序信息。

本模型基于文生图模型 CogView2 进行训练，在训练时使用了 5.4 M 视频-文本对数据。在训练时，文本条件是通过 in context learning，也就是将文本 token 直接拼接在图像 token 序列前方的方式实现的。除此之外还引入了多帧率层次化训练的训练策略，通过调整帧率来动态地调整视频的长度。在生成时，首先生成关键帧，然后用一个插值模型生成中间帧。

多帧率层次化训练

CogVideo 也采用了比较常见的方式，用 VQVAE 将视频序列转换为离散的 token 序列，再使用 transformer 对 token 序列进行学习。在训练时，token 序列的长度是固定的，也就是总共包含对应于 5 帧的 token 序列。不过和通常的方法不同的是，这里虽然序列的长度是固定的，但是实际上对应的视频的长度是可变的。

具体来说，CogVideo 在序列开始的时候加入了一个表示帧率的 tokenFrame Rate。虽然论文原文直接把这个称为frame rate，不过这个和实际上视频的帧率感觉还是有一点区别的。这个表示的是在这个序列的 5 帧中，每两帧之间相隔的视频中帧的数量。这样，对于比较长的视频，可以设置两帧之间相隔较多帧，反之亦然。有了这个设定，无论训练视频多长都可以用固定的序列长度表示，可以实现对变长视频的处理。

这样做主要有两点好处：

首先是可以处理变长视频，防止因对视频进行截断导致与文字之间的不对齐现象；
其次是在一般的视频中，相邻帧一般比较类似，如果直接对原始数据进行学习，容易让模型学到直接 copy 上一帧的 shortcut，导致模型退化。
不过这样训练之后的模型生成的两帧之间也会比较跳跃，因此需要用一个额外的插帧模型对生成的关键帧进行插帧。

在生成阶段，首先依然是需要生成最开始的五个关键帧，在生成关键帧后，CogVideo 采用了一种递归的插帧方式。具体来说就是在现有帧的基础上，将帧率减半，然后在每两帧之间再用自回归的方式生成一帧，这样每次生成之后序列的长度就会变成原来的 2 倍。（但因为整体的 token 序列长度不变，所以每次需要拆成两半插帧两次）

除此之外，CogVideo 还使用了 CogLM 的双向注意力机制，不同于 GPT 等只有单向注意力的模型，引入双向注意力可以使生成过程关注前后文的信息。

双通道注意力

相比于图像生成模型，视频生成模型需要关注时序信息。为此，CogVideo 直接在文生图模型 CogView2 上进行改进，因为后者已经能比较好地处理文本-图像的信息，所以可以作为视频生成模型的预训练。

为了使模型能够比较好地处理时序信息，CogVideo 在原有的 attention 的基础上又加入了一个 3D attention，如下图所示。图中的 Attention-base 就是原来的 CogView2 自带的 attention，其是以图像作为单位进行处理，可以理解为每次计算 attention 都是在图像内部做，主要关注的是图像级别的生成。而 Attention-plus 则是 3D 注意力，在进行 attention 计算的时候有多帧的内容都参与计算，这样可以关注时序信息。

对于 Attention-plus 的选择，原文使用了两种选择，即 3D local attention 和 3D Swin attention。两个通道的 attention 加权求和后作为双通道注意力的整体输出。