Information Bottleneck 與深度學習的關係

這篇文章想要介紹的是另一個從資訊理論角度切入來分析深度學習。我們在前文簡介了 KL Divergence 與機器學習的關係，而這篇文章要介紹的 Information Bottleneck 是用 mutual information 來分析一個網路輸入、輸出、以及中間層的關係。

Mutual Information 的定義

兩個隨機變數 X 與 Y 的 mutual information 我們用 I(X;Y) = I(Y;X) 來表示，其定義為： \[ I(X;Y) = I(Y;X) = \int_Y \int_X p(x,y) log(\frac{p(x,y)}{p(x)p(y)})\ dxdy \] 我們可以想成是 X 與 Y 共享的訊息量，更多細節請參考維基百科的介紹。另外根據定義，如果 X 與 Y 獨立時，他們的 mutual information 便為 0。

 

Data Processing Inequality

這個不等式是指當隨機變數 X, Y, Z 是 Markov Chain 時（也就是 \(X \rightarrow  Y \rightarrow  Z\)），那麼以下不等式成立： \[ I(X;Y) \geq I(X;Z) \] 也就是說當每次資料經過一個函數傳遞時，它與原先輸入會越來越不像，也就是 mutual information 越來越小。 

 

Information Plane

我們用以下圖片來說明，圖片取自於這篇 paper：Opening the black box of Deep Neural Networks via Information [1]。

上圖中 Y 是網路的輸出，而 X 是網路的輸入，拿 MNIST 來說 X 就是輸入的圖片，Y 是這個圖片對應的手寫數字。把 Y 放在 X 前面是因為 X 其實是從給定的 Y 產生的，打個比方來說是先有個數字 1，再產生手寫數字的圖片 X。由 Data Processing Inequality 我們可以得到以下式子： \[ I(X;Y) \geq I(T_1;Y) \geq I(T_2; Y) \geq \cdots \geq I(\hat{Y};Y) \\ H(X) \geq I(X;T_1) \geq I(X;T_2) \geq \cdots \geq I(X;\hat{Y}) \]而 information plane 便是以 I(X;T) 當成橫軸，I(T;Y) 當成縱軸的平面，也就是說對於每一個 hidden layer T，我們都可以算出它對應的 I(X;T) 及 I(T;Y)，並且把這個點標在此平面上。

為什麼要討論 I(X;T) 及 I(Y;T) 呢？

這邊就要先提起資訊理論中的 Rate-distortion Theory 了。這個理論大致上的想法是我們在壓縮時想要同時達到兩個目標：

• 要盡可能地壓縮輸入信號，也就是 I(X;T) 要最小
  
• 在信號重建時（也就是解壓縮時）要讓誤差（distortion）最小，也就是讓 T 與 Y 越接近越好，換句話說要讓 I(T;Y) 最大
  

這兩個目標是矛盾的，而 Rate-distortion Theory 告訴我們可以用最佳化來解決這個問題， 比如說 Lagrange multiplier： \[ min\ I(X;T) - \beta I(T;Y) \] 而這個方法就是本篇的主題 Information Bottleneck。 

下圖出於同一篇文章 [1]，它表示了在訓練過程中所有 hidden layer在 information plane 中的變化：

上圖中橘色的點為最後一個 hidden layer，我們可以看出來一開始的時候橘色的點都在最左下角，代表它完全沒有任何 X 與 Y 的資訊，而經過 400 回合以後，它同時吸收了 X 與 Y 的訊息，然後再緩慢地壓縮 X 的訊息，也就是往平面中的左上角移動。文章中的解釋是一開始 hidden layer 努力記住 Y 的訊息，但同時也得記住 X 的訊息，而到中後期時記起來的東西夠多了，便開始壓縮來自 X 的訊息。這個與張無忌學習太極拳時有點像，一開始要記住所有的招式，但又不能完全記住（不然會 overfitting），記得差不多時便要開始忘記，最後才能打出最純粹的太極拳。

運用 Information Bottleneck 於 GAN 之中（VIB-GAN）

從另一個角度看上面的 Information Bottleneck 的式子，我們可以寫成： \[ max\ I(T;Y) \\ s.t. I(X;T)\ \leq I_c \] 這個問題的解與上面的式子用 Lagrange multiplier 的解是等價的，但這種寫法我們可以解釋成 \(I_c\) 是一個 regularization term，讓 T 與 Y 不要 overfitting。 

接下來我們連結前文 Variational Inference 與 GAN 的內容。首先我們用 Z 表示隱變量，我們可以想成是利用 z 來生成資料，比如說 z 為多維度的向量，而我們利用類神經網路生成一張圖片。在 VIB-GAN 中，我們想要限制 X 與 Z 的 mutual information，也就是： \[ I(X;Z) \leq I_c \] 根據 mutual information 的定義： \[ I(X;Z) = \int p(x,z) log(\frac{p(x,z)}{p(x)p(z)})\ dxdz = \int p(x,z) log(\frac{p(z|x)}{p(z)})\ dxdz \\ = \int p(x,z) log(p(z|x))\ dxdz - \int p(z) log(p(z))\ dz \] 因為要計算出 p(z) 的複雜度太高了，我們直接用上一篇 Variational Inference 的技巧：用一個近似分布 r(z) 來代替 p(z)。由於 \(D_{KL}(p(z)||r(z)) \geq 0\)，所以我們可以得到以下式子： \[ \int p(z) log(p(z))dz \geq \int p(z) log(r(z))dz \] 把這個不等式代回 I(X;Z)： \[ I(X;Z) = \int p(x,z) log(p(z|x))\ dxdz - \int p(z) log(p(z))\ dz \\ \leq \int p(x,z) log(p(z|x))\ dxdz - \int p(z) log(r(z))dz \\ = \int p(x)p(z|x)log(\frac{p(z|x)}{r(z)})\ dxdz \\ = E_{x \sim p(x)}[D_{KL}(p(z|x) || r(z))] \] 因此，只要讓上面這一項 KL Divergence 有個上界 \(I_c\)，就可以保證 I(X;Z) 一定小於等於此上界 \(I_c\)。在 VIB-GAN [4] 中就用了這個想法，讓 X（也就是\(p_{data}\)）與 Z（也就是\(p_g\)）不要被分隔的太遠，因此 discriminator 便不會太強造成梯度都為 0 了。

參考資料

[1] Opening the black box of Deep Neural Networks via Information, Ravid Schwartz-Ziv, Naftali Tishby, 2017

[2] 如何評價 Tishby 的打開深度學習黑箱的 Information Bottleneck 理論?

[3] Deep Variational Information Bottleneck, Alexander A. Alemi, Ian Fischer, Joshua V. Dillon, Kevin Murphy, 2016

[4] Variational Discriminator Bottleneck: Improving Imitation Learning, Inverse RL, and GANs by Constraining Information Flow, Xue Bin Peng, Angjoo Kanazawa, Sam Toyer, Pieter Abbeel, Sergey Levine, 2019