关于sora的技术，看这一篇就够了

一、sora的横空出世

openai在某一天突然宣布了一款能够文生图的大模型即将进入公测阶段，并宣称能够未来能够通过技术，构建世界模型，引起外界关注。

二、sora的训练逻辑

将其中的sora模型抽象出来，整个完整的模型框架如下所示：

首先将输入视频数据通过VAE（Variational Autoencoder，变分自动编码器）中的编码器模块编码成低维的空间向量表示，经Patches网络，将其中的向量分割为指定大小的模块（例如：16*16*16），并将时空模块通过SD模型进行扩散学习，将经过学习后的数据通过多层的以U-net为backbone的transformer网络恢复成原有大小模块，合并后得到需要的向量表示，最后逆向通过VAE解码器将向量表示输出成指定格式和时长的视频。

整个模型中主要有三大模块组成：VAE、StableDiffusion、Scaling Transformers，中间作为桥梁进行连接的是：patches和multi-condition

三、三大模块和patches

3.1 VAE

VAE是一种基于变分推断（Variational Inference, Variational Bayesian methods）的概率模型（Probabilistic Model），它属于生成模型（当然也是无监督模型）。在变分推断中，除了已知的数据（观测数据，训练数据）外还存在一个隐含变量，这里已知的数据集记为 X = { x i } i = 0 N X=\{x^{i}\}_{i=0}^N X={xi}i=0N​由$N$个连续变量或者离散变量$x$组成，而未观测的随机变量记为$z$，那么数据的产生包含两个过程：

• 从一个先验分布$p_{\theta }(z)$中采样一个$z^i$；
• 根据条件分布$p_{\theta }(x|z)$，用$z^i$生成$x^i$。

这里的$\theta$指的是分布的参数，比如对于高斯分布就是均值和标准差。我们希望找到一个参数${\theta }^{\ast }$来最大化生成真实数据的概率：
$${\theta }^{\ast }=arg\underset{\theta }{\max }\prod _{i=0}^n{p}_{\theta }(x^i)$$
这里$p_{\theta }(x^i)$可以通过对$z$积分得到：
$$p_{\theta }(x^i)=\int p_{\theta }(x^i|z)p_{\theta }(z)dz$$
而实际上要根据上述积分是不现实的，一方面先验分布$p_{\theta }(z)$是未知的，而且如果分布比较复杂，对$z$穷举计算也是极其耗时的。为了解决这个难题，变分推断引入后验分布$p_{\theta }(x|z)$来联合建模，根据贝叶斯公式，后验等于：
$$P_{\theta }(z|x)=\frac{p_{\theta }(x|z)p_{\theta }(z)}{p_{\theta }(x)}$$

这样的联合建模如上图所示，实线代表的是我们想要得到的生成模型$p_{\theta }(x|z)p_{\theta }(z)$，其中先验分布$p_{\theta }(z)$往往是事先定义好的（比如标准正态分布），而$p_{\theta }(x|z)$可以用一个网络来学习，如果把$z$ 看成隐含特征，那么这个网络就可以看成一个probabilistic decoder。虚线代表的是对后验分布$p_{\theta }(x|z)$的变分估计，记为$q_{\varphi }(x|z)$，它也可以用一个网络来学习，这个网络可以看成一个probabilistic encoder。

建模已经完成，下面我们来推导一下VAE的优化目标。对于估计的后验$q_{\varphi }(x|z)$，我们希望它接近真实的后验分布$p_{\theta }(x|z)$，评估两个分布差异最常用的方式就是计算KL散度（Kullback-Leibler divergence）。对$q_{\varphi }(x|z)$和$p_{\theta }(x|z)$计算KL散度，如下所示：
$$\begin{array}{rl}{D}_{KL}(q_{\varphi }(z|x)||p_{\theta }(z|x))& =\int q_{\varphi }(z|x)log\frac{q_{\varphi }(z|x)}{p_{\theta }(z|x)}dz\\ & =\int q_{\varphi }(z|x)log\frac{q_{\varphi }(z|x)p_{\theta }(x)}{p_{\theta }(z,x)}dz\\ & =\int q_{\varphi }(z|x)(log{p}_{\theta }(x)+log\frac{q_{\varphi }(z|x)}{p_{\theta }(z,x)})dz\\ & =log{q}_{\varphi }(x)+\int q_{\varphi }(z|x)log\frac{q_{\varphi }(z|x)}{p_{\theta }(z,x)}dz\\ & =log{p}_{\theta }(x)+\int q_{\varphi }(z|x)log\frac{q_{\varphi }(z|x)}{p_{\theta }(x|z)p_{\theta }(z)}dz\\ & =log{p}_{\theta }(x)+E_{z{q}_{\varphi }(z|x)}[log\frac{q_{\varphi }(z|x)}{p_{\theta }(z)}-log{p}_{\theta }(x|z)]\\ & =log{p}_{\theta }(x)+D_{KL}(q_{\varphi }(z|x)||p_{\theta }(z))-E_{z{q}_{\varphi }(z|x)}log{p}_{\theta }(x|z)\end{array}$$
最终可以得到：
$$\begin{array}{rl}{D}_{KL}(q_{\varphi }(z|x)||p_{\theta }(z|x))& =log{p}_{\theta }(x)+D_{KL}(q_{\varphi }(z|x)||p_{\theta }(z))-E_{z{q}_{\varphi }(z|x)}log{p}_{\theta }(x|z)\end{array}$$
这里我们适当调整一下上述等式中各个项的位置，可以得到：
$$\begin{array}{rl}log{p}_{\theta }(x)-D_{KL}(q_{\varphi }(z|x)||p_{\theta }(z|x))& =E_{z{q}_{\varphi }(z|x)}log{p}_{\theta }(x|z)-D_{KL}(q_{\varphi }(z|x)||p_{\theta }(z))\end{array}$$
这里$log{p}_{\theta }(x)$是生成真实数据的对数似然，对于生成模型，我们希望最大化这个对数似然，而$D_{KL}(q_{\varphi }(z|x)||p_{\theta }(z|x))$是估计的后验分布和真实分布的KL散度，我们希望最小化该KL散度（KL散度为0时两个分布没有差异），所以上述等式的左边就是联合建模的最大化优化目标，这等价于最大化等式的右边。这个等式的右边又称为Evidence lower bound，简称为ELBO，这主要是因为$p_{\theta }(x)$一般称为evidence，而由于KL散度的非负性，所以有下述不等式：
$$log{p}_{\theta }(x)\ge E_{z{q}_{\varphi }(z|x)}log{p}_{\theta }(x|z)-D_{KL}(q_{\varphi }(z|x)||p_{\theta }(z|x))$$
所以ELBO是evidence的下限，ELBO是变分推断中经常用到的优化目标。对于VAE，ELBO取负就是其要最小化的训练目标：
$$\begin{array}{rl}{L}_{VAE}(\theta ,\varphi )& =-E_{z{q}_{\varphi }(z|x)}log{p}_{\theta }(x|z)+D_{KL}(q_{\varphi }(z|x)||p_{\theta }(z|x))\\ {\theta }^{\ast },{\varphi }^{\ast }& =arg\underset{\theta ,\varphi }{\min }{L}_{VAE}\end{array}$$
对于优化目标的第二项，即计算$q_{\varphi }(z|x)$和$p_{\theta }(z|x)$的KL散度，首先我们必须要对两个分布做一定的假设：
$$\begin{gathered} q_{\varphi }(z|x)=N(z;{\mu }^i,{\sigma }^{2i}I) \\ {p}_{\theta }=N(z,0,I) \end{gathered}$$
即$q_{\varphi }(z|x)$为各分量独立的多元高斯分布（协方差矩阵为对角矩阵），那么encoder网络预测的就是高斯分布的均值$\mu$和方差${\sigma }^2$（实际处理时预测$log{\sigma }^2$，因为该值是无约束的）。而先验$p_{\theta }(z)$为标准正态分布，这样就变成了计算两个多元高斯分布的KL散度。对于多元高斯分布，其概率密度函数为：
$$p(x)=\frac{1}{\sqrt{{2\pi }^n}det(\sum )}exp-\frac{1}{2}(x-\mu {)}^T\sum ^{-1}(x-\mu )$$
最终可以得到$q_{\varphi }(z|x)$和$p_{\theta }(z|x)$的KL散度：
$$\begin{array}{rl}KL(D_{KL}(q_{\varphi }(z|x)||p_{\theta }(z|x)))& =KL(N(z;{\mu }^i;{\sigma }^{2i}I)||N(z,0,I))\\ & =\frac{1}{2}(tr({\sigma }^{2i})I+({\mu }^i{)}^T{\mu }^i-n-log(det{\sigma }^{2i}I)\\ & =\frac{1}{2}\sum _{j=n}^n(({\sigma }^{2i}{)}^2+({\mu }^i{)}^2-1-log(({\sigma }^{2i}{)}^2))\end{array}$$

现在我们来分析优化目标的第一项$-E_{z{q}_{\varphi }(z|x)}log{p}_{\theta }(x|z)$，它一般被称为重建误差（reconstruction error），因为$log{p}_{\theta }(x|z)$正是给定$z$下生成真实数据$x$的似然（Likelihood）。对于一个给定的训练样本$x^i$，我们可以采蒙特卡洛方法（Monte Carlo method）来估计这个数学期望，即从$q_{\varphi }(z|x)$多次采样来估计：
$$

• E_{z~q_{\phi}(z|x)}logp_{\theta}(x|z) \approx -\frac{1}{L}\sum_{l=1}^{L}(log(p_{\theta}(x|z)))
  KaTeX parse error: Can't use function '$' in math mode at position 5: 这里的$̲L$为采样的总次数，实际上在具…
  z \sim q_{\phi}(z|x) = N(z;\mu,\sigma^{2i} I)\
  z = \mu + \sigma \odot \epsilon,where \epsilon \sim N(0,I)
  $$
  直观上讲，就是首先从标准正态分布随机采样一个样本，然后乘以encoder预测的标准差，再加上encoder预测的均值，这样就能计算该损失对encoder网络参数的梯度了。

根据建模的数据类型，$p_{\theta }(x|z))$分布可以是一个高斯分布也可以是一个伯努利分布，这里以更通用的高斯分布为例。假定$p_{\theta }(x|z))$分布也属于一个各分量独立的多元高斯分布：$p_{\theta }(x|z)=N(z;\mu ,{\sigma }^{2i}I)$。由于各个分量独立，所以我们可以单独计算每个分量：
$$\begin{array}{rl}log{p}_{\theta }(x|z)& =log\frac{1}{\sqrt{2\pi {\sigma }^2}}exp(-\frac{(x-\mu {)}^2}{2{\sigma }^2})\\ & =log\frac{1}{\sqrt{2\pi {\sigma }^2}}-\frac{1}{2{\sigma }^2}(x-\mu {)}^2\end{array}$$
对于这个高斯分布的标准差，我们往往假定它是一个常量，而均值是由decoder预测得出：$\mu =f_{\theta }(z)$。那么则有:
$$-log{p}_{\theta }(x^i|z^{i,l})=C_1+C_2\sum _d^D(x_d^i-(f_{\theta }(z^{i,l}))d{)}^2$$
这里$C_1$和$C_2$均是常量，而D是变量x的维度大小。如果忽略常量$C_1$的话，那么重建误差其实就是L2损失。上面我们是假定$p_{\theta }(x|z)$分布是一个高斯分布，如果$p_{\theta }(x|z)$是一个伯努利分布即0-1分布的话，此时decoder直接预测概率值（sigmoid激活函数），重建误差就是交叉熵：
$$log{p}_{\theta }(x|z)=\sum _d^D(x_{d}log(f_{\theta }(z{)}_d)+(1-x)log(1-f_{\theta }(z{)}_d))$$
根据上述分析，对给定的一个训练样本$x_i$，其训练损失（假定是高斯分布）为：
$$\begin{array}{rl}{L}_{V}AE(\theta ,\varphi ,x^i)& =-E_{z\sim q_{\varphi }(z|x^i)}log{p}_{\theta }(x^i|z)+D_{KL}(q_{\varphi }(z|x^i)||p_{\theta }(z))\\ & =\frac{C}{L}\sum _{l=0}^L\sum _{d=0}^D(x_d-(f_{\theta }(z^{i,l}))d{)}^2+\frac{1}{2}\sum _{j=0}^n(({\sigma }_j^i{)}^2+({\mu }_j^i{)}^2-1-log(({\sigma }_j^i{)}^2))\end{array}$$
如果把KL散度项看到一个正则化的话，那么VAE的损失函数就是重建误差+正则化，这样VAE就可以看成是一个加了约束的AE。VAE的整个训练流程如下所示：输入x，encoder首先计算出后验分布的均值和标准差，然后通过重采样方法采样得到隐变量z，然后送入decoder得到重建的数据x′。

训练完成后，我们就得到生成模型$p_{\theta }(x|z)p_{\theta }(z)$，其中$p_{\theta }(x|z)$就是decoder网络，而先验$ p_{\theta}(z) $为标准正态分布，我们从$p_{\theta}(z)$随机采样一个z，送入decoder网络，就能生成与训练数据X类似的样本。

3.2 StableDiffusion

3.3 Scaling Transformers

3.4 patches模块

视频由多帧的图片组成，因此，一个视频文件是本质上是包含了时间维度的三维图片数据。

因此，对数据进行patches进行打包方法其实有两种：

• 摊大饼法：从输入视频剪辑中均匀采样 n_t 个帧，使用与ViT 相同的方法独立地嵌入每个2D帧(embed each 2D frame independently using the same method as ViT)，并将所有这些 token连接在一起

• 切面包法：

四、具体的工程难点

通过上面的技术讲解，已经基本涵括了本次sora模型中最为基础的模型框架，但依旧存在一些不知道的问题工程细节问题：

视频压缩网络的压缩率是多少，Encoder中的复杂度是具体怎么设置的，时空patches中排布的方法等。

参考

生成模型之VAE：https://zhuanlan.zhihu.com/p/452743042
sora技术参考：https://github.com/datawhalechina/sora-tutorial/tree/main/docs