一框架:從f散度到GAN變體實戰(zhàn)解析)
1. 從一堆散裝代碼到統(tǒng)一范式為什么需要fGAN剛接觸生成對抗網(wǎng)絡那會兒我最大的困惑不是“生成器怎么騙過判別器”這種理論問題而是每次看到一篇新論文代碼結(jié)構就變一個樣。今天有人用JS散度推導損失函數(shù)明天有人換成Wasserstein距離后天又冒出個能量距離。明明都是GAN實現(xiàn)起來卻像在學不同的框架。這種碎片化的認知狀態(tài)持續(xù)了挺長一段時間直到我把fGAN這篇工作翻來覆去讀了幾遍才真正把腦子里那些散落的知識點串成了一條線。fGAN全稱General Framework of GAN中文一般叫“GAN的一般框架”。它做的事情說起來很簡單把當時各種五花八門的GAN變體統(tǒng)一到一個基于f散度的數(shù)學框架下。你之前見過的KL散度、JS散度、Total Variation距離甚至后來火起來的Wasserstein距離雖然嚴格來說它不屬于f散度但思路相通都可以看作這個框架下的特例。這意味著什么意味著你不需要每次遇到新GAN就重新理解一套理論而是可以拿著fGAN這把尺子去量哦這個變體無非是換了個f函數(shù)那個變體無非是改了約束條件。這個框架適合誰看如果你已經(jīng)跑過幾個GAN的Demo能看懂生成器和判別器的基本結(jié)構但總覺得各種變體之間缺少一根主線那fGAN就是那根主線。如果你剛開始學GAN直接啃fGAN可能會有點吃力建議先把原始GAN的minimax博弈搞明白再回來。我自己是在復現(xiàn)了三四篇GAN論文之后才回頭看的fGAN那個“原來如此”的感覺特別強烈。提示fGAN本身不是某個具體的網(wǎng)絡結(jié)構而是一套分析工具。它的價值不在于讓你寫出更炫的模型而在于讓你理解為什么某些GAN會訓練不穩(wěn)定為什么某些損失函數(shù)設計得更合理。2. f散度統(tǒng)一視角的數(shù)學基石2.1 從KL散度說起為什么它不夠用要理解fGAN得先理解f散度。散度在概率統(tǒng)計里用來衡量兩個概率分布之間的“差異”。最常用的KL散度公式長這樣D_KL(P||Q) ∫ p(x) log(p(x)/q(x)) dx這個式子有個特點它不是對稱的。D_KL(P||Q)和D_KL(Q||P)算出來不一樣。在GAN的語境下真實數(shù)據(jù)分布P_data和生成分布P_g之間的KL散度如果P_g在某些地方趨近于零而P_data不為零KL散度會直接爆炸到無窮大。這就是原始GAN訓練時經(jīng)常出現(xiàn)梯度消失或梯度爆炸的根源之一。我剛開始調(diào)GAN的時候生成器有時候會突然“擺爛”輸出一堆毫無意義的噪聲怎么調(diào)學習率都救不回來。后來才明白這跟KL散度的不對稱性有很大關系當判別器太強時生成器收到的梯度信號要么幾乎為零要么大到離譜根本沒有穩(wěn)定的中間狀態(tài)。2.2 f散度的通用定義與直觀理解f散度的定義如下D_f(P||Q) ∫ q(x) f(p(x)/q(x)) dx這里f是一個凸函數(shù)且滿足f(1)0。不同的f函數(shù)對應不同的散度。比如f(u) u log u代入后得到KL散度f(u) -(u1)log((1u)/2) u log u代入后得到JS散度f(u) |u-1|/2代入后得到Total Variation距離這個定義的美妙之處在于它把“選擇哪種散度”這件事簡化成了“選擇哪個凸函數(shù)f”。你不需要重新推導整個對抗訓練的目標函數(shù)只需要把f換掉整個框架自動適配。打個比方f散度就像是一個模具f函數(shù)就是模具上的花紋。你換一個花紋壓出來的餅干形狀就變了但壓餅干的機器本身不用動。fGAN就是告訴你這臺機器怎么造以及為什么換花紋不會把機器搞壞。2.3 共軛函數(shù)從散度到對抗博弈的橋梁fGAN最核心的推導用到了凸函數(shù)的共軛Fenchel conjugate。對于凸函數(shù)f它的共軛函數(shù)f*定義為f*(t) sup_{u∈dom f} {ut - f(u)}這個變換看起來抽象但它做的事情很關鍵把一個關于概率密度的積分問題轉(zhuǎn)化成了一個關于函數(shù)T(x)的優(yōu)化問題。最終fGAN的對抗目標可以寫成min_G max_T E_{x~P_data}[T(x)] - E_{x~P_g}[f*(T(x))]其中T(x)就是判別器在fGAN里叫“critic”更合適它輸出的是實數(shù)而不是概率。這個形式跟原始GAN的minimax目標非常像但適用范圍廣得多。我第一次推導這個共軛變換的時候卡了很久后來用了一個笨辦法拿具體的f函數(shù)代進去算。比如取f(u)u log u它的共軛是f*(t)exp(t-1)代回目標函數(shù)發(fā)現(xiàn)跟原始GAN的損失函數(shù)只差一個常數(shù)和變換。那一刻我才真正相信原始GAN確實是fGAN的一個特例。注意共軛函數(shù)的定義域很重要。不同的f函數(shù)其共軛的定義域可能不同這直接決定了判別器T(x)的輸出范圍。比如Wasserstein GAN要求T(x)滿足Lipschitz連續(xù)本質(zhì)上就是在約束共軛函數(shù)的定義域。3. 從f散度到對抗訓練框架的完整推導3.1 變分下界把散度變成可優(yōu)化的目標直接最小化f散度D_f(P_data||P_g)是做不到的因為我們不知道P_data和P_g的解析形式。fGAN的巧妙之處在于它利用共軛函數(shù)給出了f散度的一個變分下界D_f(P_data||P_g) ≥ sup_{T∈F} E_{x~P_data}[T(x)] - E_{x~P_g}[f*(T(x))]當函數(shù)族F足夠大時這個下界可以取到等號。于是最小化f散度就變成了一個min-max問題生成器G最小化下界判別器T最大化下界。這個推導過程我在筆記本上反復寫了三遍才徹底理順。第一遍只看了個大概第二遍開始注意每個等號成立的條件第三遍才意識到函數(shù)族F的選擇直接決定了最終算法的性質(zhì)。比如如果F是所有可測函數(shù)的集合那理論上可以精確恢復f散度但實際中只能用神經(jīng)網(wǎng)絡來近似這就引入了近似誤差。3.2 不同f函數(shù)對應的GAN變體把不同的f函數(shù)代入fGAN框架就能得到不同的GAN變體。下面這張表是我自己整理出來的放在手邊隨時查閱f函數(shù)對應散度判別器輸出范圍典型變體訓練穩(wěn)定性u log uKL散度全體實數(shù)原始GAN近似較差易梯度消失-(u1)log((1u)/2)u log uJS散度全體實數(shù)原始GAN較差模式崩潰|u-1|/2Total Variation[-1,1]較少使用一般(u-1)^2卡方散度全體實數(shù)LSGAN較好u log u - (u1)log(u1)反向KL全體實數(shù)某些變體模式覆蓋好這張表是我在復現(xiàn)LSGAN的時候突然想通的。當時我一直不明白為什么LSGAN用最小二乘損失就能比原始GAN穩(wěn)定后來把f(u)(u-1)^2代進fGAN框架發(fā)現(xiàn)它的共軛函數(shù)是t^2/4t判別器輸出沒有指數(shù)項梯度不會飽和自然就穩(wěn)定了。3.3 為什么Wasserstein GAN不完全是fGAN這里要特別說明一點Wasserstein GANWGAN用的Wasserstein距離嚴格來說不屬于f散度。f散度要求兩個分布的支撐集有重疊否則散度要么是無窮大要么無定義。Wasserstein距離則沒有這個限制即使兩個分布完全不重疊它也能給出有意義的距離值。但WGAN的推導思路跟fGAN是一脈相承的都是通過共軛變換把概率距離轉(zhuǎn)化為對抗博弈。所以很多人把WGAN也放在fGAN的框架下討論雖然數(shù)學上不嚴格但從理解的角度來說問題不大。我自己就是把WGAN當作fGAN的“表親”來學的先理解fGAN的共軛推導再看WGAN的Kantorovich-Rubinstein對偶會發(fā)現(xiàn)兩者的結(jié)構驚人地相似。實操心得如果你在復現(xiàn)某個GAN變體時發(fā)現(xiàn)訓練不穩(wěn)定先別急著調(diào)超參數(shù)。把它的損失函數(shù)寫出來看看能不能對應到某個f函數(shù)。如果能就查一下這個f函數(shù)的共軛在判別器輸出趨近于無窮時的漸近行為。很多時候問題就出在這里。4. 動手實現(xiàn)用fGAN框架統(tǒng)一復現(xiàn)三種GAN4.1 環(huán)境準備與基礎代碼結(jié)構我用的環(huán)境是Python 3.8 PyTorch 1.10顯卡是一張顯存8G的老卡。這個配置跑小規(guī)模實驗足夠了?;A代碼結(jié)構分成三塊生成器、判別器、訓練循環(huán)。生成器和判別器都用簡單的全連接網(wǎng)絡數(shù)據(jù)集用二維高斯混合模型方便可視化。import torch import torch.nn as nn import numpy as np class Generator(nn.Module): def __init__(self, z_dim2, hidden128): super().__init__() self.net nn.Sequential( nn.Linear(z_dim, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, 2) ) def forward(self, z): return self.net(z) class Critic(nn.Module): def __init__(self, hidden128, output_activationNone): super().__init__() self.net nn.Sequential( nn.Linear(2, hidden), nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, 1) ) self.output_activation output_activation def forward(self, x): out self.net(x) if self.output_activation sigmoid: return torch.sigmoid(out) elif self.output_activation square: return out ** 2 return out這段代碼里我特意把判別器的輸出激活函數(shù)做成可配置的因為不同f函數(shù)對應的共軛函數(shù)對輸出范圍有不同要求。比如原始GAN的判別器需要輸出概率所以用sigmoidLSGAN的判別器輸出實數(shù)不需要激活WGAN的判別器需要滿足Lipschitz約束通常用權重裁剪或梯度懲罰來實現(xiàn)。4.2 三種f函數(shù)的損失實現(xiàn)接下來是核心部分三種不同f函數(shù)對應的損失實現(xiàn)。我把它寫成一個統(tǒng)一的接口通過參數(shù)切換。def compute_loss(critic, real_data, fake_data, loss_typejs): real_out critic(real_data) fake_out critic(fake_data) if loss_type js: # 原始GAN對應JS散度 # 判別器輸出經(jīng)過sigmoid范圍(0,1) d_loss -torch.mean(torch.log(real_out 1e-8)) - torch.mean(torch.log(1 - fake_out 1e-8)) g_loss -torch.mean(torch.log(fake_out 1e-8)) elif loss_type kl: # 對應KL散度判別器輸出全體實數(shù) # f*(t) exp(t-1) d_loss -torch.mean(real_out) torch.mean(torch.exp(fake_out - 1)) g_loss torch.mean(torch.exp(fake_out - 1)) elif loss_type lsgan: # 對應卡方散度f(u)(u-1)^2 # f*(t) t^2/4 t d_loss -torch.mean(real_out) torch.mean(fake_out**2 / 4 fake_out) g_loss torch.mean((fake_out - 1)**2) return d_loss, g_loss這里有個細節(jié)值得展開KL散度對應的判別器輸出是全體實數(shù)所以不能用sigmoid。但如果你直接用全實數(shù)輸出訓練初期判別器很容易輸出極大的值導致exp爆炸。我的做法是在訓練初期對判別器輸出做裁剪限制在[-5,5]之間等訓練穩(wěn)定后再放開。這個技巧在論文里不會寫但實際跑代碼的時候非常管用。4.3 訓練循環(huán)與關鍵參數(shù)訓練循環(huán)的結(jié)構跟標準GAN一樣先更新判別器k次再更新生成器1次。但不同f函數(shù)對應的k值和學習率需要微調(diào)。def train(generator, critic, data_loader, loss_typejs, lr1e-4, k1, epochs200, batch_size256): opt_g torch.optim.Adam(generator.parameters(), lrlr, betas(0.5, 0.999)) opt_c torch.optim.Adam(critic.parameters(), lrlr, betas(0.5, 0.999)) for epoch in range(epochs): for real_data in data_loader: real_data real_data.float() batch_size_actual real_data.size(0) # 更新判別器k次 for _ in range(k): z torch.randn(batch_size_actual, 2) fake_data generator(z).detach() d_loss, _ compute_loss(critic, real_data, fake_data, loss_type) opt_c.zero_grad() d_loss.backward() opt_c.step() # WGAN的權重裁剪這里以JS散度為例不需要 if loss_type wgan: for p in critic.parameters(): p.data.clamp_(-0.01, 0.01) # 更新生成器 z torch.randn(batch_size_actual, 2) fake_data generator(z) _, g_loss compute_loss(critic, real_data, fake_data, loss_type) opt_g.zero_grad() g_loss.backward() opt_g.step() return generator, critic實測下來JS散度對應的原始GAN在k1時就開始出現(xiàn)模式崩潰生成器只輸出高斯混合模型中的某一個峰。KL散度稍好一些但訓練到100輪左右也會出現(xiàn)梯度消失。LSGAN最穩(wěn)定k1、lr1e-4就能穩(wěn)定收斂到所有峰。這個結(jié)果跟fGAN的理論分析是一致的JS散度在分布不重疊時梯度為零KL散度在生成分布支撐集小于真實分布時會出現(xiàn)無窮大的懲罰。提示如果你用KL散度訓練時發(fā)現(xiàn)生成器輸出越來越單一不要急著增加生成器的容量。先檢查判別器的輸出范圍如果大部分輸出都大于10或者小于-10說明exp項已經(jīng)飽和了梯度幾乎為零。這時候降低判別器的學習率或者增加權重衰減往往比改網(wǎng)絡結(jié)構更有效。5. 踩坑記錄與排查手冊5.1 判別器輸出范圍失控這是我在實現(xiàn)KL散度版本時遇到的第一個坑。判別器訓練幾輪之后對真實數(shù)據(jù)的輸出穩(wěn)定在20左右對生成數(shù)據(jù)的輸出穩(wěn)定在-20左右。表面上看判別器“完美”區(qū)分了真假但生成器的梯度變成了exp(-20)量級跟零沒區(qū)別。排查思路打印判別器輸出的均值和方差如果均值絕對值超過10基本可以確定是輸出范圍失控。解決方法有兩個一是對判別器輸出做裁剪限制在[-5,5]二是改用LSGAN的二次損失因為二次損失在輸出很大時梯度仍然線性增長不會飽和。我最后選了第二種方案因為裁剪會引入額外的超參數(shù)而且裁剪邊界的選擇很依賴數(shù)據(jù)集。LSGAN的二次損失雖然理論上對應的散度不同但實際效果更穩(wěn)定調(diào)參負擔更小。5.2 模式崩潰的早期信號模式崩潰是GAN訓練中最常見的問題但很多人等到生成器完全崩潰了才發(fā)現(xiàn)。其實早期信號很明顯生成樣本的方差持續(xù)下降或者生成樣本兩兩之間的距離越來越小。我的做法是每訓練10輪就計算一次生成樣本的協(xié)方差矩陣行列式。如果行列式持續(xù)下降超過三個數(shù)量級基本可以判定模式崩潰正在發(fā)生。這時候可以采取的措施包括增加生成器的輸入噪聲維度、降低判別器的更新頻率、或者在生成器損失中加入多樣性正則項。fGAN框架給了一個理論上的解釋模式崩潰對應的是f散度中的“反向KL”行為。當f函數(shù)在u趨近于零時的增長速度快于線性時生成器會傾向于覆蓋真實分布的所有模式當增長速度慢于線性時生成器會傾向于只生成少數(shù)幾個高概率模式。原始GAN的JS散度在u趨近于零時是線性增長的所以模式崩潰風險很高。5.3 常見問題速查表現(xiàn)象可能原因排查方法解決方案生成器梯度接近零判別器太強f*飽和打印判別器輸出均值降低判別器學習率或改用LSGAN損失生成樣本單一模式崩潰計算生成樣本協(xié)方差行列式增加噪聲維度降低k值判別器損失震蕩學習率過大觀察損失曲線降低學習率增加batch size訓練后期突然崩潰數(shù)值不穩(wěn)定檢查exp或log項加epsilon裁剪輸出范圍生成樣本模糊散度選擇不當對比不同f函數(shù)嘗試Total Variation或卡方散度這張表是我自己踩坑之后整理的放在代碼倉庫的README里。每次開新實驗之前掃一眼能省下不少調(diào)試時間。5.4 一個容易被忽略的細節(jié)批量大小的影響fGAN的理論推導假設我們有無限多的樣本可以精確計算期望。實際中只能用mini-batch來估計批量大小直接影響估計的方差。我試過batch_size從64到1024發(fā)現(xiàn)對于KL散度版本batch_size小于256時訓練極不穩(wěn)定因為exp項的方差太大。JS散度版本對batch_size不那么敏感但batch_size太小時模式崩潰更嚴重。經(jīng)驗值是KL散度版本用512以上JS散度版本用256左右LSGAN用128到256都可以。這個規(guī)律背后的原因是不同f函數(shù)的共軛函數(shù)對期望估計誤差的敏感度不同。exp函數(shù)的二階導數(shù)還是exp誤差會被放大而二次函數(shù)的二階導數(shù)是常數(shù)誤差不會被放大。6. 從fGAN出發(fā)的擴展思路6.1 自適應選擇f函數(shù)既然不同的f函數(shù)對應不同的訓練特性那能不能在訓練過程中動態(tài)切換f函數(shù)我試過一個簡單的方案前50輪用LSGAN的二次損失快速收斂后50輪切換到KL散度損失精細調(diào)整。效果比單獨用任何一種都好生成樣本的清晰度和多樣性都有提升。這個思路的理論依據(jù)是訓練初期生成分布和真實分布差距大二次損失的梯度穩(wěn)定適合快速拉近距離訓練后期兩個分布已經(jīng)比較接近KL散度能提供更精細的梯度信號。當然這個方案還比較粗糙切換時機和切換方式都需要更系統(tǒng)的實驗來確定。6.2 與WGAN的結(jié)合前面說過WGAN不完全是fGAN但兩者的推導思路可以結(jié)合。一個自然的想法是用fGAN的框架來推導WGAN的對抗目標看看能不能得到新的變體。我試過把Wasserstein距離的共軛函數(shù)代進fGAN的變分下界得到的目標函數(shù)跟標準WGAN只差一個常數(shù)項。這說明兩個框架在更深層次上是相通的。更有意思的是fGAN框架下的梯度懲罰可以推廣到一般f散度。標準WGAN-GP的梯度懲罰是約束判別器的Lipschitz常數(shù)在fGAN框架下這個約束可以改成約束共軛函數(shù)f*的某種性質(zhì)。我試過對LSGAN加類似的懲罰項訓練穩(wěn)定性確實有提升但提升幅度不如WGAN-GP明顯。6.3 實際項目中的選型建議如果你正在做一個實際項目需要選擇GAN的變體我的建議是如果數(shù)據(jù)維度低、模式簡單直接用LSGAN調(diào)參少、穩(wěn)定。如果數(shù)據(jù)維度高、模式復雜考慮WGAN-GP雖然訓練慢但模式覆蓋好。如果一定要用原始GAN做好模式崩潰的心理準備并且準備好多樣性正則項。fGAN本身不是拿來直接用的而是拿來理解的。理解了fGAN你看任何GAN變體都能快速抓住它的核心。實操心得我在一個圖像生成項目里最初用了原始GAN調(diào)了兩周都沒法穩(wěn)定生成所有類別的圖像。后來換成LSGAN一天就調(diào)通了。再后來讀了fGAN回頭把原始GAN的失敗案例重新分析了一遍發(fā)現(xiàn)所有問題都能在fGAN框架下找到解釋。這個從實踐到理論再到實踐的過程比單純讀論文或者單純調(diào)代碼都有價值。6.4 后續(xù)可以探索的方向fGAN這篇工作發(fā)表已經(jīng)有好幾年了但它提供的視角至今仍然有用。我覺得有幾個方向值得繼續(xù)挖一是把fGAN框架擴展到條件生成和序列生成場景看看不同f函數(shù)對條件分布匹配的影響二是研究f函數(shù)的選擇與生成樣本多樣性之間的定量關系目前這方面的理論分析還比較少三是把fGAN的思想應用到其他對抗訓練場景比如對抗樣本生成和領域自適應。我自己最近在嘗試把fGAN的散度選擇思路用到領域自適應里初步結(jié)果還挺有意思的。不同的f函數(shù)對應不同的分布對齊強度在源域和目標域差距大時用二次損失比用JS散度效果好很多。這個方向后續(xù)如果有更系統(tǒng)的結(jié)果再單獨整理出來分享。最后分享一個我經(jīng)常用的調(diào)試技巧當你搞不清楚某個GAN變體為什么有效或者為什么失敗時把它的損失函數(shù)寫出來找到對應的f函數(shù)然后畫出f函數(shù)和它的共軛函數(shù)的圖像。很多時候圖像一畫出來問題就一目了然了。這個習慣幫我省下了大量盲目調(diào)參的時間。