手機(jī)站首頁散文詩歌雜文隨筆日記小小說

散文網(wǎng) » 科技 »學(xué)習(xí) » VQ-VAE——離散化特征學(xué)習(xí)模型

VQ-VAE——離散化特征學(xué)習(xí)模型

2023-03-31 08:00 作者:小牛翻譯NiuTrans 0人讀過 | 我要投稿

本文首發(fā)于網(wǎng)站?機(jī)器翻譯學(xué)堂

轉(zhuǎn)載事宜請(qǐng)后臺(tái)詢問哦

譯者|黃鵬程

單位|東北大學(xué)自然語言處理實(shí)驗(yàn)室

前言

VQ-VAE是一個(gè)強(qiáng)大的無監(jiān)督表征學(xué)習(xí)模型，它學(xué)習(xí)的離散編碼具有很強(qiáng)的表征能力，最近比較火的文本轉(zhuǎn)圖像模型DALL-E也是基于VQ-VAE的。

在具體介紹VQ-VAE模型前，需要先介紹一下該模型的前身工作AutoEncoder模型以及VAE模型。

1.1 AutoEncoder

自編碼器Auto-Encoder是無監(jiān)督學(xué)習(xí)的一種方式，可以用來做降維、特征提取等。其模型結(jié)構(gòu)如圖1.1.1所示：

從模型圖中可以很明顯的看出，對(duì)于AutoEncoder模型，采取的是無監(jiān)督訓(xùn)練的方式，對(duì)于輸入的 $x$ 經(jīng)過一個(gè)Encoder層后得到一個(gè)特征向量 $z$ ，再將該向量 $z$ 通過一個(gè)Decoder層得到最終輸出 $x%20%E2%80%99$ ，通過最小化重構(gòu)模型的輸入 $x$ 和模型的輸出 $x%E2%80%99$ 的誤差作為損失函數(shù)訓(xùn)練模型得到一個(gè)較好的關(guān)于輸入 $x$ 的特征向量 $z$ ，模型設(shè)計(jì)的初衷的獲得一個(gè)對(duì)應(yīng)于源數(shù)據(jù) $x$ 的一個(gè)低維特征向量 $z$ ，在獲得此向量的基礎(chǔ)上可以應(yīng)用在很多分類任務(wù)上，但是AE模型并不適用于生成任務(wù)。

盡管AE已經(jīng)可以獲得較好的向量表示，在還原任務(wù)上可以做出較好的效果，但其并不是一個(gè)生成式模型，這是因?yàn)閷?duì)于一個(gè)生成模型而言，他一般需要滿足兩個(gè)條件限制：

1.生成模型的編碼器和解碼器是可以分離開的；

2.對(duì)于固定維度下任意采樣出的編碼，解碼器都能產(chǎn)生一張真實(shí)且清晰的圖片。

AE模型并不滿足第二點(diǎn)的條件，舉個(gè)例子來說[2]，對(duì)于輸入的全月圖和半月圖，通過對(duì)AE模型的訓(xùn)練可以很好的完成還原任務(wù)，但是我們對(duì)于二者特征向量中取一個(gè)點(diǎn)，對(duì)于一個(gè)正常的生成模型而言，應(yīng)該生成一個(gè)介于全月和半月之間的圖片。然而，對(duì)于真實(shí)的AE而言，它生成的結(jié)果要么是亂碼要么就是異常模糊的圖片。為什么會(huì)發(fā)生這種情況呢？因?yàn)槟Ｐ驮谟?xùn)練的時(shí)候并沒有顯性對(duì)中間變量 $z$ 的分布 $P(Z)$ 進(jìn)行建模，在模型訓(xùn)練時(shí)所采用的 $f(z)$ 是有限的，而對(duì)于 $Z$ 所處的空間存在大量 $f(z)$ 外的點(diǎn)而言模型是并不理解的，如果像該例子隨機(jī)在全月和半月中采樣一個(gè)點(diǎn)，大概率得到不能夠生成有效圖片的點(diǎn)。

1.2 VAE：Auto-Encoding Variational Bayes

既然介紹過了AE模型存在的一些問題，我們?yōu)槭裁床唤o定 $Z$ 一個(gè)簡(jiǎn)單的分布呢？這樣我們就可以通過神經(jīng)網(wǎng)絡(luò)的方法去學(xué)習(xí)一個(gè)模擬分布的映射，再通過這個(gè)學(xué)習(xí)到分布就可以完成生成任務(wù)了，這就是VAE所用到的思想了。我們假設(shè) $Z$ ~ $N(0%2CI)$ ，其中 $I$ 代表一個(gè)單位矩陣。也就是說，我們將 $z$ 看作是一個(gè)服從標(biāo)準(zhǔn)多元高斯分布的多維隨機(jī)變量,對(duì)于訓(xùn)練樣本中的每一個(gè) $x$ ，通過Encoder模塊獲得對(duì)于該樣本的均值和方差，即對(duì)于每一個(gè) $x_%7Bi%7D%20$ 而言，獲取其對(duì)應(yīng)的分布 $z_%7Bi%7D%20$ ，如圖1.2.1所示，將得到的每一個(gè) $z_%7Bi%7D%20$ 混合起來，即可得到數(shù)據(jù)的大致分布，如圖1.2.2所展示的一樣，那么之前存在的問題就解決了，既然獲得了 $X$ 的分布 $P(X)$ ，那么想實(shí)現(xiàn)生成任務(wù)就可以從樣本空間中隨機(jī)采樣再送給Decoder進(jìn)行生成任務(wù)，本文的重點(diǎn)并不在于VAE，因此對(duì)于其細(xì)節(jié)、公式就不一一推導(dǎo)了，感興趣的可以自己看看論文。

VQ-AVE：Neural Discrete Representation Learning

2.1 Introduction

終于來到了本文的重點(diǎn)了，VQ在這里的含義為vector quantised，即將VAE做量化，也可以稱為離散化，那為什么要使用離散化的VAE呢，根據(jù)論文所言（which are potentially a more natural fit for many of the modalities），根據(jù)我個(gè)人理解，雖然自然界的各種信號(hào)不論是圖像還是語言，他都是連續(xù)的，大部分任務(wù)都是一個(gè)回歸任務(wù)，但當(dāng)我們處理這些問題時(shí)，往往是將問題離散化后再進(jìn)行處理，如圖像劃分成像素、語音也是經(jīng)過抽樣，回歸任務(wù)也變成了分類任務(wù)。

在這里也一樣，如果使用VAE的方式，樣本的分布 $P%EF%BC%88z%EF%BC%89$ 并不是很好學(xué)，VQ-VAE就使用了一個(gè)codebook去替代了VAE中學(xué)習(xí)樣本分布的過程，我們假設(shè)codebook是 $K%5Ctimes%20D$ 維的，其中 $K$ 是指codebook的長(zhǎng)度，一般設(shè)定為8192維，而 $D$ 則是每一位向量的長(zhǎng)度，一般設(shè)為512，如圖2.1所示，codebook的長(zhǎng)度 $K$ 可以簡(jiǎn)單的理解為codebook對(duì)于 $K$ 個(gè)聚類中心。

模型經(jīng)過Encoder后，得到一個(gè) $H%5Ctimes%20W$ 的特征圖，然后我們就將特征圖的向量去與codebook中的 $K$ 個(gè)向量作比較，將最相似的 $e_%7Bi%7D%20$ 的 $index$ 存入特征矩陣 $Z$ 中，再通過 $Z$ 和codebook得到一個(gè)新的特征圖 $f_%7Bq%7D%20$ （即量化后的特征），之后的流程又和AE相似了將這個(gè)新拿到的特征圖通過Decoder獲取新的生成圖片 $x%E2%80%99$ ，盡量讓 $x%E2%80%99$ 和 $x$ 相似，這樣就完成了VQ-VAE模型Encoder和Decoder部分的訓(xùn)練，其流程可以簡(jiǎn)單的劃分為如下的幾個(gè)步驟：

1.將一張圖片送到Encoder后，會(huì)得到一個(gè) $H%5Ctimes%20W%5Ctimes%20D$ 的feature map，即圖2.1所示 $z_%7Be%7D%EF%BC%88x%EF%BC%89%20$ ；

2.將Encoder得到的 $z_%7Be%7D%EF%BC%88x%EF%BC%89%20$ 中的 $H%5Ctimes%20W$ 個(gè) $D$ 維的向量分別與codebook中的 $K$ 個(gè)向量分別計(jì)算相似度，找到最近的 $e_%7Bi%7D%20$ ，用index表示，得到 $H%5Ctimes%20W$ 的特征矩陣 $q%EF%BC%88z%7Cx%EF%BC%89$ 。

3.對(duì)于特征圖 $z$ 中的每個(gè)index用codebook對(duì)應(yīng)的隱變量 $e_%7Bi%7D%20$ 表示，得到新的feature map，即圖2.1中的 $z_%7Bq%7D%EF%BC%88x%EF%BC%89%20$ ，該特征圖作為Decoder的輸入，最終通過Decoder得到重構(gòu)后的圖片。

2.2 Training

接下來來到了模型的訓(xùn)練環(huán)節(jié)，模型的整體損失函數(shù)如下所示：

$L%3Dlogp(x%7Cz_q(x))%2B%7C%7Csg%5BZ_e(x)%5D-e%7C%7C_2%5E2%2B%5C%5C%20%5Cbeta%7C%7CZ_e(x)-sg%5Be%5D%7C%7C_2%5E2$

它由三部分組成：

1.第一部分 $%5Clog%20p%20(x%7CZ_%7Bq%7D(x)%20)$ 為重構(gòu)損失（reconstruction loss），與Autoencoder的訓(xùn)練loss一致，都是通過最小化重構(gòu)模型輸入 $x$ 以及Decoder部分。

正常來說，Encoder輸出的每一個(gè)詞向量需要從codebook中找到最相似的詞向量來代替他送入Decoder完成計(jì)算。這個(gè)“選擇”的過程自然是不可導(dǎo)的，也就是說，目前的計(jì)算方式會(huì)導(dǎo)致Encoder無法計(jì)算自己的梯度，自然無法訓(xùn)練。對(duì)于這個(gè)問題，作者采用了Straight-Through思想，直接將原本能夠正常計(jì)算出來的codebook身上的梯度 $dL%2Fdz_%7Bq%7D%20$ 直接作為Encoder的梯度，即 $dL%2Fdz_%7Be%7D%20$ 。（Straight-Through最早出自Benjio的論文《Estimating or Propagating Gradients Through Stochastic Neurous for Conditional Computation》它想說的就是前向傳播的時(shí)候可以用不可導(dǎo)的變量參與計(jì)算，而反向傳播的時(shí)候，用為它特別設(shè)計(jì)的梯度來更新它的參數(shù)。）如果有朋友對(duì)于這個(gè)操作還不太理解的話可以看看下面這段代碼（此代碼取自vq-wav 2vec）。

這段代碼中的ze代表Encoder的輸出，zq代表codebook的輸出，detach的作用為將該變量從計(jì)算圖中剝離出來，這會(huì)導(dǎo)致它只會(huì)參與前向計(jì)算，而在反向計(jì)算的過程中，它沒有任何梯度，因此它的參數(shù)不會(huì)被更新。

這段代碼的作者沒有直接把zq送給vq vae的Decoder，而是先通過一個(gè)_pass_grad函數(shù)。我們?cè)賮砜確pass_grad函數(shù)，這個(gè)函數(shù)的作用就是將原本能夠正常計(jì)算出來的codebook身上的梯度 $dL%2Fdz_%7Bq%7D%20$ 直接作為Encoder的梯度，即 $dL%2Fdz_%7Be%7D%20$ 。它的具體做法是計(jì)算了一個(gè)zq.detach()

+(ze-ze.detach())，這種方式就會(huì)使得函數(shù)返回的結(jié)果是 $z_%7Bq%7D%20$ ，而 $z_%7Be%7D%20$ 的值因?yàn)樽约簻p自己就沒了。同時(shí)，其中一個(gè)ze通過detach消除了他身上的梯度，因此梯度可以被傳遞到另一個(gè)ze身上。相反，如果我們直接計(jì)算ze - ze的話，那么ze身上就不會(huì)產(chǎn)生任何梯度。代碼的作者通過這種方式實(shí)現(xiàn)Straight-Through使得Encoder能夠被正常訓(xùn)練到。

注：有朋友可能會(huì)想x - x.detach()是合理的，那么為什么_pass_grad是y.detach() + (x - x.detach())而不是y + (x - x.detach())？這其實(shí)是vq-wav 2vec的代碼作者復(fù)現(xiàn)vq vae原作者所說“Due to the straight-through gradient estimation of mapping from? $z_%7Be%7D(x)$ to? $z_%7Bq%7D(x)%20$ ，the embeddings? $e_%7Bi%7D%20$ ?receive no gradients from the reconstruction loss? $%5Clog%20p(z%7C%20z_p(x))$ ."而設(shè)計(jì)的。但實(shí)際上，我也嘗試了y+（x-x.detach()）,發(fā)現(xiàn)這時(shí)候模型也是正常訓(xùn)練的，并且 $dL%2Fdz_e$ 與? $dL%2Fdz_q$ 完全相同，符合之前的所有設(shè)計(jì)，因此僅用這個(gè)reconstruction loss就實(shí)現(xiàn)了Encoder、codebook、Decoder三者一起進(jìn)行訓(xùn)練。但是為了更好地介vq vae，我們接下倆還是按照y.detach()+(x-x.detach())，也就是Straight-Through導(dǎo)致的 $dL%2Fdz_q$ 為零來講解。

2.第二部分 $%7C%7Csg%5BZ_e(x)%5D-e%7C%7C_2%5E2$ 為codebook loss。我們前邊說了，由于straight-through gradient estimation of mapping from ze（x）to zq（x），因此codebook并不能通過 $log%20p(x%7Cz_q(x))$ 獲得任何梯度。對(duì)于這個(gè)問題，作者使用了vector quantisation（VQ）算法，即增加一個(gè)新的loss，來最小化 $z_e(x)$ 和embedding? $e$ 之間的 $L2$ 距離，這會(huì)使得 $e$ 即codebook接近encoder的輸出，以此來訓(xùn)練codebook。公式中 $sg$ 表示stopgradient operator，即在前向計(jì)算的時(shí)候保持相應(yīng)的量不變，但在后向計(jì)算的時(shí)候使得梯度為0（也即pytorch中的detach方法）。這么說的有點(diǎn)抽象，讓我們直接看看代碼是如何實(shí)現(xiàn)的：

代碼的作者通過mseloss來計(jì)算zq和ze之間的距離，同時(shí)按照公式中的 $sg%5BZ_e(x)%5D~%2C$ 將ze使用detach從計(jì)算圖中剝離，因此這個(gè)loss產(chǎn)生的梯度只會(huì)被傳導(dǎo)到zq身上。這段代碼的含義就是只更新codebook的參數(shù)使得其與 $z_%7Be%7D%EF%BC%88x%EF%BC%89$ 接近。

3.第三部分 $%5Cbeta%7C%7CZ_e(x)-sg%5Be%5D%7C%7C_2%5E2$ 稱為commitment loss，即原文中的（To make sure the encoder commits to an embedding and its output does not grow,we add a commitment loss,the third term in equation 3）它只訓(xùn)練模型的Encoder部分，個(gè)人理解其目的是為了讓Encoder的輸出穩(wěn)定在一個(gè)codebook聚類，不在codebook里面亂跳，作者通過實(shí)驗(yàn)發(fā)現(xiàn)最后的結(jié)果對(duì)于 $%5Cbeta%20$ 并不敏感，故設(shè)置了一個(gè)0.25的值。對(duì)于這個(gè)loss的計(jì)算代碼如下：

在訓(xùn)練好VQ-VAE后，我們并不能直接使用VQ-VAE模型的Decoder用作生成式任務(wù)，還需要訓(xùn)練一個(gè)先驗(yàn)?zāi)Ｐ蛠韺?shí)現(xiàn)數(shù)據(jù)生成，拿圖像來舉例子，論文中采用的就是PixelCNN模型，與以往不同的是，PixelCNN模型的輸入并不是圖像所對(duì)應(yīng)的pixels，而是VQ-VAE模型學(xué)到的那個(gè)離散編碼。首先，我們需要使用已經(jīng)訓(xùn)練好的VQ-VAE對(duì)訓(xùn)練圖像進(jìn)行推理，得到其相對(duì)應(yīng)的特征矩陣 $z(m*m)$ 這個(gè)特征矩陣在一定程度上保留了輸入圖片的位置信息，所以我們?cè)谶@里可以用自回歸模型如PixelCNN，來對(duì)特征矩陣 $z'(m*m)$ ，再通過codebook得到相應(yīng)的特征圖 $f(m*m*D)$ ，最后經(jīng)過Decoder就可以生成圖片了。[4]

2.3 Experiments

2.3.1 lmage

作者首先展示了重構(gòu)的效果，如圖2.4.1所示：

從圖像上可以看出重構(gòu)出來的圖基本和原圖一致，印證了使用離散隱變量的方式重構(gòu)圖片還是效果不錯(cuò)的。

隨后再使用 $m%5Ctimes%20m%5Ctimes%201$ 的latent space去訓(xùn)練PixeICNN，從PixeICNN提取的樣本被VQ-VAE的Decoder映射到像素空間，可以在圖2.4.2中看到。

2.3.2 Audios

作者在VCTK上進(jìn)行實(shí)驗(yàn)，decoder使用了類WaveNet的架構(gòu)。

由于是multi-speaker dataset，decoder中還輸入了spekaer-id。

Reconstruction的結(jié)果如圖2.4.3所示：

原論文中給出了sample page，感興趣的同學(xué)可以去聽一下。

聽了之后可以發(fā)現(xiàn)雖然reconstruct audio的style會(huì)和原audio有不同，但是內(nèi)容是相同的。這說明VQ-VAE在沒有text信息進(jìn)行監(jiān)督的情況下也能學(xué)習(xí)到語音中的關(guān)鍵信息，且忽略了一些low-level的信息。

之后作者使用WaveNet重新學(xué)習(xí)了prior，然后進(jìn)行了sample，作者發(fā)現(xiàn)模型可以生成比較清晰且流暢的語音，這說明VQ-VAE實(shí)際上只使用語音數(shù)據(jù)就學(xué)習(xí)到了一個(gè)phoneme-level language model。

作者也進(jìn)行了speaker conversion的實(shí)驗(yàn)，即使用不同的speaker-id進(jìn)行decoding，作者發(fā)現(xiàn)模型可以生成內(nèi)容相同但speaker不同的語音，這說明encoder排除了speaker相關(guān)的信息。

最后作者也嘗試將每個(gè)embedding映射到最有可能的phoneme上，這里有128個(gè)latent variable，41個(gè)phoneme，最終得到的映射的accuracy為49.3%，這比起random的情況（7.2%）要好很多，這說明VQ-VAE學(xué)習(xí)到的latent variable中包含的信息和phoneme很類似。[5]

2.3.3 Videos

從圖2.4.4可以看出，該模型已經(jīng)學(xué)會(huì)了在不降低視覺質(zhì)量的情況下成功生成以給定動(dòng)作為條件的幀序列，同時(shí)保持局部幾何形狀的正確性。

2.4 Conclusion

講完了VQ-AVE的大致思路，我們會(huì)發(fā)現(xiàn)，現(xiàn)在學(xué)習(xí)到的又是一個(gè)固定的codebook，這也意味著它又沒辦法像VAE一樣，通過隨機(jī)采樣生成圖片，準(zhǔn)確的說VQ-VAE并不像一個(gè)VAE，而更像一個(gè)AE，它學(xué)習(xí)到的codebook適用于類似分類的任務(wù)而不是生成任務(wù)，如果想要VQ-VAE做生成任務(wù)，就需要像論文里提到的一樣，再訓(xùn)練一個(gè)prior網(wǎng)絡(luò)，利用codebook實(shí)現(xiàn)圖像的生成任務(wù)。

引用

[1]知乎：https://zhuanlan.zhihu.com/p/58111908

[2]知乎：https://zhuanlan.zhihu.com/p/112513743

[3]知乎：https://zhuanlan.zhihu.com/p/249296925

[4]蘇劍林. (Jun. 24, 2019). 《VQ-VAE的簡(jiǎn)明介紹：量子化自編碼器》[Blog post]. Retrieved from?https://spaces.ac.cn/archives/6760

[5]知乎：https://zhuanlan.zhihu.com/p/382305612