迁移学习:互信息的变分上下界

1 导引

在机器学习,尤其是涉及异构数据的迁移学习/联邦学习中,我们常常会涉及互信息相关的优化项,我研一下期的处女作(发在SDM'24上)也是致力于此(ArXiv论文链接:FedDCSR,GitHub源码链接:FedDCSR)。其思想虽然简单,但其具体的估计与优化手段而言却大有门道,我们今天来好好总结一下,也算是对我研一的一个收尾。

我们知道,随机变量XY的互信息定义为其联合分布(joint)p(x,y)和其边缘分布(marginal)的乘积p(x)p(y)之间的KL散度(相对熵)[1]

(1)I(X;Y)=DKL(p(x,y)p(x)p(y))=Ep(x,y)[logp(x,y)p(x)p(y)]

直观地理解,互信息表示一个随机变量包含另一个随机变量信息量(即统计依赖性)的度量;同时,互信息也是在给定另一随机变量知识的条件下,原随机变量不确定度的缩减量,即I(X;Y)=H(X)H(XY)=H(Y)H(YX)。当XY一一对应时,I(X;Y)=H(X)=H(Y);当XY相互独立时I(X;Y)=0

在机器学习的情境下,联合分布p(x,y)一般是未知的,因此我们需要用贝叶斯公式将其继续转换为如下形式:

(2)I(X;Y)=(1)Ep(x,y)[logp(xy)p(x)]=(2)Ep(x,y)[logp(yx)p(y)]

那么转换为这种形式之后,我们是否就可以开始对其进行估计了呢?答案是否定的。我们假设现在是深度表征学习场景,X是数据,Y是数据的随机表征,则对于第(1)种形式来说,条件概率分布p(x|y)=p(y|x)p(x)p(y|x)p(x)dx是难解(intractable)的(由于p(x)未知);而对于第(2)种形式而言,边缘分布p(y)也需要通过积分p(y)=p(yx)p(x)dx来进行计算,而这也是难解的(由于p(x)未知)。为了解决互信息估计的的难解性,我们的方法是不直接对互信息进行估计,而是采用变分近似的手段,来得出互信息的下界/上界做为近似,转而对互信息的下界/上界进行最大化/最小化[2]

2 互信息的变分下界(对应最大化)

我们先来看互信息的变分下界。我们常常通过最大化互信息的下界来近似地对其进行最大化。具体而言,按照是否需要解码器,我们可以将互信息的下界分为两类,分别对应变分信息瓶颈(解码项)[3][4]Deep InfoMax[5][6]这两种方法。

2.1 数据VS表征:变分信息瓶颈(解码项)

对于互信息的第(1)种表示法即I(X;Y)=Ep(x,y)[logp(xy)p(x)],我们已经知道条件分布p(x|y)是难解的,那么我们就采用变分分布q(x|y)将其转变为可解(tractable)的优化问题。这样就可以导出互信息的Barber & Agakov下界(由于KL散度的非负性):

(3)I(X;Y)=Ep(x,y)[logq(xy)p(x)]+Ep(y)[DK L(p(xy)q(xy))]Ep(x,y)[logq(xy)]+H(X)IBA,

这里H(X)X的微分熵,BA是论文[7]两位作者名字的缩写。当q(x|y)=p(x|y)时,该下界是紧的,此时上式的第一项就等于条件熵H(X|Y)

上式可不可解取决于微分熵H(X)是否已知。幸运的是,限定在 X是数据,Y是表征 的场景下,H(X)=Exp(x)logp(x)仅涉及数据生成过程,和模型无关。这意味着我们只需要最大化IBA的第一项,而这可以理解为最小化VAE中的重构误差(失真,distortion)。此时,IBA的梯度就与“编码器”p(y|x)和变分“解码器”q(x|y)相关,而这是易于计算的。因此,我们就可以使用该目标函数来学习一个最大化I(X;Y)的编码器p(y|x),这就是大名鼎鼎的变分信息瓶颈(variational information bottleneck) 的思想(对应其中的解码项部分)。

2.2 表征VS表征:Deep Infomax

我们在 2.1 中介绍的方法虽然简单好用,但是需要构建一个易于计算的解码器q(x|y),这在X是数据,Y是表征的时候非常容易,然而当 XY都是表征 的时候就直接寄了,首先是因为解码器q(x|y)是难以计算的,其次微分熵H(X)也是未知的。为了导出不需要解码器的可解下界,我们转向去思考q(x|y)变分族的的非标准化分布(unnormalized distributions)。

我们选择一个基于能量的变分族,它使用一个判别函数/网络(critic)f(x,y):X×YR,并经由数据密度p(x)缩放:

(4)q(xy)=p(x)Z(y)ef(x,y), where Z(y)=Ep(x)[ef(x,y)]

我们将该分布代入公式(3)中的IBA中,就导出了另一个互信息的下界,我们将其称为UBA下界(记作IUBA),可视为Barber & Agakov下界的非标准化版本(Unnormalized version):

(5)Ep(x,y)[f(x,y)]Ep(y)[logZ(y)]IUBA

f(x,y)=logp(y|x)+c(y)时,该上界是紧的,这里c(y)仅仅是关于y的函数(而非x)。注意在代入过程中难解的微分熵H(X)被消掉了,但我们仍然剩下一个难解的log配分函数logZ(y),它妨碍了我们计算梯度与评估。如果我们对Ep(y)[logZ(y)]这个整体应用Jensen不等式(log为凹函数),我们能进一步导出式(5)的下界,即大名鼎鼎的Donsker & Varadhan下界[7]

(6)IUBAEp(x,y)[f(x,y)]logEp(y)[Z(y)]IDV

然而,该目标函数仍然是难解的。接下来我们换个角度,我们不对Ep(y)[logZ(y)]这个整体应用Jensen不等式,而考虑对里面的logZ(y)应用Jensen不等式即logZ(y)=logEp(x)[ef(x,y)]Ep(x)[logef(x,y)]=Ep(x)[f(x,y)],那么我们就可以导出式(5)的上界来对其进行近似:

(7)IUBAEp(x,y)[f(x,y)]Ep(x)p(y)[f(x,y)]IMINE

然而式(5)本身做为互信息的下界而存在,因此IMINE严格意义上讲既不是互信息的上界也不是互信息的下界。不过这种方法可视为采用期望的蒙特卡洛近似来评估IDV,也就是作为互信息下界的无偏估计。已经有工作证明了这种嵌套蒙特卡洛估计器的收敛性和渐进一致性,但并没有给出在有限样本下的成立的界[8][9]

IMINE思想的基础之上,论文Deep Infomax[6]又向前推进了一步,认为我们无需死抱着信息的KL散度形式不放,可以大胆采用非KL散度的形式。事实上,我们主要感兴趣的是最大化互信息,而不关心它的精确值,于是采用非KL散度形式可以为我们提供有利的trade-off。比如我们就可以基于p(x,y)p(x)p(y)Jensen-Shannon散度(JSD),来定义如下的JS互信息估计器:

(8)IJSDEp(x,y)[sp(f(x,y))]Ep(x)p(y)[sp(f(x,y))]

这里x是输入样本,x是采自p(x)=p(x)的负样本,sp(z)=log(1+ex)softplus函数。这里判别网络f被优化来能够区分来自联合分布的样本对(正样本对)和来自边缘乘积分布的样本对(负样本对)。

此外,噪声对比估计(NCE)[10]做为最先被采用的互信息下界(被称为“InfoNCE”),也可以用于互信息最大化:

(9)I(X;Y)Ep(x,y)[f(x,y)Ep(x)[logxef(x,y)]]IInfoNCE

对于Deep Infomax而言,IJSDIInfoNCE形式的之间差别在于负样本分布p(x)的期望是套在正样本分布p(x,y)期望的里面还是外面,而这个差别就意味着对于DVJSD而言一个正样本只需要一个负样本,但对于InfoNCE而言就是一个正样本就需要N个负样本(N为batch size)。此外,也有论文[6]分析证明了IJSD对负样本的数量不敏感,而IInfoNCE的表现会随着负样本的减少而下降。

3 互信息的变分上界(对应最小化)

我们接下来来看互信息的变分上界。我们常常通过最小化互信息的上界来近似地对互信息进行最小化。具体而言,按照是否需要编码器,我们可以将互信息的下界分为两类,而这两个类别分别就对应了变分信息瓶颈的编码项[4]解耦表征学习[11]

3.1 数据VS表征:变分信息瓶颈(编码项)

对于互信息的第(2)种表示法即I(X;Y)=Ep(x,y)[logp(yx)p(y)],我们已经知道边缘分布p(y)=p(yx)p(x)dx是难解的。但是限定在 X是数据,Y是表征 的场景下,我们能够通过引入一个变分近似q(y)来构建一个可解的变分上界:

(10)I(X;Y)Ep(x,y)[logp(yx)p(y)]=(1)Ep(x,y)[logp(yx)q(y)q(y)p(y)]=(2)Ep(x,y)[logp(yx)q(y)]DK L(p(y)q(y))(3)Ep(x)[DK L(p(yx)q(y))]R,

注意上面的(1)是分子分母同时乘以q(y)(2)是单独配凑出KL散度;(3)是利用KL散度的非负性(证明变分上下界的常用技巧)。最后得到的这个上界我们在生成模型在常常被称为Rate[12](也就是率失真理论里的那个率),故这里记为R。当q(y)=p(y)时该上界是紧的,且该上界要求logq(y)是易于计算的。该变分上界经常在深度生成模型(如VAE)[13][14] 被用来限制随机表征的容量。在变分信息瓶颈[4]这篇论文中,该上界被用于防止表征携带更多与输入有关,但却和下游分类任务无关的信息(即对应其中的编码项部分)。

3.2 表征VS表征:解耦表征学习

上面介绍的方法需要构建一个易于计算的编码器p(y|x),但应用场景也仅限于在X是数据,Y是表征的情况下,当 XY都是表征 的时候(即对应解耦表征学习的场景)也会遇到我们在2.2中所面临的问题,从而不能够使用了。那么我们能不能效仿2.2中的做法,对导出的IJSDIInfoNCE加个负号,从而将互信息最大化转换为互信息最小化呢?当然可以但是效果不会太好。因为对于两个分布而言,拉近它们距离的结果是确定可控的,但直接推远它们距离的结果就是不可控的了——我们无法掌控这两个分布推远之后的具体形态,导致任务的整体表现受到负面影响。那么有没有更好的办法呢?

我们退一步思考:最小化互信息I(X;Y)的难点在于XY都是随机表征,那么我们可以尝试引入数据随机变量D,使得互信息I(X;Y)可以进一步拆分为DXY之间的互信息(如I(D;X)以及I(D;Y)。已知三个随机变量的互信息(称之为Interation information[1])的定义如下:

(11)I(X;Y;D)=(1)I(X;Y)I(X;YD)=(2)I(X;D)I(X;DY)=(3)I(Y;D)I(Y;DX)

联立上述的等式(1)和等式(2),我们有:

(12)I(X;Y)=I(X;D)I(X;DY)+I(X;YD)

对解耦表征学习而言,在概率图模型中的结构化假设(V型结构)中XY共同为D的潜在因子,而XY互不影响(详情可参见论文[11])。XYD对应的结构化概率关系如下图所示:

反映在后验概率上即表征后验分布q满足q(XD)=q(XD,Y),因此上述等式的最后一项就消失了:

(13)I(X;YD)=H(XD)H(XD,Y)=H(XD)H(XD)=0

这样我们就有:

(14)I(X;Y)=(1)I(D;X)I(D;XY)=(2)I(D;X)+I(D;Y)I(D;X,Y)

上述的(1)是由于I(X;YD)=0(2)是由于互信息的链式法则即I(D;X,Y)=I(D;Y)+I(D;XY)

I(X;Y)等价变换至此,真相已经逐渐浮出水面:我们可以可以通过最小化I(D;X)I(D;Y),最大化I(D;X,Y)来完成对I(X;Y)的最小化。其直观的物理意义也就是惩罚表征XY中涵盖的总信息,并使得XY共同和数据D相关联。

基于我们在3.12.1中所推导的I(D;X)I(D;Y)的变分上界与I(D;X,Y)的变分下界,我们就得到了I(X;Y)的变分上界:

(15)I(X;Y)Ep(D)[DK L(q(xD)p(x))+DK L(q(yD)p(y))]Ep(D)[Eq(x|D)q(y|D)[logp(Dx,y)]]

直观地看,上式地物理意义为使后验q(xD)q(yD)都趋近于各自的先验分布(一般取高斯分布),并减小XYD的重构误差,直觉上确实符合表征解耦的目标。

4 总结

总结起来,互信息的所有上下界可以表示为下图[2](包括我们前面提到的IBAIUBAIDVIMINEIInfoNCE等):

图中节点的代表了它们估计与优化的易处理性:绿色的界表示易估计也易于优化,黄色的界表示易于优化但不易于估计,红色的界表示既不易于优化也不易于估计。孩子节点通过引入新的近似或假设来从父亲节点导出。

参考

  • [1] Cover T M. Elements of information theory[M]. John Wiley & Sons, 1999.
  • [2] Poole B, Ozair S, Van Den Oord A, et al. On variational bounds of mutual information[C]//International Conference on Machine Learning. PMLR, 2019: 5171-5180.
  • [3] Tishby N, Pereira F C, Bialek W. The information bottleneck method[J]. arXiv preprint physics/0004057, 2000.
  • [4] Alemi A A, Fischer I, Dillon J V, et al. Deep variational information bottleneck[J]. arXiv preprint arXiv:1612.00410, 2016.
  • [5] Belghazi M I, Baratin A, Rajeshwar S, et al. Mutual information neural estimation[C]//International conference on machine learning. PMLR, 2018: 531-540.
  • [6] Hjelm R D, Fedorov A, Lavoie-Marchildon S, et al. Learning deep representations by mutual information estimation and maximization[J]. arXiv preprint arXiv:1808.06670, 2018.
  • [7] Barber D, Agakov F. The im algorithm: a variational approach to information maximization[J]. Advances in neural information processing systems, 2004, 16(320): 201.
  • [8] Rainforth T, Cornish R, Yang H, et al. On nesting monte carlo estimators[C]//International Conference on Machine Learning. PMLR, 2018: 4267-4276.
  • [9] Mathieu E, Rainforth T, Siddharth N, et al. Disentangling disentanglement in variational autoencoders[C]//International conference on machine learning. PMLR, 2019: 4402-4412.
  • [10] Oord A, Li Y, Vinyals O. Representation learning with contrastive predictive coding[J]. arXiv preprint arXiv:1807.03748, 2018.
  • [11] Hwang H J, Kim G H, Hong S, et al. Variational interaction information maximization for cross-domain disentanglement[J]. Advances in Neural Information Processing Systems, 2020, 33: 22479-22491.
  • [12] Alemi A, Poole B, Fischer I, et al. Fixing a broken ELBO[C]//International conference on machine learning. PMLR, 2018: 159-168.
  • [13] Rezende D J, Mohamed S, Wierstra D. Stochastic backpropagation and approximate inference in deep generative models[C]//International conference on machine learning. PMLR, 2014: 1278-1286.
  • [14] Kingma D P, Welling M. Auto-encoding variational bayes[J]. arXiv preprint arXiv:1312.6114, 2013.
posted @   orion-orion  阅读(2372)  评论(6编辑  收藏  举报
相关博文:
阅读排行:
· 清华大学推出第四讲使用 DeepSeek + DeepResearch 让科研像聊天一样简单!
· 推荐几款开源且免费的 .NET MAUI 组件库
· 实操Deepseek接入个人知识库
· 易语言 —— 开山篇
· Trae初体验
点击右上角即可分享
微信分享提示