ARTICLE · INTELLIGENCE

战地情报 · 详情页

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

InfoNCE与变分互信息界:对比学习背后的数学原理与选型指南

InfoNCE与变分互信息界:对比学习背后的数学原理与选型指南

第一次把 InfoNCE 损失看成互信息下界的时候,我突然意识到自己此前的理解有多浅。当时我在某跨模态检索项目里调一个对比学习目标,训练曲线漂移得让人失眠,同事看了一眼说:你这不是在优化一个靠谱的目标,而是在优化一个你管不住的下界。后来我花了两周把 “On Variational Bounds of Mutual Information” 这类材料的核心思想啃了一遍,才明白这句话的分量有多重。

这篇文章想把自己这些年对变分互信息界的理解整理成一份能直接拿去用的经验笔记。会讲清楚三件事:互信息为什么在真实数据上几乎算不出来、变分界到底在“变分”什么、以及 InfoNCE、MINE、CLUB 这三种主流界各自的推导逻辑和选型思路。适合正在做对比学习、多模态对齐、特征压缩的算法工程师,也适合刚入坑想搞懂“为什么 InfoNCE 有效”的研究生。

1. 从 InfoNCE 说起:互信息估计为什么成了算法岗的必修课

1.1 对比学习真正在优化的对象

几乎所有做深度学习表征学习的同学,都绕不开对比学习。它的训练目标看起来非常直白:从同一个样本出发,做两个不同增强,得到正样本对(anchor 和 positive),同一个 batch 里的其他样本当作负样本(negatives)。模型要做的,就是在表征空间里把正样本对拉近、把负样本对推远。

这种设计太符合直觉了,以至于很多人(包括刚开始接触时的我)都觉得它就是一个带标签的度量学习。但如果把这件事往概率方向推一步,会发现一个很关键的事实:给定一对样本 (x, y),一个足够好的打分器 f(x, y) 是可以学到联合分布 p(x, y) 和边缘乘积 p(x)p(y) 的比值,也就是密度比:

r(x, y) = p(x, y) / (p(x) * p(y))

对比学习里的分类任务,本质上就是在估计这个密度比。而密度比的期望对数,正好就是互信息:

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

所以对比学习真正在做的事情,是在用一个可训练的分类任务去逼近互信息的下界。这也解释了为什么 InfoNCE 这类损失在大量实验里都有效——它不是被拍脑袋设计出来的度量学习目标,而是有明确信息论语义的代理目标。

1.2 密度比与互信息只差一个 log 的期望

这里我想强调一个容易被忽略的点:互信息 I(X; Y) 本身很难直接算,也很难直接优化,但密度比 r(x, y) 是可以通过神经网络打分器 f_θ(x, y) 去逼近的。一旦 f 学到了密度比,我们就能用一批真实数据样本,把E[log f]当成互信息的近似值。

这个“把互信息估计转化成密度比估计”的思路,就是几乎所有变分互信息下界的源头。InfoNCE、MINE 这些听起来高大上的名字,说到底都是密度比估计的不同实现方式。理解了这一点,后面再看它们各自的公式,就不会觉得是一堆无根无据的 Trick 了。

1.3 谁真的需要互信息估计

互信息估计不是纯粹的数学玩具,它在工程里的使用频率比我之前以为的高得多。我自己整理过一张任务清单:

  • 自监督预训练:最大化样本与其增强视图之间的互信息,学到通用表征;
  • 多模态对齐:最大化图像与文本之间的互信息,让两种模态的表征落到共享空间;
  • 信息压缩:最小化表征与原始输入之间的互信息,让表征只保留任务相关信息;
  • 去偏学习:最小化表征与敏感属性之间的互信息,削弱偏见信息。

注意,这里已经出现了两类方向完全相反的需求:最大化用下界,最小化用上界。如果方向选错,你在训练里看似在降低 loss,实际上和真实目标背道而驰。这个问题我会在第三章专门展开。

2. 精确互信息的死结:未知密度与高维积分

2.1 公式看起来很美,算起来直接崩溃

互信息的经典定义是:

I(X; Y) = ∫∫ p(x, y) log[ p(x, y) / (p(x)p(y)) ] dx dy

或者写成信息熵的形式:

I(X; Y) = H(X) - H(X|Y)

意思是:知道了 Y 之后,X 的不确定性减少了多少。这个定义在信息论课堂上干净利落,可一旦放到真实数据上,立刻会遇到三个致命问题。

第一个问题:我们根本不知道联合密度 p(x, y)。我们手里只有一批观测样本,比如一堆图片和对应的文本描述,顶多能估计样本出现的频率,但拿不到背后的真实概率密度函数。

第二个问题:即使想用直方图之类的离散化方法去近似,维度一高就立刻失效。假设每个维度分成 10 个格子,100 维的特征空间就有 10^100 个格子。现实中我们的图像特征往往是 512 维、1024 维,离散化这条路基本是断的。

第三个问题:就算我们硬着头皮去做高维积分,蒙特卡洛积分的收敛速度也撑不住。高维空间里的样本点分布极其稀疏,积分误差会随维度增长急剧恶化。

2.2 “先估计密度,再算 KL”路线为什么走不通

一个很自然的替代思路是:先用核密度估计(KDE)之类的非参数方法,把 p(x, y)、p(x)、p(y) 都估计出来,然后再数值积分算 KL 散度。

理论上是可行的,但工程量劝退。KDE 在高维空间里有一个众所周知的毛病:样本需求随维度指数增长。在 2 维平面上,几百个点就能画出大致密度轮廓;到了 20 维,就需要天文数字级别的样本才能维持同样的估计质量。而我们的视觉特征、文本 embedding 往往上百维,KDE 估计出来的密度本身噪声就很大。

更麻烦的是,KL 散度对密度估计误差极其敏感。密度比的 log 是非线性变换,密度估计的微小偏差经过 log 之后可能被放大成很大的误差。我在一个模拟实验里体会过:二维高斯分布下,用 KDE 估计再算 KL,和解析解差了 30% 以上;维度再高一点,误差直接失控。所以这条路在真实任务里基本没人走。

2.3 换个思路:把“估计”改成“优化”

既然精确计算走不通,那能不能不追求精确值,而是去找一个可控的目标函数,先优化起来再说?

这恰恰就是变分界的思想:互信息可以表示成某个泛函的最优值,我们用参数化函数族去逼近这个最优值,逼近得好,得到的就是一个足够贴近真实值的“界”。代价是我们可能拿不到精确的数值,拿到的只是一个下界或上界。但在大部分深度学习场景里,我们本来就不需要精确值——我们需要的是一个方向正确、可以反向传播、还算紧的优化目标。

这个思路的转换,是我觉得理解变分互信息界最关键的一步:忘掉“我要算出一个数”,接受“我要训练出一个可靠的界”。

3. 变分界的思想内核:用优化代替估计

3.1 KL 散度的变分表示是起点

要理解变分互信息界,得先看一个经典的事实:KL 散度存在一个变分表示。

简单说,对于任意两个分布 p 和 q,KL(p||q) 可以写成某个关于函数 T 的目标函数的上确界:

KL(p||q) = sup_T E_p[T(x)] - log E_q[exp(T(x))]

这个公式的含义是:如果你构造一个打分函数 T,并且把E_p[T] - log E_q[e^T]最大化,那么最大值就等于真实的 KL 散度。

但注意,这里有几个“取决于你怎么做”的关键点。首先,T 是任意函数,我们不可能真的搜索所有函数;其次,实践中我们通常用一个神经网络 T_θ 来参数化 T,网络容量有限,大概率取不到最优的 T。于是我们得到的不再是等式,而是一个下界:

KL(p||q) ≥ E_p[T_θ] - log E_q[exp(T_θ)]

把 p 换成联合分布 p(x, y),把 q 换成边缘乘积 p(x)p(y),就有了最经典的互信息神经估计下界。MINE 就是直接走这条路线的代表。

3.2 下界和上界的分工完全不同

理解了“参数化导致下界”,就很容易看清一个工程上的关键问题:不同的任务需要不同方向的界。

如果目标函数里需要“最大化互信息”,比如对比学习、多模态对齐,你应该用一个下界来驱动优化。因为下界在训练中被不断抬高,说明真实互信息也至少在被抬高,方向是安全的。

反过来,如果任务需要“最小化互信息”,比如特征压缩、信息瓶颈、去偏,那你需要的是上界。上界被压低,真实互信息才会被压到更低水平。如果你在压缩任务里错误地用一个下界去优化,结果会非常荒谬:下界被不断抬高,你以为你在压缩,实际上模型在偷偷保留更多信息。

我用一个生活化的比喻来记这个事:下界像一把往井里探深度的尺子,尺子不够长只能告诉你“井至少这么深”;上界像一把量身高的卡尺,卡尺放高了只能告诉你“身高最多这么高”。你需要哪个方向,取决于你要“证明”什么。

3.3 为什么神经网络可以担当 T 这个角色

很多人第一次看到变分表示时都会问:凭什么神经网络可以逼近这个 T?

道理其实不复杂。理论上,只要 T 属于足够大的函数空间,它就能逼近任意可测函数。神经网络在足够宽、足够深时,在逼近能力上是相当强的。更重要的是,神经网络是可微的,我们可以把E_p[T_θ] - log E_q[e^{T_θ}]这个目标用 SGD 端到端地最大化。这正好满足深度学习训练的需求。

不过,“足够强”只是理论上的保证,实际操作中,网络容量、训练策略、批大小都会直接决定这个界到底紧不紧。这也为后面的各种坑埋下了伏笔。

4. 三种主流界的拆解与对比:InfoNCE、MINE、CLUB

4.1 InfoNCE:披着分类外壳的互信息下界

InfoNCE 是这几年自监督学习里最常用的损失之一。它的训练逻辑是:给定一个 anchor x,我们有 1 个正样本 y_pos 和 K-1 个负样本 y_neg,模型要正确地把正样本从这 K 个候选中挑出来。

用公式来表达,就是:

L_InfoNCE = E[ log( f(x, y_pos) / ( (1/K) * Σ_j f(x, y_j) ) ) ]

其中 j 遍历一个 batch 里的 K 个样本。这里的 f 通常是一个带温度系数的相似度打分器。

为什么它是互信息下界?关键在“最优分类器”的分析:当 softmax 分类器训练到最优时,它输出的概率刚好是和密度比成正比的形式,即 f(x, y) ∝ p(x, y)/(p(x)p(y))。此时 L_InfoNCE 的期望就等于E log f = I(X; Y)。因为神经网络不一定能达到这个最优,所以实际得到的期望值是互信息的一个下界。

还有一个很重要的不等式,很多做工程的人容易忽略:

I(X; Y) ≥ E[L_InfoNCE],且I(X; Y) ≤ E[L_InfoNCE] + log K

这告诉你:InfoNCE 的紧度直接受 K 控制。K 越小,这个下界离真实互信息越远;K 越大,下界越紧。这也是为什么对比学习普遍要求大 batch,因为 batch 里的样本数就是公式里的 K。

实践里一个典型的 PyTorch 实现长这样:

import torch import torch.nn.functional as F def info_nce_loss(q, k, tau=0.07): # q: [B, D] 锚点表征 # k: [B, D] 正样本表征,约定第 i 行与 q 的第 i 行构成正对 logits = torch.matmul(q, k.T) / tau # [B, B] labels = torch.arange(q.size(0), device=q.device) return F.cross_entropy(logits, labels)

这个实现就是用第 i 行和它自己对应的第 i 行当成正样本,其余均为负样本。温度 tau 会直接影响训练难度,这个我在第五章展开讲。

4.2 MINE:直接优化变分表示的神经估计器

MINE(互信息神经估计)走的是另一条路,它不去构造分类任务,而是直接优化上一章提到的变分下界:

I_MINE = E_{p(x,y)}[T_θ(x, y)] - log E_{p(x)p(y)}[exp(T_θ(x, y))]

第一项用联合分布下的正样本对来估计,第二项用边缘乘积分布下的负样本对来估计。在实现上,负样本对通常通过对 batch 内样本做随机 shuffle 来构造,让 x 和 y 来自不同的样本。

MINE 的特点是非常灵活,T 可以是任意结构的神经网络,不像 InfoNCE 那样要求打分函数具备“密度比”含义。但它的缺点也很明显:目标函数里包含 exp(T),T 的取值稍微大一点,指数就爆炸;而且第二项期望的梯度方差通常很大,SGD 训练时会剧烈抖动。

所以在实现 MINE 时,常见的做法是用指数滑动平均(EMA)来维护E_q[e^T]的基线,避免每个 batch 只靠当前样本估计带来的大方差。否则你会在训练日志里看到一个来回横跳、根本降不下来的曲线。

4.3 CLUB:当我们需要一个真正的上界

前面两个都是下界,适合“最大化”。那“最小化互信息”的任务怎么办?这时候就要请出上界了。

CLUB 的上界构造逻辑非常漂亮,它不绕弯子,直接从条件分布入手:

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

这个等式在条件分布 p(y|x) 已知时严格成立。实际中我们并不知道真实条件分布,所以用一个参数化的变分条件分布 q_ψ(y|x) 去替换:

I_CLUB = E_{p(x,y)}[ log q_ψ(y|x) ] - E_{p(x)p(y)}[ log q_ψ(y|x) ]

如果 q 恰好等于真实条件分布,那么 I_CLUB 就是真值的上界;q 拟合得越好,这个上界越紧。所以在实操里,使用 CLUB 通常需要先训练一个条件分布模型,再去估计互信息或做正则化。

这带来的隐患也很明显:如果 q 拟合得很差,I_CLUB 可能不再严格是上界。这一点我在第五章会专门讲。

4.4 一张表对比三种界

界方向核心目标函数(示意)主要用途常见风险
InfoNCE下界softmax 分类损失对比学习、互信息最大化需要大量负样本,K 决定紧度
MINE下界E_p[T] - log E_q[e^T]小规模互信息估计、表征学习梯度方差大,exp 易溢出
CLUB上界E_p[log q(yx)] - E_q[log q(yx)]

选界时我首先看方向:最大化选 InfoNCE 或 MINE,最小化选 CLUB。然后再看实现成本和稳定性,后面我会给一个更具体的流程。

5. 训练变分界时反复踩过的坑

5.1 温度系数:调不好,下界比没有还难受

InfoNCE 里的温度系数 τ 是我第一个踩的坑。表面上看,它只是在控制 softmax 的陡峭程度,实际上它决定了训练过程中“困难负样本”的权重。

τ 很小(比如 0.01)时,logits 被放大,softmax 会变成接近 one-hot 的输出。模型几乎只会关注最难区分的那几个负样本,梯度非常尖锐,batch 稍小一点就会训练崩溃。τ 很大(比如 1.0)时,所有负样本的得分被压平,模型对负样本的区分度大幅下降,学出来的表征会变得含糊。

我自己的经验是,τ 放在 0.05 到 0.2 之间通常比较稳。具体值跟特征维度、数据难度都有关系,最好用一个小验证集去扫几轮,而不是把它当成默认参数永远不调。还有一个额外收获:τ 除了影响训练难度,也会影响你最终这个下界的“紧度”,因为打分器 f 的内部就带着 τ,τ 变了最优密度比估计的尺度也会变。

5.2 负样本数量直接决定 InfoNCE 的紧度

前面提到过一个不等式:I ≤ E[L_InfoNCE] + log K。这个 log K 不是个小数字。K=64 时,log K 大约是 4.16 nats,折算成 bit 大约是 6 bit。也就是说,你的下界可能离真实互信息有 6 bit 的缺口,而且这个缺口是内禀存在的,不是网络训练不够努力造成的。

在我做过的某图像检索项目里,把 batch size 从 256 提到 1024,下游任务准确率涨了不少,当时的解释就是这个:更大的 batch 带来了更大的 K,InfoNCE 下界更紧,模型优化目标更接近真实互信息最大化。

如果显存撑不住大 batch,可以用 memory bank 或者队列结构,把之前迭代的特征缓存下来当作额外的负样本。注意要对缓存特征做动量更新,否则旧特征太陈旧,会引入很大的估计偏差。

5.3 MINE 的梯度方差问题:logS 的陷阱

用 MINE 做小规模互信息估计时,我最头疼的不是网络结构,而是目标函数的梯度方差。

MINE 的目标里有- log E_q[e^T]这一项。这个期望通常用当前 batch 的负样本来估计。问题在于:log 是凹函数,而 e^T 又非常敏感,二者叠加之后梯度会剧烈波动。更隐蔽的是,理论上这个梯度是有偏的,因为你在同一个 batch 里同时估计 T 的取值和 E_q[e^T] 的期望,两者互相干扰。

常用的缓解手段有两个。第一,用 EMA 维护E_q[e^T]的基线,而不是每个 batch 重新估计。第二,把目标函数的第二项换成“每条样本单独计算后取平均”的形式,避免在 exp 之后才做平均。具体到代码里,前者实现简单,效果立竿见影;后者能在一定程度上降低方差,但也要注意数值溢出,exp 之前要把 T 减去当前 batch 的最大值。

5.4 上界失效的隐蔽场景:q 没学好,界就不严

CLUB 这类上界在实际使用时有一个隐蔽陷阱:它要求变分条件分布 q(y|x) 和真实条件分布足够接近。

如果 q 只是一个粗糙的高斯模型,而真实 p(y|x) 是多模态分布(比如输入一张图,可能对应多种合理文本描述),那么 q 对真实分布的覆盖会不足,I_CLUB 可能会偏离上界的方向。我见过一个场景:在一份公开风格的数据集上,CLUB 估计出的“上界”竟然比 InfoNCE 估计出的“下界”还低,这明显是不自洽的。当时排查的结果就是条件分布模型容量不够,把多模态输出强行拟合成了单峰高斯。

所以在决定信任 CLUB 之前,最好先在一个能算真实互信息的合成数据上做一次标定实验。比如把真实条件分布设成已知的混合高斯,看看估计出的 CLUB 是否还能保持上界性质。如果连合成数据上都保不住方向,那直接用到真实任务里就是把错误放大。

5.5 一个容易被误读的现象:下界在涨,不代表互信息真的在涨

最后再提醒一个更隐蔽的坑:训练过程中下界上升,并不一定等于真实互信息上升。

因为变分下界的值同时受两个因素影响:一个是真实互信息本身,另一个是参数化函数族 T_θ 的逼近能力。如果你在训练过程中同时提高神经网络的容量或者增加训练步数,T 对最优函数的逼近程度变好了,下界可能上升,但真实互信息纹丝未动。这种“界本身变紧”会伪装成“互信息变大”,在看训练曲线时很容易被误导。

我的做法是分开观察两类指标:一类是下界的绝对值,另一类是下游任务的表现。如果前者上升而后者毫无变化,多半是函数族在自我膨胀,而不是学习到了更有用的表征。

6. 选界判断流程与我的个人习惯

6.1 先回答三个问题

每次接到一个和互信息相关的训练任务,我都会先回答三个问题,答案基本决定选哪个界。

第一,我是要最大化互信息,还是要最小化它?最大化选下界,Minimize 选上界。第二,我的 batch 或负样本规模能否做得很大?如果只能凑出很小的 batch,InfoNCE 的紧度会很差,这时可以考虑 MINE 这种不依赖 K 的下界,但要注意方差问题。第三,我是要一个用于优化的目标,还是要一个用于报告/分析的数值?如果是报告数值,最好同时跑一个下界和一个上界,输出一个区间,比单点估计更有说服力。

判断流程可以浓缩成一张小表:

我的场景第一选择备选方案特别注意
对比学习 / 自监督预训练InfoNCE 大 batchInfoNCE + memory bank温度系数需要调
多模态对齐InfoNCEMINE模态差异会导致温度敏感
信息压缩 / 瓶颈CLUB 上界变分上界正则项条件分布模型不能太弱
互信息数值分析同时跑 InfoNCE 和 CLUB交叉验证报告区间而非单点

6.2 一个可落地的小验证流程

我一般不会直接信任何一条估计曲线,而是先花半天时间做一次标定实验。具体步骤大概是:

第一步,造一个真实互信息已知的合成数据集,比如两个多维高斯变量,它们的相关系数决定真实互信息,可以直接用解析解。第二步,分别用 InfoNCE 和 CLUB 去估计,看看估计值和真实值之间的方向对不对、误差有多大。第三步,把合成实验里发现的问题(比如 q 容量不够、温度不合适)在真实任务里对应修正,然后再把估计值写进训练日志,当成一个诊断信号。

这个小流程看起来费时间,但它能避免你在真实任务里跑了一周之后才发现目标函数的方向根本不对。我犯过的最贵的错误,就是把应该用上界的压缩任务错配成下界,白白浪费了两周训练时间,损失曲线倒是很漂亮,下游效果一路下滑。后来所有和互信息沾边的任务,我都会先做这个标定。

6.3 我自己的几个习惯

最后分享几个已经固化成习惯的做法。

第一个习惯,是尽量用模型自带的 memory bank 或动量队列来扩大负样本规模。InfoNCE 的 K 就是生命线,大 K 带来的收益通常比换网络结构还明显。第二个习惯,是给温度系数做一个小的退火规划,训练初期用略大的 τ 保证稳定,后期调小 τ 让模型更严格地区分困难负样本。第三个习惯,是同时记录下界和上界。在报告互信息数值时,我会说“真实互信息大约在 [下界, 上界] 这个区间里”,而不是扔一个虚无缥缈的单点估计。

还有一个小技巧是我最近才养成的:用不同 K 的 InfoNCE 做自洽性检验。如果 K 从 256 翻倍到 512,infoNCE 估计值变化不大,说明这个界基本收敛了;如果估计值随 K 快速增长,说明你的负样本池远远不够,这时候去调别的超参都是白费力气。

这就回到了开头那句话:变分界不是什么高深的数学装饰品,它是一套告诉你“目标到底往哪个方向优化”的纪律。理解了它,你再看自己训练的 loss 曲线,很多迷惑就散了。

RELATED READING

延伸阅读

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