Deep Variational Information Bottleneck (VIB) 變分資訊瓶頸 論文閱讀筆記。本文利用變分推斷將資訊瓶頸框架適應到深度學習模型中，可視為一種正則化方法。
變分資訊瓶頸
假設資料輸入輸出對為$(X,Y)$，假設某個判別模型有關於$X$的中間表示$Z$，本文旨在最佳化$\theta$以最小化互資訊$I(Z;X)$ ，同時最大化互資訊$I(Z;Y)$，即：
$\max\limits_{\theta}I(Z;Y|\theta)-\beta I(Z;X|\theta)$
其中$\beta>0$為平衡係數。直覺理解，上式期望$Z$能保留更少$X$資訊的同時能較好用於預測$Y$。那麼如何構造這個模型以及相應的最佳化方案？下面推導上式的下界，使其下界變大，上式即可變大。為了簡化，下面去掉$\theta$進行推導。
下界1
首先，$I(Z;Y)$展開為：
$\displaystyle I(Z; Y) = \int \int p(y, z) \log \frac{p(y|z)}{p(y)} \, dy \, dz$
其中$p(y)$是資料的標籤分佈，已知。未知而需要進行處理的是其中的$p(y,z)$和$p(y|z)$，也就是模型需要擬合的分佈。對於$p(y|z)$，可以用一個解碼器$q(y|z)$來擬合，即文中所謂的變分估計。利用KL散度的大於零性質，有以下不等式：
\begin{align*} &\text{KL}(p(Y|Z),q(Y|Z))\geq 0\\ \implies &\int \, p(y|z) \log \frac{p(y|z)}{q(y|z)} dy\geq 0\\ \implies &\int \, \frac{p(y,z)}{p(z)} \log \frac{p(y|z)}{q(y|z)} dy\geq 0\\ \implies &\int \, p(y,z) \log p(y|z) dy\geq \int \, p(y,z) \log q(y|z)dy\\ \end{align*}
則有
\begin{align*} \displaystyle I(Z; Y) &= \int \int p(y, z) \log p(y|z) - p(y, z) \log p(y) \, dy \, dz\\ &\geq \int \int p(y, z) \log q(y|z) - p(y, z) \log p(y) \, dy \, dz\\ &= \int \int \, p(y, z) \log q(y|z) dy \, dz - \int \, p(y) \log p(y) dy \\ &= \int \int \, p(y, z) \log q(y|z)dy \, dz + H(Y) \end{align*}
對於其中的$p(y,z)$，本文基於馬爾科夫假設：$Y\leftrightarrow X\leftrightarrow Z$。這個假設表明，$Y$和$Z$在$X$的條件下獨立（那在最佳化時呢？$Z$是關於$X$和$Y$的聯合分佈進行更新的）。有：
$\displaystyle p(y,z)=\int p(x,y,z)dx=\int p(x,y)p(z|x,y)dx= \int p(x,y)p(z|x)dx$
此外，由於$H(Y)$已知且固定，可忽略。則有
$\displaystyle I(Z; Y) \geq \int \int \int \, p(x,y)p(z|x) \log q(y|z)dx \, dy \, dz$
其中，$p(x,y)$是真實資料分佈，$p(z|x)$是原始模型關於$x$對中間表示$z$的推理分佈。
上界2
$I(Z;X)$展開為：
$\displaystyle I(Z;X)=\int \int p(x,z)\log \frac{p(z|x)}{p(z)}dx\,dz$
對於其中的$p(z)$，作者用另一個變分估計$r(z)$來擬合。由於有
\begin{align*} &\text{KL}(p(Z),r(Z))\geq 0\\ \implies&\int p(z)\log p(z)dz\geq\int p(z)\log r(z)dz\\ \implies&\int\int p(x,z)\log p(z)dx\,dz\geq\int \int p(x,z)\log r(z)dx\,dz \end{align*}
則有
\begin{align*} I(Z;X)\leq \int\int p(x)p(z|x)\log \frac{p(z|x)}{r(z)}dx\,dz \end{align*}
總體下界和最佳化
結合下界1和上界2，有：
\begin{align*} &I(Z; Y) - \beta I(Z; X) \\ \geq&\int \int \int \, p(x,y)p(z|x) \log q(y|z)dx \, dy \, dz - \beta \int\int p(x)p(z|x)\log \frac{p(z|x)}{r(z)}dx\,dz = L \end{align*}
針對上式，用經驗分佈來代替真實分佈。即用$\frac{1}{N}\sum_{n=1}^N\delta_{x_n}(x)$代替$p(x)$，用$\frac{1}{N}\sum_{n=1}^N\delta_{y_n}(y)$代替$p(y)$，用$\frac{1}{N}\sum_{n=1}^N\delta_{(x_n,y_n)}(x,y)$代替$p(x,y)$。其中$\delta_{x_n}(x)$表示狄拉克函式，其空間內積分為1，且僅在$x_n$上非零。假設經驗分佈有$N$各樣本$\{(x_n,y_n)\}_{n=1}^N$。文中額外引入所謂狄拉克函式讓人看不懂，實際上直接把機率積分改成離散樣本的求和取平均即可。則上式可被估計為：
$\displaystyle L\approx \frac{1}{N}\sum\limits_{n=1}^N\int p(z|x_n)\log q(y_n|z)-\beta\, p(z|x_n)\log\frac{p(z|x_n)}{r(z)}\, dz$
文中將$z$視為隱變數，利用VAE的重引數技巧將$p(z|x_n)$實現為一個關於$x_n$的正態分佈$\mathcal{N}(f_e^\mu(x_n),f_e^{\Sigma}(x_n))$，其中$f_e^\mu(x_n),f_e^\Sigma(x_n)$分別為基於$x_n$生成的均值和協方差矩陣。將$z$抽樣表示為$f(x_n,\epsilon)= f_e^{\Sigma}(x_n))\epsilon + f_e^\mu(x_n) $，其中$\epsilon\sim \mathcal{N}(0,1)$。則最大化$L$可表示為最小化：
$\displaystyle J_{IB} = \frac{1}{N}\sum\limits_{n=1}^N \mathbb{E}_{\epsilon\sim \mathcal{N}(0,1)}\left[-\log q(y_n|f(x_n,\epsilon))\right] + \beta\,\text{KL}\left[p(Z|x_n);r(z)\right]$
其中$r(z)$利用某一特定分佈實現，文中使用標準正態分佈實現。
直覺理解
直覺上理解：模型要把每個$x_n$分別對映到特定的分佈，這些分佈既不能偏離標準正態分佈太遠，又需要讓模型後續能根據這些分佈的抽樣來預測$x_n$的標籤。那麼這種做法為什麼能從$x_n$中抽取對預測$y_n$有效的關鍵資訊而忽略無關資訊呢（即資訊瓶頸）？我的理解是，模型被懲罰以使不同$x_n$得到的$z_n$分佈靠近同一分佈，但為了有效預測$y_n$，又必須產生一定的不一致。不同$x_n$對應的$z$分佈越一致，透過$z$而流向$y$的差異性資訊將越少，導致$q$更難利用取樣的$z$預測$y$，從而促使模型忽略$x$中的冗餘資訊而保留預測$y$所需的關鍵資訊。$\beta$則用於控制$z$保留$x$資訊的程度，越大保留資訊越少。
相較於一般的判別模型：當不把$z $視為隱變數，而變成關於$x$唯一確定的中間表示時，就是一般的判別模型。這種方式隱式地假定了表示的連續性，然而無法確保所有$z$都不是被離散地分散在表示空間中。最壞的過擬合情況下，每個$(x_n,y_n)$都孤立地確定了一箇中間表示$z_n$來實現一一對映，導致無泛化。而對於使用了資訊瓶頸$z$的判別模型，由於$x$僅僅確定$z$的生成分佈，不同的$x_i,x_j$依然可能抽樣出同一個$z$，這種模式強制這個抽樣出的$z $必須共享這兩個樣本的相似資訊並忽略不同的資訊，從而表示語義的相似性被強制由線性距離控制，實現表示語義的連續性，從而顯式地確定了模型的泛化。
實驗
表1：資訊瓶頸加成的模型和各種正則化後模型的對比。
圖1：不同$\beta$、$z$維度$K$下VIB模型在MNIST上的錯誤率，以及兩個互資訊的平衡。
圖2：$z$維度$K=2$時，1000張圖片的$z$分佈的視覺化。
後續是一些對抗魯棒的實驗，不記錄