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 和 qKL(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、CLUB4.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, tau0.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), deviceq.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 不是个小数字。K64 时log K 大约是 4.16 nats折算成 bit 大约是 6 bit。也就是说你的下界可能离真实互信息有 6 bit 的缺口而且这个缺口是内禀存在的不是网络训练不够努力造成的。在我做过的某图像检索项目里把 batch size 从 256 提到 1024下游任务准确率涨了不少当时的解释就是这个更大的 batch 带来了更大的 KInfoNCE 下界更紧模型优化目标更接近真实互信息最大化。如果显存撑不住大 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 或负样本规模能否做得很大如果只能凑出很小的 batchInfoNCE 的紧度会很差这时可以考虑 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 翻倍到 512infoNCE 估计值变化不大说明这个界基本收敛了如果估计值随 K 快速增长说明你的负样本池远远不够这时候去调别的超参都是白费力气。这就回到了开头那句话变分界不是什么高深的数学装饰品它是一套告诉你“目标到底往哪个方向优化”的纪律。理解了它你再看自己训练的 loss 曲线很多迷惑就散了。