KL vanishing 問題的一般解

技術 · 發表 2019-04-22 16:46:04

摘要：本文主要針對阿里最近發表於 ICASSP2019 上的一篇論文 “ Improve Diverse Text Generation by Self Labeling Conditional Variational Auto Encoder” [1] ，通過博弈論的角度分析這篇論文存在的一些問...

本文主要針對阿里最近發表於 ICASSP2019 上的一篇論文 “ Improve Diverse Text Generation by Self Labeling Conditional Variational Auto Encoder” ^[1] ，通過博弈論的角度分析這篇論文存在的一些問題，並給出了 KL vanishing 問題的一般解。因此，本文的貢獻主要包含了兩方面：

嘗試從博弈論的角度分析 KL vanishing problem 的納什均衡解，並從這個角度得出阿里的這篇論文並不能很好地解決 KL vanishing problem
從原始的 VAE 優化函數出發，通過引入互資訊的限制條件來解決 KL vanishing problem ，並從博弈論角度給出文字提出的優化演算法的納什均衡解不再會遇到 KL vanishing 問題

KL vanishing problem

在 VAE 裡面， KL vanishing problem 主要是指當 KL(q(z|x) || p(z)) 消失時，並且當隱變數 z 和 x 獨立時，此時的 ELBO=log(p(x)) 【推導過程見 Supplemental Materials 】。由 KL(q(z|x) || p(z|x)) = log(p(x)) - ELBO ，得到此時 KL(q(z|x) || p(z|x))=0 ，因此得出 q(z|x)==p(z|x) ，但這個解並不是我們需要的理想解，因為此時的互資訊 I(x, z)=0 。

從博弈論觀點看待 KL vanishing problem

通過對 KL vanishing problem 的分析，假設當 KL(q(z|x) || p(z))=0 時，存在著兩種情況，一個是 z 與 x 獨立，另一個是 z 和 x 不獨立。我們想得到的解是當 KL(q(z|x) || p(z))=0 時， z 和 x 不獨立。現在我們引入博弈論的思想，假設模型的兩個 modules 之間進行博弈，它們分別是決定 z 和 x 獨立與否的 module A ，以及決定 q(z|x) 於 p(z) 是否相等的 module B 。而它們的收益函式都是 ELBO 或 ELBO 的相關變種。

首先對傳統 VAE 進行分析，得到 module A 和 module B 的收益矩陣【計算過程見 Supplemental Materials 】：

Module A \ Module B

q(z|x) == p(z)

q(z|x) ≠ p(z)

z 與 x 獨立

(logp(x), logp(x))

(logp(x)KL(q(z),p(z)), logp(x)-KL(q(z),p(z))

z 與 x 不獨立

(E _p(z) [logp(x|z)],

E _p(z) [logp(x|z)])

(ELBO, ELBO)

由於 E _p(z) [logp(x|z)] <= logp(x) ，當且僅當 x 和 z 獨立時等式成立，因此可以得到納什均衡解是 (z 與 x 獨立， q(z|x) == p(z)) ，此時就會造成所謂的 KL vanishing problem 。

而 “Improve Diverse Text Generation by Self Labeling Conditional Variational Auto Encoder” ^[1] 這篇論文提出通過引入新的損失函式：

可以避免出現 KL vanishing problem 。他們的基本思想是，通過使隱變數 z 包含 x 的資訊，從而避免模型收斂到 (z 與 x 獨立， q(z|x) == p(z)) 這個解。但這種方法並不能完全避免模型收斂到 (z 與 x 獨立， q(z|x) == p(z)) 。我們先一般化這個思路，即定義新的，具體的，這裡的，此處的就是 module A 和 module B 採取的策略，即 module A 選擇 z 是否與 x 獨立， module B 選擇 q(z|x) 是否等於 p(z) 。由此，我們得到新的收益矩陣：

Module A \ Module B

q(z|x) == p(z)

q(z|x) ≠ p(z)

z 與 x 獨立

( ,

)

(

)

z 與 x 不獨立

( ,

)

( ,

)

因為要想讓模型收斂到納什均衡解 (z 與 x 不獨立， q(z|x) == p(z)) ，必須滿足：

但是當時，由於，而當 z 與 x 不獨立時，，故條件 (1) 不成立。因此此時的納什均衡解依然是 (z 與 x 獨立， q(z|x) == p(z) ) 。所以論文裡提出的演算法並不能解決 KL vanishing problem 。

新方法

我們發現如果嘗試去通過修改中間優化函式，比如 ELBO ，往往會導致最終得到的新的 objective function 不再和最原始的 objective function ： min KL(q(z|x), p(z|x)) 等價，這樣的 objective function 往往是一個四不像，雖然從區域性分析往往對模型有利，但如果從全域性進行分析就會發現這樣做會顧此失彼。因此，能否在原始的 objective function 上做修改，比如引入一個限制條件，使得隱變數 z 和 x 的互資訊大於某個常數，這樣我們能夠保證求解得到的解與求解原始 objective function 等價。

定義新的 objective function

由於 VAE 最初的目標是優化 KL(q(z|x), p(z|x)) ，而我們只要在優化過程中引入一些限制條件，比如使得 x 和 z 的互資訊 I(x,z) >= c ，因此得到下面新的 objective function ：

通過引入拉格朗日乘子，得到不帶約束的優化問題：

我們可以看出，這個優化目標函式可以使得當 I(x, z) < c 時，將得到一個懲罰項，只有當 I(x, z) >= c 時，懲罰項才為 0 。

由

因此，若 I(x, z) >= c ，則

而當 I(x, z) < c 時，由

進一步得到：

在實際訓練的時候，我們讓：

因此新的求解演算法具體過程如下：

1. Choose hyper-parameters ：， and initialize model parameters randomly

While not stopped:

Sample a mini-batch data: {x ₁ , x ₂ , ..., x _n }

Calculate mini-batch mean KL(q(z|x), p(z)):

Calculate gradients:

Update model parameters

演算法分析

文字提出的優化演算法的收益矩陣為【計算過程見 Supplemental Materials 】：

Module A \ Module B

q(z|x) == p(z)

q(z|x) ≠ p(z)

I(x, z) < c

( , )

(

)

I(x, z) >= c

( ,

)

( ,

)

特別的，當 I(x, z)=0 ，並且 q(z|x) == p(z) 時的值取最大：

根據收益矩陣可以得出，我們只要滿足條件：

而這兩項條件只要我們設定合理的值就可以滿足。此時的納什均衡解為（ I(x, z) >= c, q(z|x) == p(z) ），而這正是我們希望得到的解，因為此時在保證 q(z|x) == p(z) 的同時， x 與 z 不獨立。

與β-VAE對比

我們知道 β-VAE可以學到比較好的解耦表示，即使得每個維度的隱變數z _i 之間的相關度很小，但同時它會減小隱變數 z與x之間的互資訊 ^[2] 。因此會導致 reconstruction網路效能的下降。由於β-VAE 只是簡單地增大對 KL(q(z|x), p(z))的懲罰，隨著訓練的進行，有可能會push隱變數的分佈朝著跟x獨立的分佈的方向去優化。而本文提出的演算法，可以動態去調整KL(q(z|x), p(z))的懲罰項，從而避免減小隱變數 z和x之間的互資訊。同時我們發現，在優化目標函式 (3)和(4)裡引入β(β>1) 後，收益矩陣變為：