从CogView到CogVideoX：CogView1&2

前言

  从本文开始我们计划将经典的开源工作系列Cog-X比较系统地介绍一下。我们将从最早的CogView开始介绍到最近的CogVideoX。

CogView

  CogView出自NIPS 21的《CogView: Mastering Text-to-Image Generation via Transformers》，是在diffusion爆发之前一篇经典工作，基于VQ-VAE和causal transformer (GPT)。虽然以现在的视角来看比较naive了，而且文中还花了大量篇幅介绍本质上和VQ没啥关系的VAE，但是作为和DALL-E同时期的两阶段T2I工作还是比较经典的。框架如下：

  其中有几个细节值得关注。首先是图像tokenize那里讨论了几种常见的优化方法：

• $argmin+STE$：这也是原始VQ-VAE提出的方法，基本思路和常见的如坍缩问题我们在之前的文章中也已介绍过了。

• Gumbel sampling+STE：是DALL-E中使用的方法，具体：

  $$ z_{q,ij}=\arg \max_k \left(g_k-\frac{\Vert e_k-z_{ij} \Vert^2}{\tau} \right),g_k\sim\text{Gumbel}(0,1) $$

  缺点是比较麻烦，而且用到了$argmax$还是不可微需要STE，训练不稳定。

• nearest-neighbor+EMA：即简单将最近邻的平均值作为量化，并用EMA更新。
• nearest-neighbor+fixed：和第三种方法类似也用最近邻量化，不过codebook就定死不更新了。

  作者给出了四种方法训练时的重建L2损失曲线如下。比较意外的是四种方法表现基本相当，效果区别不大。最后作者选择了第三种方法。

  其次还有在FP16训练下提升稳定性的一些工程细节。第一个是Precision Bottleneck Relaxation (PB-Relax)：作者发现训练loss溢出问题常发生在两个瓶颈操作中：LayerNorm和attention计算。改进如下：

• layernorm的输入值放缩：$LN(x)=LayerNorm(x/\max[x])$。

• attention计算顺序调整：

  $$ softmax(\frac{Q^\top K}{\sqrt d})=softmax\Bigg(\left(\frac{Q^\top}{\alpha\sqrt d}K\right)-\alpha\cdot \max\left[\frac{Q^\top}{\alpha\sqrt d}K\right] \Bigg) $$

  最后是layernorm的顺序。一般将残差和norm的顺序不同可以分为post-norm和pre-norm；这里作者在pre的基础上前后分别做一次layer normalization，作者将之称为Sandwich-LN。从下右图可以发现上述做法对提升训练稳定性还是很重要的。