VQGAN是一个可以在多种图像生成任务上（如无条件图像生成、图像补全、条件引导图像生成等）均有着良好性能的生成模型，其中最大的亮点在于超像素级别的图像生成（百万级像素）。
模型整体流程图如下：
整体上来说与VQVAE相近，但是将VQVAE中的PixelCNN替换为Transformer，同时在训练过程中加入PatchGAN的判别器做对抗损失。
损失函数
在VQVAE中，其损失函数可以表示如下：
\[ \begin{aligned} L_{VQ}(E,G,Z) &= L_{rec} + \vert\vert\text{sg}[E(x)] - z_q\vert\vert_2^2 + \vert\vert\text{sg}[z_q] - E(x)\vert\vert_2^2 \\ &= \vert\vert x - \hat{x}\vert\vert_2^2 + \vert\vert\text{sg}[E(x)] - z_q\vert\vert_2^2 + \vert\vert\text{sg}[z_q] - E(x)\vert\vert_2^2 \end{aligned} \]
其中\(E,G,Z\)分别代表encoder、decoder和codebook。\(\text{sg}\)是stop gradient操作。\(L_{rec}\)是重建损失。
在VQGAN中，作者将\(L_{rec}\)替换为感知损失（perceptual loss），同时还加入了GAN中的对抗loss。在实际应用中，为了在某一任务上实现更好的效果，可能会出现L2损失与感知损失一同使用的情况，甚至是将L2损失替换为L1损失，需要按照具体任务通过实验确定。
Perceptual loss
在一些图像风格迁移、超分辨率等的任务中，传统的损失函数pixel-loss（如L1、L2或cross entropy等）都是追求模型输出与目标在像素级上相等，但像素值上一个很小的偏差，在实际效果上基本没有区别，而pixel-loss的值却会很大。感知损失主要关注图像的风格、颜色以及纹理之类的特征，可以更好地反映模型优化目标。
该损失函数的计算需要借助额外的网络来提取特征，具体计算流程大致如下：
上图中的\(f_W\)对应的是我们需要训练的模型。VGG-16是在Imagenet上预训练好的模型，主要用来提取特征，不参与权重更新。感知损失分为两部分：content loss和style loss。
content loss就是在feature层上的L2损失。选择在feature层上而不是原图上去做是为了避免output与target过分相似，同时还保留有指导output中content更新的能力。数学表达如下：
\[ L^{\phi, j}_{feat}(\hat{y}, y) = \frac{1}{C_jH_jW_j}\vert\vert \phi_j(\hat{y}) - \phi_j(y)\vert\vert_2^2 \]
其中\(\phi_j(\hat{y})\)和\(\phi_j(y)\)分别代表模型输出\(\hat{y}\)和目标\(y\)在VGG-16网络中第\(j\)层的输出特征，且有\(\phi_j(\hat{y}), \phi_j(y)\in\mathbb{R}^{H_j\times W_j\times C_j}\)。由上图可以看出，论文中提取的特征层为relu3_3的输出。
style loss用来计算图像在风格、颜色、纹理等方面的不同。假设\(\phi_j(x)\in\mathbb{R}^{H_j\times W_j\times C_j}\)表示输入\(x\)在模型\(\phi\)中第\(j\)层网络的输出特征。这里先定义一个格拉姆矩阵（Gram matrix）\(G_j^\phi(x)\)，大小为\(C_j\times C_j\)：
\[ G_j^\phi(x)_{c,c^\prime} = \frac{1}{C_jH_jW_j}\sum^{H_j}_{h=1}\sum^{W_j}_{w=1}\phi_j(x)_{h,w,c}\phi_j(x)_{h,w,c^\prime} \]
实际计算中我们可以将\(\phi_j(x)\)reshape为\(\psi\in\mathbb{R}^{C_j\times H_jW_j}\)，则格拉姆矩阵的计算可以简化如下：
\[ G^\phi_j(x) = \frac{\psi\psi^T}{C_jH_jW_j} \]
style loss则是定义在格拉姆矩阵上的Frobenius范数距离：
\[ L_{style}^{\phi,j}(\hat{y},y) = \vert\vert G_j^\phi(\hat{y}) - G_j^\phi(y)\vert\vert_F^2 \]
格拉姆矩阵计算输出的大小只与输入的通道数有关（输出大小为\(C_j\times C_j\)），所以只要保证通道数相同，即使模型输出与目标的尺寸不一致也是可以计算损失的。
在特征提取网络中，不同深度的卷积层提取的特征是不一样的（低层网络会提取一些点线特征，高层网络会提取一些抽象特征），所以为了能综合评估模型的style能力，我们需要提取各个不同的网络层输出来计算style损失。依据上图。作者分别提取了relu1_2、relu2_2、relu3_3和relu4_3这四层的输出特征去计算style损失，最终求和作为style loss的值。
关于Frobenius范数，其计算公式如下（\(A\in\mathbb{R}^{m\times n}\)）：
\[ \vert\vert A\vert\vert_F = \sqrt{\sum^m_{i=1}\sum^n_{j=1}\vert a_{ij}\vert^2} = \sqrt{\text{trace}(A^\star A)} = \sqrt{\sum_{i=1}^{\min(m,n)}\sigma_i^2} \]
上面给出了Frobenius范数的三种计算方式。其中\(A^\star\)表示\(A\)的共轭转置矩阵（共轭转置矩阵在将行与列对换后还要将每个元素共轭一下，实数矩阵的共轭转置矩阵就是转置矩阵），\(\sigma_i\)则是\(A\)的奇异值，且使用了迹函数。
Adversarial loss
关于对抗损失部分，作者选用的是适用于GAN的Hinge loss。一般的Hinge loss通常用于最大边际分类器（maximal margin classifer）如支持向量机中，通过尽可能地拉大模型正负样本之间的得分差距来提高模型性能。数学表示如下：
\[ L = \max(0, 1-\hat{y}y) \]
其中\(\hat{y}\)表示模型输出的样本得分，\(y\)表示样本真实标签（1或-1）。
- 当\(\hat{y}\)与\(y\)同号，即模型分类类别与真实类别相同。此时若\(\vert\hat{y}\vert\le 1\)，则表示此时的间距还不够大，损失函数不为0，还会继续更新，直到\(\vert\hat{y}\vert\gt 1\)损失函数为0
- 当\(\hat{y}\)与\(y\)异号，即模型分类类别与真实类别不同，损失函数不为0
在Geometric GAN中，引入了Hinge loss且变换为以下形式：
\[ \begin{aligned} L(G,D) &= L(G) + L(D) \\ &= -E(D(G(z))) + E(\max(0, 1-D(x))) + E(\max(0, 1+D(G(z)))) \end{aligned} \]
Hinge loss主要应用在了判别器\(D\)上。只有当\(D(x)\lt 1\)的真实样本和\(D(G(z))\gt -1\)的生成样本会产生损失，即只有一些被不合理划分的样本会有贡献，这样GAN模型的训练会更加稳定。
在VQGAN中，对抗损失笼统地表示如下：
\[ L_{GAN}(\{E,G,Z\}, D) = [\log(D(x)) + \log(1-D(\hat{x}))] \]
我们的目标是求得最优的\(E, G, Z\)，最后总结的优化目标如下：
\[ Q^\star = argmin_{E,G,Z}\max_D\mathbb{E}_{x\sim p(x)}[L_{VQ}(E,G,Z) + \lambda L_{GAN}(\{E,G,Z\}, D)] \]
其中\(\lambda\)是自适应权重，表达式如下：
\[ \lambda = \frac{\nabla_{G_L}[L_{rec}]}{\nabla_{G_L}[L_{GAN}]+\delta} \]
\(L_{rec}\)就是感知损失，\(\nabla_{G_L}[\cdot]\)是一个函数，表示该函数的输入相对于模型decoder最后一层\(L\)的梯度，\(\delta=10^{-6}\)是为了防止除0错误。
VQGAN中新加的判别器\(D\)结构是PatchGAN的判别器，它是一个全卷积的模型结构。通常的GAN判别器输出的只是一个值，而PatchGAN判别器输出的是一个\(n\times n\)的矩阵，矩阵中的每个值都代表输入模型在图像中对应的每个小区域（patch）的得分。这种设计可以让模型更关注图像的细节，也不再是用单单一个值来评价图像的真实性。
随机图像生成
上面的优化目标其实已经可以完成模型各个模块的训练了，但在随机图像生成任务中，我们是需要在没有encoder的情况下，采样出\(z_q\)给到decoder的。VQVAE的方法通过PixelCNN来预测\(z_q\)。考虑到Transformer具有长距离注意力，并且这个采样过程也是可以类比文本序列预测问题，所以作者采用了Transformer结构的模型来预测\(z_q\)（具体结构是GPT-2）。
将采样过程视为自回归预测，即让模型去预测分布\(p(s)=\prod_i p(s_i\vert s_{\lt i})\)，相应的损失函数是最大化对数似然函数：
\[ L_{Transformer} = \mathbb{E}_{x\sim p(x)}[-\log p(s)] \]
Transformer模块的训练是需要在其它模块训练好后进行的，具体的训练过程如下：
- 输入经过encoder和codebook后得到\(z_q\in \mathbb{R}^{h\times w\times n}\)并展开为二维形式\(z_q^\prime\in \mathbb{R}^{hw\times n}\)，即得到\(hw\)个维度为\(n\)的编码。定义此时\(z_q^\prime\)中的编码组合对应到codebook中的索引为unmodified_indices
- 将\(z_q^\prime\)中的部分编码替换为随机生成的编码（维度依然为\(n\)），定义修改后的编码组合对应的索引为modified_indices。替换为随机编码这一步相当于是在特征中添加了噪声，以此来提升Transformer的泛化性
- 将modified_indices输入到Transformer模型中，通过交叉熵损失训练模型，使其重构出unmodified_indices
借助于Transformer，VQGAN实现了超像素级的图像生成任务。
参考
Taming Transformers for High-Resolution Image Synthesis
Perceptual Losses for Real-Time Style Transfer and Super-Resolution
Image-to-Image Translation with Conditional Adversarial Networks