强大的NVAE：以后再也不能说VAE生成的图像模糊了
?PaperWeekly 原創 ·?作者｜蘇劍林
單位｜追一科技
研究方向｜NLP、神經網絡
前幾天筆者在日常刷 arixv 的時候，然后被一篇新出來的論文震驚了！論文名字叫做 NVAE: A Deep Hierarchical Variational Autoencoder，顧名思義是做 VAE 的改進工作的，提出了一個叫 NVAE 的新模型。
說實話，筆者點進去的時候是不抱什么希望的，因為筆者也算是對 VAE 有一定的了解，覺得 VAE 在生成模型方面的能力終究是有限的。結果，論文打開了，呈現出來的畫風是這樣的：
▲ NVAE的人臉生成效果
然后筆者的第一感覺是這樣的：W!T!F! 這真的是 VAE 生成的效果？這還是我認識的 VAE 么？看來我對 VAE 的認識還是太膚淺了啊，以后再也不能說 VAE 生成的圖像模糊了...
不過再看了看作者機構，原來是 NVIDIA，這也大概能接受了。最近幾年可能大家都留意到 NVIDIA 通常都在年底發布個生成模型的突破，2017 年底是 PGGAN [1] ，2018 年底是 StyleGAN [2] ，2019 年底是 StyleGAN2 [3] 。
今年貌似早了些，而且動作也多了些，因為上個月才發了個叫 ADA 的方法，將 Cifar-10 的生成效果提到了一個新高度，現在又來了個 NVAE。
論文標題：NVAE: A Deep Hierarchical Variational Autoencoder
論文鏈接：https://arxiv.org/abs/2007.03898
那這個 NVAE 究竟有什么特別的地方，可以實現 VAE 生成效果的突飛猛進呢？
VAE回顧
可能讀者認真觀察后會說：好像還是有點假呀，那臉部也太光滑了，好像磨過皮一樣，還比不上 StyleGAN 呀～
是的，這樣評價并沒有錯，生成痕跡還是挺明顯的。但如果你沒感覺到震驚，那估計是因為你沒看過之前的 VAE 生成效果，一般的 VAE 生成畫風是這樣的：
▲ 一般的VAE的隨機生成效果
所以，你還覺得這不是一個突破嗎？
那么，是什么限制了（以前的）VAE 的表達能力呢？這一次的突破又是因為改進了哪里呢？讓我們繼續看下去。
1.1 基本介紹
VAE，即變分自編碼器（Variational Auto-Encoder），本人已經有不少文章介紹過了，在公眾號后臺搜索“變分自編碼器”就能搜到很多相關文章。這里做個簡單的回顧和分析。
在筆者對 VAE 的推導里邊，我們是先有一批樣本，這批樣本代表著一個真實的（但不知道形式的）分布 ，然后我們構建一個帶參數的后驗分布 ，兩者就組成一個聯合分布 。
接著，我們再定義一個先驗分布 q(z)，已經定義一個生成分布 ，這樣構成另一個聯合分布 。最后，我們的目的就是讓 相互接近起來，所以我們去優化兩者之間的 KL 散度：
這就是 VAE 的優化目標。
1.2 困難分析
對 的要求是：1. 能寫出解析表達式；2. 方便采樣。
然而連續型分布的世界里這樣的分布并不多，最常用的也就是高斯分布了，而這其中又以“各分量獨立的高斯分布”最為簡單，所以在一般的 VAE 里邊， 都被設為各分量獨立的高斯分布：
以及
問題是，“各分量獨立的高斯分布”不能擬合任意復雜的分布，當我們選定 的形式后，有可能不管我們怎么調它的參數， 和 都不能成為高斯分布，這就意味著 從理論上來說就不可能為 0，所以讓 相互逼近的話，只能得到一個大致、平均的結果，這也就是常規 VAE 生成的圖像偏模糊的原因。
1.3 相關改進
改進 VAE 的一個經典方向是將 VAE 與 GAN 結合起來，比如 CVAE-GAN [4] 、AGE [5] 等，目前這個方向最先進結果大概是 IntroVAE [6] 。從理論上來講，這類工作相當于隱式地放棄了 q(x|z) 是高斯分布的假設，換成了更一般的分布，所以能提升生成效果。
不過筆者覺得，將 GAN 引入到VAE中有點像“與虎謀皮”，借助 GAN 提升了性能，但也引入了 GAN 的缺點（訓練不穩定等），而且提升了性能的 VAE 生成效果依然不如純 GAN 的。另外一個方向是將 VAE 跟 flow 模型結合起來，比如 IAF-VAE [7] 以及筆者之前做的 f-VAE，這類工作則是通過 flow 模型來增強 或 的表達能力。
還有一個方向是引入離散的隱變量，典型代表就是 VQ-VAE，其介紹可以看筆者的《VQ-VAE 的簡明介紹：量子化自編碼器》[8]?。VQ-VAE 通過特定的編碼技巧將圖片編碼為一個離散型序列，然后 PixelCNN 來建模對應的先驗分布 q(z)。
前面說到，當 z 為連續變量時，可選的 p(z|x),q(z) 都不多，從而逼近精度有限；但如果z是離散序列的話，p(z|x),q(z) 對應離散型分布，而利用自回歸模型（NLP 中稱為語言模型，CV 中稱為 PixelRNN/PixelCNN 等）我們可以逼近任意的離散型分布，因此整體可以逼近得更精確，從而改善生成效果。
其后的升級版 VQ-VAE-2 [9] 進一步肯定了這條路的有效性，但整體而言，VQ-VAE 的流程已經與常規 VAE 有很大出入了，有時候不大好將它視為 VAE 的變體。
NVAE梳理
鋪墊了這么久，總算能談到 NVAE 了。NVAE 全稱是?Nouveau VAE（難道不是 Nvidia VAE？），它包含了很多當前 CV 領域的新成果，其中包括多尺度架構、可分離卷積、swish 激活函數、flow 模型等，可謂融百家之所長，遂成當前最強 VAE～
提醒，本文的記號與原論文、常見的 VAE 介紹均有所不同，但與本人其他相關文章是一致的，望讀者不要死記符號，而是根據符號的實際含義來理解文章。
自回歸分布
前面我們已經分析了，VAE 的困難源于 不夠強，所以改進的思路都是要增強它們。首先，NVAE 不改變 ，這主要是為了保持生成的并行性，然后是通過自回歸模型增強了先驗分布 和后驗分布 。
具體來說，它將隱變量分組為 ，其中各個 還是一個向量（而非一個數），然后讓：
而各個組的 依舊建立為高斯分布，所以總的來說 就被建立為自回歸高斯模型。這時候的后驗分布的 KL 散度項變為：
當然，這個做法只是很樸素的推廣，并非 NVAE 的首創，它可以追溯到 2015 年的 DRAW [10] 、HVM [11] 等模型。NVAE 的貢獻是給式 (2) 提出了一種“相對式”的設計：
也就是說，沒有直接去后驗分布 的均值方差，而是去建模的是它與先驗分布的均值方差的相對值，這時候我們有（簡單起見省去了自變量記號，但不難對應理解）。
原論文指出這樣做能使得訓練更加穩定。
3.1 多尺度設計
現在隱變量分成了 L 組 ，那么問題就來了：1. 編碼器如何一一生成 ？2. 解碼器如何一一利用 ？也就是說，編碼器和解碼器如何設計？
▲ NVAE 中的編碼器和解碼器架構。其中 r 代表殘差模塊，h 代表可訓練參數，藍色部分是參數共享的NVAE 巧妙地設計了多尺度的編碼器和解碼器，如上圖所示。首先，編碼器經過層層編碼，得到最頂層的編碼向量 ，然后再慢慢地從頂層往下走，逐步得到底層的特征 。
至于解碼器，自然也是一個自上往下地利用 的過程，而這部分剛好也是與編碼器生成 的過程有共同之處，所有 NVAE 直接讓對應的部分參數共享，這樣既省了參數量，也能通過兩者間的相互約束提高泛化性能。
這種多尺度設計在當前最先進的生成模型都有體現，比如 StyleGAN [2] 、BigGAN、VQ-VAE-2 [9] 等，這說明多尺度設計的有效性已經得到比較充分的驗證。此外，為了保證性能，NVAE 還對殘差模塊的設計做了仔細的篩選，最后才確定了如下的殘差模塊，這煉丹不可謂不充分極致了：
▲ NVAE中的殘差模塊
3.2 其他提升技巧
除了以上兩點比較明顯的特征外，其實 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 增強了模型對分布的擬合能力。不過這個自回歸只是對組間進行的，對于組內的單個分布 和 ，依然假設為各分量獨立的高斯分布，這說明擬合能力依然有提升空間。
更徹底的方案是，對于組內的每個分量也假設為自回歸分布，但是這樣一來在采樣的時候就巨慢無比了（所有的分量串聯遞歸采樣）。
NVAE 提供了一個備選的方案，通過將組內分布建立為 flow模型來增強模型的表達能力，同時保持組內采樣的并行性。實驗結果顯示這是有提升的，但筆者認為引入 flow 模型會大大增加模型的復雜度，而且提升也不是特別明顯，感覺能不用就不用為好。
節省顯存的技巧。盡管 NVIDIA 應該不缺顯卡，但 NVAE 在實現上還是為省顯存下了點功夫。一方面，它采用了混合精度訓練，還順帶在論文推了一波自家的 APEX 庫? [12] 。
另一方面，它在 BN 那里啟用了 gradient check-point [13]（也就是重計算）技術，據說能在幾乎不影響速度的情況下節省 18% 的顯存。總之，比你多卡的團隊也比你更會省顯存～
3.3 更多效果圖
到這里，NVAE 的技術要點基本上已經介紹完畢了。如果大家還覺得意猶未盡的話，那就多放幾張效果圖吧，讓大家更深刻地體會 NVAE 的驚艷之處。
▲?NVAE 在 CelebA HQ 和 FFHQ 上的生成效果。值得注意的是，NVAE 是第一個 FFHQ 數據集上做實驗的 VAE 類模型，而且第一次就如此驚艷了
▲ 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 組，每組隱變量大小還不一樣，分別是 ，如此算來，要生成一個 的 FFHQ 圖像，需要一個總大小有：
維的隨機向量，也就是說，采樣一個 600 萬維的向量，生成一個 （不到 20 萬）維的向量。這跟傳統的 VAE 很不一樣，傳統的 VAE 一般只是將圖片編碼為單個（幾百維）的向量，而這里的編碼向量則多得多，有點全卷積自編碼器的味道了，所以清晰度提升也是情理之中。
Nouveau是啥？
最后，筆者饒有性質地搜索了一下 Nouveau 的含義，以下是來自維基百科的解釋 [14] ：
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。
參考文獻
[1] https://arxiv.org/abs/1710.10196
[2] https://arxiv.org/abs/1812.04948
[3] https://arxiv.org/abs/1912.04958
[4] https://arxiv.org/abs/1703.10155
[5] https://arxiv.org/abs/1704.02304
[6] https://arxiv.org/abs/1807.06358
[7] https://arxiv.org/abs/1606.04934
[8] https://kexue.fm/archives/6760
[9] https://arxiv.org/abs/1906.00446
[10] https://arxiv.org/abs/1502.04623
[11] https://arxiv.org/abs/1511.02386
[12] https://github.com/NVIDIA/apex
[13] https://arxiv.org/abs/1604.06174
[14] https://zh.wikipedia.org/wiki/Nouveau
更多閱讀
#投 稿?通 道#
?讓你的論文被更多人看到?
如何才能讓更多的優質內容以更短路徑到達讀者群體，縮短讀者尋找優質內容的成本呢？答案就是：你不認識的人。
總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學者和學術靈感相互碰撞，迸發出更多的可能性。?
PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優質內容，可以是最新論文解讀，也可以是學習心得或技術干貨。我們的目的只有一個，讓知識真正流動起來。
?????來稿標準：
? 稿件確系個人原創作品，來稿需注明作者個人信息（姓名+學校/工作單位+學歷/職位+研究方向）?
? 如果文章并非首發，請在投稿時提醒并附上所有已發布鏈接?
? PaperWeekly 默認每篇文章都是首發，均會添加“原創”標志
?????投稿郵箱：
? 投稿郵箱：hr@paperweekly.site?
? 所有文章配圖，請單獨在附件中發送?
? 請留下即時聯系方式（微信或手機），以便我們在編輯發布時和作者溝通
????
現在，在「知乎」也能找到我們了
進入知乎首頁搜索「PaperWeekly」
點擊「關注」訂閱我們的專欄吧
關于PaperWeekly
PaperWeekly 是一個推薦、解讀、討論、報道人工智能前沿論文成果的學術平臺。如果你研究或從事 AI 領域，歡迎在公眾號后臺點擊「交流群」，小助手將把你帶入 PaperWeekly 的交流群里。
與50位技術專家面對面20年技術見證，附贈技術全景圖總結
以上是生活随笔為你收集整理的强大的NVAE：以后再也不能说VAE生成的图像模糊了的全部內容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： 苹果最薄 Apple Watch 拆解：
- 下一篇： 实录 | DSTC 8“基于Schema