论文笔记-Discrete Latent Variables Based Generation

  • VQ-VAE: Neural Discrete Representation Learning (NIPS2017)
  • VQ-VAE2: Generating Diverse High-Resolution Images with VQ-VAE-2
  • DALL-E: Zero-Shot Text-to-Image Generation
  • VideoGPT: Video Generation using VQ-VAE and Transformers
  • LVT: Latent Video Transformer
  • Feature Quantization Improves GAN Training (ICML2020)
  • DVT-NAT: Fast Decoding in Sequence Models Using Discrete Latent Variables (ICML2018)
  • NWT: Towards natural audio-to-video generation with representation learning

VQ-VAE

在认识VQ-VAE之前,先回顾一下AE、VAE。很早以前看李宏毅老师的视频学过一次,这里就直接使用之前整理的PPT笔记了. 以及对应的代码可以看这里 (vae.ipynb - Colaboratory (google.com))

Auto-Encoder

我们在重构一个图像时,通常输入和输出都是图像本身。为了保证神经网络不是直接的copy,很自然会想到这种降维再升维的方式。这里的latent vector/codings 就是我们希望学到的一个原图像压缩后的低维特征表示。

模型代码如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class Encoder(nn.Module):
def __init__(self, latent_dims):
super(Encoder, self).__init__()
self.linear1 = nn.Linear(784, 512)
self.linear2 = nn.Linear(512, latent_dims)

def forward(self, x):
z = torch.flatten(x, start_dim=1)
z = F.relu(self.linear1(z))
return self.linear2(z)

class Decoder(nn.Module):
def __init__(self, latent_dims):
super(Decoder, self).__init__()
self.linear1 = nn.Linear(latent_dims, 512)
self.linear2 = nn.Linear(512, 784)

def forward(self, z):
x_hat = F.relu(self.linear1(z))
x_hat = torch.sigmoid(self.linear2(x_hat))
return x_hat.reshape((-1, 1, 28, 28))

class Autoencoder(nn.Module):
def __init__(self, latent_dims):
super(Autoencoder, self).__init__()
self.encoder = Encoder(latent_dims)
self.decoder = Decoder(latent_dims)

def forward(self, x):
z = self.encoder(x)
return self.decoder(z)

看源代码可以直观的发现,对于每一个样本,我们学到了对应的压缩后的特征表示的维度是 latent_dims=2.(潜空间的维度当然也可以更大,这里为了可视化更方便,我们设置为2)。

可视化代码:

1
2
3
4
5
6
7
8
def plot_latent(autoencoder, data, num_batches=100):
for i, (x, y) in enumerate(data):
z = autoencoder.encoder(x.to(device))
z = z.to('cpu').detach().numpy()
plt.scatter(z[:, 0], z[:, 1], c=y, cmap='tab10')
if i > num_batches:
plt.colorbar()
break

将得到的这些2 维的 latent vectors可视化如下图:

通过上图我们可以发现,我们训练的每一个样本对应着潜空间中的一个点,这些点是离散的。这意味着,我们并不能生成新的图像,而只能“复制”原有的图像(也可以说,如果在潜空间采样得到的点在这些离散的点之间/外时,生成的样本会比较模糊)。但也有有趣的发现,降维确实是起到了一个聚类的效果,只是效果一般吧,类与类之间并没有完全的分开。这个应该是可以人为控制的。

我们可以试着去生成一些潜空间的vector,然后在此基础上去重构,看看会发现什么?

重构代码:

1
2
3
4
5
6
7
8
9
10
def plot_reconstructed(autoencoder, r0=(-10, 5), r1=(-5, 10), n=12):
w = 28
img = np.zeros((n*w, n*w))
for i, y in enumerate(np.linspace(*r1, n)):
for j, x in enumerate(np.linspace(*r0, n)):
z = torch.Tensor([[x, y]]).to(device)
x_hat = autoencoder.decoder(z)
x_hat = x_hat.reshape(28, 28).to('cpu').detach().numpy()
img[(n-1-i)*w:(n-1-i+1)*w, j*w:(j+1)*w] = x_hat
plt.imshow(img, extent=[*r0, *r1])

代码中我们选取的潜空间的范围是 $x\in [-5, 10]$ , $y\in [10, 5]$ ,在结合可视化的图,可以看到主要是集中在数字 “0” 的区域。右上角是 (10, 5) 可以从可视化的图中也能看到,有部分接近“1”的区域,但生成的样本很模糊。

自编码除了潜空间是离散的,还有个缺点,就是当神经网络太强时,会overfitting。去噪自编码可以缓解这个问题。

VAE

在前面提到AE的缺点是,其得到的潜空间是非连续的,导致采样发生在离散点之间/外时,会生成比较模糊的图片。而VAE就是我们假设先验 $p(z)$ 的每个维度就是一个连续的distribution. AE是用潜空间中的一个点(fixed vector)来对应一个训练样本,VAE则是用一个连续的分布来对应一个训练样本。通过代码来理解这句话:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
class VariationalEncoder(nn.Module):
def __init__(self, latent_dims):
super(VariationalEncoder, self).__init__()
self.linear1 = nn.Linear(784, 512)
self.linear2 = nn.Linear(512, latent_dims)
self.linear3 = nn.Linear(512, latent_dims)

self.kl = 0

def forward(self, x):
x = torch.flatten(x, start_dim=1)
x = F.relu(self.linear1(x))
mu = self.linear2(x)
sigma = torch.exp(self.linear3(x))

z = mu + sigma*torch.randn_like(sigma)
self.kl = 0.5*(sigma**2 + mu**2 - torch.log(sigma) - 1).sum()

return z

class Decoder(nn.Module):
def __init__(self, latent_dims):
super(Decoder, self).__init__()
self.linear1 = nn.Linear(latent_dims, 512)
self.linear2 = nn.Linear(512, 784)

def forward(self, z):
x_hat = F.relu(self.linear1(z))
x_hat = torch.sigmoid(self.linear2(x_hat))
return x_hat.reshape((-1, 1, 28, 28))

class VariationalAutoencoder(nn.Module):
def __init__(self, latent_dims):
super(VariationalAutoencoder, self).__init__()
self.encoder = VariationalEncoder(latent_dims)
self.decoder = Decoder(latent_dims)

def forward(self, x):
z = self.encoder(x)
return self.decoder(z)

第16行代码 z = mu + sigma*torch.randn_like(sigma) 中得到的是潜空间的多维特征表示 z, 每个维度都是一个正态分布(均值为 mu, 方差为 sigma). 第17行代码的kl散度看着很疑惑,这里需要推导得到,先暂放一边。

  • 先验 $p_{\theta}(z)$

  • 似然 $p_{\theta}(x|z)$

  • 后验 $p_{\theta}(z|x)$

右图中实线时生成模型(先验*似然),我们的优化目标可以有两种选择,一个是最大化似然概率(MLH),另一个是最大化后验概率(MAP). 但是对于 $p(x|z)p(z)$ 这个的优化是很难的,我们不可能去采样所有的 $p(z)$。

于是为了解决这个问题,我们可以采用变分近似的方法,也就是我们通过另一个判别模型来代替这个后验概率。也就是 $q_{\phi}(z|x)\sim p_{\theta}(z|x)$ . 这个判别模型就是下图中的encoder。

如代码中13-14行显示那样,我们通过encoder直接生成每一个维度的均值和方差,得到一个多维的正态分布,然后在此基础上去解码。

接下来我们再回过头来说说第17行代码的kl散度是咋回事。我们了解的kl散度通常作为一个loss来衡量两个分布是否接近,那这里我们也确实有两个分布需要拉近,就是下式中的第一项,先验和后验我们需要尽量一致:

img

后验$q_{\phi}(z|x)$就是我们的encoder,先验$p_{\theta}(z|x)$ 就是我们的先验。除此之外,还有第二项极大似然(也就是我们希望最大化观测的样本的概率)。

对于如何通过极大似然推导得到我们的目标函数,李宏毅老师PPT中讲的我觉得没有下面这个清楚。这个就很直观了。

假设之前的我们懂了,我们总归是得到了这样的一个目标损失函数。一个是重构loss,一个是kl散度。

对于第一项kl散度,我们假设先验 $p(z)$ 是正态分布,通过推导可以得到:

image-20210914092720814

这样就跟第17行中的代码一样了。也就是 lower bound的第一项。

对于第二项重构损失:

这里用到的就是一种在参数技巧。直接生成均值和方差,把后验分布这个转换成一个多维的正态分布。

至此,loss中的两项就讲完了。训练代码:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
def train(model, data, epochs=20):
opt = torch.optim.Adam(model.parameters())
for epoch in range(epochs):
print(f"epoch: {epoch+1}/{epochs}")
for x, y in data:
x = x.to(device) # GPU
opt.zero_grad()
x_hat = model(x)
loss = ((x - x_hat)**2).sum() + model.encoder.kl
loss.backward()
opt.step()
return model

latent_dims = 2
vae = VariationalAutoencoder(latent_dims).to(device) # GPU
vae = train(vae, data)

同样为了可视化方便,我们把latent_dims=2. 可视化代码与之前一致:

边缘的位置还是看起来并不连续,可能是可视化样本比较少。

重构代码与之前一致:

效果确实好了一些。。那么VAE的“后验崩塌”在这个图里有体现吗,好像确实更倾向于生成“1”和“7” ?

VQ-VAE

VAE生成的是连续的潜空间,AE生成的是离散的潜空间。那么问题来了,这不是回到原始AE了吗。确实,VQ-VAE本质上更像AE,而不是VAE。但是不同于AE的是,对于vanilla AE,一个训练样本/图像对应一个fixed vector(维度为latent_dims),也就是潜空间中的一个点。而VQ-VAE则是将图像中的一个patch映射到一个codebook中的一个embedding vector(可视化这些codes可能会有有趣的发现)。可以这么说,一个样本/图像是由潜空间中的多个点构成的,这个潜空间就我们需要学习/维护的codebook。

我的理解,不同的潜空间中的点组合就可以生成新的样本/图像了?而且相比AE,潜空间明显更强大了,所以可以避免“后验崩塌”(decoder直接忽略潜向量去生成样本/图像)

“Introducing the VQ-VAE model, which is simple, uses discrete latents, does not suffer from “posterior collapse” and has no variance issues. “ 这是论文中的一句话。

下面是reddit网友/大佬的解释:

Unlike VAEs with continuous latent spaces, VQVAEs use a discrete space by having a codebook containing a large number of continuous vectors. The VQVAE encoding process maps an image to a continuous space then for each spatial location changes each vector to the closest one in the codebook.

Brief recap of posterior collapse in case you’re not sure: my understanding is that VAE models struggle with posterior collapse when (a) the latents contain little information about the input data and (b) a powerful generative model (e.g. an autoregressive decoder) is used that can model the data distribution without any latents. At the start of training the latents often contain little information about the data so the generative model can ignore them and focus on modelling the data distribution on its own. This results in a lack of gradients to the encoder amplifying the problem and so the latents are never used (i.e. posterior collapse).

Specifically with VQVAEs the latents (although discrete) are pretty high dimensional so can store a LOT of information about the input, so this helps with (a). As for (b), the decoder tends to be a fairly simple conv net so the latents are definitely needed to reconstruct the input.

As for the variance issues since VAE encoders are probabilistic, training requires sampling the encoder outputs which can have high variance. Gaussian VAEs bypass this using the reparameterisation trick (here is a great discussion on this https://ermongroup.github.io/cs228-notes/extras/vae/). While VQVAEs have deterministic encoders, the discretisation process can introduce variance, however, they use the straight through estimator which is biased but has low variance, I believe. This leads to other issues such as codebook collapse, where some codes are never used. DALL-E on the other hand uses Gumbel Softmax.

这段话解释了VQ-VAE为啥能解决“后验崩塌”的问题,但是也带来了新的问题,“codebook collapse”,也就是通过encoder压缩之后的vector映射到codebook中的向量时,只会使用到codebook中的很小一部分vectors。

模型前向的框架如下:

很好理解的一个过程,我们可以把 $z_e \rightarrow z_q$ 看作是一个聚类的过程(也确实可以用kmeans方法),甚至也可以看作是一个再参数化的过程,只不过这个过程不可导。

因为存在argmin所以不可导,那么这个就是我们需要解决的一个问题。

怎么解决这个问题呢,把 $z_q$ 的梯度传递到 $z_e$,这听起来太玄乎了。让我们来看看代码吧:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
class Encoder(nn.Module):
def __init__(self, latent_dims, pic_channels=1):
super(Encoder, self).__init__()
self.conv1 = nn.Conv2d(in_channels=pic_channels, out_channels=latent_dims//2, kernel_size=4)
self.conv2 = nn.Conv2d(in_channels=latent_dims//2, out_channels=latent_dims, kernel_size=4)

def forward(self, x):
x = self.conv1(x)
x = F.relu(x)
x = self.conv2(x)
#print(x)
return x

class Decoder(nn.Module):
def __init__(self, latent_dims, pic_channels=1):
super(Decoder, self).__init__()
self.conv_trans1 = nn.ConvTranspose2d(
in_channels=latent_dims, out_channels=latent_dims//2, kernel_size=4)
self.conv_trans2 = nn.ConvTranspose2d(
in_channels=latent_dims//2, out_channels=pic_channels, kernel_size=4)

def forward(self, x):
x = self.conv_trans1(x)
x = F.relu(x)
x = self.conv_trans2(x)
return x


class VectorQuantizer(nn.Module):
def __init__(self, latent_dims, num_codes=32, beta=0.25):
super(VectorQuantizer, self).__init__()
self.K = num_codes
self.D = latent_dims
self.beta = beta

self.codebook = nn.Embedding(self.K, self.D)
self.codebook.weight.data.uniform_(-1 / self.K, 1 / self.K)

self.vq_loss = 0

def forward(self, latents):
'''
latents: (batch, dim, height, width)
codebook: (K, dim)
'''
# convert latents from BCHW -> BHWC
latents = latents.permute(0, 2, 3, 1).contiguous() # (B, H, W, dim)
latents_shape = latents.shape

# Flatten latent
flat_latent = latents.view(-1, self.D) # (BHW, dim)

# Compute L2 distance between latents and codes in codebook
dist = (flat_latent.unsqueeze(1) - self.codebook.weight.unsqueeze(0)) ** 2 # (BHW, 1, dim) - (1, K, dim) -> (BHW, K, dim)
dist = dist.sum(-1) # (BHW, K)

# Get the code index that has the min distance
nearest_idxs = torch.argmin(dist, dim=1).unsqueeze(1) # (BHW, 1)

# Convert to one-hot
nearest_one_hot = torch.zeros(nearest_idxs.size(0), self.K, device=latents.device) # (BHW, K)
nearest_one_hot.scatter_(1, nearest_idxs, 1) # .scatter(dim,index,src)

# Quantize the latents
quantized_latents = torch.matmul(nearest_one_hot, self.codebook.weight).view(latents_shape) # (BHW, K) * (K, dim) = (BHW, dim) -> (B, H, W, dim)

# Compute the VQ Losses
commitment_loss = F.mse_loss(quantized_latents.detach(), latents)
codebook_loss = F.mse_loss(quantized_latents, latents.detach())

self.vq_loss = commitment_loss * self.beta + codebook_loss

# convert quantized from BHWC -> BCHW
return quantized_latents.permute(0, 3, 1, 2).contiguous()

class VQVariationalAutoencoder(nn.Module):
def __init__(self, latent_dims, ema=True):
super(VQVariationalAutoencoder, self).__init__()
self.encoder = Encoder(latent_dims)
if ema:
self.vector_quantizer = VectorQuantizerEMA(latent_dims)
else:
self.vector_quantizer = VectorQuantizer(latent_dims)
self.decoder = Decoder(latent_dims)

def forward(self, x):
z_e = self.encoder(x)
z_q = self.vector_quantizer(z_e) # (batch, dim, 22, 22)
return self.decoder(z_q)

看代码其实也挺简单的:

  • 52-65行计算 l2 距离,然后选择距离最小的index,再通过这个index去codebook中选取对应的vector

  • 这里的 quantized_latents 也就是我们想到得到的 $z_q$

  • 68-69行代码对应下式中后两个loss

这里有个不太懂的问题,$z_q$的计算过程因为第58行argmin肯定是不可导的。可是代码中也没有见到做任何处理。在[苏剑林的博客](VQ-VAE的简明介绍:量子化自编码器 - 科学空间|Scientific Spaces)中讲解了这一点:

按照这个思路,我觉得第74行代码应该改成下面这样才对,这样确确实实做到了把 $z_q$ 的梯度传递给 $z_e$.

1
2
quantized_latents = latents + (quantized_latents - latents).detach()
return quantized_latents.permute(0, 3, 1, 2).contiguous()

提供的代码链接确实是有问题的,改成上述两行代码之后就可以跑出效果了~

接下来我们训练VQ-VAE,并可视化训练过程中的codebook和重构情况。为了方便可视化,我们同样设置latent_dims=2

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def plot_codebook(autoencoder):
codes = autoencoder.vector_quantizer.codebook.weight.detach().cpu().numpy()
for i in range(codes.shape[0]):
plt.scatter(codes[i][0], codes[i][1])
plt.show()

def plot_recon(autoencoder, data, n=10):
x = next(iter(data))[0][:n]
x_hat = autoencoder(x.to(device))
x_hat = x_hat.to('cpu').detach().numpy().squeeze(1)

w = x_hat.shape[1]
img = np.zeros((w, n*w))

print("original:")
for i in range(x_hat.shape[0]):
img[:, i*w:(i+1)*w] = x[i]
plt.imshow(img)
plt.show()

print("reconstructed:")
for i in range(x_hat.shape[0]):
img[:, i*w:(i+1)*w] = x_hat[i]
plt.imshow(img)
plt.show()

def train(autoencoder, data, epochs=20):
opt = torch.optim.Adam(autoencoder.parameters())
for epoch in range(epochs):
print(f"epoch: {epoch+1}/{epochs}")
for x, y in data:
x = x.to(device) # GPU
opt.zero_grad()
x_hat = autoencoder(x)
loss = ((x - x_hat)**2).sum() + autoencoder.vector_quantizer.vq_loss
loss.backward()
opt.step()
plot_codebook(autoencoder)
plot_recon(autoencoder, data)
return autoencoder

训练的可视化过程如下:

这效果有点玄学,第一次跑出来效果挺好,之后跑出来就很模糊。。

把维度扩大到 latent_dims=32后效果会变好,效果如下:

EMA update codebook:

对于loss function中的第二项 $||sg(E(x))- e_{k}||$ 只是用来更新codebook,这个可以用 EMA的方法来更新。也就是对于某一个code $e_i$,它是被 $z_{i,1},z_{i,2},…,z_{i,n}$ 选中的code vector,因此它应该于这些encoder output vector的均值 $\frac{1}{n}\sum_{j}^{n}z_{i,j}$ 更接近。

当使用小批量(minibatches) 训练时,由于数据量不足,直接用这个均值来更新 $e_i$是不准确的。所以用 指数滑动平均(exponential moving average) 来更新 $e_i$:

img

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
class VectorQuantizerEMA(nn.Module):
def __init__(self, latent_dims, num_codes=32, beta=0.25, gamma=0.99, epsilon=1e-5):
super(VectorQuantizerEMA, self).__init__()
self.K = num_codes
self.D = latent_dims
self.beta = beta
self.gamma = gamma
self.epsilon = 1e-5

self.codebook = nn.Embedding(self.K, self.D)
self.codebook.weight.data.normal_()

self.N = None
self.m = None

self.vq_loss = 0

def forward(self, latents):
'''
latents: (batch, dim, height, width)
codebook: (K, dim)
'''
# convert latents from BCHW -> BHWC
latents = latents.permute(0, 2, 3, 1).contiguous() # (B, H, W, dim)
latents_shape = latents.shape

# Flatten latent
flat_latent = latents.view(-1, self.D) # (BHW, dim)

# Compute L2 distance between latents and codes in codebook
dist = (flat_latent.unsqueeze(1) - self.codebook.weight.unsqueeze(0)) ** 2 # (BHW, 1, dim) - (1, K, dim) -> (BHW, K, dim)
dist = dist.sum(-1) # (BHW, K)

# Get the code index that has the min distance
nearest_idxs = torch.argmin(dist, dim=1).unsqueeze(1) # (BHW, 1)

# Convert to one-hot
nearest_one_hot = torch.zeros(nearest_idxs.size(0), self.K, device=latents.device) # (BHW, K)
nearest_one_hot.scatter_(1, nearest_idxs, 1) # .scatter(dim,index,src)

# Quantize the latents
quantized_latents = torch.matmul(nearest_one_hot, self.codebook.weight).view(latents_shape) # (BHW, K) * (K, dim) = (BHW, dim) -> (B, H, W, dim)

# Compute the VQ Losses
commitment_loss = F.mse_loss(quantized_latents.detach(), latents)
self.vq_loss = commitment_loss * self.beta

# EMA update cookbook
n = torch.sum(nearest_one_hot, 0) # (K)
self.N = self.N * self.gamma + (1 - self.gamma) * n if self.N is not None else n
N_ = torch.sum(self.N.data) # Laplace smoothing of the cluster size
self.N = (self.N + self.epsilon) / (N_ + self.K * self.epsilon) * N_
z = torch.matmul(nearest_one_hot.T, flat_latent) # (K, BHW) * (BHW, dim) = (K, dim)
self.m = nn.Parameter(self.m * self.gamma + (1 - self.gamma) * z) if self.m is not None else nn.Parameter(z)
self.codebook.weight = nn.Parameter(self.m / self.N.unsqueeze(1))

# convert quantized from BHWC -> BCHW
quantized_latents = latents + (quantized_latents - latents).detach()
return quantized_latents.permute(0, 3, 1, 2).contiguous()

潜空间维度 latent_dims=2, 可视化如下:

潜空间维度为32, 可视化如下:

前面这些可视化只是在训练过程中的重构,那么我们是否可以像VAE那样,丢掉encoder,直接从 latent space中采样去生成新的样本/图像呢? 但是随机采样出来的codes组合起来的image大概乱七八糟的。。这个时候就在这个低维空间去训练一个自回归模型,来正确采样出有效有意义的codes

在训练VQ-VAE完之后。我们并不能去生成新的样本/图像。怎么去从潜空间采样出这些离散的codes呢?现在我们需要训一个自回归的模型,让这些codes的排列变得有意义。我们可以把这些codes的序列看成一句话,然后训练这样一个自回归模型,也就是stage2 prior training。下图是VQ-VAE-2的伪代码(2.0版本可以看作是一个层次化的VQ-VAE).

简单解释下:

  • $e_{top}\leftarrow Quantize(h_{top})$ 中 $h_{top}$ 就是encoder的输出,通过量化(映射)之后得到 $e_{top}$
  • 对于bottom也是一样的

prior training:

  • $T_{top}$ 就是 $e_{top}$ 的序列

  • $p_{top}=TrainPixelCNN(T_{top})$ 训练自回归模型

VideoGPT

image-20210914210857358

左边就是原始的VQ-VAE. 右侧是训练一个自回归网络。这么看这篇paper确实没啥创新点啊。。就是把vq-vae应用在了video上,不过代码其实复杂了很多。

分两个阶段:

  • VQ-VAE training
  • VideoGPT training

看完代码,第二阶段forward过程如下:

需要注意的是,经过vqvae.encoder之后的维度变成了 [t/4, h/4,w/4], 这就是 latent vectors的个数,后续就是在这个上面做self-attention以及自回归 (注意自回归需要decoder input right-shift).

LVT

这篇跟 VideoGPT很像,都是在潜空间内做自回归,但是这篇paper在VideoGPT之前发表。而且看了openviewer的审稿意见,reviewers 就怼VideoGPT跟这篇很像。。但是两者的模型结构是有区别的

DVT-NAT

Fast Decoding in Sequence Models Using Discrete Latent Variables, (ICML2018)

这篇paper好呀好呀,是真好!!关键还是18年就发表的,是真的牛逼!反思一下为啥别人就能在紧跟前沿,而且还能从CV领域follow出这么厉害的NLP的工作呢?

Abstract

  • 提出了一种对目标序列进行离散建模的方法,能有效提高自回归建模的效率,从而提升解码速度

    To overcome this limitation, we propose to introduce a sequence of discrete latent variables $l_1 . . . l_m$, with $m < n$, that summarizes the relevant information from the sequence $y_1 . . . y_n$. We will still generate $l_1 . . . l_m$ autoregressively, but it will be much faster as $m < n$ (in our experiments we mostly use $m = n/8$ ). Then, we reconstruct each position in the sequence $y_1 . . . y_n$ from $l_1 . . . l_m$ in parallel.

    以前的seq2seq方式是 $(x_1,…,x_L) \rightarrow (y_1,…,y_n)$ ,现在是 $(x_1,…,x_L) \rightarrow (l_1,..,l_m)\rightarrow (y_1,…,y_n)$ ,其中 $l_1,…,l_m$ 的生成依然是自回归的,但是长度更短, 所以提高了效率。

  • 将 latent transformer 应用在机器翻译上,提升翻译效率。但是在BLEU上仍然比自回归的方法要差很多。

Contribution

  • 提出了一个基于离散潜变量的快速解码框架
  • 提出新的离散化技术,能有效缓解VQ-VAE中的index collapse问题
  • latent transofmer 应用在机器翻译上,提升了翻译速度

Discretization Techniques

  • Gumbel-Softmax

    • 将argmax/argmin可微化。将采样过程用可导的softmax代替,同时加上gumbel noise使得最终的结果跟
    • PyTorch 32.Gumbel-Softmax Trick - 科技猛兽的文章 - 知乎 https://zhuanlan.zhihu.com/p/166632315

    image-20210916151621009 image-20210916151649926

  • Improved Semantic Hashing

  • Vector Quantization

Decomposed Vector Quantization

motivation

index collapse, where only a few of the embedding vectors get trained due to a rich getting richer phenomena

Sliced Vector Quantization
  • break up the encoder output enc(y) into n_d smaller slices, like multi-head attention in transformer (eq. 11)

Latent Transformer

  • Main Steps

    • VAE encoder encodes target sentence $y$ into shorter discrete latent variables $l$ (parallel)
    • Latent prediction model - Transformer, is trained to predict $l$ from source sentence $x$ (autoregressive)
    • VAE decoder decodes predicted $l$ back to sequence $y$ (parallel)
  • Loss Function

    • reconstruction loss $l_r$ from VAE
    • latent prediction loss $l_{lp}$ from Latent Transformer
    • in first 10k steps, true targets $y$ is given to transformer-decoder instead of decompressed latents $l$ which ensures self-attention part has reasonable gradients to train the whole architecture
  • Architectures of VAE

    • Encoder

      • conv residual blocks + attention + conv to scale down the dimension
      • $C = n/m, C=2^c$, in the setting, $C=8, c=3$
    • Decoder

      • conv residual blocks + attention + up-conv to scale up the dimension
      • Transformer Decoder
作者

Xie Pan

发布于

2021-09-12

更新于

2021-09-17

许可协议

评论