PyTorch Probability Distributions: 変換による確率密度関数の変化量を計算する「log_abs_det_jacobian」の解説と詳細なコード例

理解を深めるためのポイント

• 対数: 絶対値の対数を取ることで、ヤコビアン行列式のスケールを調整します。
• 絶対値: ヤコビアン行列式の絶対値は、行列の各要素の絶対値を用いて算出されます。
• ヤコビアン行列式: 変換関数 f(x) がベクトル値出力を持つ場合、ヤコビアン行列式は f(x) の各出力成分に対する入力変数の偏導関数の行列を表します。

メソッドの役割

log_abs_det_jacobian() メソッドは、確率分布を変換する際に、変換関数による確率密度関数の変化量を計算するために用いられます。具体的には、以下の役割を果たします。

1. 変換関数 f(x) の絶対ヤコビアン行列式 |J(f(x))| を計算します。
2. |J(f(x))| の対数 log(|J(f(x))|) を計算します。
3. 計算結果を返します。

import torch
from torch.distributions import Normal, TransformedDistribution
from torch.distributions.transforms import AffineTransform

# 元となる確率分布を定義
base_distribution = Normal(torch.zeros(1), torch.ones(1))

# アフィン変換を定義
affine_transform = AffineTransform(loc=1.0, scale=2.0)

# 変換された確率分布を定義
transformed_distribution = TransformedDistribution(base_distribution, affine_transform)

# サンプルデータを作成
x = torch.randn(10)

# 密度関数を計算
log_prob = transformed_distribution.log_prob(x)

# log_abs_det_jacobian() を用いて変換による密度関数の変化量を計算
jacobian = transformed_distribution.transforms.log_abs_det_jacobian(x)

# 密度関数と変換による変化量を足し合わせる
adjusted_log_prob = log_prob + jacobian

print(adjusted_log_prob)

上記の例では、まず元となる確率分布 base_distribution とアフィン変換 affine_transform を定義します。その後、TransformedDistribution クラスを用いて変換された確率分布 transformed_distribution を定義します。サンプルデータ x を作成し、log_prob() メソッドを用いて密度関数を計算します。

続いて、log_abs_det_jacobian() メソッドを用いて変換による密度関数の変化量 jacobian を計算します。最後に、密度関数 log_prob と変化量 jacobian を足し合わせて、変換後の密度関数 adjusted_log_prob を求めます。

---

import torch
import torch.nn as nn
import torch.distributions as distributions
from torch.distributions.transforms import AffineTransform, Compose

# データ生成
data = torch.randn(1000, 2)

# 標準正規分布を定義
base_distribution = distributions.Normal(torch.zeros(2), torch.ones(2))

# スケーリングとシフトを行うアフィン変換を定義
affine_transform1 = AffineTransform(scale=2.0, loc=torch.tensor([1.0, 0.0]))
affine_transform2 = AffineTransform(scale=0.5, loc=torch.tensor([0.0, 1.0]))

# 変換を合成
transform = Compose([affine_transform1, affine_transform2])

# 変換付き確率分布を定義
transformed_distribution = distributions.TransformedDistribution(base_distribution, transform)

# 密度関数と確率質量関数を計算
pdf = transformed_distribution.pdf(data)
cdf = transformed_distribution.cdf(data)

# 結果を可視化
import matplotlib.pyplot as plt

plt.figure(figsize=(10, 6))

# 密度関数
plt.subplot(1, 2, 1)
plt.contourf(data[:, 0], data[:, 1], pdf.detach().numpy())
plt.colorbar(label='PDF')
plt.title('PDF')

# 確率質量関数
plt.subplot(1, 2, 2)
plt.contourf(data[:, 0], data[:, 1], cdf.detach().numpy())
plt.colorbar(label='CDF')
plt.title('CDF')

plt.show()

1. 2次元標準正規分布に従うデータ data を生成します。
2. 標準正規分布 base_distribution を定義します。
3. スケーリングとシフトを行うアフィン変換 affine_transform1 と affine_transform2 を定義します。
4. Compose クラスを用いて、2つのアフィン変換を合成した変換 transform を定義します。
5. 変換付き確率分布 transformed_distribution を定義します。
6. 生成データ data に対する密度関数 pdf と確率質量関数 cdf を計算します。
7. 密度関数と確率質量関数の等高線を可視化します。

このコード例は、torch.distributions.transforms モジュールと TransformedDistribution クラスを用いて、確率分布をどのように変換し、その結果を可視化できるかを示しています。

このコード例はあくまでも一例であり、状況に応じて様々な改造が可能です。例えば、以下のような変更が考えられます。

• 可視化の方法を変更する
• 変換の種類を変更する
• 使用する確率分布を変更する

---

以下では、log_abs_det_jacobian() の代替方法として検討できるいくつかの方法を紹介します。

数値微分

数値微分を用いて、ヤコビアン行列式を直接計算する方法です。この方法は、比較的シンプルですが、計算精度と計算量とのバランスを調整する必要があります。

import torch

def numerical_jacobian(f, x):
    n = x.shape[1]
    jac = torch.zeros(x.shape[0], n, n)
    eps = torch.eye(n, device=x.device) * 1e-5
    for i in range(n):
        jac[:, i, :] = (f(x + eps[:, i, :]) - f(x - eps[:, i, :])) / (2 * eps[i, i])
    return jac.det().log()

# ... (データ生成 and 変換定義) ...

# 数値微分を用いて log_abs_det_jacobian を計算
jacobian = numerical_jacobian(transformed_distribution.transforms.forward, x)

解析的な計算

変換関数が単純な場合、解析的にヤコビアン行列式と絶対値を計算することができます。この方法は、数値微分よりも精度が高く、計算量も少ない場合があります。

import torch

class MyTransform(Transform):
    def forward(self, x):
        # ... (変換関数) ...

    def log_abs_det_jacobian(self, x):
        # ... (解析的な計算) ...

# ... (データ生成 and 変換定義) ...

# 解析的な計算を用いて log_abs_det_jacobian を計算
jacobian = transformed_distribution.transforms.log_abs_det_jacobian(x)

近似手法

ヤコビアン行列式の近似計算を行う手法です。代表的な手法として、以下のようなものがあります。

• ランダムサンプリング: ランダムな方向ベクトルを用いて、ヤコビアン行列式の積分を近似的に計算します。
• 線形近似: 変換関数を局所的に線形関数で近似し、ヤコビアン行列式を計算します。

これらの手法は、数値微分よりも精度が低くなる場合がありますが、計算量を大幅に削減することができます。

GPU 活用

計算量が多い場合は、GPU を活用することで計算時間を短縮することができます。PyTorch は GPU をサポートしており、log_abs_det_jacobian() メソッドも GPU 上で実行することができます。

import torch.cuda.amp as amp

# ... (データ生成 and 変換定義) ...

# GPU 上で計算
with amp.autocast():
    # ... (計算処理) ...

どの代替方法が適切かは、状況によって異なります。考慮すべき要素は以下の通りです。

• 変換関数の複雑さ: 解析的な計算は、変換関数が単純な場合のみ可能です。
• 計算時間: GPU 活用は計算時間を短縮できますが、GPU 環境が必要となります。
• 計算量: 数値微分は最も計算量が多く、近似手法は最も計算量が少ない
• 精度: 数値微分は最も精度が高く、近似手法は最も精度が低くなります。

これらの要素を考慮し、状況に合った代替方法を選択することが重要です。

上記以外にも、以下のような選択肢が考えられます。

• 専門家への相談: 複雑な変換関数の場合は、確率論や統計学の専門家に相談することを検討しても良いでしょう。
• ライブラリの活用: torch-geometry などのライブラリには、ヤコビアン行列式を計算する関数などが用意されている場合があります。