Mutual Information Neural Estimation
原文
参考:https://ruihongqiu.github.io/posts/2020/07/mine/
背景
互信息可以衡量两个随机变量之间的相关性:
I ( X ; Z ) = H ( X ) − H ( X ∣ Z ) = H ( Z ) − H ( Z ∣ X ) = H ( X ) + H ( Z ) − H ( X , Z ) I(X;Z)=H(X)-H(X|Z)=H(Z)-H(Z|X)=H(X)+H(Z)-H(X,Z) I(X;Z)=H(X)−H(X∣Z)=H(Z)−H(Z∣X)=H(X)+H(Z)−H(X,Z)
互信息量和KL散度的关系如下:
I ( X ; Z ) = ∑ x ∈ X ∑ z ∈ Z p ( x , z ) l o g p ( x , y ) p ( x ) p ( y ) = D ( p ( x , y ) ∣ ∣ p ( x ) p ( y ) ) I(X;Z)=\sum_{x\in \mathcal{X}}\sum_{z\in \mathcal{Z}}p(x,z)log\frac{p(x,y)}{p(x)p(y)}=D(p(x,y)||p(x)p(y)) I(X;Z)=x∈X∑z∈Z∑p(x,z)logp(x)p(y)p(x,y)=D(p(x,y)∣∣p(x)p(y))
但实际计算中,特别是对于高维空间来说,其边缘熵 H ( X ) H(X) H(X)、 H ( Z ) H(Z) H(Z)和条件熵 H ( X ∣ Z ) H(X|Z) H(X∣Z)难以计算。
解决方案
作者给出了两种利用梯度下降算法逼近的互信息估计,分别是The Donsker-Varadhan representation和The f-divergence representation。
The Donsker-Varadhan representation
D K L ( P ∣ ∣ Q ) = sup T : Ω → R E P [ T ] − l o g ( E Q [ e T ] ) D_{KL}(\mathbb{P}||\mathbb{Q})=\sup\limits_{T:\Omega \rightarrow \mathbb{R}}\mathbb{E_{\mathbb{P}}}[T]-log(\mathbb{E_{\mathbb{Q}}}[e^T]) DKL(P∣∣Q)=T:Ω→RsupEP[T]−log(EQ[eT])
其中 P \mathbb{P} P和 Q \mathbb{Q} Q是两个任意分布, T T T是从样本空间 Ω \Omega Ω映射到实数 R \mathbb{R} R的任意函数。
证明见大佬Ruihong Qiu中2.2节
The f-divergence representation
The f-divergence representation可以看做是The Donsker-Varadhan representation的弱化版本,由2.1和不等式 x e > l o g x \frac{x}{e}> log\mathcal{x} ex>logx易得。
最终形式
I ( X ; Z ) ≥ I Θ ( X , Z ) = sup θ ∈ Θ E P X Z [ T θ ] − l o g ( E P X P Z [ e T θ ] ) I(X;Z)\geq I_{\Theta}(X,Z)=\sup\limits_{\theta\in\Theta}\mathbb{E}_{\mathbb{P}_{XZ}}[T_{\theta}]-log(\mathbb{E}_{\mathbb{P}_X\mathbb{P}_Z}[e^{T_{\mathbb{\theta}}}]) I(X;Z)≥IΘ(X,Z)=θ∈ΘsupEPXZ[Tθ]−log(EPXPZ[eTθ])
我们希望用一个可以利用梯度更新的神经网络模型来计算上式,则有:
I ( X ; Z ) n ^ = sup θ ∈ Θ E P X Z ( n ) [ T θ ] − l o g ( E P X ( n ) P ^ Z ( n ) [ e T θ ] ) \hat{I(X;Z)_n}=\sup\limits_{\theta\in\Theta}\mathbb{E}_{\mathbb{P}^{(n)}_{XZ}}[T_{\theta}]-log(\mathbb{E}_{\mathbb{P}^{(n)}_X\mathbb{\hat{P}}^{(n)}_Z}[e^{T_{\mathbb{\theta}}}]) I(X;Z)n^=θ∈ΘsupEPXZ(n)[Tθ]−log(EPX(n)P^Z(n)[eTθ])
其中 T T T是一个神经网络; X X X、 Z Z Z是两个样本集。得到估计的梯度为:
G ^ B = E B [ ∇ θ T θ ] − E B [ ∇ θ T θ e T θ ] E B [ e T θ ] \hat{G}_B=\mathbb{E}_B[\nabla_{\theta}T_{\theta}]-\frac{\mathbb{E}_B[\nabla_{\theta}T_{\theta}e^{T_{\theta}}]}{\mathbb{E}_B[e^{T_{\theta}}]} G^B=EB[∇θTθ]−EB[eTθ]EB[∇θTθeTθ]
但是这种方式是有偏的。可以通过滑动平均来估计 E B [ e T θ ] \mathbb{E}_B[e^{T_\theta}] EB[eTθ]
完整的过程如下: