【论文笔记——DIM】Learning Deep Representations By Mutual Information Estimation and Maximization

arxiv：https://arxiv.org/abs/1808.06670
code：https://github.com/rdevon/DIM
report video：https://www.youtube.com/watch?v=o1HIkn8LEsw&t=256s

1 Introduction

无监督学习一个重要的问题就是学习有用的 representation，本文的目的就是训练一个 representation learning 函数（即编码器encoder） ，其通过最大编码器输入和输出之间的互信息(MI)来学习对下游任务有用的 representation，而互信息可以通过 MINE ${}^{[1]}$的方法进行估算。

作者还提到，直接最大化全部输入和编码器的输出（即全局的MI）更适合于重建性的任务，而在分类的下游任务上效果不太好。而最大化输入的局部区域（例如不是完整的图片，而是图片中的一块）和输出（即学到的representation）的平均互信息在下游任务（如图像分类）上的效果更好。

因为这种方法类似于对抗自动编码器(AEE, adversarial autoencoders)${}^{[2]}$的，将MI最大化和先验匹配结合起来，根据期望的统计特性约束 representation，并且还与 infomax 的优化规则密切相关${}^{[3]}$，因此作者称其为 Deep InfoMax(DIM).

因此，文章的主要贡献如下：

• 提出了 Deep InfoMax(DIM)，可以同时估算和最大化输入数据和高级representation之间的互信息(MI)
• 作者提出的最大化互信息的方法，可以根据下游任务是分类还是重建，来对优化全局还是局部的信息进行调整。
• 使用对抗学习约束representation，来使其具有特定于先验的期望统计特征
• 介绍了两种测量 representation 质量的新方法，一个基于Mutual Information Neural Estimation(MINE)${}^{[1]}$，一个基于neural dependency measure(NDM)${}^{[4]}$，并用它们将DIM与其他无监督学习方法进行比较

2 Objectives of Deep InfoMax

首先，先来看一下传统的图像 encoder 结构 ，如下图：

$x$ 表示原始图像，$f_{\psi }(x)$ 表示经过卷积层输出的 $M\times M$（注意，DIM 里的 $M\times M$ 指的是 $M\times M$ 块，而不是像素）的 feature map，而 $y=h_{\psi }\left(f_{\psi }(x)\right)$ 则为经过全连接层之后最终的 feature vector（也可以理解为图像的一种 representation），再将这个 representation 应用于一些下游任务，如分类。

而 DIM 的整体思路如下图，也就是“用图片 encode 后得到的 representation（$y$），与同一张图片的 feature map 组成正样本对，和另一张图片的 feature map 组成负样本对，然后利用 GAN 的思想训练一个 Discriminator 来区分这两中样本对”。

既然要训练 Discriminator ，就不得不提一下 DIM 的损失函数，那么就得先说一下其优化目标。

首先，我们定义一些数学符号：

• $x\in X$为 image 数据集；编码器 encoder 为$E_{\psi }=h_{\psi }\circ f_{\psi }$（$\circ$表示是由后面那两个函数组成 ）；$f_{\psi }(x)$ 表示由将原始图像 $x$ 映射到 feature map；$h_{\psi }(f_{\psi }(x))$ 表示将 feature map 再映射到最后的输出 representation，即 $y=h_{\psi }(f_{\psi }(x))=E_{\psi }(x)\in Y$。

• 将样本从 $f_{\psi }(X)$ 经过 $h_{\psi }$ 得到的 $y\in Y$ 的分布，也叫做 ”push-forward distribution“（目前还不是很理解这个名词），记作$Q_{\psi ,X}$

而 DIM 最终的优化目标也就是以下两点：

• 最大化 $f_{\psi }(X)$ 和 $Y$ 之间的互信息。（作者根据具体任务的不同，又分为了 global 和 local 两种，后面再提）
• 在最终的 representation 中再加入一个统计性的约束，使得到的 $y$ 的分布（push-forward distribution）$Q_{\psi ,X}$尽量与先验分布 $Q_{prior}$ 相匹配，有点类似对抗自编码器（AAE）。

3 Loss function

Global Infomax

知道了优化目标，下面就是如何去建立和推导该目标的数学表达方式。

我们首先解决第一个优化目标，也就是最大化互信息。互信息就是衡量两个变量的相关性，如果 $X$ 和 $Y$ 相互独立，即$P_{XY}=P_X{P}_Y$，那么它们的互信息 $I(X;Y)$ 也就为0。互信息的定义如下：
$$\begin{array}{rl}\mathcal{I}(X;Y)& =\sum _{X,Y}P(X,Y)\log \frac{P(X\mid Y)}{P(X)}\end{array}\text{\hspace{1em}(1)}$$
KL散度则是衡量同一个随机变量两个分布之间的差异，其值越小，说明分布越接近，定义如下：
$$\begin{array}{rl}{\mathcal{D}}_{KL}(X||Y)& ={\mathbb{E}}_{P_{XY}}[\log \frac{P(X)}{P(Y)}]\end{array}\text{\hspace{1em}(2)}$$
MI和KL散度之间的关系也不难推导，如下：
$$\begin{array}{rl}\mathcal{I}(X;Y)& =\sum _{X,Y}P(X,Y)\log \frac{P(X\mid Y)}{P(X)}\\ & =\iint P(X,Y)log\frac{P(X|Y)}{P(X)}\\ & ={\mathbb{E}}_{P_{XY}}[log\frac{P(X|Y)}{P(X)}]\\ & ={\mathbb{E}}_{P_{XY}}[log\frac{P(X,Y)}{P(X)P(Y)}]\\ & ={\mathcal{D}}_{KL}(P_{XY}||P(X)P(Y))\end{array}\text{\hspace{1em}(3)}$$
因为互信息没办法精确的计算，但是已经有一些算法可以对其进行估计，作者首先使用了Mutual Information Neural Estimation（MINE）${}^{[1]}$方法，其基于KL散度的 Donsker-Varadhan representation 给出了互信息的下限，如下公式：
$$\mathcal{I}(X;Y):={\mathcal{D}}_{KL}(\mathbb{J}||\mathbb{M})\ge {\hat{\mathcal{I}}}_{\omega }^{(DV)}(X;Y):={\mathbb{E}}_{\mathbb{J}}\left[T_{\omega }(x,y)\right]-\log {\mathbb{E}}_{\mathbb{M}}\left[e^{T_{\omega }(x,y)}\right]\text{\hspace{1em}(4)}$$
上式中的 $\mathbb{J}$ 表示 $X$ 和 $Y$ 的联合分布 $P_{XY}$，$\mathbb{M}$ 表示 $X$ 和 $Y$ 的边缘分布的乘积 $P_X{P}_y$，$T_{\omega }$ 表示参数为 $\omega$ 的鉴别器。

更具体来说，这里的 $X$ 为图像的 feature map（$f_{\psi }(X)$），$Y$ 为图像编码后的得到的 representation（$E_{\psi }(X)$）。 $P_{XY}$ 表示图片的 $f_{\psi }(X)$ 和其 $E_{\psi }(X)$ 之间的联合分布（因为来自同一张图片），而 $P_{\tilde{X}}{P}_Y$ 表示另一张图片 $f_{\psi }(\tilde{X})$ 和同一个 $E_{\psi }(X)$ 之间这两个边缘分布的乘积（因为来自不同的图片）。前者需要最大化，而后者需要最小化。因为鉴别器 $T_{\psi }$ 接收两个参数 $f_{\psi }(X)/f_{\psi }(\tilde{X})$ 和 E_{\psi}(X)$，然后输出一个分数，表示有多大的把握确信这个 feature map 是对应于这个 representation的（即来自同一张图片），显然来自同一张图片的输出越大越好，来自不同图片的输出越小越好。

因此最大化 $\mathcal{I}(X;Y)$ 的目标变成了最大化其的下限 ${\hat{\mathcal{I}}}_{\omega }^{(DV)}(X;Y)$，而又转化为了一个对抗训练问题。这样，我们就可以写出这一目标的损失函数了，如下（$\psi$ 为编码器 $E$ 的参数，$\omega$ 为鉴别器的参数）：
$$\begin{array}{rl}los{s}_G& =\underset{\omega ,\psi }{\max }{\hat{\mathcal{I}}}_{\omega }^{(DV)}(f_{\psi }(X);E_{\psi }(X))\\ & =\underset{\omega ,\psi }{\max }{\mathbb{E}}_{P_{XY}}\left[T_{\omega }(f_{\psi }(X),E_{\psi }(X))\right]-\log {\mathbb{E}}_{P_X\times P_{\tilde{X}}}\left[e^{T_{\omega }(f_{\psi }(X),E_{\psi }(X))}\right]\\ & =\underset{\omega ,\psi }{\min }-({\mathbb{E}}_{P_{XY}}\left[T_{\omega }(f_{\psi }(X),E_{\psi }(X))\right]-\log {\mathbb{E}}_{P_X\times P_{\tilde{X}}}\left[e^{T_{\omega }(f_{\psi }(X),E_{\psi }(X))}\right])\end{array}\text{\hspace{1em}(5)}$$
但由于我们的优化目标并不需要知道 MI 的具体估计值，而只需要能够将其最大化即可，所以我们不一定需要使用 KL 散度（也就是基于 DV representation 的方法），因此将上述损失函数可以写成更广泛的形式，如下：
$$\begin{array}{rl}los{s}_G& =\underset{\omega ,\psi }{\max }{\hat{\mathcal{I}}}_{\omega }\left(f_{\psi }(X);E_{\psi }(X)\right)\\ & =\underset{\omega ,\psi }{\min }-{\hat{\mathcal{I}}}_{\omega }\left(f_{\psi }(X);E_{\psi }(X)\right)\end{array}\text{\hspace{1em}(6)}$$
式子中的 ${\hat{\mathcal{I}}}_{\omega }$ 可以换成其他的 MI 估计器，只要能给出 MI 的边界将其进行最大化即可。在文章中，作者尝试了 Jensen-Shannon MI estimator ${}^{[4]}$ 和 infoNCE ${}^{[5]}$ 两种方法，如下，并对其各自的优势与性能进行了比较分析：
$$\begin{gathered} {\hat{\mathcal{I}}}_{\omega ,\psi }^{(\mathrm{J}\mathrm{S}\mathrm{D})}\left(X;E_{\psi }(X)\right):={\mathbb{E}}_{\mathbb{P}}\left[-\mathrm{sp}\left(-T_{\psi ,\omega }\left(x,E_{\psi }(x)\right)\right)\right]-{\mathbb{E}}_{\mathbb{P}\times \tilde{\mathbb{P}}}\left[\mathrm{sp}\left(T_{\psi ,\omega }\left(x^{\prime },E_{\psi }(x)\right)\right)\right] \\ ,wheresp(z)=log(1+e^z)\text{\hspace{1em}(7)} \end{gathered}$$

$${\hat{\mathcal{I}}}_{\omega ,\psi }^{(\mathrm{i}\mathrm{n}\mathrm{f}\mathrm{o}\mathrm{N}\mathrm{C}\mathrm{E})}\left(X;E_{\psi }(X)\right):={\mathbb{E}}_{\mathbb{P}}\left[T_{\psi ,\omega }\left(x,E_{\psi }(x)\right)-{\mathbb{E}}_{\tilde{\mathbb{P}}}\left[\log \sum _{x^{\prime }}{e}^{T_{\psi ,\omega }\left(x^{\prime },E_{\psi }(x)\right)}\right]\right]\text{\hspace{1em}(8)}$$

可以看到，我们上面得到的损失函数有一个下标 $G$，实际上就是代表这是 Global 情况下的目标。所谓 Gloabl ，也就是说我们是直接最大化整个图像的输入和输出之间的互信息的，在说什么是 Local Infomax 之前，再总结一下 Global Infomax 的流程：

（1）从数据集中取样原始图像 $x_{+}^{(1)},\dots ,x_{+}^{(n)}\sim P_X$ ，然后计算 feature map $f_{\psi }(x_{+}^{(i)}){\forall }_i$

（2）计算图像的 representation $y^{(i)}=h_{\psi }(f_{\psi }(x_{+}^{(i)}))$

（3）将 $(f_{\psi }(x_{+}^{(i)}),y^i)$ 组成正样本对

（4）从数据集中取样不同的图像 $x_{-}^{(1)},\dots ,x_{-}^{(n)}\sim P_X$，然后计算 feature map $f_{\psi }(x_{-}^{(i)}){\forall }_i$

（5）将 $(f_{\psi }(x_{-}^{(i)}),y^i)$ 组成负样本对

（6）通过优化编码器的参数 $\psi$ 和鉴别器的参数 $\omega$ 来最小化$-(\frac{1}{n}{\sum }_{i=1}^n{T}_{\psi }(f_{\psi }(x_{+}^{(i)}),y^{(i)}))-log\frac{1}{n}{\sum }_{i=1}^n{e}^{T_{\psi }(f_{\psi }(x_{-}^{(i)}),y^{(i)}))})$

Local Infomax

先说一下为什么要提出 Local Infomax？直觉上来看，对于一张图片来说，如果我们的下游任务不是重建类的任务，只是对图片进行分类，那么就没有必要对一些琐碎或者对分类任务无关紧要的像素。而如果我们的目标是最大化整张输入图片的 feature map 与 representation，那么编码器为了符合最后的全局最优情况，就有可能会选择到这些对下游任务并无实际作用的部分进行编码，这样得到的 representation 就肯定不会是针对下游任务最优的 representation。

而 Local Infoamx 的思想就是，我们并不将整张图片的 feature map 一次性输入损失函数来进行 MI 最大化，而是将其分为 $M\times M$块（$M$ 不是指像素，而是指被分成了 $M^2$ 个块），一次输入一个块和同一个 representation，最终目标是使这 $M^2$ 个块和整张图片的 representation 的平均 MI 达到最大。这样就可以使给每个块之间共享的一些信息进行编码。

文章中也用实验证明了，根据下游任务的不同，Local Inofmax 在图像分类等一些下游任务中确实具有更好的效果。

因此，也就很容易写出 Local Infomax 的损失函数（与 $los{s}_G$ 的 $\omega$ 不是同一个）：
$$\begin{array}{rl}los{s}_L& =\underset{\omega ,\psi }{\max }\frac{1}{M^2}\sum _{i=1}^{M^2}{\hat{\mathcal{I}}}_{\omega ,\psi }\left(f_{\psi }^{(i)}(X);E_{\psi }(X)\right)\\ & =\underset{\omega ,\psi }{\min }\frac{1}{M^2}-\sum _{i=1}^{M^2}{\hat{\mathcal{I}}}_{\omega ,\psi }\left(f_{\psi }^{(i)}(X);E_{\psi }(X)\right)\end{array}\text{\hspace{1em}(9)}$$
总结一下 Local Infomax 的流程：

（1）从数据集中取样原始图像 $x_{+}^{(1)},\dots ,x_{+}^{(n)}\sim P_X$ ，然后计算 feature map $f_{\psi }(x_{+}^{(i)}){\forall }_i$

（2）计算图像的 representation $y^{(i)}=h_{\psi }(f_{\psi }(x_{+}^{(i)})){\forall }_i$

（3）将 $(f_{\psi }(x_{+}^{(i)}{)}_j,y^i){\forall }_j$ 组成正样本对

（4）从数据集中取样不同的图像 $x_{-}^{(1)},\dots ,x_{-}^{(n)}\sim P_X$，然后计算 feature map $f_{\psi }(x_{-}^{(i)}){\forall }_i$

（5）将 $(f_{\psi }(x_{-}^{(i)}{)}_j,y^i){\forall }_j$ 组成负样本对

（6）通过优化编码器的参数 $\psi$ 和鉴别器的参数 $\omega$ 来最小化$-\frac{1}{M^2}{\sum }_{j=1}^{M^2}(\frac{1}{n}{\sum }_{i=1}^n{T}_{\psi }(f_{\psi }(x_{+}^{(i)}{)}_j,y^{(i)}))-log\frac{1}{n}{\sum }_{i=1}^n{e}^{T_{\psi }(f_{\psi }(x_{-}^{(i)}{)}_j,y^{(i)}))})$

看起来和 Global Infomax 的流程差不多，主要是在第（3）（5）步中，现在需要对每个 patch 执行此操作（也就是Global的 $M^2$ 倍） ，而之所以只用 $j$ 来对 patch 进行索引，是因为 patch 在哪一行哪一列并不重要。第（5）步中的每个patch $j$ 则是来自另一张不同的图像。

Local Infomax的原理图如下：

至于怎么计算这个平均互信息，文章的附录中给了两种方法，一个是直接将最后的 representation（$Y$）复制然后接到 feature map(+) 的块后面形成real，接到 feature map(-) 后面形成 fake，然后再利用一个 $1\times 1$ 的卷积鉴别器对对进行评分，如下图：

另一种方法则是利用点击运算，即对 representation（$Y$）进行全连接网络的编码，每个块用一个 $1\times 1$ 的网络进行编码，最后得到的两个结果是想同纬度的，然后进行点乘操作得到对应块的分数：

Matching representations to prior

若学习到的隐变量服从标准正态分布的先验分布，这有利于使得编码空间更加规整，甚至有利于解耦特征，便于后续学习。因此，在 DIM 中，我们同样希望加上这个约束，作者利用对抗自编码器（AAE）的思路引入对抗来加入这个约束，即训练一个新的鉴别器，而将编码器当做生成器。鉴别器的目标是区分 representation 分布的真伪（即是否符合先验分布），而编码器则是尽量欺骗判别器，输出更符合先验分布的 representation，如下图：

具体做法是，训练另一个鉴别器$D_{\varphi }$，我们需要学习到一种 representation 来让这个鉴别器 $D_{\varphi }$ 确信其来自先验分布 $Q_{prior}$，这不就是一种对抗的思想。

训练该鉴别器的损失函数如下：
$$\begin{array}{rl}los{s}_{matching}& =\underset{\psi }{\min }\underset{\varphi }{\max }{\hat{\mathcal{D}}}_{\varphi }\left(Q_{prior}\Vert Q_{\psi ,X}\right)\\ & =\underset{\psi }{\min }\underset{\varphi }{\max }{\mathbb{E}}_{Q_{prior}}\left[\log D_{\varphi }(y)\right]+{\mathbb{E}}_{Q_{\psi ,X}}\left[\log \left(1-D_{\varphi }\left(E_{\psi }(x)\right)\right)\right]\end{array}\text{\hspace{1em}(10)}$$

Final loss

将上面三个目标整合后，也就得到了 DIM 的最终目标，如下：
$$\begin{array}{rl}& \underset{{\omega }_1,\psi }{\max }\alpha {\hat{\mathcal{I}}}_{{\omega }_1}(f_{\psi }(X);E_{\psi }(X))(globaInfoMax)\\ & +\underset{{\omega }_2,\psi }{\max }\frac{\beta }{M^2}\sum _{i=1}^{M^2}{\hat{\mathcal{I}}}_{{\omega }_2,\psi }\left(f_{\psi }^{(i)}(X);E_{\psi }(X)\right)(localInfoMax)\\ & +\underset{\psi }{\min }\underset{\varphi }{\max }\gamma {\hat{\mathcal{D}}}_{\varphi }\left(Q_{prior}\Vert Q_{\psi ,X}\right)(priormatching)\\ & =\underset{{\omega }_1,{\omega }_2,\psi }{\min }(-\alpha {\hat{\mathcal{I}}}_{{\omega }_1}(f_{\psi }(X);E_{\psi }(X))-\frac{\beta }{M^2}\sum _{i=1}^{M^2}{\hat{\mathcal{I}}}_{{\omega }_2,\psi }\left(f_{\psi }^{(i)}(X);E_{\psi }(X)\right))+\underset{\psi }{\min }\underset{\varphi }{\max }\gamma {\hat{\mathcal{D}}}_{\varphi }\left(Q_{prior}\Vert Q_{\psi ,X}\right)\end{array}$$

之所以加上$\alpha 、\beta 、\gamma$三个参数，是因为有时候我们只想使用 global InfoMax (如重建类下游任务)，就可以将 $\beta$ 设置为0；而有时候只想使用 Iocal InfoMax (如分类任务)，就可以将 $\alpha$ 设置为0；但这两种情况下，最佳的$\gamma$是不同的，所以也需要 $\gamma$ 来进行调节。

References

[1] Ishmael Belghazi, Aristide Baratin, Sai Rajeswar, Sherjil Ozair, Y oshua Bengio, Aaron Courville, and R Devon Hjelm. Mine: mutual information neural estimation. arXiv preprint arXiv:1801.04062, ICML’2018, 2018.

[2] Alireza Makhzani, Jonathon Shlens, Navdeep Jaitly, Ian Goodfellow, and Brendan Frey. Adversarial autoencoders. arXiv preprint arXiv:1511.05644, 2015.

[3] Anthony J Bell and Terrence J Sejnowski. An information-maximization approach to blind separation and blind deconvolution. Neural computation, 7(6):1129–1159, 1995.

[4] Philemon Brakel and Y oshua Bengio. Learning independent features with adversarial nets for non-linear ica. arXiv preprint arXiv:1710.05050, 2017.

[5] Sebastian Nowozin, Botond Cseke, and Ryota Tomioka. f-gan: Training generative neural samplers using variational divergence minimization. In Advances in Neural Information Processing Systems, pp. 271–279, 2016.

[6] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807.03748, 2018.