InfoNCE損失とは — NCEからの系譜と相互情報量の下界を導出で理解する

SimCLRもCLIPもMoCoも、自己教師あり学習の主要手法はみな同じひとつの損失関数を最小化しています。それが InfoNCE損失 です。「正例を当てる分類問題」という見かけの裏に、「相互情報量を最大化する」という情報理論の目標が隠れている——この二面性こそがInfoNCEを学ぶ価値です。

たとえば次の2つの場面を考えてみましょう。

  • 画像の事前学習(SimCLR): 1枚の画像から2通りの拡張ビューを作り、「この2つは同じ画像由来だ」と当てさせます。ラベルは一切使いません。
  • 画像とテキストの結びつけ(CLIP): 大量の (画像, キャプション) のペアから、「この画像に合うキャプションはどれか」を当てさせます。これだけでゼロショット分類ができる表現が育ちます。

どちらも「正しい相方(正例)を、たくさんのダミー(負例)の中から見分ける」という同じ形をしています。この識別タスクを支える損失がInfoNCEです。本記事では、対照学習の入門記事では「損失の一部」として軽く触れられがちなInfoNCEそのものを主役に据え、

  1. なぜ NCE(Noise Contrastive Estimation) という「データ対ノイズの二値分類」が生まれたのか
  2. InfoNCEが $(K{+}1)$クラスのsoftmax交差エントロピー と等価であること
  3. InfoNCEが 相互情報量の下界 $I(X;C) \geq \log N – L_{\text{InfoNCE}}$ を最大化していること(1行ずつ導出)
  4. 温度 $\tau$負例数 $N$ が下界と勾配にどう効くか

を、導出とPyTorch実験で徹底的に掘り下げます。

画像なし
対照学習(Contrastive Learning)の理論と損失関数の導出
対照学習の基本原理とInfoNCE損失の概要。本記事の出発点になる記事です。
画像なし
対照学習の理論と実装(SimCLR)
InfoNCEを使ったSimCLRのアーキテクチャとPyTorch実装。

前提知識

この記事を読む前に、以下を押さえておくと理解がスムーズです。

対照学習そのものの問題設定(正例ペア・負例ペア・データ拡張)は上の2記事で詳しく扱っています。本記事はその「損失関数の中身」に絞って深掘りします。

そもそもの困りごと: 正規化定数が計算できない

InfoNCEの出発点を理解するには、確率モデルの学習でよく出てくる「正規化定数の壁」を知る必要があります。直感から入りましょう。

データの分布 $p_{\text{data}}(x)$ を、パラメータ $\theta$ を持つモデルで近似したいとします。多くのモデルは次の形をしています。

$$ p_\theta(x) = \frac{1}{Z(\theta)} \, \tilde{p}_\theta(x), \qquad Z(\theta) = \int \tilde{p}_\theta(x)\, dx $$

ここで $\tilde{p}_\theta(x)$ は「スコア」や「エネルギー」のような 正規化前の量 で、計算は簡単です。問題は分母の $Z(\theta)$ です。これは全空間にわたる積分(または巨大な語彙にわたる和)で、高次元では事実上計算できません。最尤推定 $\max_\theta \sum_i \log p_\theta(x_i)$ をやろうとすると、この $Z(\theta)$ が必ず顔を出して詰まります。

ここで発想を変えます。「$p_\theta(x)$ の値そのもの」を求めるのではなく、目の前の点が『本物のデータ』なのか『でっち上げのノイズ』なのかを当てる分類器 を作ればよいのではないか——これが次に見るNCEのアイデアです。

NCE(Noise Contrastive Estimation): 推定を分類に置き換える

直感: 偽物を混ぜて見分けさせる

贋作を見抜く目利きを育てたいとします。本物の絵だけを延々と見せても「本物らしさ」の基準は育ちません。本物と精巧な贋作を混ぜて「どちらが本物か」を当てさせ続ける と、目利きは本物と偽物を分ける特徴を自然に学びます。

NCE(Gutmann & Hyvärinen, 2010)はこれを確率モデルの学習に持ち込みました。自分で用意した既知の ノイズ分布 $q(x)$(一様分布やガウス分布など、サンプルが簡単に作れて密度も計算できるもの)から偽サンプルを生成し、本物のデータと混ぜて「データかノイズか」の二値分類を解かせます。

NCEのデータ対ノイズ二値分類の概念図

上の図のように、緑が本物のデータ分布 $p_{\text{data}}$、赤が自分で用意したノイズ分布 $q$ です。分類器は「いま見ている点はどちらの壺から出てきたか」を判定します。重なり合う領域でうまく判定するには、分類器は本物の密度の形を内部で表現せざるを得ません。ここに学習が宿ります。

式: 二値分類のロジットが対数密度比になる

データを 1 個、ノイズを $\nu$ 個混ぜた集合を考えます。ある点 $x$ が「データ(ラベル $C=1$)」である事前確率は $\frac{1}{1+\nu}$、「ノイズ($C=0$)」である事前確率は $\frac{\nu}{1+\nu}$ です。ベイズの定理で、点 $x$ がデータである事後確率を書き下します。

$$ P(C=1 \mid x) = \frac{\tfrac{1}{1+\nu}\, p_\theta(x)}{\tfrac{1}{1+\nu}\, p_\theta(x) + \tfrac{\nu}{1+\nu}\, q(x)} $$

分子分母を $\frac{1}{1+\nu}$ で約分して整理すると、

$$ P(C=1 \mid x) = \frac{p_\theta(x)}{p_\theta(x) + \nu\, q(x)} $$

となります。ここがNCEの肝です。この事後確率をロジスティック関数 $\sigma(u)=1/(1+e^{-u})$ の形に書き直すと、

$$ P(C=1 \mid x) = \sigma\!\left( \log \frac{p_\theta(x)}{q(x)} – \log \nu \right) $$

つまり、二値分類器のロジット(決定境界からの距離)が対数密度比 $\log \frac{p_\theta(x)}{q(x)}$ そのもの になります。分類がうまくできるようにロジットを学習すれば、密度比、ひいてはデータの密度の形が手に入る、というわけです。

NCEのうれしい点は、この分類の損失(ロジスティック回帰の交差エントロピー)には 正規化定数 $Z(\theta)$ を陽に計算する必要がない ことです。$Z$ を $\theta$ から切り離して「もう1つの学習可能な定数」とみなしても、分類問題として整合的に学習できることがGutmann & Hyvärinenによって示されています。困りごとだった積分が、分類器の学習に化けたのです。

ここまでで「推定を分類に置き換える」という発想を手に入れました。InfoNCEは、この二値分類を 多値分類(1個の正例 vs 複数の負例) に拡張し、さらに表現学習に最適化した形だと考えると、すんなり理解できます。

InfoNCE損失の定義

直感: 正しい相方を当てる多肢選択クイズ

NCEは「データかノイズか」の2択でした。InfoNCEは選択肢を増やします。1個の正例と $K$ 個の負例、合わせて $K{+}1$ 個の候補の中から、正例がどれかを当てる多肢選択クイズ だと思ってください。

アンカー(基準)$x$ に対して、本来の相方が正例 $x^+$、無関係な候補が負例 $x_1, \dots, x_K$ です。エンコーダ $f$ で全部をベクトルに直し、アンカーとの類似度(内積)をスコアにします。正解(正例)のスコアが一番高くなるように学習する——これがInfoNCEの全体像です。

定義式

$N = K{+}1$ 個の候補のうち1つが正例であるとき、InfoNCE損失は次のように定義されます。

$$ \begin{equation} L_{\text{InfoNCE}} = -\,\mathbb{E}\!\left[\, \log \frac{\exp\!\big(f(x)^{\top} f(x^+)/\tau\big)}{\displaystyle\sum_{j=1}^{N} \exp\!\big(f(x)^{\top} f(x_j)/\tau\big)} \,\right] \end{equation} $$

ここで、

  • $f(\cdot)$ はエンコーダ(多くの場合、出力を $\|f(\cdot)\|=1$ に正規化)
  • 分子は 正例 とのスコア、分母は正例を含む 全候補 のスコアの和($x_j$ の和は正例 $x^+$ を含む)
  • $\tau > 0$ は 温度パラメータ

です。スコアに内積 $f(x)^{\top} f(x_j)$ を使い、ベクトルを正規化していればこれはコサイン類似度になります。

$(K{+}1)$クラスのsoftmax交差エントロピーとの等価性

定義式 (1) の中身は、見覚えのある形をしています。$C$ クラス分類のsoftmax交差エントロピーは、正解クラス $y$ に対して

$$ L_{\text{CE}} = -\log \frac{\exp(\ell_y)}{\sum_{c=1}^{C} \exp(\ell_c)} $$

でした($\ell_c$ はクラス $c$ のロジット)。両者を見比べると、対応関係は明らかです。

  • ロジット $\ell_j \;\longleftrightarrow\; f(x)^{\top} f(x_j)/\tau$(アンカーと候補 $j$ の類似度を温度でスケール)
  • クラス数 $C \;\longleftrightarrow\; N = K{+}1$(候補の数)
  • 正解クラス $y \;\longleftrightarrow\;$ 正例 $x^+$ のインデックス

つまり InfoNCEは「$K{+}1$クラスのsoftmax交差エントロピー」そのもの です。「候補のうち正例はどれか」という分類確率を

$$ p(\text{正例}=x^+ \mid x) = \frac{\exp\!\big(f(x)^{\top} f(x^+)/\tau\big)}{\sum_{j=1}^{N} \exp\!\big(f(x)^{\top} f(x_j)/\tau\big)} $$

と置き、その対数尤度を最大化していると読めます。

InfoNCEをK+1クラスのsoftmax分類として見る図

図は、6個の候補(正例1個+負例5個)に対するsoftmax確率の例です。正例(緑)のスコアを高く、負例(灰)のスコアを低く押し下げるほど、正例の確率が1に近づき損失が下がります。候補数 $N$ が増えるほどクイズの選択肢が増え、難しい分類になる——この「難しさ」が後で効いてきます。

この「分類としての顔」はSimCLRやCLIPの実装(F.cross_entropy をそのまま使う)と直結します。では、なぜこの単純な分類が「良い表現」を生むのでしょうか。その答えが次の 相互情報量の下界 です。

相互情報量の下界としてのInfoNCE

ゴールの宣言

ここでは、Oord et al. (2018, “Representation Learning with Contrastive Predictive Coding”) に従って、次の不等式を導出します。

$$ \begin{equation} I(X; C) \;\geq\; \log N – L_{\text{InfoNCE}} \end{equation} $$

$I(X;C)$ はアンカー(文脈)$C$ と正例 $X$ の 相互情報量、$N$ は候補数です。この式が言いたいのは「InfoNCEを小さくすると、相互情報量の下界が押し上がる」ということ。導出の途中で、最適なクリティック(スコア関数)が密度比 $p(x \mid c)/p(x)$ に比例する という重要な事実も出てきます。1行ずつ進みます。

ステップ1: 問題を「正例当て」の確率として書く

候補集合 $X = \{x_1, \dots, x_N\}$ を考えます。このうち1つは文脈 $c$ と整合する正例で、条件付き分布 $p(x \mid c)$ から引かれます。残り $N{-}1$ 個は文脈と無関係な負例で、周辺分布 $p(x)$ から独立に引かれます。

「候補 $i$ こそが正例である」という事後確率を、ベイズの定理で求めます。$N$ 個のうちどれが正例かは事前に等確率とすると、

$$ p(\text{正例}=i \mid X, c) = \frac{p(x_i \mid c)\displaystyle\prod_{l \neq i} p(x_l)}{\displaystyle\sum_{j=1}^{N} \left[ p(x_j \mid c)\prod_{l \neq j} p(x_l)\right]} $$

と書けます。分子は「$x_i$ が正例($p(x_i\mid c)$ から)で、残りが負例($p(x_l)$ から)」という同時確率です。

ステップ2: 約分して密度比だけを残す

分子・分母に共通して $\prod_{l} p(x_l)$ という因子が(ほぼ)含まれます。各項を $\prod_{l} p(x_l)$ で割って整理しましょう。たとえば分子は

$$ p(x_i \mid c)\prod_{l \neq i} p(x_l) = \Big[\textstyle\prod_{l} p(x_l)\Big]\cdot \frac{p(x_i \mid c)}{p(x_i)} $$

と変形できます($x_i$ の分だけ $p(x_i \mid c)$ にして、その代わり $p(x_i)$ で割って帳尻を合わせた)。分母の各項も同様です。共通因子 $\prod_l p(x_l)$ が分子分母でキャンセルして、

$$ \begin{equation} p(\text{正例}=i \mid X, c) = \frac{\dfrac{p(x_i \mid c)}{p(x_i)}}{\displaystyle\sum_{j=1}^{N} \dfrac{p(x_j \mid c)}{p(x_j)}} \end{equation} $$

が得られます。これは美しい結果です。正解を当てる事後確率は、密度比 $\frac{p(x \mid c)}{p(x)}$ をスコアとしたsoftmaxの形 をしています。

ステップ3: 最適クリティックは密度比に比例する

式 (3) を、InfoNCEのsoftmax(式の分子分母)と見比べます。InfoNCEのスコアは $\exp\!\big(f(x)^{\top}f(c)/\tau\big)$ という学習可能なクリティック(critic)$h(x,c)$ です。式 (3) と一致するのは、クリティックが

$$ \begin{equation} h^\star(x, c) \;\propto\; \frac{p(x \mid c)}{p(x)} \end{equation} $$

を満たすときです。つまり InfoNCEを最小化する最適なスコア関数は、密度比 $p(x\mid c)/p(x)$ に比例する。比例係数(正例当てのsoftmaxでは消える定数)は学習で吸収されます。

最適クリティックが密度比 p(x|c)/p(x) に比例する図

図のように、文脈と整合する分布 $p(x\mid c)$(緑)が、無関係な分布 $p(x)$(赤破線)より相対的に大きい領域でスコア(青)が高くなります。クリティックが密度比を学ぶ、というのはまさにこの「文脈と整合する点を、無関係な背景に対して相対的に高く評価する」ことに対応します。NCEで見た「ロジット=対数密度比」が、多値版でもそのまま生きていることに注目してください。

ステップ4: 損失を相互情報量に結びつける

最適クリティック (4) を、InfoNCE損失 (1) の定義に代入します。$h^\star(x_i,c) = \frac{p(x_i\mid c)}{p(x_i)}$ とおくと、正例(インデックスを便宜上 $i{=}1$ とする)に対する損失は

$$ L_{\text{InfoNCE}}^\star = -\,\mathbb{E}\!\left[\, \log \frac{\dfrac{p(x_1 \mid c)}{p(x_1)}}{\displaystyle\sum_{j=1}^{N} \dfrac{p(x_j \mid c)}{p(x_j)}} \,\right] $$

です。分母の和を「正例の項 + 負例の項」に分けます。

$$ L_{\text{InfoNCE}}^\star = \mathbb{E}\!\left[\, \log\!\left( 1 + \frac{p(x_1)}{p(x_1 \mid c)} \sum_{j=2}^{N} \frac{p(x_j \mid c)}{p(x_j)} \right) \right] $$

(分子分母を正例の密度比で割りました。)負例 $x_j$ は周辺分布 $p(x)$ から引かれているので、負例の項の期待値は $\mathbb{E}_{x_j \sim p(x)}\big[\frac{p(x_j\mid c)}{p(x_j)}\big] = \int p(x_j)\frac{p(x_j\mid c)}{p(x_j)}dx_j = \int p(x_j \mid c)\,dx_j = 1$ になります。負例は $N{-}1$ 個あるので、和はおおよそ $N{-}1$ に置き換わります(大数の法則)。

$$ L_{\text{InfoNCE}}^\star \;\approx\; \mathbb{E}\!\left[\, \log\!\left( 1 + \frac{p(x_1)}{p(x_1 \mid c)}\,(N-1) \right) \right] $$

中身を $\frac{p(x_1)}{p(x_1\mid c)}(N-1) \le \frac{p(x_1)}{p(x_1\mid c)}\,N$ で上から押さえ、さらに $1 + \frac{p(x_1)}{p(x_1\mid c)}N \ge \frac{p(x_1)}{p(x_1\mid c)}N$ を使うと($\log$ は単調増加なので不等号の向きは保たれます)、

$$ L_{\text{InfoNCE}}^\star \;\gtrsim\; \mathbb{E}\!\left[\, \log\!\left( \frac{p(x_1)}{p(x_1 \mid c)}\,N \right) \right] = \log N – \mathbb{E}\!\left[\, \log \frac{p(x_1 \mid c)}{p(x_1)} \right] $$

最後の期待値こそが、相互情報量の定義 $I(X;C) = \mathbb{E}\big[\log \frac{p(x\mid c)}{p(x)}\big]$ そのものです。代入して移項すると、

$$ \boxed{\; I(X; C) \;\geq\; \log N – L_{\text{InfoNCE}} \;} $$

が得られます。目標の式 (2) です。導出を振り返ると、(i) 正例当ての事後確率が密度比のsoftmaxになり、(ii) 最適クリティックが密度比に比例し、(iii) それを損失に戻すと相互情報量が顔を出す、という流れでした。InfoNCEを下げる=相互情報量の下界を上げる という二面性が、これで腑に落ちます。

相互情報量の下界 logN-L_InfoNCE と logN の頭打ち

この不等式の重要な含意が図に表れています。下界 $\log N – L$(緑)は、天井である $\log N$(青)を超えられません。つまり どれだけ学習しても、InfoNCEで捉えられる相互情報量は高々 $\log N$ nats です。相互情報量が大きいデータ($I \gg \log N$)を扱うには、$N$(=候補数、実質的には負例の数)を増やすしかありません。これが「対照学習は大バッチが効く」理由の理論的な核心です。

ここまでが理論の山場です。次に、実装上のつまみである 温度 $\tau$負例数 $N$ が、この描像とどう噛み合うかを見ます。

温度パラメータ $\tau$ の役割

直感: 採点の厳しさを決めるつまみ

温度 $\tau$ は、softmaxに入れる前にスコアを $1/\tau$ 倍するパラメータです。$\tau$ を小さくするとスコアの差が拡大され、softmaxが尖ります。$\tau$ を大きくすると差が圧縮され、分布が平らになります。「採点をどれだけ厳しくするか」のつまみだと思ってください。

温度τによるsoftmax分布の鋭さの変化

同じ類似度の並びでも、$\tau=0.05$(赤)では正例にほぼ全確率が集中し、最も紛らわしい候補だけが残ります。$\tau=1.0$(青)ではどの候補もそれなりの確率を持ち、区別が緩くなります。$\tau$ は「正例だけを際立たせるか、全候補をなだらかに見るか」を決めています。

ハード負例への重み付け

なぜ $\tau$ が表現の質に効くのか。鍵は 勾配 です。InfoNCE損失(1サンプル分)をスコア $s_j = f(x)^{\top}f(x_j)$ で微分すると、負例 $x_j$ に対する勾配の大きさは softmax 確率 $p_j = \frac{\exp(s_j/\tau)}{\sum_k \exp(s_k/\tau)}$ に比例します。

$\tau$ が小さいと、$p_j$ は 最も類似度の高い負例(ハード負例)に集中 します。つまり「正例と紛らわしい、際どい負例」を重点的に押し下げる学習になります。一方 $\tau$ が大きいと、すべての負例をほぼ均等に少しずつ押し下げます。

小さい $\tau$ はハード負例を鋭く罰するので強い表現が育ちやすい反面、勾配が一点に集中して不安定になりがちです(紛らわしい正例ペアを過剰に引き離す「均一性のペナルティ」が強くなりすぎる)。実務ではおおむね $\tau \in [0.05, 0.5]$ が使われ、SimCLRは $0.5$、CLIPは学習可能な $\tau$(初期値 $0.07$ 相当)を採用しています。

この「ハード負例への集中」を、実際に数値で確かめてみましょう。次のコードは固定された(学習していない)埋め込みに対して $\tau$ を振り、正例への確率集中と、負例側の「実効的に注目される個数」を測ります。

import torch
import torch.nn.functional as F
import numpy as np

def make_pairs(B, N, d, gap, seed):
    """B個のアンカーと、各々N個の候補(候補0が正例)を作る。
    正例はアンカー方向にgapだけ寄せて類似度を上げる。すべて単位ベクトル。"""
    g = torch.Generator().manual_seed(seed)
    anchor = F.normalize(torch.randn(B, d, generator=g), dim=1)
    cand = F.normalize(torch.randn(B, N, d, generator=g), dim=2)
    pos = F.normalize((1 - gap) * cand[:, 0, :] + gap * anchor, dim=1)
    cand[:, 0, :] = pos
    pos_idx = torch.zeros(B, dtype=torch.long)  # 正例は常にインデックス0
    return anchor, cand, pos_idx

B, N, d = 512, 64, 32
taus = [0.02, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0]
gap = 0.5

for tau in taus:
    anchor, cand, pos_idx = make_pairs(B, N, d, gap, seed=42)
    logits = torch.einsum("bd,bnd->bn", anchor, cand) / tau   # (B, N)
    p = F.softmax(logits, dim=1)
    posprob = p[torch.arange(B), pos_idx].mean().item()       # 正例への平均確率
    # 負例だけを正規化し、その分布の実効サイズ exp(エントロピー) を測る
    mask = torch.ones_like(p, dtype=torch.bool)
    mask[torch.arange(B), pos_idx] = False
    pneg = p[mask].reshape(B, N - 1)
    pneg = pneg / pneg.sum(dim=1, keepdim=True)
    ent = -(pneg * pneg.clamp_min(1e-12).log()).sum(dim=1)
    neg_eff = torch.exp(ent).mean().item()                    # 実効的な負例数
    print(f"tau={tau}: 正例確率={posprob:.4f}, 実効負例数={neg_eff:.2f}")

実行すると、次の出力が得られます。

tau=0.02: 正例確率=0.9971, 実効負例数=1.97
tau=0.05: 正例確率=0.9790, 実効負例数=5.79
tau=0.1: 正例確率=0.7882, 実効負例数=19.66
tau=0.2: 正例確率=0.2744, 実効負例数=43.80
tau=0.5: 正例確率=0.0580, 実効負例数=59.26
tau=1.0: 正例確率=0.0307, 実効負例数=62.04
tau=2.0: 正例確率=0.0220, 実効負例数=62.76

温度スイープ:正例への確率集中とハード負例への注目

数値とグラフから2点が読み取れます。第一に、$\tau$ を $2.0 \to 0.02$ と下げると、正例に割り当てる平均確率が $0.022 \to 0.997$ へ単調に上がります。$\tau$ が小さいほど確率が一点に集中するわけです。第二に、負例側の「実効的に注目される個数」は $\tau=2.0$ で約63(ほぼ全負例を均等に見ている)だったのが、$\tau=0.02$ では約2まで激減します。つまり $\tau$ が小さいほど、ごく少数のハード負例に勾配が集中 します。これが「小さい $\tau$ はハード負例を重視する」という主張の定量的な裏付けです。

温度が「採点の厳しさ」を決めるのに対し、次に見る負例数 $N$ は「捉えられる情報の天井」を決めます。

負例数 $N$ の依存性と大バッチの動機

$\log N$ の天井をどう上げるか

下界の式 (2) $I(X;C) \ge \log N – L$ をもう一度見ます。$N$ は「候補数」で、実質的には 負例の数 $+1$ です。下界の天井が $\log N$ である以上、

  • 負例が少ない($N$ が小さい)と、いくら学習しても捉えられる相互情報量は $\log N$ で頭打ち
  • 負例を増やす($N$ を大きくする)と、天井が上がり、より多くの相互情報量を捉えられる

ことになります。SimCLRが数千のバッチサイズを必要としたり、MoCoが メモリバンク(キュー) で過去の表現を負例として溜め込んだり、CLIPが数万規模のバッチで学習したりするのは、すべて「$\log N$ の天井を上げる」ための工夫です。

画像なし
CLIP: 画像とテキストの対照学習
画像・テキストペアにInfoNCEを適用し、巨大バッチでゼロショット表現を学習する。

固定埋め込みで $N$ を振り、損失と下界 $\log N – L$ の挙動を確かめます。

import torch
import torch.nn.functional as F
import numpy as np

def make_pairs(B, N, d, gap, seed):
    g = torch.Generator().manual_seed(seed)
    anchor = F.normalize(torch.randn(B, d, generator=g), dim=1)
    cand = F.normalize(torch.randn(B, N, d, generator=g), dim=2)
    pos = F.normalize((1 - gap) * cand[:, 0, :] + gap * anchor, dim=1)
    cand[:, 0, :] = pos
    pos_idx = torch.zeros(B, dtype=torch.long)
    return anchor, cand, pos_idx

d, tau, gap = 32, 0.2, 0.5
Ns = [2, 4, 8, 16, 32, 64, 128, 256, 512]

for N in Ns:
    anchor, cand, pos_idx = make_pairs(1024, N, d, gap, seed=7)
    logits = torch.einsum("bd,bnd->bn", anchor, cand) / tau
    loss = F.cross_entropy(logits, pos_idx).item()   # InfoNCE = 交差エントロピー
    bound = np.log(N) - loss                          # MIの下界
    print(f"N={N:4d}: 損失={loss:.4f}, logN={np.log(N):.4f}, 下界(logN-L)={bound:.4f}")

出力は次の通りです。

N=   2: 損失=0.0431, logN=0.6931, 下界(logN-L)=0.6500
N=   4: 損失=0.1256, logN=1.3863, 下界(logN-L)=1.2607
N=   8: 損失=0.2740, logN=2.0794, 下界(logN-L)=1.8055
N=  16: 損失=0.5034, logN=2.7726, 下界(logN-L)=2.2692
N=  32: 損失=0.8591, logN=3.4657, 下界(logN-L)=2.6067
N=  64: 損失=1.3352, logN=4.1589, 下界(logN-L)=2.8237
N= 128: 損失=1.8779, logN=4.8520, 下界(logN-L)=2.9741
N= 256: 損失=2.5086, logN=5.5452, 下界(logN-L)=3.0365
N= 512: 損失=3.1554, logN=6.2383, 下界(logN-L)=3.0829

負例数スイープ:損失とMI下界 logN-L の挙動

グラフ左から、$N$ が増えると損失 $L$(オレンジ)は増えますが、その増え方は $\log N$(灰破線)より緩やかです。これは「選択肢が増えてクイズが難しくなる」分の損失増加です。グラフ右の下界 $\log N – L$(緑)は、$N$ を増やすほど単調に上がり、$0.65 \to 3.08$ nats まで改善しています。ただし増分は次第に小さくなり、飽和の兆しが見えます。負例を増やすほど捉えられる情報の下界が上がるが、収穫は逓減する という、対照学習の経験則がそのまま数値に現れています。これが大バッチ・メモリバンクの動機であり、同時にその限界でもあります。

ここまでで $\tau$ と $N$ の役割を実験で確かめました。最後に、InfoNCEを相互情報量推定として使うときの限界を整理しておきます。

限界と注意点

InfoNCEは強力ですが、万能ではありません。理論と実用の両面で、押さえておくべき注意点があります。

1. 下界が緩い(MI推定としての限界)。式 (2) は下界であって等号ではありません。導出のステップ4で複数の不等式を使ったように、実際の相互情報量との間にはギャップがあります。とくに相互情報量が大きいデータでは、$\log N$ の天井に阻まれて下界が大きく過小評価になります。「InfoNCEの値からMIを正確に読む」ことはできません。InfoNCEはMIの 推定器 というより、MIを増やす方向に表現を動かす 代理目的 だと捉えるのが健全です。

2. $\log N$ の天井。前節で見た通り、捉えられる情報量は候補数 $N$ で頭打ちになります。バッチサイズやメモリバンクの大きさが、そのまま表現の上限を縛ります。

3. 大バッチ要求とコスト。$N$ を増やすほど良いという性質は、裏を返せば「大きな計算資源を要求する」ということです。バッチサイズに比例してメモリと計算が増え、類似度行列は $O(N^2)$ になります。MoCo(メモリバンク)やBYOL/SimSiam(負例なし)は、この要求を回避するための異なる設計です。

4. 密度比推定の難しさ。最適クリティックは密度比 $p(x\mid c)/p(x)$ に比例しますが、高次元で密度比を正確に推定するのは本質的に困難です。負例のサンプリング方法(どの分布から、いくつ引くか)が学習を大きく左右し、ハード負例の質が結果を決めます。

5. 表現の質はMIだけでは決まらない。後続研究(Wang & Isola, 2020 ほか)は、InfoNCEで学習される表現の良し悪しが「相互情報量の大きさ」よりも、alignment(正例同士の近さ)uniformity(表現が球面上に広がる度合い) という2つの幾何的性質で説明できることを示しました。MI下界はInfoNCEを理解する強力な視点ですが、唯一の視点ではありません。

画像なし
対照学習のAlignmentとUniformity
InfoNCEが最適化する表現の幾何(正例の近さと球面上の一様性)を分解して理解する。

これらの注意点を踏まえると、InfoNCEは「相互情報量の最大化」という美しい看板と、「正例を当てる分類」という実用的な実装の両方を持つ、絶妙なバランスの損失だと分かります。

補足: 類似度行列と埋め込み幾何で全体像をつかむ

最後に、InfoNCEが表現空間で何をしているかを2枚の図で直感的にまとめます。

類似度行列のヒートマップ(対角が正例)

実装では、バッチ内の全アンカーと全候補の類似度を一括計算した 類似度行列 を作り、各行で「対角成分(正例)が最大になる」ようにcross_entropyをかけます。図の対角が赤(高類似度)、非対角が青(低類似度)になっているのが理想形です。これは「各行で正解インデックス=対角を当てる分類」という、InfoNCEの実装上の顔そのものです。

単位球面上で正例を引き寄せ負例を押し離す埋め込み幾何

幾何的には、表現を単位球面上の点とみなし、アンカー(青)に対して正例(緑)を引き寄せ、負例(赤)を押し離します。学習前(左)はばらばらだった正例とアンカーが、学習後(右)では重なるほど近づいています。これがInfoNCEの最小化で起きていることの全体像です。CPC(次図)のように「文脈から未来を予測する」設定でも、本質は同じ「正例を負例の中から見分ける」操作です。

CPCの文脈から未来を予測する概念図

InfoNCEの名付け親であるCPC(Contrastive Predictive Coding)は、過去の系列を集約器 $g_{ar}$ で文脈 $c_t$ にまとめ、その文脈から未来 $x_{t+k}$ を「他系列(負例)の中から当てる」枠組みでした。SimCLRの「同じ画像の別ビューを当てる」も、CLIPの「画像に合うテキストを当てる」も、すべてこの同じInfoNCEの上に乗っています。

まとめ

本記事では、対照学習の心臓部であるInfoNCE損失を主役に据えて解説しました。

  • NCEからの系譜: 正規化定数 $Z(\theta)$ が計算できないという困りごとを、「データ対ノイズの二値分類」に置き換えたのがNCE。二値分類器のロジットが対数密度比になる。
  • $(K{+}1)$クラスsoftmaxとの等価性: InfoNCEは「$N=K{+}1$個の候補から正例を当てるsoftmax交差エントロピー」そのもの。実装で F.cross_entropy をそのまま使える理由がここにある。
  • 相互情報量の下界: $I(X;C) \ge \log N – L_{\text{InfoNCE}}$ を導出した。途中で最適クリティックが密度比 $p(x\mid c)/p(x)$ に比例することが分かる。InfoNCEを下げる=MIの下界を上げる。
  • 温度 $\tau$: 小さいほど分布が尖り、ハード負例に勾配が集中する(実験で実効負例数が63→2へ激減)。
  • 負例数 $N$: 下界は $\log N$ で頭打ち。$N$ を増やすほど下界が上がる(収穫逓減)。これが大バッチ・メモリバンクの動機。
  • 限界: 下界は緩く、$\log N$ の天井があり、大バッチを要求する。表現の質はMIだけでなくalignment/uniformityでも語れる。

InfoNCEは「分類」と「情報理論」の二面性を持つ、対照学習の屋台骨です。次に読む記事として、この損失が育てる表現の幾何そのものを分解する以下をおすすめします。