【最適輸送】ワッサースタイン距離でドメインシフトを定量化する
はじめに
この記事は、サイバーエージェント26卒内定者 Advent Calendarの21日目の記事です。
「学習データでは高精度だったのに、いざ本番環境（推論時）に投入すると精度がガクッと落ちてしまった……」と思ったことはありませんか？
深層学習において、学習時と推論時でデータの統計的な性質が変化してしまう現象は避けて通れない課題です。本記事では、この「ドメインシフト」という現象に対して、分布の「幾何学的な形」を考慮して便利に定量化できるワッサースタイン距離という概念について説明し、そのユースケースを紹介します。
ドメインシフトとワッサースタイン距離
ここでは、本記事の核となる「ドメインシフト」の定義と、それを測るための「ワッサースタイン距離」の基礎について整理します。
ドメインとは？
機械学習におけるドメイン（Domain）とは、一言で言えば「データが従っている分布」のことです。
数学的には、入力空間
ドメインシフトとは？
ドメインシフト（Domain Shift）とは、学習時（ソースドメイン）と運用時（ターゲットドメイン）で、データの確率分布が変化してしまう現象を指します。
例えば、
-
画像認識： 晴天の日に撮った写真で学習したモデルを、雨の日の画像に適用する。
-
音声認識： 静かなスタジオの音声で学習し、騒がしい街頭の音声に適用する。
-
LLM： 2023年までのデータで学習し、2025年の最新トレンドについて回答させる。
このように分布がズレると、モデルが学習時に獲得した「境界線」が適合しなくなり、精度が著しく低下します。これを防ぐには、まずどれくらいズレているかを正しく測る必要があります。
ワッサースタイン距離とは？
ワッサースタイン距離（Wasserstein Distance）は、別名土運搬距離（Earth Mover's Distance; EMD）とも呼ばれます。
ある分布
ここで
KL-divergenceとの違い
分布の距離を測る指標として一般的な「KLダイバージェンス」と比較したとき、ワッサースタイン距離には以下の3つの大きな利点があります。
対称性があり直感的に理解できる
KLダイバージェンスは
分布が重なっていなくても計算できる
これが最大のメリットです。KLダイバージェンスやJSダイバージェンスは、2つの分布のサポート（値が存在する範囲）が重なっていない場合、値が無限大に発散したり、勾配が消失したりします。 ワッサースタイン距離は、分布がどれだけ離れていても「移動距離」として有限の値を返します。
点同士の対応関係が得られる
ワッサースタイン距離を計算する過程で「どのサンプルをどこに移動させるのが最適か」という最適な輸送計画（Optimal Transport Map）が得られます。これにより、「ドメインAのこのデータは、ドメインBではこれに相当する」といった解釈が可能になります。
音声基盤モデルのドメインシフトを定量化する
ここでは音声基盤モデルのノイズ耐性を例に、ドメインシフトの定量化を試します。
WavLMやHuBERT、wav2vec 2.0などの音声基盤モデルは音声認識や話者識別、感情認識など様々なタスクで使用されています。これらのモデルは自己教師あり学習 (Self-Supervised Learning; SSL) で学習されていることが多いことから音声SSLとも呼ばれることがあります。
音声基盤モデルで推論を行う上で、よく課題になるのがノイズに対する耐性です。音声基盤モデルに入力する音声にノイズが含まれている場合、モデルの性能が低下してしまうことが知られています。この問題を避けるための1つの方法として音声基盤モデルの出力に変化がないかを確認することが挙げられます。
ここでは、入力音声に意図的にノイズを加えて、どの程度出力が変化するかを定量化することを目的とします。
WavLM特徴量の抽出
WavLMの使い方は簡単で、wavファイルなどに含まれる音声波形をモデルに入力すれば、一定時間ごとに高次元の音声特徴量を抽出することができます。この特徴量を使って様々な機械学習モデルを構築します。
例えば、wavlm-baseから特徴量を抽出する場合には以下のようになります。ここでは時間方向に平均プーリングを行うことで(サンプル数, 次元数=768)の特徴量を抽出しています。
import torch
from torch import nn
from transformers import WavLMModel
class WavLMEmbedding(nn.Module):
def __init__(self, model_name: str = "microsoft/wavlm-base"):
super().__init__()
self.model = WavLMModel.from_pretrained(model_name)
self.model.eval()
def forward(self, waveform: torch.Tensor) -> torch.Tensor:
outputs = self.model(waveform).last_hidden_state # (batch, seq, dim)
return outputs.mean(dim=1) # (batch, dim)
t-SNEで定性的に確認してみる
次元削減アルゴリズムであるt-SNEを使って、まずは定性的にどのような変化が生じているかを確認してみます。今回使うデータは英語の音声感情コーパスであるIEMOCAPです。このコーパスからランダムに100件ほど音声を選択して可視化していきます。
まずはノイズがない綺麗な音声をwavlm-baseに入力し、音声特徴量を抽出します。
続いてこれらの音声に対して、一定の強さの ノイズ(ホワイトノイズ) を加えます。加えるノイズの強さについては SNR（信号対雑音比、Signal-to-Noise Ratio） という指標を使いますが、ざっくりと「SNRが大きいほどノイズが小さい（信号がクリア）」と理解していただければと思います。SNRが0dBの場合にはノイズと元音声の比率が同じになります。
このノイズが含まれる音声も同様にwavlm-baseに入力し、音声特徴量を抽出します。ノイズはSNR 10dB・20dB・30dBのパターンを試します。ここまでで、(100件 * 4パターン, 768次元)の音声特徴量を抽出することができました。
これらの音声特徴量をt-SNEで2次元まで次元削減し可視化してみます。t-SNEはscikit-learnから使用することができます。
可視化した結果は以下の通りです。SNRが0dBから30dBと大きくなるにつれて、音声特徴量の分布がシフトしていることが分かります。
t-SNEによる可視化
分かりやすいようにSNR 0dBと30dBの場合のみの結果を可視化してみます。この2つの点群の間ではかなり大きな差があることが分かります。
t-SNEによる可視化 (SNR 0dB・30dB)
定性的にはノイズが大きいほど音声特徴量が変化することが分かりました。では、この現象をどのように定量化すれば良いのでしょうか？ここでワッサースタイン距離の出番です。
ワッサースタイン距離による定量化
ワッサースタイン距離はライブラリPOTを使って、以下のように簡単に計算することができます。
def calc_wasserstein_distance(
x1: np.ndarray, # [samples, dim]
x2: np.ndarray, # [samples, dim]
) -> float: # type: ignore
matrix = ot.dist(x1, x2, metric="euclidean", p=1)
matrix = np.array(matrix)
a = np.ones(x1.shape[0]) / x1.shape[0]
b = np.ones(x2.shape[0]) / x2.shape[0]
wasserstein_distance = ot.emd2(a, b, matrix)
return float(wasserstein_distance)
例えば、SNR 0dBの場合の音声特徴量(100, 768)と、SNR 30dBの場合の音声特徴量(100, 768)の間でワッサースタイン距離を計算するには、x1とx2にそれぞれ入力すれば良いです。
これら全ての組み合わせに対してワッサースタイン距離を計算した結果は以下の通りです。
ワッサースタイン距離
この結果からSNRが近い音声特徴量は距離が近くなっていることが分かります。特にSNR 30dBとの距離を比較してみると、ノイズが大きい(SNR 0dBに近い)ほどワッサースタイン距離が大きくなっています。
この結果は先ほどのt-SNEの「ノイズが大きいほど音声特徴量が変化する」という知見と一致しており、ノイズによる変化の定量化に成功しました。
また、KLダイバージェンスで定量化した結果と比較してみます。KLダイバージェンスは以下のようなコードを使って計算することができます。KLダイバージェンスの入力は確率分布となるので、softmaxを使うことでモデル出力を確率に変換しています。
def softmax(x: np.ndarray) -> np.ndarray:
x_max = np.max(x, axis=-1, keepdims=True)
e_x = np.exp(x - x_max)
probs = e_x / np.sum(e_x, axis=-1, keepdims=True)
return probs
def calc_kl_divergence(
x1: np.ndarray, # [samples, 768]
x2: np.ndarray, # [samples, 768]
eps: float = 1e-10,
) -> float:
p1 = np.mean(softmax(x1), axis=0)
p2 = np.mean(softmax(x2), axis=0)
p1 = np.clip(p1, eps, 1.0)
p2 = np.clip(p2, eps, 1.0)
kl_divergence = np.sum(p1 * np.log(p1 / p2))
return float(kl_divergence)
結果は以下の図の通りです。今回の場合は2つの分布のサポートが重なっていたからか、ワッサースタイン距離とほぼ同様の結果が得られています。ただしワッサースタイン距離は対称で直感的に理解しやすかったのに対して、KLダイバージェンスは非対称になっているという違いがあります。
KLダイバージェンス
LLMの違いを定量化する
次にLLMの違いをワッサースタイン距離で定量化することを試みます。
大規模言語モデル（LLM）は日進月歩で進化しており、OpenAI, Anthropic, Google, Meta、あるいは各社の独自ファインチューニングモデルなど、選択肢が無数に存在します。ここでどのモデルが、どのモデルの出力に近いのかを定量化できれば、モデルのリプレイスを行う際に客観的に挙動変化を予測できるなどいくつかのメリットがあります。
本来であれば、ICA (独立成分分析)やコサイン類似度、その他研究されている指標などを使うのが正しいと思われますが、ここでは敢えてワッサースタイン距離を使ってみようと思います。
LLMからの出力を集める
まずは、各種LLMに多様な入力文を与え、出力文を収集します。入力文は合計で50種類用意し、API経由で出力文を取得します。入力文の例は以下の通りです。
量子もつれ（Quantum Entanglement）について、10歳の子供でもわかるように比喩を使って3行で説明してください。
この入力文に対して、合計で6種類のLLM (gpt-5.2-mini・gpt-5-mini・gemini-3-pro-preview・gemini-2.5-flash・claude-opus-4-5・claude-haiku-4-5)からの出力を取得しました。gpt-5.2-miniの場合は以下のような結果になりました。
量子もつれは、2つの「ふしぎなサイコロ」をペアにして作ると、どんなに遠くに離してもつながったままになることだよ。 \n片方をふって目が決まった瞬間、もう片方の目も同時に決まって（あらかじめ決まっていたみたいに）わかるんだ。 \nまるで「離れていても同時に色が変わる手袋」みたいな、不思議なペアなんだよ。
文埋め込みモデルRuriで出力文を埋め込み表現にする
テキストのままでは何もすることができないので、この出力文を文埋め込み表現(ベクトル表現) に変換します。文埋め込みモデルにはcl-nagoya/ruri-v3-310mを使用します。文章が長い場合は256文字で打ち切りました。
この処理により(モデル数 * 出力文, 次元数=768)の文埋め込み表現を取得できました。
from sentence_transformers import SentenceTransformer
model = SentenceTransformer(model_name)
embeddings = model.encode(texts)
t-SNEでLLMの違いを見てみる
この文埋め込み表現を先ほどと同様にt-SNEで可視化してみます。結果は図の通りで、大まかには入力文ごとにクラスターを形成している傾向が見られます。各クラスターの中身をよく見てみるとモデルによって微妙に違いがあるが見て取れます。
t-SNEによる可視化
より拡大して見てみます。良く見ると同じ組織が提供しているモデルは位置が近いようにも見えます。例えばGPTはクラスターの中心から離れた位置にありがちで、Gemini・Claudeは固まっているようにも見えます。
しかしこの図だけではまだ結論づけることはできません。そこでワッサースタイン距離でLLMごとの違いを定量化してみます。
t-SNEによる可視化(拡大)
ワッサースタイン距離による定量化
先ほどと同様にワッサースタイン距離を計算しました。結果は以下の図の通りです。
t-SNEの図でも見られたように同じ組織が提供しているモデルはワッサースタイン距離が近いようです。特にGeminiとClaudeは顕著です。
またGPTは他のモデルとの距離が遠いようです。こちらもt-SNEの図の傾向と一致しています。学習データや学習方法が他とは異なるのでしょうか...？
ワッサースタイン距離
KLダイバージェンスでも計算してみます。KLダイバージェンスは非対称であることに加え、Claudeの2つのモデル間の差が大きいようです。ワッサースタイン距離の場合にはこの2つのモデルの距離は短かった(8.71)ですが、KLダイバージェンスではGPTの2つのモデルよりも大きな値(4.72e-4または4.38e-4)となっています。
今回の場合はどちらが正しいと結論づけることは難しいですが、ワッサースタイン距離とKLダイバージェンスの計算結果はしばしば異なることには留意した方が良いと思います。
KLダイバージェンス
最後に
本記事ではワッサースタイン距離について紹介し、2つのテーマに沿って実用例を紹介しました。
個人的に最適輸送は数学・統計チックな理論が展開されていることが好きです。機械学習の研究は帰納的なアプローチを取ることが多いのに対して、演繹的な導出があると安心できます。
より詳しく知りたい方は佐藤竜馬先生著「最適輸送の理論とアルゴリズム」および松井孝太先生/熊谷亘先生著「転移学習」が分かりやすいので、是非見てみてください。
ご覧いただきありがとうございました😀
参考資料
Discussion