昨天早上，筆者在日常刷arixv的時候，然後被一篇新出來的論文震驚了！論文名字叫做《NVAE: A Deep Hierarchical Variational Autoencoder》，顧名思義是做VAE的改進工作的，提出了一個叫NVAE的新模型。說實話，筆者點進去的時候是不抱什麼希望的，因為筆者也算是對VAE有一定的瞭解，覺得VAE在生成模型方面的能力終究是有限的。結果，論文開啟了，呈現出來的畫風是這樣的：
NVAE的人臉生成效果
然後筆者的第一感覺是這樣的：
W!T!F! 這真的是VAE生成的效果？這還是我認識的VAE麼？看來我對VAE的認識還是太膚淺了啊，以後再也不能說VAE生成的影像模糊了...
不過再看了看作者機構，原來是NVIDIA，這也大概能接受了。最近幾年可能大家都留意到NVIDIA通常都在年底釋出個生成模型的突破，2017年底是PGGAN，2018年底是StyleGAN，2019年底是StyleGAN2，今年貌似早了些，而且動作也多了些，因為上個月才發了個叫ADA的方法，將Cifar-10的生成效果提到了一個新高度，現在又來了個NVAE。
那這個NVAE究竟有什麼特別的地方，可以實現VAE生成效果的突飛猛進呢？
VAE回顧 #
可能讀者認真觀察後會說：
好像還是有點假呀，那臉部也太光滑了，好像磨過皮一樣，還比不上StyleGAN呀～
是的，這樣評價並沒有錯，生成痕跡還是挺明顯的。但如果你沒感覺到震驚，那估計是因為你沒看過之前的VAE生成效果，一般的VAE生成畫風是這樣的：
一般的VAE的隨機生成效果
所以，你還覺得這不是一個突破嗎？
那麼，是什麼限制了（以前的）VAE的表達能力呢？這一次的突破又是因為改進了哪裡呢？讓我們繼續看下去。
基本介紹 #
VAE，即變分自編碼器（Variational Auto-Encoder），本空間已經有不少文章介紹過了，在右側搜尋欄搜尋“變分自編碼器”就能搜到很多相關博文。這裡做個簡單的回顧和分析。
在筆者對VAE的推導裡邊，我們是先有一批樣本，這批樣本代表著一個真實的（但不知道形式的）分佈$\tilde{p}(x)$，然後我們構建一個帶引數的後驗分佈$p(z|x)$，兩者就組成一個聯合分佈$p(x,z)=\tilde{p}(x)p(z|x)$。接著，我們再定義一個先驗分佈$q(z)$，已經定義一個生成分佈$q(x|z)$，這樣構成另一個聯合分佈$q(x,z)=q(z)q(x|z)$。最後，我們的目的就是讓$p(x,z),q(x,z)$相互接近起來，所以我們去最佳化兩者之間的KL散度：
\begin{equation}\begin{aligned}
KL\big(p(x,z)\big\Vert q(x,z)\big)=&\iint p(x,z)\log \frac{p(x,z)}{q(x,z)} dzdx\\
=&\mathbb{E}_{x\sim \tilde{p}(x)} \Big[\mathbb{E}_{z\sim p(z|x)}\big[-\log q(x|z)\big]+KL\big(p(z|x)\big\Vert q(z)\big)\Big] + \text{常數}
\end{aligned}\end{equation}
這就是VAE的最佳化目標。
困難分析 #
對$p(z|x),q(z),q(x|z)$的要求是：1、能寫出解析表示式；2、方便取樣。然而連續型分佈的世界裡這樣的分佈並不多，最常用的也就是高斯分佈了，而這其中又以“各分量獨立的高斯分佈”最為簡單，所以在一般的VAE裡邊，$p(z|x),q(z),q(x|z)$都被設為各分量獨立的高斯分佈：$p(z|x)=\mathcal{N}(z;\mu_1(x),\sigma_1^2(x))$、$q(z)=\mathcal{N}(z;0,1)$以及$q(x|z)=\mathcal{N}(x;\mu_2(z),\sigma_2^2(z))$。
問題是，“各分量獨立的高斯分佈”不能擬合任意複雜的分佈，當我們選定$p(z|x)$的形式後，有可能不管我們怎麼調它的引數，$\int \tilde{p}(x)p(z|x)dx$和$\frac{\tilde{p}(x)p(z|x)}{\int \tilde{p}(x)p(z|x)dx}$都不能成為高斯分佈，這就意味著$KL\big(p(x,z)\big\Vert q(x,z)\big)$從理論上來說就不可能為0，所以讓$p(x,z),q(x,z)$相互逼近的話，只能得到一個大致、平均的結果，這也就是常規VAE生成的影像偏模糊的原因。
相關改進 #
改進VAE的一個經典方向是將VAE與GAN結合起來，比如CVAE-GAN、AGE等，目前這個方向最先進結果大概是IntroVAE。從理論上來講，這類工作相當於隱式地放棄了$q(x|z)$是高斯分佈的假設，換成了更一般的分佈，所以能提升生成效果。不過筆者覺得，將GAN引入到VAE中有點像“與虎謀皮”，藉助GAN提升了效能，但也引入了GAN的缺點（訓練不穩定等），而且提升了效能的VAE生成效果依然不如純GAN的。另外一個方向是將VAE跟flow模型結合起來，比如IAF-VAE以及筆者之前做的f-VAE，這類工作則是透過flow模型來增強$p(z|x)$或$q(x|z)$的表達能力。
還有一個方向是引入離散的隱變數，典型代表就是VQ-VAE，其介紹可以看筆者的《VQ-VAE的簡明介紹：量子化自編碼器》。VQ-VAE透過特定的編碼技巧將圖片編碼為一個離散型序列，然後PixelCNN來建模對應的先驗分佈$q(z)$。前面說到，當$z$為連續變數時，可選的$p(z|x),q(z)$都不多，從而逼近精度有限；但如果$z$是離散序列的話，$p(z|x),q(z)$對應離散型分佈，而利用自迴歸模型（NLP中稱為語言模型，CV中稱為PixelRNN/PixelCNN等）我們可以逼近任意的離散型分佈，因此整體可以逼近得更精確，從而改善生成效果。其後的升級版VQ-VAE-2進一步肯定了這條路的有效性，但整體而言，VQ-VAE的流程已經與常規VAE有很大出入了，有時候不大好將它視為VAE的變體。
NVAE梳理 #
鋪墊了這麼久，總算能談到NVAE了。NVAE全稱是Nouveau VAE（難道不是Nvidia VAE？），它包含了很多當前CV領域的新成果，其中包括多尺度架構、可分離卷積、swish啟用函式、flow模型等，可謂融百家之所長，遂成當前最強VAE～
（提醒，本文的記號與原論文、常見的VAE介紹均有所不同，但與本部落格其他相關文章是一致的，望讀者不要死記符號，而是根據符號的實際含義來理解文章。）
自迴歸分佈 #
前面我們已經分析了，VAE的困難源於$p(z|x),q(z),q(x|z)$不夠強，所以改進的思路都是要增強它們。首先，NVAE不改變$q(x|z)$，這主要是為了保持生成的並行性，然後是透過自迴歸模型增強了先驗分佈$q(z)$和後驗分佈$p(z|x)$。具體來說，它將隱變數分組為$z=\{z_1,z_2,\dots,z_L\}$，其中各個$z_l$還是一個向量（而非一個數），然後讓
\begin{equation}q(z)=\prod_{l=1}^L q(z_l|z_{< l}),\quad p(z|x)=\prod_{l=1}^L p(z_l|z_{< l},x)\label{eq:arpq}\end{equation}
而各個組的$q(z_l|z_{< l}),p(z_l|z_{< l},x)$依舊建立為高斯分佈，所以總的來說$q(z),p(z|x)$就被建立為自迴歸高斯模型。這時候的後驗分佈的KL散度項變為
\begin{equation}KL\big(p(z|x)\big\Vert q(z)\big)=KL\big(p(z_1|x)\big\Vert q(z_1)\big)+\sum_{l=2}^L \mathbb{E}_{p(z_{< l}|x)}\Big[KL\big(p(z_l|z_{< l}, x)\big\Vert q(z_l|z_{< l})\big)\Big]\end{equation}
當然，這個做法只是很樸素的推廣，並非NVAE的首創，它可以追溯到2015年的DRAW、HVM等模型。NVAE的貢獻是給式$\eqref{eq:arpq}$提出了一種“相對式”的設計：
\begin{equation}\begin{aligned}&q(z_l|z_{< l})=\mathcal{N}\left(z_l;\mu(z_{< l}),\sigma^2(z_{< l})\right)\\
&p(z_l|z_{< l},x)=\mathcal{N}\left(z_l;\mu(z_{< l})+\Delta\mu(z_{< l},x),\sigma^2(z_{< l})\otimes \Delta\sigma^2(z_{< l}, x)\right)
\end{aligned}\end{equation}
也就是說，沒有直接去後驗分佈$p(z_l|z_{< l},x)$的均值方差，而是去建模的是它與先驗分佈的均值方差的相對值，這時候我們有（簡單起見省去了自變數記號，但不難對應理解）
\begin{equation}KL\big(p(z_l|z_{< l}, x)\big\Vert q(z_l|z_{< l})\big)=\frac{1}{2} \sum_{i=1}^{|z_l|} \left(\frac{\Delta\mu_{(i)}^2}{\sigma_{(i)}^2} + \Delta\sigma_{(i)}^2 - \log \Delta\sigma_{(i)}^2 - 1\right)\end{equation}
原論文指出這樣做能使得訓練更加穩定。
多尺度設計 #
現在隱變數分成了$L$組$z=\{z_1,z_2,\dots,z_L\}$，那麼問題就來了：1、編碼器如何一一生成$z_1,z_2,\dots,z_L$？2、解碼器如何一一利用$z_1,z_2,\dots,z_L$？也就是說，編碼器和解碼器如何設計？
NVAE中的編碼器和解碼器架構。其中r代表殘差模組，h代表可訓練引數，藍色部分是引數共享的
NVAE巧妙地設計了多尺度的編碼器和解碼器，如上圖所示。首先，編碼器經過層層編碼，得到最頂層的編碼向量$z_1$，然後再慢慢地從頂層往下走，逐步得到底層的特徵$z_2,\dots,z_L$；至於解碼器，自然也是一個自上往下地利用$z_1,z_2,\dots,z_L$的過程，而這部分剛好也是與編碼器生成$z_1,z_2,\dots,z_L$的過程有共同之處，所有NVAE直接讓對應的部分引數共享，這樣既省了引數量，也能透過兩者間的相互約束提高泛化效能。
這種多尺度設計在當前最先進的生成模型都有體現，比如StyleGAN、BigGAN、VQ-VAE-2等，這說明多尺度設計的有效性已經得到比較充分的驗證。此外，為了保證效能，NVAE還對殘差模組的設計做了仔細的篩選，最後才確定瞭如下的殘差模組，這煉丹不可謂不充分極致了：
NVAE中的殘差模組
其他提升技巧 #
除了以上兩點比較明顯的特徵外，其實NVAE還包含了很多對效能有一定提升的技巧，這裡簡單列舉一些。
BN層的改進。當前很多生成模型已經棄用BN（Batch Normalization）了，多數會改用IN（Instance Normalization）或WN（Weight Normalization），因為發現用BN會損失效能。NVAE透過實驗發現，其實BN對訓練還是有幫助的，但對預測有害，原因是預測階段所使用的滑動平均得來的均值方差不夠好，所以NVAE在模型訓練完後，透過多次取樣同樣batch_size的樣本來重新估算均值方差，從而保證了BN的預測效能。此外，為了保證訓練的穩定性，NVAE還給BN的$\gamma$的模長加了個正則項。
譜正則化的應用。我們知道，任意兩個分佈的KL散度是無上界的，所以VAE裡邊的KL散度項也是無上界的，而最佳化這種無上界的目標是很“危險”的，說不準啥時候就發散了。所以同樣是為了穩定訓練，NVAE給每一個卷積層都加了譜正則化，其概念可以參考筆者之前的《深度學習中的Lipschitz約束：泛化與生成模型》。加譜歸一化可以使得模型的Lipschitz常數變小，從而使得整個模型的Landscape更為光滑，更利於模型穩定訓練。
flow模型增強分佈。透過自迴歸模型，NVAE增強了模型對分佈的擬合能力。不過這個自迴歸只是對組間進行的，對於組內的單個分佈$p(z_l|z_{< l}, x)$和$q(z_l|z_{< l})$，依然假設為各分量獨立的高斯分佈，這說明擬合能力依然有提升空間。更徹底的方案是，對於組內的每個分量也假設為自迴歸分佈，但是這樣一來在取樣的時候就巨慢無比了（所有的分量串聯遞迴取樣）。NVAE提供了一個備選的方案，透過將組內分佈建立為flow模型來增強模型的表達能力，同時保持組內取樣的並行性。實驗結果顯示這是有提升的，但筆者認為引入flow模型會大大增加模型的複雜度，而且提升也不是特別明顯，感覺能不用就不用為好。
節省視訊記憶體的技巧。儘管NVIDIA應該不缺顯示卡，但NVAE在實現上還是為省視訊記憶體下了點功夫。一方面，它採用了混合精度訓練，還順帶在論文推了一波自家的APEX庫。另一方面，它在BN那裡啟用了gradient check-point（也就是重計算）技術，據說能在幾乎不影響速度的情況下節省18%的視訊記憶體。總之，比你多卡的團隊也比你更會省視訊記憶體～
更多效果圖 #
到這裡，NVAE的技術要點基本上已經介紹完畢了。如果大家還覺得意猶未盡的話，那就多放幾張效果圖吧，讓大家更深刻地體會NVAE的驚豔之處。
NVAE在CelebA HQ和FFHQ上的生成效果。值得注意的是，NVAE是第一個FFHQ資料集上做實驗的VAE類模型，而且第一次就如此驚豔了
基於NVAE的影像檢索實驗。左邊是隨機生成的樣本，右邊是訓練集的最相似樣本，主要是檢驗模型是否只是單純地記住了訓練集
CelebA HQ的更多生成效果
個人收穫 #
從下述訓練表格來看，我們可以看到訓練成本還是蠻大的，比同樣解析度的StyleGAN都要大，並且縱觀整篇論文，可以發現有很多大大小小的訓練trick（估計還有不少是沒寫在論文裡邊的，當然，其實在StyleGAN和BigGAN裡邊也包含了很多類似的trick，所以這不算是NVAE的缺點），因此對於個人來說，估計是不容易復現NVAE的。那麼，對於僅僅是有點愛好生成模型的平民百姓（比如筆者）來說，從NVAE中可以收穫到什麼呢？
NVAE的訓練引數與成本
對於筆者來說，NVAE帶來的思想衝擊主要有兩個。
第一，就是自迴歸的高斯模型可以很有力地擬合複雜的連續型分佈。以前筆者以為只有離散分佈才能用自迴歸模型來擬合，所以筆者覺得在編碼時，也需要保持編碼空間的離散型，也就是VQ-VAE那一條路。而NVAE證明了，哪怕隱變數是連續型的，自迴歸高斯分佈也能很好地擬合，所以不一定要走VQ-VAE的離散化道路了，畢竟連續的隱變數比離散的隱變數更容易訓練。
第二，VAE的隱變數可以不止一個，可以有多個的、分層次的。我們再次留意上表，比如FFHQ那一列，關於隱變數$z$的幾個資料，他一共有$4+4+4+8+16=36$組，每組隱變數大小還不一樣，分別是$\left\{8^2,16^2,32^2,64^2,128^2\right\}\times 20$，如此算來，要生成一個$256\times 256$的FFHQ影像，需要一個總大小有
\begin{equation}\left(4\times 8^2 + 4\times 16^2 + 4\times 32^2 + 8\times 64^2 + 16\times 128^2 \right)\times 20=6005760\end{equation}
維的隨機向量，也就是說，取樣一個600萬維的向量，生成一個$256\times 256\times 3 = 196608$（不到20萬）維的向量。這跟傳統的VAE很不一樣，傳統的VAE一般只是將圖片編碼為單個（幾百維）的向量，而這裡的編碼向量則多得多，有點全卷積自編碼器的味道了，所以清晰度提升也是情理之中。
Nouveau是啥？ #
最後，筆者饒有性質地搜尋了一下Nouveau的含義，以下是來自維基百科的解釋：
nouveau (/nuːˈvoʊ/) 是一個自由及開放原始碼顯示卡驅動程式，是為Nvidia的顯示卡所編寫，也可用於屬於系統晶片的NVIDIA Tegra系列，此驅動程式是由一群獨立的軟體工程師所編寫，Nvidia的員工也提供了少許幫助。 該專案的目標為利用逆向工程Nvidia的專有Linux驅動程式來創造一個開放原始碼的驅動程式。由讓freedesktop.org託管的X.Org基金會所管理，並以Mesa 3D的一部分進行散佈，該專案最初是基於只有2D繪圖能力的“nv”自由與開放原始碼驅動程式所開發的，但紅帽公司的開發者Matthew Garrett及其他人表示原先的程式碼被混淆處理過了。nouveau以MIT許可證許可。 專案的名稱是從法文的“nouveau”而來，意思是“新的”。這個名字是由原作者的的IRC客戶端的自動取代功能所建議的，當他鍵入“nv”時就被建議改為“nouveau”。
這是不是說，其實 Nouveau VAE 跟 Nvidia VAE 算是同義詞了呢？原來咱們開始的理解也並沒有錯呀～
文章小結 #
本文介紹了NVIDIA新發表的一個稱之為NVAE的升級版VAE，它將VAE的生成效果推向了一個新的高度。從文章可以看出，NVAE透過自迴歸形式的隱變數分佈提升了理論上限，設計了巧妙的編碼-解碼結構，並且幾乎融合了當前所有生成模型的最先進技術，打造成了當前最強的VAE～
轉載到請包括本文地址：https://spaces.ac.cn/archives/7574
更詳細的轉載事宜請參考：《科學空間FAQ》