PyTorchで確率分布を可視化：RelaxedOneHotCategorical分布の温度パラメータの影響を分析

torch.distributions.relaxed_categorical.RelaxedOneHotCategorical.temperature は、RelaxedOneHotCategorical 分布の重要なパラメータであり、サンプリングされるカテゴリの集中度を制御します。温度が高いほど、サンプリングされるカテゴリは一様になり、温度が低いほど、サンプリングされるカテゴリは集中します。

数学的な定義

RelaxedOneHotCategorical 分布の確率密度関数は以下の式で表されます。

p(y | logits, temperature) = exp((y - logits) / temperature) / Z

ここで、

Z は、正規化定数であり、以下の式で計算されます。
temperature は、温度と呼ばれるパラメータです。
logits は、各カテゴリの非正規化対数確率を表すベクトルです。
y は、K 個のカテゴリのうち 1 つのカテゴリを表す one-hot ベクトルです。

Z = sum(exp((y - logits) / temperature))

温度の影響

温度が低い場合:
- 確率の高いカテゴリが 優先的に サンプリングされます。
- 分布は集中します。
- 予測されるカテゴリの 確実性 が高くなります。
温度が高い場合:
- 各カテゴリがサンプリングされる確率はほぼ 等しくなります。
- 分布は一様に近づきます。
- 予測されるカテゴリの 不確実性 が高くなります。

例

以下のコード例は、RelaxedOneHotCategorical 分布を使用してカテゴリをサンプリングする方法を示します。

import torch
from torch.distributions import RelaxedOneHotCategorical

# ロジッ トを定義
logits = torch.tensor([1.0, 2.0, 3.0])

# 温度を定義
temperature = 0.5

# RelaxedOneHotCategorical 分布を作成
distribution = RelaxedOneHotCategorical(logits=logits, temperature=temperature)

# 10 個のサンプルをサンプリング
samples = distribution.rsample(sample_shape=(10,))

# 結果を出力
print(samples)

このコードを実行すると、以下の出力が得られます。

tensor([[ 0.0332,  0.9268,  0.0399],
        [ 0.0144,  0.0456,  0.9400],
        [ 0.9517,  0.0433,  0.0050],
        [ 0.0340,  0.8781,  0.0879],
        [ 0.0101,  0.9880,  0.0019],
        [ 0.0041,  0.0159,  0.9799],
        [ 0.8852,  0.1098,  0.0050],
        [ 0.0297,  0.9024,  0.0679],
        [ 0.0012,  0.0018,  0.9969],
        [ 0.0475,  0.8894,  0.0631]])

温度が 0.5 の場合、各カテゴリがほぼ等しい確率でサンプリングされていることがわかります。

import matplotlib.pyplot as plt
import numpy as np
import torch
from torch.distributions import RelaxedOneHotCategorical

# ロジッ トを定義
logits = torch.tensor([1.0, 2.0, 3.0])

# 温度のリストを定義
temperatures = np.linspace(0.0, 10.0, 20)

# 各温度におけるサンプリング結果を格納するリスト
samples_list = []

# 各温度についてサンプリングを実行
for temperature in temperatures:
    # RelaxedOneHotCategorical 分布を作成
    distribution = RelaxedOneHotCategorical(logits=logits, temperature=temperature)

    # 100 個のサンプルをサンプリング
    samples = distribution.rsample(sample_shape=(100,))

    # サンプリング結果をリストに追加
    samples_list.append(samples.numpy())

# 結果を可視化
fig, axs = plt.subplots(4, 5, figsize=(15, 10))

for i, temperature in enumerate(temperatures):
    samples = samples_list[i]

    # カテゴリごとにサンプリング数をカウント
    counts, _ = np.histogram(samples, bins=3, range=(0, 3))

    # 棒グラフを作成
    axs[i // 5, i % 5].bar(np.arange(3), counts)

    # タイトルと軸ラベルを設定
    axs[i // 5, i % 5].set_title(f"Temperature: {temperature:.2f}")
    axs[i // 5, i % 5].set_xlabel("Category")
    axs[i // 5, i % 5].set_ylabel("Count")

# 全体的なタイトルと軸ラベルを設定
fig.suptitle("RelaxedOneHotCategorical Distribution: Temperature vs. Sampling Distribution")
plt.xlabel("Temperature")
plt.tight_layout()
plt.show()

図中の棒グラフは、各温度におけるサンプリング結果の分布を示しています。温度が低いほど、特定のカテゴリがサンプリングされる確率が高くなり、分布が集中します。逆に、温度が高いほど、各カテゴリがサンプリングされる確率がほぼ等しくなり、分布が均一になります。

この可視化コードは、RelaxedOneHotCategorical 分布における temperature パラメータの影響を理解するのに役立ちます。

コードを実行するには、PyTorch 1.10 以降が必要です。
このコードは、matplotlib と numpy ライブラリを使用しています。これらのライブラリがインストールされていない場合は、事前にインストールする必要があります。

このパラメータの代替方法として、以下の方法が考えられます。

温度スケーリング

temperature パラメータを直接設定する代わりに、スケーリング因子を使用して調整することができます。例えば、以下の式のように logits をスケーリングすることができます。

scaled_logits = logits / temperature

この方法により、温度が低い場合でも、より高い確率で高確率なカテゴリがサンプリングされるようになります。

混合分布

RelaxedOneHotCategorical 分布と Categorical 分布を混合することで、温度を制御することができます。混合比は、温度に基づいて決定することができます。例えば、以下の式のように混合比を定義することができます。

mixing_prob = 1.0 / (1.0 + temperature)

この方法により、温度が低い場合、Categorical 分布から直接サンプリングされるようになり、より集中した分布を得ることができます。

カスタム分布

独自の分布を定義することで、温度以外の要因も考慮したサンプリングを行うことができます。例えば、以下の式のように、温度とカテゴリのエンベディングに基づいて確率を計算する分布を定義することができます。

def custom_distribution(logits, temperature, embeddings):
    # 温度に基づいて確率を計算
    probs = torch.exp((logits - embeddings) / temperature)

    # 正規化
    probs = probs / probs.sum(dim=-1, keepdim=True)

    # サンプリング
    return Categorical(probs=probs).rsample()

この方法は、より柔軟な制御が可能ですが、実装が複雑になるというデメリットがあります。

最適な方法の選択

最適な方法は、具体的なタスクや要件によって異なります。以下の点を考慮して選択する必要があります。

実装の容易さ: 温度スケーリングは最も簡単ですが、混合分布やカスタム分布は実装が複雑になる可能性があります。
計算コスト: カスタム分布は、他の方法よりも計算コストが高くなる可能性があります。
必要な集中度: 温度が高いほど、サンプリングされるカテゴリは一様になります。

上記以外にも、様々な代替方法が考えられます。最適な方法は、具体的な状況に合わせて検討してください。