InfoNCE

PyTorch 实现了来自 "Representation Learning with Contrastive Predictive Coding" <https://arxiv.org/abs/1807.03748>__ 的 InfoNCE 损失函数。在对比学习中，我们希望学习如何将高维数据映射到低维嵌入空间。这种映射应使语义相似的样本在嵌入空间中彼此靠近，而语义不同的样本则相距较远。InfoNCE 损失函数可用于实现对比学习的目的。

该包已在 PyPI 上发布，地址为：info-nce-pytorch <https://pypi.org/project/info-nce-pytorch/>__，可通过以下命令安装：

.. code::

pip install info-nce-pytorch

使用示例

首先导入本包。

.. code:: python

from info_nce import InfoNCE, info_nce

可以在不显式提供负样本键的情况下使用，此时每个样本都会与批次中的其他样本进行比较。

.. code:: python

loss = InfoNCE()
batch_size, embedding_size = 32, 128
query = torch.randn(batch_size, embedding_size)
positive_key = torch.randn(batch_size, embedding_size)
output = loss(query, positive_key)

也可以使用负样本键，此时查询样本与每个负样本键的所有组合都会被比较。

.. code:: python

loss = InfoNCE(negative_mode='unpaired') # negative_mode='unpaired' 是默认值
batch_size, num_negative, embedding_size = 32, 48, 128
query = torch.randn(batch_size, embedding_size)
positive_key = torch.randn(batch_size, embedding_size)
negative_keys = torch.randn(num_negative, embedding_size)
output = loss(query, positive_key, negative_keys)

还可以使用负样本键，但此时每个查询样本仅与其配对的负样本键进行比较。

.. code:: python

loss = InfoNCE(negative_mode='paired')
batch_size, num_negative, embedding_size = 32, 6, 128
query = torch.randn(batch_size, embedding_size)
positive_key = torch.randn(batch_size, embedding_size)
negative_keys = torch.randn(batch_size, num_negative, embedding_size)
output = loss(query, positive_key, negative_keys)

损失函数图

假设我们有初始的均值向量 µ_q、µ_p、µ_n，以及协方差矩阵 Σ = I/10，那么我们可以通过从具有插值均值向量的分布中采样来绘制 InfoNCE 损失的值。给定插值权重 α 和 β，我们定义查询样本的分布 Q ~ N(µ_q, Σ)，正样本的分布 P_α ~ N(αµ_q + (1-α)µ_p, Σ)，以及负样本的分布 N_β ~ N(βµ_q + (1-β)µ_n, Σ)。下图展示了在不同 α 和 β 值下，从上述分布中采样得到的输入所对应的损失值。

.. image:: https://raw.githubusercontent.com/RElbers/info-nce-pytorch/main/imgs/loss.png

info-nce-pytorch 快速上手指南

环境准备

• 系统要求：支持 Linux、macOS 或 Windows。
• 前置依赖：
  • Python 3.6+
  • PyTorch (建议安装最新稳定版)
  • pip 包管理工具

安装步骤

推荐使用国内镜像源加速安装（如清华大学开源软件镜像站）：

pip install info-nce-pytorch -i https://pypi.tuna.tsinghua.edu.cn/simple

或者使用官方源安装：

pip install info-nce-pytorch

基本使用

以下是最简单的用法示例：将批次内的每个样本与其他样本进行对比（无需显式提供负样本键）。

import torch
from info_nce import InfoNCE

# 初始化损失函数
loss_fn = InfoNCE()

# 模拟数据：批次大小 32，嵌入维度 128
batch_size, embedding_size = 32, 128
query = torch.randn(batch_size, embedding_size)
positive_key = torch.randn(batch_size, embedding_size)

# 计算 InfoNCE Loss
output = loss_fn(query, positive_key)

print(f"Loss value: {output.item()}")

该模式会自动将当前 batch 中的其他样本视为负样本，适用于大多数标准的对比学习场景。