基于Wasserstein距離的雙向學習推理

花強,劉軼功,張峰,董春茹

(河北大學 河北省機器學習與計算智能重點實驗室， 河北 保定 071002)

近年來，基于無監督學習的深度學習技術受到了越來越多學者的關注[1-4]，其中生成對抗網絡(generative adversarial networks，GAN)[5-9]和變分自編碼器(variational auto-encoder，VAE)[10]是高維復雜數據建模的2類最重要的深度生成模型. GAN模型的優勢在于不需要對生成分布進行顯式表達，避免了VAE模型中復雜的馬爾可夫鏈采樣和推理計算，可以訓練得到生成高品質樣本的生成模型.然而，原始GAN及其眾多變種不包含從數據空間到隱變量空間的映射，缺少有效的推理機制，并且缺乏完備的理論保障，從而使得GAN的訓練需要謹慎地選擇各個超參數.進一步的研究發現，GAN對生成樣本多樣性和準確性的懲罰不平衡，有可能導致生成器傾向于重復生成少數幾種樣本，出現模式坍塌(mode collapse)[11-12]問題.因此，Donahue等[13]受VAE的編碼器原理啟發，提出了一種雙向生成對抗網絡(adversarial feature learning，BiGAN).BiGAN在原始GAN框架的基礎上，引入編碼器實現了隱變量特征空間的學習和推理機制.同一時期，Dumoulin等[14]獨立的提出了與BiGAN類似的模型——對抗學習推理(adversarially learned inference，ALI)，ALI將一個稱為推理器的編碼網絡和一個深度定向生成網絡集成在GAN的框架下共同訓練.模型結合了GAN和VAE的部分優點，具有良好的學習性能，與當前流行的自監督和弱監督特征學習方法[15-16]相比也具有一定的競爭力.由于BiGAN和ALI模型在優化目標函數時，需要最小化真實數據分布和生成樣本分布的差異，通常利用分布散度作為基本度量，如Jensen-Shannon(JS)散度或Kullback-Leibler(KL)散度.然而，真實數據分布和生成樣本分布的支撐集是高維空間中的低維流形時，2個分布重疊部分的測度為0，這將導致在訓練模型迭代過程中，出現梯度為零或者無窮大的情況，從而使得生成器無法接收有效的梯度信息，導致訓練失敗，影響模型魯棒性.基于此原因，Arjovsky等[17]提出了Wasserstein GAN(WGAN)模型.WGAN使用Wasserstein距離代替GAN模型損失函數中的KL散度作為衡量2個概率分布之間的相似性.理論和實驗表明，該方法能夠在一定程度上緩解GAN模型訓練過程中的梯度消失和爆炸現象，但是WGAN不具備隱含特征學習功能，容易出現模型坍塌的問題.

為了同時解決上述生成模型存在的模式坍塌和梯度消失爆炸問題，本文提出了一種基于Wasserstein距離的雙向學習推理模型(Wasserstein bidirectional learned inference，WBLI).WBLI使用Wasserstein距離代替BiGAN中的KL散度作為衡量概率分布差異的度量；同時，WBLI模型由生成器、編碼器和判別器3個網絡模塊構成，其中生成器和編碼器在數據特征空間和與之對應的隱變量空間的聯合分布之間建立雙向聯系，而判別器度量了2個聯合分布的Wasserstein距離；最后，WBLI采用了交替迭代算法對網絡參數進行訓練.在MNIST和Fashion MNIST數據集上的實驗結果表明，WBLI模型可以有效緩解基于KL散度的模型在訓練過程中梯度消失或梯度爆炸的缺陷；此外，WBLI通過引入有學習數據樣本內在特征的逆映射的編碼器結構，具有類似于BiGAN和ALI模型具有的隱式正則化、模式覆蓋等優點.

本文的主要貢獻：1)提出了一種基于Wasserstein距離的雙向學習推理模型(WBLI)，該模型緩解了基于KL散度的BiGAN及ALI訓練過程中梯度消失或梯度爆炸問題，從而提高了模型對于樣本分布的魯棒性；同時WBLI一定程度可以緩解WGAN中的模式坍塌問題.2)從結構及實驗2方面將WBLI與BiGAN和WGAN進行了深度比較，結果表明WBLI從模型功能和圖像生成效果上都有一定的提高.

1 相關模型

首先介紹雙向生產對抗網絡模型及Wasserstein距離，并引入本文工作所需的主要數學符號和基礎概念.

1.1 雙向生成對抗網絡

GAN最早于2014年由Goodfellow等[5]提出，是一種實現復雜數據分布學習的無監督生成模型.該模型主要由生成器網絡G和判別器網絡D兩部分構成，其中生成器將輸入的隨機噪聲映射為生成樣本，而判別器同時接收真實樣本和生成樣本，并判別輸入樣本的真偽(即判別樣本是真實樣本還是生成樣本).在GAN模型的訓練過程中，通過構建目標函數引入競爭機制讓這2個網絡同時得到優化，最終使得生成器生成與真實樣本數據分布足夠相似的新數據分布.GAN模型的結構如圖1所示.

圖1 GAN 結構Fig.1 Structure of GAN

設q(x)為真實數據分布，其中x∈ΩX，設p(z)為一個固定的隱編碼分布，其中z∈ΩZ，通常定義為簡單分布，例如標準正態分布p(z)=N(0,1)，生成器G∶ΩZ→ΩX:可以將隱編碼分布映射到數據分布，D(x)代表x來自于真實數據分布q(x)而不是生成樣本分布的概率.據此GAN網絡的優化目標函數如下：

(1)

圖2 BiGAN結構Fig.2 Structure of BiGAN

(2)

BiGAN采用與GAN相同的基于交替梯度的EM優化算法來優化目標函數[18].理論上，在BiGAN達到最優解時，即KL散度收斂達到最小的情況下，可認為所有邊緣分布和所有條件分布都已達到匹配.然而，如引言部分所述，BiGAN目標函數中采用KL散度衡量數據分布間的差異，在某些情況下會出現梯度爆炸的情形，從而導致訓練失敗，影響模型魯棒性[19].

1.2 Wasserstein距離

生成模型的傳統設計方法依靠最大似然估計，或者最小化未知的真實數據分布q(x)和生成樣本分布pG(x)之間的KL散度

(3)

文獻[19]中證明當處理2個由低維流形支持的分布時，那么這2個低維流形將會具有極小重疊甚至沒有重疊，這意味著KL散度在大部分區域是無意義的，即KL(q(x)‖pG(x))=∞，并且JS散度將變為常數log2，這將導致判別器損失函數的梯度為無窮或零，從而導致模型訓練失敗.因此，Arjovsky等[17]通過全面的理論分析，把Wasserstein距離與其他廣受關注的度量概率分布的距離和散度相比，用Wasserstein距離替換原始GAN中的KL散度，提出了WGAN模型，其采用的Wasserstein距離定義為

(4)

其中，∏(pr,pg)是以pr和pg為邊緣分布的所有可能的聯合概率分布的集合.對于每個聯合分布γ(x,y)，都可以通過采樣的方法獲得(x,y)～γ.計算(x,y)的范數‖x-y‖，這樣就可以計算每個聯合分布γ(x,y)的期望值E(x,y)～γ[‖x-y‖].W(pr,pg)為γ(x,y)期望的下確界，更直觀的說，它表示為了將pr移動到pg需要將x移動到y的最小距離或能量.

Wasserstein距離相對KL散度與JS散度具有相對平滑特性，即使2個分布之間沒有交集，Wasserstein距離亦能夠正確度量它們之間的差異，進而產生有意義的梯度.因此，WGAN能有效緩解基于KL散度或JS散度的GAN模型的梯度消失或梯度爆炸問題.然而，理論上目前的深度神經網絡只能夠逼近連續映射，而GAN訓練過程中，目標映射是具有間斷點的非連續映射，不在深度神經網絡的可表示泛函空間之中，這導致了收斂困難，從而產生了模式坍塌，可見WGAN并沒有完全克服GAN中的模式坍塌問題.

2 基于Wasserstein距離的雙向學習推理

基于以上分析，本文將Wasserstein距離引入到BiGAN中，提出了一種基于Wasserstein距離的雙向學習推理模型(即WBLI)，以綜合BiGAN和Wasserstein距離的優點，從而獲得更加穩定的學習模型.

(5)

其中，θG,θE,θD分別表示生成器G、推理器E和判別器D的模型參數.WBLI使用對抗方式聯合訓練一個生成器與一個推理器，生成器G將服從簡單分布的隱變量映射到數據空間，而推理器E將訓練樣本從數據空間反映射回隱變量空間.因此，對抗博弈在G，E與D之間展開.因為WBLI具有逆映射結構，所以推理器在編碼的過程中將相似樣本的隱變量聚在一起，使得流形連續，達到隱式正則化的效果，從而可以提高模型泛化能力.

γ([x1,x2],[y1,y2])∈∏(q(x,z),p(x,z)),

此時Wasserstein距離的計算公式如(6)式

(6)

(7)

其中

由于映射函數f的參數可調節，故f要滿足在γ?π時使整體附加項趨于無窮，使得supf無解，從而達到類似(6)式的約束效果.而當γ∈π時，s和x都是從同一個邊緣分布中采樣，s和x兩個隨機變量分布的期望值相等，故Ef(s)-Ef(x)=0，t和y同理，從而整體附加項等于零.這樣就成功地去掉了γ∈π的約束.將(7)式2項合并得到(8)式

(8)

由于sup為凸函數且inf為凹函數，根據極小極大原理[20]，得到 (9) 式

(9)

(10)

在具體算法實現中，可將函數f(x)用一簇參數為w的神經網絡參數化，并采用權重裁剪方法使得函數滿足Lipschitz連續，此時求上界的問題便可轉化為如下(11)式所表示的優化問題

(11)

接下來生成器要近似地最小化Wasserstein距離，可由此設置最小化損失函數L，由于Wasserstein距離的優良性質，可有效避免判別器的梯度消失問題.再考慮到L的第1項與生成器無關，得到WBLI最終的2個損失函數

LG=-Ex～p(z)[fw(x)]，

(12)

(13)

LD是(11)式的相反數，可以指示訓練進程，其數值越小，表示真實分布與生成分布的Wasserstein距離越小，WBLI訓練得越好.在訓練生成器時，首先固定判別器的參數，從正態分布N(0,1)中采樣m個樣本作為一個批次的訓練數據輸入判別器，然后根據(13)式計算生成器損失，同樣采用RMSProp[21]算法更新其參數.由于更優的判別器可以反向傳播給生成器更準確的梯度信息，因此從訓練開始，在每次更新生成器之前，均需更新判別器n次，以使判別器D更快收斂.完整的訓練過程如算法1所示.

算法1 WBLI輸入 z: 隱變量;T: 訓練數據集;m: 批次大小;α: 學習率;c: 判別器的梯度剪裁數.n: 生成器和推理器優化過程中的判別器更新次數.輸出 判別器參數θD;生成器參數θG;推理器參數θE1)隨機初始化θD和θG2)重復3) for t=0,…,n do 4) z1,…,zm～p(z)5) x1,…,xm～q(x)6) z^1,…,z^m～qE(z^|x=xi),i=1,…,m7) x^1,…,x^m～pG(x^|z=zj),j=1,…,m8) LD←1m∑mi=1fθD(D(xi,z^i))-1m∑mi=1fθD(D(x^i,zi))9) θD←θD+α×RMSProp(θD,?θDLD)10)剪裁θD,將其限制在[-c,c]范圍內11)結束重復12)z1,…,zi～p(z)13) LG←1m∑mi=1fθD(D(xi,z^i))14) θG←θG+α×RMSProp(θG,?θGLG),θE←θE+α×RMSProp(θE,?θELG)15) 結束至判別器收斂

3 實驗結果與分析

3.1 實驗設置

本文實驗運行操作系統為64位Windows10，編程語言為Python3.6，基于TensorFlow開源框架，使用的數據集為MNIST[22]和FashionMNIST[23].MNIST是手寫體數據集，分為60 000個訓練樣本和10 000個測試樣本.所有數字圖像都標準化為28X28像素的固定大小.每個像素由0到255之間的值表示，其中0為黑色，255為白色，介于兩者之間的數值代表不同的灰色影像.FashionMNIST由德國Zalando科技公司的研究部門提供，FashionMNIST的大小、格式和訓練集、測試集劃分與原始的MNIST完全一致，涵蓋了來自10種類別的共70 000個不同商品的正面圖片.

3.2 生成樣本質量與Wasserstein距離

眾所周知，在訓練GAN和BiGAN過程中，生成器G的目標是盡量生成真實的圖片去欺騙判別器D；而D的目標是盡量把G生成的圖片和真實的圖片區分開.這樣，G和D就構成了一個動態的“博弈過程”.然而GAN和BiGAN在訓練過程中沒有任何指示訓練進度的指標，只能基于經驗和生成樣本的效果來判斷模型是否收斂.本文提出的WBLI模型在引入Wasserstein距離度量后，自動建立了生成模型訓練的進程監視指標.本節實驗就是驗證生成樣本質量同Wasserstein距離的正比關系.

在本節實驗，WBLI模型中生成器G、推理器E、判別器D均由3層神經元網絡實現.設定3個網絡的隱含層神經元個數均取128，則生成網絡G輸入、隱含和輸出層神經元個數依次為10-128-784，推理網絡E與G的神經元結構鏡像對稱為784-128-10，判別網絡D接收G和E的聯合數據進行真偽判斷，因此網絡結構為794-128-1.模型學習率設為0.000 1.根據算法1，每訓練5次判別網絡，則更新訓練1次生成網絡和推理網絡，輸出損失函數并記錄生成樣本.抽取了5張不同迭代階段的生成樣本和對應Wasserstein距離值，關系如圖3所示，可以直觀地看到，判別器所輸出的Wasserstein距離與生成器的生成圖片的質量高度相關(更多隨迭代次數生成的樣本序列見圖4).隨著Wasserstein距離的不斷減小，生成樣本的質量逐漸提高.因此，Wasserstein距離可作為訓練階段進程評判指標，這也是WBLI相比于其他基于KL散度模型的一個優勢.

圖3 生成樣本質量和Wasserstein距離的關系Fig.3 Relationship between the quality of generated images and Wasserstein distance

圖4 生成樣本序列Fig.4 Generating sample sequence graphs

3.3 WBLI模型生成樣本多樣性測試

如前所述，GAN和WGAN不具備隱含特征學習功能，易出現模式坍塌問題，即無論輸入模型的簡單隨機分布如何變化，生成器生成的樣本都缺乏多樣性，不能支撐數據空間.WBLI模型能否生成多樣性樣本，是評估該模型優劣的重要指標.

在計算機圖像領域，定量評價圖像相似度的方法有很多，例如常用方法之一是基于單尺度的結構相似性指標SSIM[24-25]

SSIM(x,y)=[l(x,y)]α·[c(x,y)]β·[x(x,y)]γ,

(14)

其中x和y分別表示2個視窗圖像，l(x,y),c(x,y)和s(x,y)分別表示2個視窗圖像的亮度、對比度和結構相似度度量，其具體計算公式見文獻[24-25]；而參數α,β,γ用于控制3部分度量在SSIM中所占比例，一般設α=β=γ=1.當計算2張影像的結構相似性值時，會創建一個局部性視窗，并按式(14)計算視窗內圖像的結構相似性值，每次以像素為單位移動視窗，直到整張影像每個位置的SSIM都計算完畢并求取均值，作為2張影像的結構相似性指標.而為構建更貼近主觀的圖像相似性質量評價方法，文獻[26]在SSIM的基礎上提出了多層級結構相似性指標MS-SSIM[26]，其基本思路為同時考慮多個尺度對圖像相似度進行度量，具體定義為

(15)

其中M為尺度層數，lM(x,y)為視窗圖像x和y在M層上的亮度相似性度量，而cj(x,y)和sj(x,y)分別為視窗圖像x和y在第j層尺度上的對比度和結構相似性度量；而參數αM用于控制第M層亮度相似性所占比例，βj和γj為第j層尺度上對比度相似性和結構相似性度量所占的比例.本文采用多層級結構相似性指標MS-SSIM衡量圖片集的相似性，其值越小，代表圖像集多樣性越好.

圖5為FashionMNIST訓練集和測試集各10類數據的MS-SSIM平均值，可看到大部分類的MS-SSIM值都小于0.25，因此選擇0.25作為判斷生成樣本是否達到真實數據集多樣性標準的閾值.同時，通過統計生成樣本MS-SSIM值的變化也可監控模式坍塌情況.

圖5 數據集的MS-SSIM平均值Fig.5 MS-SSIM average graphs of data sets

圖6表明隨著訓練迭代增加，生成樣本的MS-SSIM值逐步減少，最后降低至FashionMNIST數據集的MS-SSIM平均值之下.由此可見，WBLI模型的訓練過程穩定，沒有產生模式坍塌現象.作為對照，圖7給出了BiGAN訓練過程中MS-SSIM值的變化趨勢.由圖7可見，在BiGAN訓練過程中生成的樣本MS-SSIM值停留在較高水平，并隨迭代次數增加有增加趨勢，這表明生成樣本出現了相似度過高現象，樣本趨于單調，從而發生模式坍塌.

圖7 BiGAN模式坍塌下MS-SSIM值Fig.7 Mode collapse of BiGAN

圖6 WBLI生成樣本的MS-SSIM值Fig.6 MS-SSIM graphs for generating samples

為更加直觀地體現WBLI生成樣本的多樣性，圖8給出了部分WBLI生成樣本的示例.

圖8 WBLI生成樣本示例Fig.8 Generated samples by WBLI

3.4 生成樣本分類識別率

評估生成樣本的質量，只是通過圖像觀察是不規范的.基于同一Le Net-5[25]卷積分類器，分別對WBLI、BiGAN和WGAN模型生成的樣本數據在給定類別信息后，進行分類測試.

首先為了對BiGAN，WGAN和WBLI有更全面的分析，表1中對這3種模型從是否包含隱變量編碼、梯度穩定性以及判別器輸入變量構成3個方面進行了比較.由表1可看出，BiGAN模型可以通過編碼器從數據空間學習到隱含特征，并且判別器的輸入是包含原始數據和隱變量數據的高維向量，使得網絡能夠對低層信息x和高層信息z共同進行判別，從而提高了判別能力.WGAN憑借Wasserstein距離的優點和權重裁剪技術避免梯度消失，且其值保持平滑穩定.本文提出的WBLI模型正是同時集成了2類模型的優勢.

表1 3種模型功能的比較

分類識別率實驗以MNIST手寫數字數據集為初始訓練集.實驗設定3種模型分別以MNIST手寫數據集進行訓練.將訓練好的3種模型分別生成1 000個生成樣本，然后輸入到同樣已經預先訓練好的Le Net-5卷積分類器中，3種模型生成樣本的分類測試正確率如表2所示，同時表2也列出部分生成樣本圖像示例.從生成的示例樣本可看到，WBLI模型生成的樣本更清晰，各類數字間特征更明顯，辨識度更高.而識別正確率結果也說明WBLI模型生成的樣本更具有真實樣本的特征，而這正是因為WBLI綜合了WGAN的梯度穩定性與BiGAN的推理結構.

表2 3種生成樣本分類識別率

4 總結

在BiGAN模型結構基礎上，引入Wasserstein距離代替KL散度用于計算分布間的差異性，建立了基于Wasserstein距離的魯棒無監督生成式學習模型WBLI，并將基于Wasserstein距離的多維優化問題轉化為可求解形式，得到了模型的生成器和判別器的對抗損失函數.一方面，由于WBLI采用的Wasserstein距離具有整體平滑的特性，理論上解決了當前基于KL散度或JS散度的無監督生成模型(如GAN，BiGAN)的梯度爆炸或梯度消失問題；另一方面，借鑒BiGAN中引入推理器E從而使得模型可以有效緩解模式坍塌問題.WBLI解決了原模型訓練不穩定的問題，建立了一個可靠的與生成樣本的質量高度相關的訓練進程指標，實驗結果驗證了上述優點.