Taming Transformers for High-Resolution Image Synthesis 论文阅读

论文地址 2012.09841.pdf (arxiv.org)

摘要

结合CNN的归纳偏置有效性和transformer的表达能力去合成高分辨率图像。
展示了(i)用CNNs去学习一个语义构成的图像的特征(ii)通过这些语义生成图像

实现

Figure2:

学习一个有效地Codebook（具体可以看VQVAE，我也是一知半解，推一篇文VQ-VAE解读）

应该用一个序列来表现图像的构成( 这里的constitution of images 不太会翻译 )而不是独立的像素，需要使用一种离散码本（codebook）的方法来解决。
一个图像RH×W×3\mathbb{R}^{H\times W\times 3}RH×W×3 能被码本的条目的空间集合表示zq∈Rh×w×nzz_q\in \mathbb{R}^{h\times w\times n_z}zq∈Rh×w×nz , nzn_znz表示一段码的维数。
先学习一个由一个编码器EEE和一个解码器GGG组成的卷积网络，它将去学习用码（code）来代表图像。离散的码本Z=zkk=1K⊂Rnz\mathbb{Z}={z_k}^K_{k=1}\subset \mathbb{R}^{n_z}Z=zkk=1K⊂Rnz 。
图像xxx通过编码器EEE，z^=E(x)⊂Rh×w×nz\hat{z}=E(x) \subset \mathbb{R}^{h\times w\times n_z}z^=E(x)⊂Rh×w×nz
然后量化z^\hat{z}z^，表示为q(z^)q(\hat{z})q(z^)，将每段码量化到最相近的密码本中的条目。zq=q(z^):=arg min⁡zk∈Z∥z^ij−zk∥∈Rh×w×nzz_q=q(\hat{z}):=\underset {z_k\in \mathbb{Z}}{\operatorname {arg\,min}}\|\hat{z}_{ij}-z_k\| \in \mathbb{R}^{h\times w\times n_z}zq=q(z^):=zk∈Zargmin∥z^ij−zk∥∈Rh×w×nz
然后用解码器GGG重建图像，使得x^≈x\hat x \approx xx^≈x，x^=G(zq)=G(q(E(x)))\hat x = G(z_q)=G(q(E(x)))x^=G(zq)=G(q(E(x)))
不可微的量化操作的反向传播是通过直接的梯度估计（gradient estimator）实现的，简单的将梯度从编码器复制到解码器，所以这个模型能被端到端训练，通过这个loss函数(VQ指Vector Quantised矢量量化)：
LVQ(E,G,Z)=∥x−x^∥+∥sg[E(x)]−zq∥22+β∥sg[zq]−E(x)∥22\mathcal L_{VQ}(E,G,Z)=\|x-\hat x\|+\|sg[E(x)]-z_q\|^2_2+\beta \|sg[z_q]-E(x)\|^2_2LVQ(E,G,Z)=∥x−x^∥+∥sg[E(x)]−zq∥22+β∥sg[zq]−E(x)∥22
其中Lrec=∥x−x^∥2\mathcal L_{rec}=\|x-\hat x\|^2Lrec=∥x−x^∥2是重建loss，sg[]sg[]sg[]表示停止梯度计算（stop-gradient operation，即gradient backpropagation到此为止, 不再往前传）
然后∥sg[zq]−E(x)∥22\|sg[z_q]-E(x)\|^2_2∥sg[zq]−E(x)∥22就是所谓的承诺损失（commitment loss，不知道咋翻译），β\betaβ为权重。

用transformer去学习一个感知丰富（意译为特征丰富可能会好一些）的码本，来代表隐藏的图像分布，需要我们突破压缩的限制。为了完成它，我们提出了VQGAN，一个VQVAE的变体。用一个判别器（discriminatior）和感知损失（perceptual loss），在提高压缩率的情况下保持良好的感知质量。
需要注意的是，以前的工作仅在轻量模型的基础上应用基于像素和变换的方法。
在GAN部分，我们这里将Lrec\mathcal L_{rec}Lrec中的L2L_2L2替换为判别器D，来区分真实和重建的图像：
LGAN({E,G,Z},D)=[logD(x)+log(1−D(x^))]\mathcal L_{GAN}(\{E,G,Z\},D)=[log D(x)+log(1-D(\hat x))]LGAN({E,G,Z},D)=[logD(x)+log(1−D(x^))]
去寻找最优压缩模型Q∗=arg min⁡E,G,Zmax⁡DEx∼p(x)[LVQ(E,G,Z)+λLGAN({E,G,Z},D)]\mathcal Q^*=\underset {E,G,Z}{\operatorname {arg\,min}} \underset {D}{\operatorname {max}} \mathbb E_{x\sim p(x)}[\mathcal L_{VQ}(E,G,Z)+\lambda \mathcal L_{GAN}(\{E,G,Z\},D) ]Q∗=E,G,ZargminDmaxEx∼p(x)[LVQ(E,G,Z)+λLGAN({E,G,Z},D)]
计算自适应权重λ=∇GL[Lrec]∇GL[LGAN]+δ\lambda=\frac{\nabla_{GL}[\mathcal L_{rec}]}{\nabla_{G_L}[\mathcal L_{GAN}]+\delta}λ=∇GL[LGAN]+δ∇GL[Lrec]，∇GL\nabla_{GL}∇GL表示解码器最后一个layer的梯度，δ=10−6\delta=10^{-6}δ=10−6用来维持除法的有意义。
为了将这些离散的context聚合，我们在最低分辨率处应用了一个single attention layer。当展开隐藏的码（code）的时候，这道工序显著减少了序列长度。因此能够使用更强大的transformer模型。

用transformer学习图像的结构（composition）

有了E和G，我们能依据codebook的code的索引来表示图像。更准确的说，图像xxx的量化编码由zq=q(E(x))∈Rh×w×nzz_q=q(E(x)) \in \mathbb R^{h\times w\times n_z}zq=q(E(x))∈Rh×w×nz给出，等价于一个序列s∈{0,...,∣Z∣−1}h×ws\in \{0,..., |Z|-1\}^{h\times w}s∈{0,...,∣Z∣−1}h×w，其中sss是codebook的索引序列，因此有sij=k=>(zq)ij=zk]s_{ij}=k~=>~(z_q)_{ij}=z_k]sij=k => (zq)ij=zk] (根据Figure2能看出s与Zk的关系，文字表达不是很清晰，大致就是说s是Zq上对应codebook条目的索引，知道了s就可以知道Zq)。

因此，有了sss序列，我们就可以用p(si∣s<i)p(s_i|s_{<i})p(si∣s<i)来后验下一个索引，那么数据整体的关联性（这里不知道怎么表达，文中说的是likelihood of the full representation）就是p(s)=∏ip(si∣s<i)p(s)=\prod_{i}p(s_i|s_{<i})p(s)=∏ip(si∣s<i)，这让我们能直接最大化损失函数:
LTransformer=Ex∼p(x)[−logp(x)]\mathcal L_{Transformer}=\mathbb E_{x\sim p(x)}[-log ~p(x)]LTransformer=Ex∼p(x)[−log p(x)]

条件合成 Conditioned Synthesis

一些任务可能会根据用户需求来生成，将这些生成的条件设为ccc，它可以是一个标签描述全局图像，也可以是另一幅图像：
p(s∣c)=∏ip(si∣s<i,c)p(s|c)=\prod_i p(s_i|s_{<i},c)p(s∣c)=∏ip(si∣s<i,c)
后面都是一些VQGAN的具体应用，暂时不翻译。