VQ-VAE

admin管理员组

文章数量:1530018

特点

1. 用于压缩编码，编码出离散隐向量（有用的表征）
2. 先验分布不再固定，可学习，离散的均匀分布。不像VAE中，先验分布是一个标准的高斯分布。
3. 避免后验坍塌，即输入任意隐变量解码器输出都一样。

结构

稀疏的高维数据->压缩到很小的空间->解码器还原为高维对象。
后验分布 q ( z ∣ x ) q(z|x) q(z∣x)，先验分布 p ( z ) p(z) p(z)，解码器 p ( x ∣ z ) p(x|z) p(x∣z)。

VQ-VAE的三部分：

1. 图像x输入编码器，得到 z e ( x ) z_e(x) ze​(x)
2. latent embedding space: KxD，K个类别的向量，每个向量维度D。
   编码器的输出 z e ( x ) z_e(x) ze​(x)通过最近邻查找，找到距离最近的离散向量 e j e_j ej​的索引 k k k。
   q ( z = k ∣ x ) = { 1 f o r   k = a r g m i n j ∣ ∣ z e ( x ) − e j ∣ ∣ 2 0 o t h e r w i s e q(z=k|x)=\left\{\begin{matrix} 1 &for \ k=argmin_j||z_e(x)-e_j||_2\\ 0&otherwise \end{matrix}\right. q(z=k∣x)={10​for k=argminj​∣∣ze​(x)−ej​∣∣2​otherwise​
   后验分布概率 q ( z ∣ x ) q(z|x) q(z∣x)是独热编码，隐变量z就是一个索引。解码器输入是索引k对应的向量 e k e_k ek​。
3. 解码器对 e k e_k ek​解码，得到重构的x。

损失

编码器，解码器，embedding space都需要训练。

1. 重构loss
2. KL散度loss，后验分布和先验分布的距离，但这里先验分布是均匀分布，后验分布也是类别分布，两个类别分布之间的KL散度是常数，所以VQ-VAE忽略KL散度loss。
3. codebook loss：为了让embedding space训练学习，让embedding vector e e e去接近编码器输出 z e ( x ) z_e(x) ze​(x)
4. commitment loss->embedding space没有约束，如果 e e e学得不够快，可能会无限膨胀，所以commitment loss让 z e ( x ) z_e(x) ze​(x) 接近 e e e，让 z e ( x ) z_e(x) ze​(x)走的慢一点。commitment loss前面一般会有一个系数 β \beta β，比如0.25.

注意

• 编码器到解码器中间的梯度传导断掉了，因为argmin不可导。解码器输入的梯度复制给编码器输出。
• 我们用 z e ( x ) z_e(x) ze​(x)找到 e k e_k ek​，所以二者维度需要一致，即编码器输出和解码器输入维度要一致，都是D维。
• stop gradient相当于看作常量

其他

• 如何评估VQ-VAE的性能？拿一张图片编码-离散化-解码，对比重构后的图片和原图差距，看重建效果。
• 训练过程，如何判断模型是否在有效训练？把编码器的类别分布信息熵打印出来。-plogp，均匀分布熵最大，熵很小说明类别变化很小，可能输出只有一个类别->后验坍塌。
• 训练好VQ-VAE后，如何生成新的图片？原本的VQ-VAE其实只是一个压缩器。无监督随机生成：大规模图片预训练VQ-VAE模型，某张图片输入VQ-VAE，那就有这张图片的离散化向量，用另一个模型对其自回归建模，用自回归生成的压缩图片，解码生成新的图片。（让自回归模型学习一个离散化分布，这样可以通过该模型采样离散向量，解码就可以产生不同的图片）
  条件生成：大规模图片训练VQ-VAE模型，得到一个有效的后验编码器；用RNN以提示prompt作为RNN的条件输入，对离散向量自回归建模。训练好的RNN模型，给定prompt自回归生成离散向量，解码即可条件生成图片。

本文标签： VQVAE