ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

互信息变分估计:从InfoNCE到CLUB的界与实战

互信息变分估计:从InfoNCE到CLUB的界与实战

做自监督对比学习的同学,对 InfoNCE 这串公式应该都不陌生,但很多人未必清楚,它的本名其实叫“互信息的变分下界”。我当初第一次在论文里看到“On Variational Bounds of Mutual Information”这个标题时,心里想的是:这不就是算两个变量之间“共享了多少信息”吗,至于要单独立一个研究方向来琢磨?后来真正上手做特征解耦和表征学习,被“互信息到底怎么算”这个看似简单的问题卡了好几天,才意识到这件事水有多深。

这篇文章不是文献综述,而是我根据自己的实践对“互信息的变分界”这一主题做的一次完整梳理。我会从“为什么精确计算互信息基本不可能”讲起,然后逐一拆解 DV 界、InfoNCE、NWJ 和 CLUB 这几类常见的变分界,解释它们背后的思路、适用的场景,以及我在实际代码里踩过的那些坑。如果你是做自监督学习、表征解耦、因果发现,或者只是想搞清楚对比学习损失函数和互信息到底什么关系,这篇内容应该能在你调参的时候帮上忙。

1. 为什么“精确算”互信息是一件基本没指望的事

想理解变分界,先得搞明白精确计算互信息为什么这么让人绝望。这不是实现上的问题,而是本质上的困难。

1.1 从定义出发:I(X;Y) 到底在算什么

互信息的定义用一句话说就是:知道随机变量 Y 之后,X 的不确定性减少了多少。写成信息论的标准形式是:

I(X;Y) = E_{p(x,y)}[ log( p(x,y) / (p(x)p(y)) ) ]

其中 p(x,y) 是 X 和 Y 的联合分布,p(x)p(y) 是两者独立时的分布。这个比值刻画的就是“联合分布”和“独立分布”之间的差异。

我一开始理解这个概念时觉得挺直观:两个人的朋友圈如果几乎完全重合,那他们之间的互信息就高;如果一个人聊篮球另一个人聊烘焙,那互信息就低。这没错。但问题在于,日常我们能拿到的只有一堆观测样本,比如一批图片和它们对应的特征向量,我们根本不知道产生这些样本的联合分布到底长什么样。

1.2 精确计算的三个瓶颈

第一个瓶颈是分布未知。我们手里只有样本,没有 p(x,y) 的解析表达式。传统的做法是用核密度估计或者直方图去近似联合分布,但一旦 X 向量的维度超过 20,直方图需要的样本数量就呈指数级增长。你想想看,如果每个维度只有 10 个分箱,20 维就需要 10 的 20 次方个格子,这个数量级在现实世界不存在。

第二个瓶颈是积分不可行。就算你假设了一个参量化联合分布,算互信息依然要处理高维积分。比如 X 是 128 维的图像特征,Y 是 64 维的表示向量,你要计算 log(p(x,y)/(p(x)p(y))) 在高维空间里的期望。这个积分没有解析解,数值积分在高维度下也彻底失效,因为体积基本上都集中在边界上,采样效率极低。

第三个瓶颈,也是最微妙的一点:很多场景下我们关心的并不是互信息的具体数值,而是怎么让模型通过“最大化/最小化”互信息来学到一个好的表示。这意味着我们真正需要的是一个可以放进梯度计算流程的目标函数,而不是一个只能在离线状态下算出来的数字。

这三个瓶颈叠加在一起,基本宣判了“精确计算”这条路在高维场景下走不通。于是学界把目光转向了变分界——也就是通过优化一个“替代目标”去逼近真实的互信息。

2. 变分界的基本思路:把“算不出”换成“优化得来”

2.1 变分法的核心:找到一个容易优化的替代目标

“变分”(variational)这个词听起来高深,但核心思想其实很简单:当你要求一个函数的最值而直接求不了的时候,就把它转换成一个优化问题,让某个可调的函数族去逼近最优值。

举一个大家更熟悉的例子:贝叶斯推断里算后验概率 p(z|x) 通常要积分 p(x) = ∫p(x,z)dz,这个积分往往算不出来,于是我们用另一个容易计算的分布 q(z) 去近似后验,然后最大化证据下界 ELBO。之所以叫“变分”,是因为我们不是在解一个固定函数,而是在一个分布族里面“变”出一个最接近目标的分布。

互信息的变分估计走的是同样的路:不直接算互信息,而是找一个带参数的函数 T,通过优化 T 使得某个表达式的值逼向真实的互信息。这个方法之所以可行,靠的是互信息本身可以写成某种“最大值”的形式——这类形式在信息论里有一套现成的理论工具。

2.2 从 KL 散度到 Donsker-Varadhan 表示

互信息本质上就是联合分布 p(x,y) 对乘积分布 p(x)p(y) 的 KL 散度:

I(X;Y) = KL( p(x,y) || p(x)p(y) )

而 KL 散度有一个非常有名的变分表示:对任意函数 T(x,y),以下不等式恒成立:

KL(p || q) ≥ sup_T [ E_p[T] - log E_q[ e^T ] ]

这个关系式叫 Donsker-Varadhan 变分表示,它的核心价值在于:把“算 KL 散度”换成了“优化一个关于 T 的目标”。如果允许 T 遍历所有可能的函数,这个上确界就能精确等于 KL 散度;如果 T 只能用某个函数族(比如神经网络)来表示,那我们就只能得到 KL 散度的一个下界。

我用一个更生活的类比来理解这个公式:假设你要评估两堆数据“有多不一样”,E_p[T] 可以理解为“我在真实数据上挑选证据的能力”,log E_q[e^T] 是“在独立分布上同一个证据产生的代价”。你调 T 的过程就像是一个侦探在真实现场找线索,同时要确保这个线索在无关场景里不容易出现。最理想的 T 就是直接指出两堆数据本质差异的那个函数。

2.3 用神经网络当变分函数

为什么大家普遍用神经网络实现变分函数 T?原因有三个,且每个都是实际工程里的硬需求。

第一,函数族表达能力要足够大。如果 T 只能用线性函数,那 DV 表示给出的界会非常松,估计出来的互信息严重低估。神经网络的泛函表达能力足以逼近很多复杂的判别函数,能保证界比较紧。第二,我们需要梯度。T 的参数需要靠梯度上升来优化,神经网络天然就支持端到端反传。第三,不同变分界实际上可以理解为对 T 施加了不同的约束,或者对 log E_q[e^T] 这一配分项做了不同的近似处理——神经网络提供了一个统一的载体,切换界的时候只需要改损失函数,不需要改整个架构。

3. 四种常见界:原理、行为与选型

变分界大致可以分成两类:下界和上界。下界用于“最大化互信息”场景,比如对比学习和表征学习;上界用于“最小化互信息”场景,比如解耦表示、信息瓶颈。下面逐个拆解。

3.1 DV 界:理论上最漂亮的下界

DV 界就是前面提到的 Donsker-Varadhan 表示直接套用到互信息上:

I(X;Y) ≥ E_{p(x,y)}[ T(x,y) ] - log E_{p(x)p(y)}[ e^{T(x,y)} ]

这个界的优点在于:它是严格下界,只要 T 函数族足够大,最终能收敛到真实的互信息,而且理论上不需要任何额外的正则化项。在学界,很多关于互信息估计误差的理论分析都是以 DV 界为基准展开的。

但在我的实际测试中,DV 界的训练体验相当折磨。问题出在第二项 log E[e^T]:当 T 输出比较大的正值时,e^T 会指数级放大,导致梯度的方差非常惊人。训练初期模型还稳定,跑几百步之后 loss 开始剧烈震荡,甚至直接 NaN。我试过把学习率降到 1e-5,也只是把崩溃时间往后拖了一点。后来我查了很多人的经验,发现要稳定训练 DV 界,通常得做额外处理,比如对 T 做梯度裁剪、谱归一化,或者用一个单独的网络作为 critic 并配合很大的 batch size。这些手段都可行,但代价不小。

因此我的结论是:DV 界非常适合作为理论研究工具,分析互信息估计的收敛性;如果要在大规模模型训练里直接拿它当损失函数,除非你有足够精力去调稳定手段,否则不建议上来就用它。

3.2 InfoNCE:表现力与实用性的平衡

InfoNCE 是当前实践中最常用的一类互信息下界。它的估计形式是:

I(X;Y) ≥ E_p[ log( exp(T(x,y)) / ( (1/K) ∑_{j=1}^K exp(T(x,y_j)) ) ) ]

这里的 y_j 包括一个正样本 y(来自联合分布)和 K-1 个负样本(来自边缘分布 p(y))。这个式子本质上是在做一个 softmax 分类:在 K 个候选样本里找出与 x 真正匹配的那个 y。

从理论角度看,InfoNCE 具备几个非常好的性质:它是互信息的下界,当负样本数量 K 逼近无穷时,这个下界会逐渐趋紧;它对 T 的约束比 DV 界温和,不会有 e^T 剧烈放大的梯度方差问题;它和对比学习之间几乎是直接映射——所谓“对比损失”,就是在最大化这个互信息下界。实际上很多经典的自监督模型用的正是这一类的损失函数,只是披了一层别的外衣。

从工程角度看,InfoNCE 几乎是为 GPU 训练量身定做的。你只需要在一个 batch 里把每个样本的增强视图当作正样本,把 batch 里的其他样本当作负样本,一次前向就能算出损失。这个设计非常自然地利用了 batch 内的构造性负样本,省去了显式的负样本采样,所以我在做一些大规模特征学习实验时,默认选 InfoNCE。

需要特别注意的一点是:InfoNCE 估计出来的数值通常比真实互信息低,尤其是负样本数不够大的时候,低估会很严重。也就是说,它“能用来排名”不代表“能用来精确度量”。如果只是把它当损失函数优化模型,那完全没问题;可要是你想报告实验里算出的“互信息绝对值”,那就得多留个心眼。

3.3 NWJ(MINE):另一个视角的下界

NWJ 界(也叫 MINE 界)在公式上和 DV 界只差一个 log 的位置。它用的是:

I(X;Y) ≥ E_{p(x,y)}[ T(x,y) ] - e^{-1} · E_{p(x)p(y)}[ e^{T(x,y)} ]

注意,DV 界是对 e^T 的期望取 log,而 NWJ 界直接在期望外面乘以一个常数项。这看起来只是小小的改动,实际影响却很大:log 运算会带来很强的压缩效应,去掉 log 之后,第二项的梯度方差会变小,更容易训练,但代价是偏置变大——特别是在 T 的函数族表达能力不够强、或者 T 和最优解差距较大时。

我个人在实际使用中发现,NWJ 界在维度不高的场景(比如几十维的特征)下是可行的,但一旦特征维度上百,第二项的期望方差还是比较大,需要额外的方差削减技巧。它更多出现在生成模型和对抗式训练框架里,因为那个框架里的鉴别器结构和 NWJ 界天然契合。

3.4 CLUB:少数能给出上界的形式

前面几个都是下界,专门用来最大化互信息。但有些任务需要的是上界,比如信息瓶颈(把无关信息压掉)和解耦表示学习(让不同因子之间尽量独立)。在这些场景里,如果你用一个下界去做“最小化”,你优化出的结果和真实互信息之间的差完全不可控,边界会崩。

CLUB 给出了一种基于条件分布的互信息上界估计:

I(X;Y) ≤ E_{p(x,y)}[ log q(y|x) ] - E_{p(x)p(y)}[ log q(y|x) ]

其中 q(y|x) 是你用神经网络拟合的条件分布。它的思想是:通过优化 q(y|x) 去逼近真实条件分布 p(y|x),然后用这个近似分布构造出互信息的上界。如果 q 拟合得越准,这个上界就越紧;如果 q 拟合得差,那这个上界就变得很松,甚至失去参考价值。

这个界在我的实验里表现相当稳定,因为它的每一项都是 log q,不会出现指数爆炸。我主要把它用在特征解耦任务上:先训练一个条件网络拟合数据之间的分布,再把 CLUB 作为正则项去惩罚不同特征维度之间的依赖。这样一来,既能利用互信息的理论支撑,又不会让训练过程爆炸。

3.5 界与界之间的横向对比

下面用一张表总结这几种界在我实践中的表现差异。注意,方差大小是相对于同一个数据规模和同一种 critic 架构而言的。

界名称类型梯度方差偏置情况对负样本依赖典型使用场景
DV下界很大渐近无偏不依赖显式负样本理论分析、小规模验证
InfoNCE下界较小K 越小低估越严重依赖负样本数量 K对比学习、表征学习
NWJ/MINE下界中等偏置较大不依赖显式负样本生成模型、对抗式训练
CLUB上界较小依赖 q 拟合程度不依赖显式负样本解耦表示、信息瓶颈

一句话选型经验:做学到表示,默认用 InfoNCE;要更紧的界且能承受调参成本,试 DV;要控制互信息不过高,用 CLUB。NWJ 在我这边使用频率不高,但如果你在做对抗式结构,它反而是合拍的。

4. 实操中的细节、调参与常见问题

理论讲再多,最后还是要落到代码上。这一部分我分享一下我在项目里跑互信息变分估计时的实际操作。

4.1 一个最小可跑的 InfoNCE 估计器

先给一个完整的 PyTorch 风格的实现,它是我最常用的估计器骨架。这个函数输入一个 batch 的特征和对应的增强视图,输出 InfoNCE 下界估计值。同时我加了一段温度参数,这在后面会讲。

import torch import torch.nn.functional as F def info_nce_loss(z_x, z_y, temperature=0.1): """ z_x, z_y: (batch_size, feature_dim) 特征矩阵 返回 InfoNCE 下界估计标量 """ batch_size = z_x.size(0) z_x = F.normalize(z_x, dim=-1) z_y = F.normalize(z_y, dim=-1) # 相似度矩阵: logits[i][j] = z_x[i] 与 z_y[j] 的余弦相似度 / temperature logits = torch.mm(z_x, z_y.t()) / temperature # 对角线元素是正样本对 positives = torch.diagonal(logits) # 对每一行,计算 log(exp(pos) / sum_j exp(logit_j)) # 等价于 pos - logsumexp(row) loss = -torch.diagonal( logits - torch.logsumexp(logits, dim=1) ).mean() # 返回的是互信息下界的估计(注意是负损失) # 若想要 InfoNCE 下界的正值,取 -loss 即可 return -loss

这段代码里最容易被忽略的是 logsumexp 那一步。很多人初次写 InfoNCE 直接用softmax拿概率然后取 log,这在数值上是等价的,但logsumexp更稳定,尤其是温度参数很小时,logits 里的值会变得很大,直接 softmax 可能出现精度问题。我在实际调参过程中确实遇到过这个坑,换掉之后就稳了。

另外注意,负样本来自同一个 batch 里的其他样本。如果 batch size 只有 16,那么每行只有 16 个候选,InfoNCE 下界会严重低估互信息;如果 batch size 拉到 512 甚至 1024,估计值会明显“变紧”。这也是为什么要强调 batch size 在对比学习里的重要性——它不只是影响训练稳定性,还直接决定了互信息下界的松紧程度。

4.2 负样本数量 K 的影响

关于 K 的选择,我在多次实验里观察到几个规律。K 越大,InfoNCE 下界越紧,训练出的线性评估指标通常也越好,但收益是递减的。从 128 到 256,提升很明显;从 256 到 512,还能看到一点;到 1024 以上,指标变化就很小了,而显存开销却大幅增加。

这里有个理论原因:InfoNCE 的下界里隐含了一个 log K 项,负样本的数量以对数尺度影响界值,所以想靠暴力加大 K 来获得线性收益是不可能的。

实践中解决 K 不够大的办法,是维护一个队列或者记忆库。把历史 batch 的特征放进一个很大的缓冲池里,当计算当前 batch 的 InfoNCE 时,从池子里取负样本。这样可以做到“batch 内样本配对 + 缓冲池里的额外负样本”,把有效负样本数提升到几千甚至上万,而不需要把 batch size 开得巨大。我用的方案是维护一个固定大小的特征队列,每次更新时用当前 batch 的特征替换最旧的那批,保存的样本数量一般在 4096 到 16384 之间。

注意一个副作用:缓冲池里的特征是不同迭代步的模型输出的,一致性会差一些。如果模型变化很快,老特征和新特征之间的分布会有偏移,负样本的“质量”会降低。解决方法是让编码器更新得更慢一点,比如使用动量更新的方式,也就是让用来生成缓冲池特征的那个“旧编码器”的权重以很小的比例跟随当前编码器更新。听起来复杂,实际就是在反向传播时额外维护一组“慢变”参数,代码量不多,但在我的实验里效果提升是肉眼可见的。

4.3 温度系数:一个“必要的不完美”

InfoNCE 的标准形式里没有可调温度,它默认在 logits 上不乘任何系数。但实际工程里几乎没人这么用,大家都会加一个 temperature 参数。我在实验里最常用的温度范围是 0.05 到 0.2,0.1 附近最为稳妥。

温度系数加上去之后,严格意义上训练的就不再是原始的互信息下界了,而是一个对概率分布做锐化/平滑后的目标。温度太高,softmax 分布趋于均匀,每个负样本都得到几乎相同的权重,模型学不到精细区分;温度太低,softmax 分布集中在最困难的负样本上,梯度开始震荡,训练不稳定甚至直接崩掉。

我在一次实验里把温度从 0.1 调到 0.01,准确率暴跌了十几个百分点,原因是模型被少数难样本带偏了。反过来,从 0.1 调到 0.5,模型的表现也逐渐下滑,因为负样本的区分度变模糊了。所以如果项目里用到了温度系数,强烈建议把它纳入超参搜索范围,不要沿用上一个项目的温度值。

4.4 常见问题速查表与避坑技巧

下面这个表格整理了我从多个项目里总结出的问题和排查方法,有不少是我自己踩过的坑。

现象可能原因处理建议
loss 剧烈震荡不收敛T 网络学习率过大 / 没有使用 log 稳定技巧降低学习率到 1e-4 左右,检查是否有 logsumexp 或梯度裁剪
互信息估计出现负数估计器偏置过大、负样本太少、模型还未收敛增大 K 或使用更大 batch size,确认不是在训练初期
加了大缓冲池但指标没提升缓冲池里的负样本特征太旧,分布不一致给编码器加动量更新,或降低缓冲池更新频率
温度调低后指标反而崩掉温度太小导致模型被困难负样本主导把温度调回 0.1 附近,观察梯度范数是否异常
训练后互信息数值远小于理论预期下界本身的低估效应,特征维度与 batch 不匹配如果目标是精确度量,改用 CLUB 或专门做校正,不要死磕 InfoNCE 的绝对值
用 DV 界训练 NaNe^T 项溢出给 T 加谱归一化,或用 InfoNCE 替代

除了表里的这些,还有一个我自己反复强调的原则:不要用“训练目标的好坏”来评估“互信息估计的准确性”。一个驱动的模型表现很好的 InfoNCE 损失,对应的互信息估计值照样可能明显低于真实值。反过来说,CLUB 给出的上界虽然在数值上更接近真实互信息,但直接用梯度下降去最小化它,并不能自动保证模型特征一定会解耦得很好。这是两个层次的问题,混在一起看往往会把调参方向带偏。

最后的一点个人体会

这几个界我都实际写过、跑过、崩过。如果让我给一个最简单的起步建议:做对比学习或者表征预训练,InfoNCE 是你最可靠的起点,它兼顾了效果与稳定;做解耦表示学习强调因子独立性,CLUB 作为正则项更合适;如果你的研究方向是理论分析或验证某种估计的收敛性,DV 界那套框架更规范。重要的是先想清楚你要的是“一个能优化模型的目标”还是“一个能给出可信数值的量”,这两种需求对应的选择完全不同。

在我个人的项目里,性价比最高的一套配置是很大 batch size 的 InfoNCE 加上一个小权重的 CLUB 正则。前者负责把表征空间撑开,后者负责压掉不必要的维度冗余。两者互补,训练也稳定。当然,这只是我基于自己数据实验得到的经验,换一个领域和数据集可能还需要重新调校,但至少可以作为一个起点。

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进