埋め込みのアライメントとは — Alignment と Uniformity で対照学習の良さを測る

「良い埋め込み」とは何でしょうか。対照学習で得たベクトルが「良い」とき、私たちは何が起きていると言えるのでしょうか。ふだんは「下流タスクの精度が高ければ良い表現だ」と片付けがちですが、それでは精度を測るまで表現の良し悪しがわかりません。学習の途中で「この表現は健全か、それとも崩壊しかけているか」を、ラベルも下流タスクも使わずに直接覗き込めたら便利です。

Wang & Isola は ICML 2020 の論文「Understanding Contrastive Representation Learning through Alignment and Uniformity on the Hypersphere」で、この問いに驚くほど明快な答えを与えました。良い埋め込みは、単位超球面の上で 測れる2つの性質 を満たしている、というのです。1つは Alignment(アライメント)——意味的に同じデータ(正例ペア)が球面上で近くに来ること。もう1つは Uniformity(一様性)——埋め込み全体が球面いっぱいに、まんべんなく広がっていること。

この見方が強力なのは、対照学習の成功を「ブラックボックスの中で何かいいことが起きている」ではなく、幾何学的に定量化できる2つの量 に翻訳してくれる点にあります。実際の応用は2つあります。

  • 表現品質の診断: 下流タスクを走らせなくても、$\mathcal{L}_{\text{align}}$ と $\mathcal{L}_{\text{uniform}}$ を測るだけで表現の健康状態がわかります。学習中にこの2値をモニタすれば、いつ表現が良くなり、いつ崩れ始めたかが一目でわかります。
  • 表現崩壊の検知: 自己教師あり学習では、すべての入力を同じベクトルに潰してしまう「崩壊(collapse)」が起こりがちです。崩壊は uniformity が悪化する形で必ず現れます。uniformity を監視することは崩壊の早期警報になります。

本記事では、以下を扱います。

  • 埋め込みを単位超球面に置くという発想と、その上での alignment / uniformity の直感
  • $\mathcal{L}_{\text{align}}$ と $\mathcal{L}_{\text{uniform}}$ の定義と、なぜその形なのかの導出
  • InfoNCE が漸近的に alignment + uniformity に分解できること(Wang & Isola の議論)
  • 次元崩壊(dimensional collapse)との関係——alignment だけ追うと定数写像に潰れる話
  • 「アライメント」という言葉の多義性(表現内 vs クロスモーダル空間整合)の整理
  • PyTorch/numpy による実装と、良い表現 vs 崩壊表現の球面可視化

対照学習そのものの仕組み(正例・負例、InfoNCE 損失の導出)は既存記事に譲り、本記事は alignment と uniformity という「ものさし」 に焦点を当てます。

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。

画像なし
対照学習(Contrastive Learning)の理論と損失関数の導出
正例ペア・負例ペア・InfoNCE損失の基礎。本記事の出発点です。
画像なし
対照学習(Contrastive Learning)の理論と実装
InfoNCE損失の情報理論的な意味とSimCLRの実装。
画像なし
InfoNCE損失を理解する
本記事で分解の対象とするInfoNCE損失そのものの詳細。
KLダイバージェンスとは
分布の散らばりを測る道具。uniformityの理解に役立ちます。

埋め込みを単位超球面の上に置く

alignment と uniformity の話に入る前に、まず舞台を整えましょう。対照学習では、エンコーダ $f$ が出した特徴ベクトルを L2正規化 して使うのが定石です。

$$ \bm{z} = \frac{f(\bm{x})}{\|f(\bm{x})\|_2} $$

正規化すると $\|\bm{z}\|_2 = 1$ になります。つまり、すべての埋め込みは原点からの距離が1の点、すなわち 単位超球面 $S^{d-1}$ の上に乗ります($d$ は埋め込み次元、$S^{d-1}$ は $d$ 次元空間の中の $(d-1)$ 次元球面)。

なぜわざわざ球面に乗せるのでしょうか。理由は、対照学習で使うコサイン類似度が「向き」だけを見る量だからです。ベクトルの長さは類似度に効きません。長さの自由度を最初に潰して向きだけを残せば、似ているかどうかは球面上の2点がどれだけ近いか(測地距離、あるいはユークリッド距離)で測れます。下の図で、正規化の前後のイメージを掴んでください。

L2正規化で埋め込みを単位超球面上に置く

左は正規化前で、ベクトルの長さがバラバラです。右では各ベクトルを長さ1に揃え、原点から伸ばした先端だけを残しています。すべての点が単位円 $S^1$(2次元の場合の球面)の上にきれいに並んでいるのがわかります。3次元なら $S^2$、高次元ならその一般化です。以後、埋め込みはすべてこの球面の上にあると考えます。

舞台が球面だとわかったところで、「良い埋め込み」が球面上でどう見えるべきかを考えましょう。Wang & Isola の主張は、たった2枚の絵にまとまります。

単位超球面上のAlignmentとUniformityの概念図

左の絵が Alignment です。同じ意味のデータ(正例ペア、同じ色の2点)が球面上で互いに近くに来ています。右の絵が Uniformity です。埋め込み全体が球面のあちこちに散らばり、特定の場所に固まっていません。良い埋め込みとは、この 両方 を同時に満たすものだ——これが論文の核心です。次節から、この2つの性質を1つずつ数式にしていきます。

Alignment:正例ペアの近さを測る

まず alignment(アライメント)です。直感はシンプルで、「正例ペアは球面上で近くにあってほしい」というだけです。日常のたとえで言えば、同じ猫を別角度から撮った2枚の写真は、埋め込み空間でも隣同士にいてほしい、ということです。同じものの2つの見え方が遠く離れていたら、その埋め込みは「同じものを同じと認識できていない」ことになります。

この「近さ」を1つの数にしたものが alignment ロスです。正例ペア $(\bm{x}, \bm{x}^+)$(同じデータから作った2つのビュー)の分布を $p_{\text{pos}}$ とすると、

$$ \begin{equation} \mathcal{L}_{\text{align}} = \mathbb{E}_{(\bm{x}, \bm{x}^+) \sim p_{\text{pos}}}\!\left[\, \|f(\bm{x}) – f(\bm{x}^+)\|_2^{\alpha} \,\right] \end{equation} $$

と定義されます。中身は「正例ペアの埋め込み間のユークリッド距離を $\alpha$ 乗したものの期待値」です。$\alpha > 0$ は指数で、論文では $\alpha = 2$(距離の2乗)がよく使われます。距離が小さいほど $\mathcal{L}_{\text{align}}$ は小さくなるので、この値を小さくする=正例ペアを揃える ことになります。下の図で、揃っているペアと揃っていないペアを見比べてください。

Alignmentは正例ペアの距離を小さくする

緑のペアは球面上で近く、距離が小さいので $\mathcal{L}_{\text{align}}$ への寄与が小さい「良い」ペアです。赤のペアは大きく離れていて距離が大きく、$\mathcal{L}_{\text{align}}$ を押し上げる「悪い」ペアです。alignment ロスは、こうしたペアごとの距離を全正例ペアで平均したものだと捉えてください。

ここで1つ注意があります。$\mathcal{L}_{\text{align}}$ は 正例ペアしか見ていません。負例(異なるデータ同士)には一切触れていないのです。すると当然の疑問が湧きます——正例だけを近づけ続けたら、いったいどこまで近づくのでしょうか。すべての点を1か所に集めてしまえば $\mathcal{L}_{\text{align}} = 0$ になりますが、それは本当に良い表現でしょうか。この疑問が、次の uniformity を必要とする理由です。

Uniformity:球面に一様に広がる性質

alignment だけを最小化する最も簡単な方法は、すべての入力を球面上の同じ1点に写す ことです。どんな $\bm{x}$ も同じ $\bm{z}_0$ に潰せば、正例ペアの距離は常に0、$\mathcal{L}_{\text{align}} = 0$ の完璧なスコアです。しかしこれは「定数写像」、つまり何の情報も持たない最悪の表現です。これが対照学習で恐れられる 崩壊(collapse) です。

崩壊を防ぐには、「近づけすぎるな、全体は広く散らばっていろ」という逆向きの圧力が要ります。その圧力を定量化するのが uniformity(一様性)です。直感は「埋め込みが球面のあちこちに均等にばらまかれているほど良い」というものです。なぜ均等が良いかというと、均等に散らばっている表現ほど、各点が他の点と区別できる情報を最大限に保持しているからです(情報を1か所に固めず、空間をフルに使っている状態)。

では「球面上で一様に散らばっている度合い」をどう数式にするか。Wang & Isola は ガウシアンポテンシャル(Gaussian potential kernel) を使いました。

$$ \begin{equation} \mathcal{L}_{\text{uniform}} = \log \, \mathbb{E}_{\bm{x},\bm{y} \sim p_{\text{data}}}\!\left[\, e^{-t\,\|f(\bm{x}) – f(\bm{y})\|_2^2} \,\right] \end{equation} $$

ここで $\bm{x}, \bm{y}$ はデータから独立に取った2点、$t > 0$ は温度に相当する定数(論文では $t = 2$)です。中身の $e^{-t\|\bm{z}_i – \bm{z}_j\|^2}$ が「2点が近いほど大きく、遠いほど小さい」反発エネルギーで、これを全ペアで平均して対数を取っています。

なぜこの形が「一様」を測るのか

この式の気持ちを掴むために、まずガウシアンポテンシャルそのものを見てみましょう。

ガウシアンポテンシャルの直感

横軸は2点間の距離、縦軸はポテンシャル $e^{-t\|\bm{z}_i – \bm{z}_j\|^2}$ です。距離が0に近いと値は1に近づき(高コスト)、距離が大きくなると急速に0へ落ちます(低コスト)。$t$ を大きくすると、近い点へのペナルティがより鋭くなります。つまりこのポテンシャルは、「近くにいる点のペア」だけを高コストとして数え上げる 仕掛けです。

$\mathcal{L}_{\text{uniform}}$ はこのポテンシャルを全ペアで平均して対数を取った量なので、「どこかに点が密集していると大きくなり、点がまんべんなく散らばっていると小さくなる」性質を持ちます。点が固まっている=近接ペアが多い=ポテンシャルの和が大きい、というわけです。したがって、

$$ \mathcal{L}_{\text{uniform}} \text{ を最小化する} \iff \text{点を互いにできるだけ離して散らす} $$

となります。直感的には、各点が「お互いを押し離す電荷」のように振る舞い、球面上で最も均等な配置に落ち着くイメージです。実際、Wang & Isola はこのポテンシャルが球面上で最小化されるとき、点の分布が 球面一様分布(uniform distribution on the sphere) に収束することを、ポテンシャル理論を使って厳密に証明しています(論文 Theorem 1)。ガウシアンポテンシャルは、球面上の点を「最も均等にばらまく」エネルギーとして数学的に正当化されているのです。

下の図で、配置による $\mathcal{L}_{\text{uniform}}$ の違いを確かめましょう。

一様分布がUniformityを最小にする

左から「球面に一様」「半球に偏り」「一点に集中」の3配置で、それぞれ $\mathcal{L}_{\text{uniform}}$ を計算しています。値は左から約 $-2.07$、$-1.33$、$-0.20$ で、一様な配置が最小(最も小さい値) になっています。偏るほど値が大きくなり、一点に集中した崩壊状態では最大(最悪)です。$\mathcal{L}_{\text{uniform}}$ が小さいことは「球面を広く使えている」ことの定量的な証拠だとわかります。

ここまでで、良い埋め込みの2つの条件——正例を近づける alignment と、全体を散らす uniformity——が数式になりました。すると次に知りたいのは、対照学習で実際に使う InfoNCE 損失が、この2つとどう関係するのかです。

InfoNCE は alignment + uniformity に分解できる

対照学習の標準的な損失は、本記事冒頭で挙げた既存記事で導出した InfoNCE 損失 です。$\bm{z}$ を正規化済み、$\bm{z}^+$ を正例、$\{\bm{z}^-\}$ を負例として、温度 $\tau$ のもとで

$$ \mathcal{L}_{\text{InfoNCE}} = \mathbb{E}\!\left[ -\log \frac{e^{\,\bm{z}^\top \bm{z}^+ / \tau}}{e^{\,\bm{z}^\top \bm{z}^+ / \tau} + \sum_{\bm{z}^-} e^{\,\bm{z}^\top \bm{z}^- / \tau}} \right] $$

と書けます。Wang & Isola が示したのは、負例の数 $M \to \infty$ の極限で、この損失が alignment 項と uniformity 項にきれいに分かれる ことです。その流れを追いましょう。

分子と分母を分ける

まず $-\log$ の中の分数を、分子と分母の対数の差に開きます。

$$ \mathcal{L}_{\text{InfoNCE}} = \mathbb{E}\!\left[ -\frac{\bm{z}^\top \bm{z}^+}{\tau} + \log\!\left( e^{\,\bm{z}^\top \bm{z}^+/\tau} + \sum_{\bm{z}^-} e^{\,\bm{z}^\top \bm{z}^-/\tau} \right) \right] $$

第1項は 正例との内積を大きくする 力、第2項は すべてのペア(正例+負例)との内積を小さくする 力です。第1項に注目しましょう。$\bm{z}, \bm{z}^+$ がともに単位ベクトルなら、内積と距離の間には

$$ \|\bm{z} – \bm{z}^+\|^2 = \|\bm{z}\|^2 – 2\,\bm{z}^\top\bm{z}^+ + \|\bm{z}^+\|^2 = 2 – 2\,\bm{z}^\top\bm{z}^+ $$

という関係があります($\|\bm{z}\| = \|\bm{z}^+\| = 1$ を代入しました)。これを $\bm{z}^\top\bm{z}^+$ について解くと $\bm{z}^\top\bm{z}^+ = 1 – \tfrac{1}{2}\|\bm{z} – \bm{z}^+\|^2$ です。第1項に入れると、

$$ -\frac{\bm{z}^\top \bm{z}^+}{\tau} = -\frac{1}{\tau} + \frac{1}{2\tau}\|\bm{z} – \bm{z}^+\|^2 $$

定数 $-1/\tau$ を除けば、これは まさに $\mathcal{L}_{\text{align}}$ そのもの($\alpha = 2$)です。正例との内積を大きくすることは、正例ペアの距離を縮めること、つまり alignment の最小化に等しいのです。

負例の項が uniformity になる

次に第2項です。負例の数を $M$ とし、$M \to \infty$ の極限を取ります。和 $\frac{1}{M}\sum_{\bm{z}^-} e^{\bm{z}^\top\bm{z}^-/\tau}$ は大数の法則で期待値に収束します。

$$ \frac{1}{M}\sum_{\bm{z}^-} e^{\,\bm{z}^\top \bm{z}^-/\tau} \;\xrightarrow{M \to \infty}\; \mathbb{E}_{\bm{z}^-}\!\left[ e^{\,\bm{z}^\top \bm{z}^-/\tau} \right] $$

正例項は負例項の和に比べて無視できるようになり($M$ が巨大なので分母は負例の和が支配的)、第2項の対数の中身は本質的に $\mathbb{E}_{\bm{z}^-}[e^{\bm{z}^\top\bm{z}^-/\tau}]$ に比例します。ここでもう一度 $\bm{z}^\top\bm{z}^- = 1 – \tfrac{1}{2}\|\bm{z} – \bm{z}^-\|^2$ を代入すると、

$$ \mathbb{E}\!\left[ \log \mathbb{E}_{\bm{z}^-}\!\left[ e^{\,\bm{z}^\top \bm{z}^-/\tau} \right] \right] = \text{const} + \mathbb{E}\!\left[ \log \mathbb{E}_{\bm{z}^-}\!\left[ e^{-\frac{1}{2\tau}\|\bm{z} – \bm{z}^-\|^2} \right] \right] $$

右辺の第2項は、$t = \tfrac{1}{2\tau}$ と置けば まさに $\mathcal{L}_{\text{uniform}}$ の形(ガウシアンポテンシャルの対数期待値)です。負例を遠ざける力は、埋め込み全体を一様に散らす uniformity の最小化に等しかったのです。

まとめると

定数項を無視して整理すると、負例数無限大の極限で

$$ \begin{equation} \mathcal{L}_{\text{InfoNCE}} \;\xrightarrow{M \to \infty}\; \frac{1}{2\tau}\,\mathcal{L}_{\text{align}} \;+\; \mathcal{L}_{\text{uniform}} \;+\; \text{const} \end{equation} $$

という関係が得られます。下の概念図がこの分解を表しています。

InfoNCEがAlignmentとUniformityに分解される

InfoNCE の 分子(正例を引き寄せる力)が alignment に、分母(負例を押し離す力)が uniformity に対応します。つまり対照学習がうまくいくのは、たまたまではなく、InfoNCE がこの2つの「良い表現の条件」を同時に最適化しているからだ、というのが Wang & Isola の説明です。この分解こそが、対照学習の成功を幾何学的に腑に落とす鍵になります。

この理解には実用的な含みもあります。InfoNCE を経由せず、$\mathcal{L}_{\text{align}}$ と $\mathcal{L}_{\text{uniform}}$ を 直接 重み付き和として最小化しても、同等以上の表現が学習できるのです(論文の実験でも確認されています)。後半の実装でこれを実際に試します。その前に、alignment と uniformity が崩れるとどうなるか——崩壊の話を見ておきましょう。

次元崩壊との関係:alignment だけでは潰れる

すでに触れたように、alignment だけを最小化すると、すべての点が1か所に集まる 完全崩壊(complete collapse) が起きます。これは定数写像で、$\mathcal{L}_{\text{align}} = 0$ という見かけ上は完璧なスコアを出しながら、表現としては全く役に立ちません。uniformity の項がこれを防いでいるわけです。

崩壊にはもっと巧妙な形もあります。次元崩壊(dimensional collapse) です。これは、埋め込みが球面の全次元を使わず、低次元の部分空間(たとえば大円や、いくつかの軸が張る平面)に潰れてしまう現象です。完全崩壊ほど派手ではないので気づきにくいのですが、表現の表現力を大きく損ないます。

次元崩壊の可視化

左は健全な表現で、点が球面($S^2$)の全体に広がり、3次元すべてを使っています。中央は次元崩壊で、点が大円(1次元の輪)に潰れ、球面の表面という2次元の自由度を捨ててしまっています。右は完全崩壊で、すべての点が1か所に集まった定数写像です。中央・右はいずれも uniformity が悪化しており、uniformity の劣化が崩壊のシグナルになる ことが見て取れます。

次元崩壊は、特異値スペクトルを見ると定量的に診断できます。埋め込み行列の特異値を大きい順に並べたとき、健全な表現では特異値がなだらかに続きますが、崩壊した表現では少数の特異値だけが大きく、残りがほぼ0に落ちます(=実効的な次元が少ない)。

特異値スペクトルで崩壊を診断する

緑(健全な表現)は特異値が緩やかに減衰し、多くの次元が有効に使われています。赤(崩壊した表現)は数本目以降の特異値が急落しており、実質的に低次元の部分空間しか使っていません。学習中にこのスペクトルや uniformity をモニタすれば、表現が崩れ始めた瞬間を捉えられます。これが冒頭で述べた「崩壊の早期警報」の実態です。

崩壊の議論からわかるのは、alignment と uniformity が トレードオフではなく、両立すべき2つの条件 だということです。次の図で、片方だけを追ったときに何が起きるかを並べてみましょう。

良い表現と崩壊表現の球面散布

左は両方を満たす良い表現で、2クラス(シアンとアンバー)がそれぞれまとまりつつ、全体は球面に広がっています。中央は alignment だけを最適化した結果で、両クラスが同じ点に潰れ、クラスの区別が消えています(完全崩壊)。右は uniformity だけを最適化した結果で、点は広がっていますが正例ペアがバラバラに散らばり、同じクラスがまとまっていません。どちらか一方だけでは良い表現にならない ことが視覚的に確認できます。

理論はここまでです。次は、これらの量を実際に Python で計算し、良い表現と崩壊表現の違いを数値で確かめましょう。

Pythonでの実装

Alignment と Uniformity を計算する

まず、$\mathcal{L}_{\text{align}}$ と $\mathcal{L}_{\text{uniform}}$ を計算する関数を書きます。論文の定義そのままで、数行で実装できます。良い表現・崩壊表現・ランダム表現の3つで値を比べ、ものさしとして機能するかを確かめます。

import numpy as np
import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
import matplotlib.font_manager as _fm
for _c in ['Hiragino Sans','Yu Gothic','Noto Sans CJK JP','IPAexGothic','Meiryo']:
    if any(_c==_f.name for _f in _fm.fontManager.ttflist):
        plt.rcParams['font.family']=_c; break
plt.rcParams['axes.unicode_minus']=False

def lalign(z1, z2, alpha=2):
    """Alignment ロス: 正例ペアの距離の alpha 乗の平均"""
    return (z1 - z2).norm(dim=1).pow(alpha).mean()

def luniform(z, t=2):
    """Uniformity ロス: ガウシアンポテンシャルの対数期待値"""
    sq_pdist = torch.pdist(z, p=2).pow(2)   # 全ペアの距離の2乗
    return sq_pdist.mul(-t).exp().mean().log()

torch.manual_seed(0)
n, d = 200, 16

# (1) 良い表現: 正例が近く、全体は球面に広がる
anchors = F.normalize(torch.randn(n, d), dim=1)
z1 = F.normalize(anchors + 0.1*torch.randn(n, d), dim=1)
z2 = F.normalize(anchors + 0.1*torch.randn(n, d), dim=1)
print("良い表現  : L_align=%.3f  L_uniform=%.3f"
      % (lalign(z1, z2), luniform(torch.cat([z1, z2]))))

# (2) 崩壊表現: すべて1点付近に集中
c = F.normalize(torch.ones(1, d), dim=1)
zc1 = F.normalize(c + 0.1*torch.randn(n, d), dim=1)
zc2 = F.normalize(c + 0.1*torch.randn(n, d), dim=1)
print("崩壊表現  : L_align=%.3f  L_uniform=%.3f"
      % (lalign(zc1, zc2), luniform(torch.cat([zc1, zc2]))))

# (3) ランダム表現: 正例ペアの対応すらない
zr1 = F.normalize(torch.randn(n, d), dim=1)
zr2 = F.normalize(torch.randn(n, d), dim=1)
print("ランダム  : L_align=%.3f  L_uniform=%.3f"
      % (lalign(zr1, zr2), luniform(torch.cat([zr1, zr2]))))

実行すると、良い表現は L_align=0.265, L_uniform=-3.460、崩壊表現は L_align=0.273, L_uniform=-0.517、ランダムは L_align=1.982, L_uniform=-3.512 となります。ここから3つのことが読み取れます。第1に、良い表現は alignment も uniformity も小さく、まさに「左下」の理想状態です。第2に、崩壊表現は alignment こそ小さいものの uniformity が極端に悪い($-0.5$ 付近)——alignment だけ見ていたら崩壊を見逃すが、uniformity が崩壊を暴いています。第3に、ランダム表現は uniformity は小さい(球面に散らばっている)のに alignment が大きく、正例を揃えられていません。2つの量を両方見ることで初めて表現の良し悪しが判別できる という主張が、数値で裏付けられました。

Alignment–Uniformity 平面にプロットする

3つだけでなく、もっと多様な表現を同じ平面にプロットすると、各表現の「位置」がよくわかります。横軸に $\mathcal{L}_{\text{align}}$、縦軸に $\mathcal{L}_{\text{uniform}}$ を取り、複数の表現を散布させます。

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

torch.manual_seed(51)
d, n = 3, 300

def make_pairs(collapse, noise):
    """collapse: 1点に潰す度合い(1=完全崩壊). noise: 正例ペアのずれ"""
    spread = 1.0 - collapse
    base = torch.randn(n, d) * spread + torch.tensor([0., 0., collapse*3])
    anchors = F.normalize(base, dim=1)
    z1 = F.normalize(anchors + noise*torch.randn(n, d), dim=1)
    z2 = F.normalize(anchors + noise*torch.randn(n, d), dim=1)
    return z1, z2

def lalign(z1, z2, a=2): return (z1-z2).norm(dim=1).pow(a).mean().item()
def luniform(z, t=2): return torch.pdist(z,2).pow(2).mul(-t).exp().mean().log().item()

cases = [("理想(両立)", 0.0, 0.08), ("align過剰(崩壊)", 0.95, 0.08),
         ("uniform過剰", 0.0, 0.8), ("中間", 0.4, 0.35), ("ランダム", 0.0, 1.5)]
for name, cl, ns in cases:
    z1, z2 = make_pairs(cl, ns)
    la = lalign(z1, z2)
    lu = luniform(torch.cat([z1, z2])[:200])
    print("%-16s L_align=%.3f  L_uniform=%.3f" % (name, la, lu))

実行結果は、理想 L_align=0.023, L_uniform=-2.093、align過剰(崩壊) L_align=0.027, L_uniform=-0.050、uniform過剰 L_align=1.373, L_uniform=-2.063、中間 L_align=0.456, L_uniform=-1.253、ランダム L_align=1.698, L_uniform=-2.084 です。これを平面にプロットすると下の図になります(下図は同じ傾向を numpy で描いたもので、各点の相対位置は一致します)。

Alignment-Uniformity平面の散布図

良い表現は 左下(両方が小さい理想領域)に来ます。崩壊(align過剰)は左上で、alignment は小さいのに uniformity が悪い——「正例を揃えすぎて全部1点に潰れた」状態がこの位置に現れます。uniform過剰とランダムは右側で、散らばってはいるが正例を揃えられていません。この平面は、表現の「どこが悪いのか」まで教えてくれる診断図として使えます。学習中にエポックごとにこの平面に点を打てば、表現が理想領域に向かっているか、崩壊に向かっているかが軌跡として見えます。

Alignment + Uniformity を直接最適化する

最後に、InfoNCE を使わずに $\mathcal{L}_{\text{align}} + \mathcal{L}_{\text{uniform}}$ を 直接 最小化して、本当に良い表現が学習できるかを確かめます。同時に、alignment だけを最適化すると崩壊することも実演します。

import torch
import torch.nn.functional as F

torch.manual_seed(0)
def lalign(z1, z2, a=2): return (z1-z2).norm(dim=1).pow(a).mean()
def luniform(z, t=2): return torch.pdist(z, 2).pow(2).mul(-t).exp().mean().log()

# 同じ潜在から2ビューを作る合成データ(正例ペア)
n, d_in, d = 256, 32, 3
latent = torch.randn(n, 8); W = torch.randn(8, d_in)
x1 = latent @ W + 0.3*torch.randn(n, d_in)
x2 = latent @ W + 0.3*torch.randn(n, d_in)

def train(objective, steps=500):
    torch.manual_seed(0)
    enc = torch.nn.Sequential(torch.nn.Linear(d_in,64), torch.nn.ReLU(),
                              torch.nn.Linear(64,d))
    opt = torch.optim.Adam(enc.parameters(), lr=1e-2)
    emb = lambda x: F.normalize(enc(x), dim=1)
    for _ in range(steps):
        z1, z2 = emb(x1), emb(x2)
        loss = objective(z1, z2)
        opt.zero_grad(); loss.backward(); opt.step()
    z1, z2 = emb(x1), emb(x2)
    return lalign(z1, z2).item(), luniform(torch.cat([z1, z2])).item()

# (A) alignment + uniformity を両方最適化
la, lu = train(lambda z1, z2: lalign(z1, z2) + luniform(torch.cat([z1, z2])))
print("両方最適化   : L_align=%.3f  L_uniform=%.3f" % (la, lu))

# (B) alignment だけを最適化 → 崩壊するはず
la2, lu2 = train(lambda z1, z2: lalign(z1, z2))
print("alignのみ    : L_align=%.3f  L_uniform=%.3f (崩壊)" % (la2, lu2))

実行すると、両方最適化では L_align=0.001, L_uniform=-2.093 と、alignment と uniformity がともに小さい理想的な表現が得られます。一方、alignment だけを最適化すると L_align=0.000, L_uniform=-0.381 となり、alignment は完璧(距離ゼロ)でも uniformity が極端に悪化——つまり すべての点が1か所に潰れた完全崩壊 が起きています。この対比は、uniformity の項が崩壊を防ぐ「つっかえ棒」であることを明快に示しています。alignment は引き寄せる力、uniformity は押し離す力で、両者の綱引きが球面上に良い表現を作るのです。

参考までに、InfoNCE 損失そのもので同じデータを学習すると L_align=0.000, L_uniform=-2.094 となり、両方を直接最適化したときとほぼ同じ理想領域に着地します。これは前節で導出した「InfoNCE が漸近的に alignment + uniformity に分解される」という主張の、ささやかな実験的裏付けになっています。InfoNCE は2つの良い性質を暗黙のうちに同時に追っていたわけです。

ここまでで、表現内の alignment という主役を一通り見てきました。最後に、この記事で言う「アライメント」と、世間でよく使われるもう一つの「アライメント」を区別しておきましょう。

「アライメント」の多義性:表現内 vs 空間整合

表現学習の文脈で「アライメント」という言葉を聞いたとき、実は2つの異なる意味があり得ます。混同しやすいので整理しておきます。

アライメントの多義性とモダリティギャップ

1つ目が 表現内アライメント——本記事の主役です。これは1つの埋め込み空間の中で、正例ペア $f(\bm{x}), f(\bm{x}^+)$ を近づけることを指します。Wang & Isola の $\mathcal{L}_{\text{align}}$ はこちらです。

2つ目が 空間アライメント(クロスモーダル/多言語の整合) です。これは異なる2つの埋め込み空間——たとえば画像の空間とテキストの空間——を互いに整合させることを指します。CLIP のように画像とテキストを共通の空間に埋め込む手法では、しばしば モダリティギャップ(modality gap) が観測されます。これは、画像の埋め込みとテキストの埋め込みが、同じ意味を表していても球面上の異なる「錐(コーン)」に分かれて固まってしまう現象です(図の左)。

画像なし
CLIP:画像とテキストの対照学習
マルチモーダルな埋め込み空間。モダリティギャップの舞台です。

空間アライメントでは、2つの空間を回転・反射で重ね合わせる Procrustes 整合(直交変換による最小二乗整合) のような手法が使われます。これは多言語の単語埋め込みを揃える際にも定番の道具です。同じ「揃える」でも、表現内アライメントが「1つの空間の中で正例を近づける」のに対し、空間アライメントは「2つの空間そのものを重ねる」点で対象が異なります。

本記事で扱ったのは前者、表現内のアライメントです。後者のクロスモーダル整合は、それ自体が大きなテーマなので、CLIP の記事などで別途掘り下げるのがよいでしょう。言葉が同じでも中身が違う点だけ、頭の片隅に置いておいてください。

まとめ

本記事では、「埋め込みのアライメントとは何か」を出発点に、Wang & Isola (ICML 2020) の AlignmentUniformity という2つの測れる性質で、対照学習がなぜ良い表現を作るのかを解説しました。

  • 舞台は単位超球面: 埋め込みを L2 正規化して $S^{d-1}$ 上に置くと、似ているかどうかは球面上の近さで測れる。
  • Alignment: 正例ペアの距離 $\mathcal{L}_{\text{align}} = \mathbb{E}\|f(x) – f(x^+)\|^\alpha$ を小さくする=同じものを近づける。ただし alignment だけだと崩壊する。
  • Uniformity: ガウシアンポテンシャルの対数期待値 $\mathcal{L}_{\text{uniform}} = \log\mathbb{E}\,e^{-t\|f(x)-f(y)\|^2}$ を小さくする=球面に一様に散らす。球面一様分布で最小になることが理論的に保証される。
  • InfoNCE の分解: 負例数無限大の極限で、InfoNCE は $\frac{1}{2\tau}\mathcal{L}_{\text{align}} + \mathcal{L}_{\text{uniform}}$ に分解できる。分子が alignment、分母が uniformity に対応する。
  • 崩壊の診断: 完全崩壊(定数写像)も次元崩壊(低次元部分空間への潰れ)も、uniformity の悪化として現れる。2つの量を両方モニタすることが表現品質の診断と崩壊検知の鍵。
  • 多義性の整理: 表現内アライメント(本記事)と、CLIP のモダリティギャップに代表される空間アライメント(Procrustes 整合など)は別物。

alignment と uniformity は、対照学習の成功を「精度が出たから良い」ではなく「球面上の2つの幾何学的条件を満たしているから良い」と言い換える、強力な視点です。これを身につけると、新しい自己教師あり手法を見たときに「この損失は alignment を担保しているか、uniformity を担保しているか」という軸で素早く評価できるようになります。

次のステップとして、以下の記事も参考にしてください。

画像なし
対照学習(Contrastive Learning)の理論と実装
SimCLR・MoCo・BYOLなど、alignmentとuniformityを実現する具体的手法。
画像なし
InfoNCE損失を理解する
本記事で分解した損失そのものの詳細な導出。
画像なし
CLIP:画像とテキストの対照学習
クロスモーダルな空間アライメントとモダリティギャップ。