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 加权求和后作为双通道注意力的整体输出。