要追溯VQ-VAE的思想就不得不谈到自回归模型。可以说VQ-VAE做生成模型的思路源于PixelRNN、PixelCNN之类的自回归模型这类模型留意到我们要生成的图像实际上是离散的而不是连续的。以cifar10的图像为例它是32×32大小的3通道图像换言之它是一个32×32×3的矩阵矩阵的每个元素是0255的任意一个整数这样一来我们可以将它看成是一个长度为32×32×33072的句子而词表的大小是256从而用语言模型的方法来逐像素地、递归地生成一张图片传入前面的所有像素来预测下一个像素这就是所谓的自回归方法()(1)(2|1)…(32|1,2,…,32−1)其中(1),(2|1),…,(32|1,2,…,32−1)每一个都是256分类问题只不过所依赖的条件有所不同。VQ-VAE针对自回归模型的固有毛病VQ-VAE提出的解决方案是先降维然后再对编码向量用PixelCNN建模。降维离散化因为PixelCNN生成的离散序列你想用PixelCNN建模编码向量那就意味着编码向量也是离散的才行。而我们常见的降维手段比如自编码器生成的编码向量都是连续性变量无法直接生成离散变量。同时生成离散型变量往往还意味着存在梯度消失的问题。最邻近重构在VQ-VAE中一张n×n×3的图片x先被传入一个encoder中得到连续的编码向量z()这里的z是一个大小为d的向量。另外VQ-VAE还维护一个Embedding层我们也可以称为编码表记为[1,2,…,]这里每个都是一个大小为d的向量。接着VQ-VAE通过最邻近搜索将z映射为这K个向量之一→,arg⁡min‖−‖2我们可以将z对应的编码表向量记为我们认为才是最后的编码结果。最后将传入一个decoder希望重构原图^()。整个流程是最邻近最邻近→→最邻近→^这样一来因为是编码表E中的向量之一所以它实际上就等价于1,2,…,这K个整数之一因此这整个流程相当于将整张图片编码为一个整数。当然上述过程是比较简化的如果只编码为一个向量重构时难免失真而且泛化性难以得到保证。所以实际编码时直接用多层卷积将x编码为m×m个大小为d的向量(1112…12122…2⋮⋮⋱⋮12…)也就是说z的总大小为m×m×d它依然保留着位置结构然后每个向量都用前述方法映射为编码表中的一个就得到一个同样大小的然后再用它来重构。这样一来也等价于一个m×m的整数矩阵这就实现了离散型编码。自行设计梯度我们知道如果是普通的自编码器直接用下述loss进行训练即可‖−()‖22但是在VQ-VAE中我们用来重构的是而不是z那么似乎应该用这个loss才对‖−()‖22但问题是的构建过程包含了argmin这个操作是没梯度的所以如果用第二个loss的话我们没法更新encoder。换言之我们的目标其实是‖−()‖22最小但是却不好优化而‖−()‖22容易优化但却不是我们的优化目标。VQ-VAE使用了一个很精巧也很直接的方法称为Straight-Through Estimator你也可以称之为直通估计。Straight-Through的思想很简单就是前向传播的时候可以用想要的变量哪怕不可导而反向传播的时候用你自己为它所设计的梯度。根据这个思想我们设计的目标函数是‖−([−])‖22其中sg是stop gradient的意思就是不要它的梯度。这样一来前向传播计算求loss的时候就直接等价于(−)()然后反向传播求梯度的时候由于−不提供梯度所以它也等价于()这个就允许我们对encoder进行优化了。维护编码表要注意根据VQ-VAE的最邻近搜索的设计我们应该期望和z是很接近的事实上编码表E的每个向量类似各个z的聚类中心出现但事实上未必如此即使‖−()‖22和‖−()‖22都很小也不意味着和z差别很小即(1)(2)不意味着12。所以为了让和z更接近我们可以直接地将‖−‖22加入到loss中‖−([−])‖22‖−‖22除此之外还可以做得更仔细一些。由于编码表相对是比较自由的而z要尽力保证重构效果所以我们应当尽量让去靠近z而不是让z去靠近而因为‖−‖22的梯度等于对的梯度加上对z的梯度所以我们将它等价地分解为‖[]−‖22‖−[]‖22第一项相等于固定z让靠近z第二项则反过来固定让z靠近。注意这个等价是对于反向传播求梯度来说的对于前向传播求loss它是原来的两倍。根据我们刚才的讨论我们希望让去靠近z多于让z去靠近所以可以调一下最终的loss比例‖−([−])‖22‖[]−‖22‖−[]‖22其中在原论文中使用的是0.25。拟合编码分布经过上述一大通设计之后我们终于将图片编码为m×m的整数矩阵了由于这个m×m的矩阵一定程度上也保留了原来输入图片的位置信息所以我们可以用自回归模型比如PixelCNN来对编码矩阵进行拟合即建模先验分布。通过PixelCNN得到编码分布后就可以随机生成一个新的编码矩阵然后通过编码表E映射为3维的实数矩阵行×列×编码维度最后经过decoder得到一张图片。CogVideoCogVideo 是基于大规模预训练 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 加权求和后作为双通道注意力的整体输出。生成阶段的滑窗注意力正常的 3D Swin attention 是不支持自回归生成过程的因此 CogVideo 为了让这个注意力机制能够支持自回归生成加入了一个自回归 mask 机制。并且这样可以让不同的帧能够在一定程度上并行生成。具体的做法如下图所示这里展示的是窗口大小为 2 的情况。图中的 ti、ti1、ti2 表示相邻的三帧后边的帧可以看到的前帧的范围就只有不是灰色的部分因为窗口有一定的大小所以可以看到的前帧中 token 的范围比当前已经生成的更多。这样就相当于第 i 帧的深绿色部分生成完的时候第 i1 帧就可以生成浅绿色的部分第 i2 帧就可以生成红框圈住的 token从而实现并行。CogVideoX和上一个工作 CogVideo 不同这个方法是基于扩散模型实现的。从框架图来看感觉 CogVideoX 同时吸取了 Sora 和 Stable Diffusion 3 的优势不仅使用了 3D VAE还引入了双路 DiT 的架构。具体来说CogVideoX 主要进行了以下几个方面的工作使用 3D VAE 编码视频有效地压缩视频维度、保证视频的连续性引入双路 DiT 分别对文本和视频进行编码并用 3D attention 进行信息交换开发了一个视频标注的 pipeline用于对视频给出准确的文本标注提出了一种渐进式训练方法和一种均匀采样方法。CogVideoX 的整体架构如下图所示文本和视频分别经过文本编码器这里是 T5和 3D VAE 编码后输入主干网络。文本和视频分别经过一条支路并在注意力部分进行交互。3D Causal VAE由于视频相比图像多了时序信息所以需要对多出来的时间维度进行处理。先前的视频生成模型都采用 2D VAE这样会导致生成的视频在时间上连续性比较差并且出现闪烁的情况。和通常的 VAE 相同3D Causal VAE 包括一个编码器、一个解码器以及一个 KL 约束。在编码的前两个阶段分别在时间和空间维度上进行 2 倍下采样在最后一个阶段只在空间维度上进行下采样。因此最后的下采样倍数是 488时间维度倍数为 4空间为 8 倍。为了防止未来的时序信息泄漏到当前或更早的时间中这里采取了一种特殊的 padding 方式也就是只在前方进行 padding这样卷积时就不会把后续 token 的信息泄露到当前 token。并且这种卷积还可以在不同 GPU 之间进行并行只需要在不同的 GPU 之间拷贝少量的数据图 b 中的粉色 token即可。Expert TransformerCogVideoX 采取和 SD3 类似的 MMDiT 架构下面来依次介绍这种架构中的各组成部分。PatchifyCogVideoX 的分块策略和 DiT 的相同同时为了使模型能够同时在视频和图像数据上进行训练这部分会在训练策略部分介绍并不在时间维度上进行分块。也就是说对于一个大小为 ××× 的输入会分成长度为 ×/×/ 的序列。3D 旋转位置编码视频经过 patchify 后每个位置可以用一个三维坐标 (,,) 来表示CogVideoX 的做法是对每一个坐标分别进行旋转位置编码再沿通道直接拼接到一起。其中表示空间位置的坐标分别占 3/8表示时间的坐标占 2/8。专家自适应层归一化我们在输入阶段将文本和视频的嵌入向量进行拼接以更好地对齐视觉信息和语义信息。然而这两种模态的特征空间差异显著其嵌入向量甚至可能具有不同的数值尺度。为了在同一序列中更好地处理它们我们采用专家自适应层归一化对每种模态进行独立处理。如图3所示参考DiT的方法我们将扩散过程的时间步长t作为调制模块的输入。随后视觉专家自适应层归一化和文本专家自适应层归一化分别将该调制作用于视觉隐藏状态和文本隐藏状态。该策略在最小化额外参数的同时促进了两种模态间特征空间的对齐。3D Full Attention先前的方法为了降低计算量而在时序和空间上分别计算 attention这样会导致当视频的变化比较快速的时候出现前后不一致的情况。因此这里使用了 Full Attention应该是对整体的所有 token 计算 attention而不是时间和空间分开。CogVideoX 的训练Frame PackCogVideoX 采用了图像与视频混合训练的方式在进行训练时将图像视为长度只有一帧的视频。并且 CogVideoX 并没有采用和其他方法相同的定长视频训练而是采用了一种打包训练的方法通过把不同长度的视频都打包在一个 batch 中来确保不同 batch 维度相同。Resolution Progressive TrainingCogVideoX 也采取了多分辨率训练的策略一方面是为了充分利用从互联网得到的带有多种分辨率的数据另一方面是为了使模型能够渐进式地学习从粗糙到精细的多种信息。由于模型最开始的训练阶段是在低分辨率数据上训练的所以在高分辨数据上训练时需要将位置编码拓展到高分辨率上。在拓展时有两种策略其一是将位置编码进行外推这样可以比较好地维持不同像素之间的相对位置关系其二是将位置编码进行插值这样可以更好地维持像素在整个图像上的全局位置。经过测试可以发现前者可以更好地生成细节而后者生成的结果比较模糊。最终 CogVideoX 使用的是前者。在最后一个训练阶段使用了高质量数据进行微调。主要包括移除了字幕以及水印的数据这部分数据占全部数据的比例大概为 20%。Explicit Uniform Sampling在扩散模型进行训练时会对时间步进行均匀采样。然而不同时间步对应的损失的尺度可能不一致因此虽然对时间步的采样是均匀的但最后得到的损失却不够均匀。为此CogVideoX 使用了一种显式均匀采样的方法具体来说在并行训练时会为每个 rank 分配一个时间步的区间然后每个 rank 都只在这个区间里进行均匀采样这样能够得到比较均匀的 loss有助于模型更好地收敛。CogVideoX 条件注入图像预处理操作形状目的PIL 读取(H, W, 3)加载原始图像ToTensor unsqueeze(0)(1, 3, H, W)转为 batch 张量resize center crop(1, 3, 480, 720)对齐模型训练分辨率* 2 - 1归一化(1, 3, 480, 720)像素值映射到 VAE 期望的 [-1, 1] 范围unsqueeze(2) 插入时间维(1, 3, 1, 480, 720)构造单帧视频适配 3D VAE 接口VAE encode空间下采样 8×(1, 16, 1, 60, 90)压缩到 latent 空间降低后续计算量permute(0,2,1,3,4)(1, 1, 16, 60, 90)调整为 (B,T,C,H,W)统一视频张量格式concat 零填充后续 12 帧(1, 13, 16, 60, 90)→c[concat]第 0 帧携带图像信息后续帧置零表示待生成给模型明确的已知/未知信号Channel ConcatOpenAIWrapper张量形状目的噪声 latentx(1, 13, 16, 60, 90)去噪目标携带当前扩散步的噪声信息图像条件c[concat](1, 13, 16, 60, 90)参考图像的 latent逐位置对齐噪声 latentcat(dim2)后(1, 13, 32, 60, 90)拼接后每个空间位置同时可见噪声与图像DiffusionTransformer操作形状目的输入x(1, 13, 32, 60, 90)拼接后的 noiseimage latenttimestep_embedding(1, 3072)将标量时间步编码为高维向量time_embedMLP(1, 512)→emb压缩到 time_embed_dim用于后续 adaLN 调制T5 文本特征context(1, 226, 4096)文本条件将与 video token 拼接做全序列 attentionPatch EmbeddingConv2d 32→3072, k2, s2操作形状目的view 展平时间(13, 32, 60, 90)批量处理所有帧Conv2d(13, 3072, 30, 45)空间再下采样 2×同时映射到 hidden_size完成 patch embeddingview 恢复时间(1, 13, 3072, 30, 45)恢复 batch 维度flatten(3)(1, 13, 3072, 1350)展平空间维度为 patch 序列transpose(2,3)(1, 13, 1350, 3072)调整为 (B,T,N,D) 格式rearrangeb t n d → b (t n) d(1, 17550, 3072)展平时空为单一序列13×135017550 个 video tokentext_proj T5 4096→3072(1, 226, 3072)将 T5 文本特征投影到与 video 相同的 hidden_sizecat([text, video], dim1)(1, 17776, 3072)拼接文本与视频 token后续做全序列 self-attention3D RoPE仅 video tokens操作形状目的频率张量(T,H,W,head_dim)(13, 30, 45, 64)构建时间高度宽度三轴的旋转频率rearrange →(T*H*W, head_dim)(17550, 64)展平为与 video token 序列对齐的位置编码应用到q/k[:, :, 226:](1, 48, 17550, 64)只对 video token ([226:]) 施加位置感知text token 不加保持文本位置无关性AdaLN Layer × 42操作形状目的emb→ adaLN Linear(512→12×3072)12 ×(1, 3072)生成 12 组调制参数text/video × attn/mlp × shift/scale/gate让时间步信息调制每层特征分离 text / video(1,226,3072)/(1,17550,3072)分别施加不同的 adaLN 参数cat → full self-attention(1, 17776, 3072)text 与 video token 互相 attend实现文本引导视频生成分离 残差 gate(1,226,3072)/(1,17550,3072)gate 控制每层更新幅度稳定深层训练cat 合并(1, 17776, 3072)输出到下一层Final Layer