CogVideo与CogVideoX模型结构

VQ-VAE（CogVideo使用的编码器）

VQ-VAE其实就是一个AE（自编码器）而不是VAE（变分自编码器）

PixelCNN

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

\[p(x)=p(x_1)p(x_2|x_1)\dots p(x_{3n^2}|x_1,x_2,\dots,x_{3n^2-1}) \]

其中\(p(x_1),p(x_2|x_1),\dots,p(x_{3n^2}|x_1,x_2,\dots,x_{3n^2-1})\)每一个都是256分类问题，只不过所依赖的条件有所不同。

VQ-VAE

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

降维离散化

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

最邻近重构

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

\[z=encoder(x) \]

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

\[E=[e_1,e_2,\dots,e_K] \]

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

\[z\to e_k,\quad k=\arg\min_j\|z-e_j\|_2 \]

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

整个流程是：

\[x\xrightarrow{encoder}z\xrightarrow{最邻近}z_q\xrightarrow{decoder}\hat{x} \]

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

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

\[z=\begin{pmatrix} z_{11}&z_{12}&\dots&z_{1m}\\ z_{21}&z_{22}&\dots&z_{2m}\\ \vdots&\vdots&\ddots&\vdots\\ z_{m1}&z_{m2}&\dots&z_{mm} \end{pmatrix}\]

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

自行设计梯度

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

\[\|x-decoder(z)\|_2^2 \]

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

\[\|x-decoder(z_q)\|_2^2 \]

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

换言之，我们的目标其实是\(\|x-decoder(z_q)\|_2^2\)最小，但是却不好优化，而\(\|x-decoder(z)\|_2^2\)容易优化，但却不是我们的优化目标。

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

\[\|x-decoder(z+sg[z_q-z])\|_2^2 \]

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

维护编码表

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

所以，为了让\(z_q\)和z更接近，我们可以直接地将\(\|z-z_q\|_2^2\)加入到loss中：

\[\|x-decoder(z+sg[z_q-z])\|_2^2+\beta\|z-z_q\|_2^2 \]

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

\[\|sg[z]-z_q\|_2^2+\|z-sg[z_q]\|_2^2 \]

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

其中\(\gamma<\beta\)，在原论文中使用的是\(\gamma=0.25\beta\)。

拟合编码分布

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

CogVideo

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

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

多帧率层次化训练

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

具体来说，CogVideo 在序列开始的时候加入了一个表示帧率的 token Frame 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 的情况。图中的 t=i、t=i+1、t=i+2 表示相邻的三帧，后边的帧可以看到的前帧的范围就只有不是灰色的部分，因为窗口有一定的大小，所以可以看到的前帧中 token 的范围比当前已经生成的更多。这样就相当于第 i 帧的深绿色部分生成完的时候，第 i+1 帧就可以生成浅绿色的部分，第 i+2 帧就可以生成红框圈住的 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 Transformer

CogVideoX 采取和 SD3 类似的 MMDiT 架构，下面来依次介绍这种架构中的各组成部分。

Patchify

CogVideoX 的分块策略和 DiT 的相同，同时为了使模型能够同时在视频和图像数据上进行训练（这部分会在训练策略部分介绍），并不在时间维度上进行分块。也就是说对于一个大小为 \(T\times H\times W\times C\) 的输入，会分成长度为 \(T\times H/p\times W/p\) 的序列。

3D 旋转位置编码

视频经过 patchify 后，每个位置可以用一个三维坐标 \((x,y,t)\) 来表示，CogVideoX 的做法是对每一个坐标分别进行旋转位置编码，再沿通道直接拼接到一起。其中，表示空间位置的坐标分别占 \(3/8\)，表示时间的坐标占 \(2/8\)。

专家自适应层归一化

我们在输入阶段将文本和视频的嵌入向量进行拼接，以更好地对齐视觉信息和语义信息。然而，这两种模态的特征空间差异显著，其嵌入向量甚至可能具有不同的数值尺度。为了在同一序列中更好地处理它们，我们采用专家自适应层归一化对每种模态进行独立处理。如图3所示，参考DiT的方法，我们将扩散过程的时间步长t作为调制模块的输入。随后，视觉专家自适应层归一化和文本专家自适应层归一化分别将该调制作用于视觉隐藏状态和文本隐藏状态。该策略在最小化额外参数的同时，促进了两种模态间特征空间的对齐。

3D Full Attention

先前的方法为了降低计算量而在时序和空间上分别计算 attention，这样会导致当视频的变化比较快速的时候出现前后不一致的情况。因此这里使用了 Full Attention，应该是对整体的所有 token 计算 attention，而不是时间和空间分开。

CogVideoX 的训练

Frame Pack

CogVideoX 采用了图像与视频混合训练的方式，在进行训练时，将图像视为长度只有一帧的视频。并且 CogVideoX 并没有采用和其他方法相同的定长视频训练，而是采用了一种打包训练的方法，通过把不同长度的视频都打包在一个 batch 中，来确保不同 batch 维度相同。

Resolution Progressive Training

CogVideoX 也采取了多分辨率训练的策略，一方面是为了充分利用从互联网得到的带有多种分辨率的数据，另一方面是为了使模型能够渐进式地学习从粗糙到精细的多种信息。

由于模型最开始的训练阶段是在低分辨率数据上训练的，所以在高分辨数据上训练时，需要将位置编码拓展到高分辨率上。在拓展时，有两种策略，其一是将位置编码进行外推，这样可以比较好地维持不同像素之间的相对位置关系；其二是将位置编码进行插值，这样可以更好地维持像素在整个图像上的全局位置。经过测试可以发现前者可以更好地生成细节，而后者生成的结果比较模糊。最终 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 Concat（OpenAIWrapper）

张量	形状	目的
噪声 latent `x`	`(1, 13, 16, 60, 90)`	去噪目标，携带当前扩散步的噪声信息
图像条件 `c["concat"]`	`(1, 13, 16, 60, 90)`	参考图像的 latent，逐位置对齐噪声 latent
`cat(dim=2)` 后	`(1, 13, 32, 60, 90)`	拼接后每个空间位置同时可见噪声与图像

DiffusionTransformer

操作	形状	目的
输入 `x`	`(1, 13, 32, 60, 90)`	拼接后的 noise+image latent
`timestep_embedding`	`(1, 3072)`	将标量时间步编码为高维向量
`time_embed` MLP	`(1, 512)` → `emb`	压缩到 time_embed_dim，用于后续 adaLN 调制
T5 文本特征 `context`	`(1, 226, 4096)`	文本条件，将与 video token 拼接做全序列 attention

Patch Embedding（Conv2d 32→3072, k=2, s=2）

操作	形状	目的
view 展平时间	`(13, 32, 60, 90)`	批量处理所有帧
Conv2d	`(13, 3072, 30, 45)`	空间再下采样 2×，同时映射到 hidden_size，完成 patch embedding
view 恢复时间	`(1, 13, 3072, 30, 45)`	恢复 batch 维度
flatten(3)	`(1, 13, 3072, 1350)`	展平空间维度为 patch 序列
transpose(2,3)	`(1, 13, 1350, 3072)`	调整为 (B,T,N,D) 格式
rearrange `b t n d → b (t n) d`	`(1, 17550, 3072)`	展平时空为单一序列，13×1350=17550 个 video token
text_proj T5 4096→3072	`(1, 226, 3072)`	将 T5 文本特征投影到与 video 相同的 hidden_size
cat([text, video], dim=1)	`(1, 17776, 3072)`	拼接文本与视频 token，后续做全序列 self-attention

3D RoPE（仅 video tokens）

操作	形状	目的
频率张量 `(T,H,W,head_dim)`	`(13, 30, 45, 64)`	构建时间+高度+宽度三轴的旋转频率
rearrange → `(THW, 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

操作	形状	目的
丢弃 text tokens	`(1, 17550, 3072)`	只需预测视频 latent，文本 token 不参与输出
Linear(3072 → 64)	`(1, 17550, 64)`	将每个 token 映射到 patch_size²×out_channels = 2²×16 = 64
unpatchify `b(thw)(cpq)→btc(hp)(wq)`	`(1, 13, 16, 60, 90)`	将 patch 序列还原为时空 latent，即模型预测的噪声/速度场

去噪 → VAE 解码

操作	形状	目的
DDPM/DDIM N 步去噪后 latent	`(1, 13, 16, 60, 90)`	迭代去噪得到干净的 latent
`/ scale_factor`	`(1, 13, 16, 60, 90)`	还原 VAE 编码时的缩放（scale_factor=0.7）
permute(0,2,1,3,4)	`(1, 16, 13, 60, 90)`	调整为 VAE 期望的 (B,C,T,H,W) 格式
VAE decode 空间上采样 8×，时间还原	`(1, 3, 49, 480, 720)`	3D VAE 解码，空间恢复到原分辨率，时间从 13 帧插值到 49 帧
clamp + 转 uint8	生成视频 49 帧	像素值裁剪到 [0,1] 并量化为 8bit，输出最终视频

CogVideo与CogVideoX笔记