NumPyro：各分布に関してまとめ

連載している記事の１つです。以前までの記事を読んでいる前提で書いているので、必要であればNumPyroの記事一覧から各記事を参考にしてください。

はじめに

ベイズを扱う際に避けて通れないのが、数多くある確率分布です。今回はこれまで扱ってきた分布をまとめていきます。

分布

分布の形状を可視化するための準備を行います。

from jax import random
from jax import numpy as jnp
import numpyro.distributions as dist
import seaborn as sns
import matplotlib.pyplot as plt

def plot_dist(list_dist, num_sample=10000, kde=True):
    plt.figure(figsize=(5, 4))
    for i, dist in enumerate(list_dist):
        samples = dist.sample(random.PRNGKey(0), (num_sample,))
        sns.histplot(samples, kde=kde, stat="density", label=str(i))
    plt.xlabel("X")
    plt.legend()

参考

• https://en.wikipedia.org/wiki/Probability_distribution
• StanとRでベイズ統計モデリング
• https://mc-stan.org/docs/functions-reference/discrete-distributions.html#discrete-distributions

連続確率分布

連続型一様分布

NumPyroでの実装

d = [dist.Uniform(low=0, high=10)]
plot_dist(d)

Probability density function（PDF）

$$Uniform(x | a, b) = \left\{ \begin{array}{ll} \frac{1}{b - a} & (a \le x \le b) \\ 0 & (otherwise) \end{array} \right.$$

support

$$[a, b]の範囲の実数$$

parameters

$$-\inf < a < b < \inf$$

平均

$$\frac{a + b}{2}$$

分散

$$\frac{(b - a)^2}{12}$$

使用例

無情報事前分布として使用されます。Stanでは事前分布を指定しない限り、とりうる値の範囲内で一様な一様分布が設定されます。NumPyroでは何かしらの事前分布を指定する必要があるので、極端に広いパラメータを指定しておけば同じような結果になります。

ベータ分布

NumPyroでの実装

# concentration1：alpha
# concentration0：beta
list_dist = [
    dist.Beta(concentration1=0.5, concentration0=0.5),
    dist.Beta(concentration1=5, concentration0=1.0),
    dist.Beta(concentration1=2, concentration0=2),
    dist.Beta(concentration1=1, concentration0=1),
]
plot_dist(list_dist, kde=True)

Probability density function（PDF）

$$Beta(x | \alpha, \beta) = \frac{1}{B(\alpha, \beta)}x^{\alpha-1}(1-x)^{\beta-1}$$

support

$$0 \le x \le 1の範囲の実数$$

parameters

$$\alpha：正の実数 \\ \beta：正の実数$$

平均

$$\frac{\alpha}{\alpha+\beta}$$

分散

$$\frac{\alpha\beta}{(\alpha+\beta)^2(\alpha+\beta+1)}$$

使用例
$0 \le x \le 1$の範囲の実数を出力するので確率を生成する分布として使用できる。そのため、ベルヌーイ分布、二項分布、負の二項分布、幾何学分布のパラメータpを生成する分布として使用されることが多い。それ以外にも故障率などの割合を表すことも可能。

その他
事前分布として何も情報を持たせたくない場合は、\alpha=\beta=1とすることで0~1の範囲の一様分布と同じになる。
ベータ分布は平均値μ（0＜μ＜1）と2つの形状パラメータの和ν＝α＋β＞0で再パラメータ化することができ、平均値が重要な統計量の場合に有効。

ディリクレ分布

NumPyroでの実装

# 三角ダイアグラムはサボりました。。。
samples = dist.Dirichlet(concentration=jnp.ones(2)).sample(random.PRNGKey(0), (100,))
plt.scatter(samples[:, 0], samples[:, 1])

Probability density function（PDF）

$$Dirichlet(\vec{x} | \vec{\alpha}) = \frac{1}{B(\alpha)}\prod_{i=1}^{K}x_i^{\alpha_i-1}$$

support

$$長さKのベクトル. 各要素は0 \sim 1の範囲の実数で合計が１.$$

parameters

$$\vec{\alpha}：正の実数を要素にもつ長さKのベクトル$$

平均

$$x_iの平均：\frac{\alpha_k}{\alpha_{sum}}$$

分散

$$x_iの分散：\frac{\alpha_k(\alpha_{sum}-\alpha_k)}{\alpha_{sum}^2(\alpha_{sum}+1)}$$

使用例
ベータ分布の多変量版。カテゴリカル分布や多項分布のパラメータ$\vec{p}$を生成するのに使われることが多い。他にも３点以上の重みづけ平均を求める際の重みや長さLの文字列をK分割する際の割合などにも使用可能。

その他
$\vec{\alpha}$の値を全て同じ値で固定した分布として対称ディリクレ分布がある。どの確率も等しくなるようにしたいという事前知識がある場合は使える。

指数分布

NumPyroでの実装

list_dist = [
    dist.Exponential(rate=0.5),
    dist.Exponential(rate=1.0),
    dist.Exponential(rate=1.5),
]
plot_dist(list_dist, kde=False)

Probability density function（PDF）

$$Exponential(x | \lambda) = \lambda \exp(-\lambda x)$$

support

$$0以上の実数$$

parameters

$$\lambda：正の実数$$

平均

$$\frac{1}{\lambda}$$

分散

$$\frac{1}{\lambda^2}$$

使用例
機械が故障するまでの時間などイベントが起こるまでの時間を表す際に使用される。また時間だけでなく、DNA鎖の突然変異の間の距離など、あるイベントが単位長さあたり一定の確率で発生する状況をモデル化する際にも使用される。また、スケールなど正の実数のパラメータを生成する分布として使用される。

その他
無記憶性を持つ唯一の連続型の確率分布で、各時間単位でイベントが起こる確率が等しいとしているモデル。そのため、12-13時の昼休憩時と13-14時の間では電話がかかってくる確率が異なるので、13-14時だけに限定してモデル化する際にはよく表現できる。
また、あるイベントが起こる時間や長さが複数の独立した作業の連続と考えられるプロセスの場合は、Erlang分布（独立した複数の指数分布変数の和の分布）に従う。

ガンマ分布

NumPyroでの実装

list_dist = [
    dist.Gamma(concentration=1, rate=1),
    dist.Gamma(concentration=1, rate=3),
    dist.Gamma(concentration=3, rate=1),
]
plot_dist(list_dist, kde=False)

Probability density function（PDF）

$$Gamma(x | \alpha, \beta) = \frac{\beta^\alpha}{\Gamma(\alpha)}x^{\alpha-1}\exp(-\beta x)$$

support

$$正の実数$$

parameters

$$\alpha：正の実数（shape parameter） \\ \beta：正の実数（rate parameter）$$

平均

$$\frac{\alpha}{\beta}$$

分散

$$\frac{\alpha}{\beta^2}$$

使用例
待ち時間をモデル化する際に使用される。また、正の実数値を取るので正のパラメータを生成する際に使用される。

その他
$\alpha$が十分大きい場合は、$Normal(\frac{\alpha}{\beta}, \frac{\sqrt{\alpha}}{\beta})$に近似できる。

1次元正規分布

NumPyroでの実装

list_dist = [
    dist.Normal(loc=0, scale=1),
]
plot_dist(list_dist, kde=False)

Probability density function（PDF）

$$Normal(x | \mu, \sigma) = \frac{1}{\sqrt{2\pi}\sigma}\exp(-\frac{1}{2}(\frac{y-\mu}{\sigma})^2)$$

support

$$実数$$

parameters

$$\mu：実数 \\ \sigma：正の実数$$

平均

$$\mu$$

分散

$$\sigma^2$$

使用例
多くのデータに当てはめることができる。モデリングする際も階層モデルで個人差やグループ差を考慮する際に使用される。

その他
裾野が短く外れ値に弱いのでコーシー分布やStudentのt分布が代わりに使用される。正規分布に従うxに対して新しい確率変数$exp(x)$は対数正規分布に従う。対数正規分布は裾野が長く、指数分布やガンマ分布より長く、コーシーより短くなる。

多次元正規分布

NumPyroでの実装

samples = dist.MultivariateNormal(loc=jnp.zeros(2), covariance_matrix=jnp.eye(2)).sample(random.PRNGKey(0), (1000,))
plt.scatter(samples[:, 0], samples[:, 1])

Probability density function（PDF）

$$MultivariateNormal(\vec{x} | \vec{\mu}, \Sigma) = \frac{1}{(2\pi)^{K/2}}\frac{1}{\sqrt{|\Sigma|}}\exp(-\frac{1}{2}(\vec{x}-\vec{\mu})^T\Sigma^{-1}(\vec{x}-\vec{\mu}))$$

support

$$\vec{x}：各要素が実数の長さKのベクトル$$

parameters

$$\vec{\mu}：各要素が実数の長さKのベクトル \\ \Sigma：正定値行列$$

平均

$$y_kの平均：\mu_k$$

分散

$$y_kの分散：\Sigma_{k,k}$$

使用例
相関がある２つの変数の分布をモデル化する際に使用する。ガウス過程を扱う際にも使用される。

その他
$\Sigma$の対角成分以外が0の場合はお互いに独立なK個の１次元正規分布から$\vec{x}$を生成しているのと同じ。

コーシー分布

NumPyroでの実装

list_dist = [
    dist.Cauchy(0, 1),
]
plot_dist(list_dist, kde=True)

Probability density function（PDF）

$$Cauchy(x | \mu, \sigma) = \frac{1}{\pi\sigma}\frac{1}{1 + ((x - \mu)/\sigma)^2}$$

support

$$実数$$

parameters

$$\mu：実数 \\ \sigma：正の実数$$

平均

$$undefined$$

分散

$$undefined$$

使用例
かなり裾が長い分布なので稀に大きな値を生成するため、外れ値を含むモデルや変化点の検出に使用される。他にも年間最大1日降雨量や河川流出量などの極端な事象に適用されます。

その他
独立な2つの正規分布Normal(0, 1)の比率はCauchy(0, 1)に従う。コーシー分布は、自由度1のスチューデントのt分布と一致する。

Studentのt分布

NumPyroでの実装

list_dist = [
    dist.StudentT(df=2, loc=0, scale=1),
    dist.StudentT(df=4, loc=0, scale=1),
    dist.StudentT(df=100, loc=0, scale=1),
]
plot_dist(list_dist, kde=True)
plt.xlim([-5, 5])

Probability density function（PDF）

$$StudentT(x | \nu, \mu, \sigma) = \frac{\Gamma((\nu+1)/2)}{\Gamma(\nu/2)\sqrt{\pi\nu}\sigma}(1+\frac{1}{\nu}(\frac{x-\mu}{\sigma})^2)^{-(\nu+1)/2}$$

support

$$実数$$

parameters

$$\nu：正の実数 \\ \mu：実数 \\ \sigma：正の実数$$

平均

$$\nu > 1の場合：0 \\ それ以外：undefined$$

分散

$$\nu > 2の場合：\sigma^2\nu/(\nu-2) \\ 1 < \nu \le 2の場合：∞ \\ それ以外：undefined$$

使用例
外れ値を含むモデルに使用される。

その他
自由度1のスチューデントのt分布はコーシー分布と一致する。自由度∞のスチューデントのt分布は正規分布と一致する。そのため自由度$\nu$は正規性パラメータとも呼ばれ、自由度が大きくなると裾が短くなるので、コーシー分布より裾を短くしたい場合は自由度が2~8の分布を使うことがある。
また、自由度はスチューデントのt分布の尖度を制御しスケールパラメータと相関がある。そのため、尤度は複数の局所最大値を持つことがあるため、自由度をかなり低い値（3~9, または5）に固定した上で他のパラメータを推定することがしばしば必要になる。

ラプラス分布

NumPyroでの実装

list_dist = [
    dist.Laplace(loc=0, scale=1),
    dist.Laplace(loc=0, scale=2),
]
plot_dist(list_dist, kde=True)

Probability density function（PDF）

$$Laplace(x | \mu, \sigma) = \frac{1}{2\sigma}\exp(-\frac{|x-\mu|}{\sigma})$$

support

$$実数$$

parameters

$$\mu：実数 \\ \sigma：正の実数$$

平均

$$\mu$$

分散

$$2\sigma^2$$

使用例
$\mu$を中心に鋭いピークがあるので回帰係数の事前分布に使用すると変数選択ができる。また、裾が正規分布より長いのでラプラス分布は年間最大1日降雨量や河川流量のような極端な事象に適用されることがある。

その他
2つの指数分布をつなぎ合わせたものと考えることができるため、二重指数分布と呼ばれることもある。指数分布と同様に裾が長い。$X \sim Laplace(0, \sigma)$の時、$|X| \sim Exponential(\sigma^{-1})$。
NumPyroではラプラス分布に似た形状であるがどこでも微分可能なdist.SoftLaplaceが実装されている。微分可能なためHMCやラプラス近似に適しています。また、２つの異なる指数分布をつなぎ合わせた分布として非対称ラプラス分布もある。

ワイブル分布

NumPyroでの実装

list_dist = [
    dist.Weibull(scale=1, concentration=1),
    dist.Weibull(scale=1, concentration=1.5),
    dist.Weibull(scale=1, concentration=5),
]
plot_dist(list_dist, kde=True)

Probability density function（PDF）

$$Weibull(x | \lambda, k) = \frac{k}{\lambda}(\frac{x}{\lambda})^{k-1}\exp(-(x/\lambda)^k)$$

support

$$正の実数$$

parameters

$$\lambda：正の実数 \\ k：正の実数$$

平均

$$\lambda\Gamma(1+1/k)$$

分散

$$\lambda^2[\Gamma(1+\frac{2}{k})-(\Gamma(1+\frac{1}{k}))^2]$$

使用例
k=1の時は指数分布とk=2の時はレイリー分布を表している。Xが故障までの時間とした時に、kの値に応じて初期不良、ランダムな外的要因による故障、摩耗などによる劣化のための故障などを表す分布と捉えることができ、幅広い分野で使用されている。

その他
故障率はワイブル分布がよく使われるが、化学反応や腐食が原因の製品故障は、通常、対数正規分布でモデル化される。ワイブル分布は極値分布の１つで、他にはガンベル分布とフレシェ分布がある。特に、ガンベル分布はさまざまな分布に従う確率変数の最大値（または最小値）が漸近的に従う分布であり、稀にしか発生しない地震や洪水などの自然災害の発生する確率をモデル化する際に使用される。

LKJ分布

Probability density function（PDF）

$$p(R|\eta) = C [det(R)]^{\eta-1} \\ C=2^{\sum_{k=1}^{d}(2\eta -2+d-k)(d-k)}\prod_{k=1}^{d-1}(B(\eta +(d-k-1)/2, \eta + (d-k-1)/2))^{d-k}$$

parameters

$$\eta：正の実数$$

使用例
共分散行列を生成する際に使用される。同様に、共分散行列を生成する分布として逆ウィシャール分布があるが問題が生じるためLKJ分布から生成した相関行列を使用して共分散行列を生成することが推奨されている。

その他
パラメータ$\eta$は、相関の強さを調整するもの。もし$\eta=1$の場合は 密度はすべての相関行列で一様で、$\eta > 1$の場合は、対角線が強い（つまり相関が小さい）行列が有利になる。もし$\eta < 1$の場合、対角線は弱く、相関が有利になる。
LKJ分布を使用した分散共分散行列の事前分布に関してはこちらを参照。
また、逆ウィシャール分布の課題やLKJ分布に関してはこちらがわかりやすかったです。

離散確率分布

離散型一様分布

NumPyroでの実装

d = [dist.DiscreteUniform(low=0, high=10)]
plot_dist(d, kde=False)

Probability density function（PDF）

$$Uniform(x | a, b) = \left\{ \begin{array}{ll} \frac{1}{b - a + 1} & (a \le x \le b) \\ 0 & (otherwise) \end{array} \right.$$

support

$$[a, b]の範囲の整数$$

parameters

$$-\inf < a < b < \infの整数$$

平均

$$\frac{a + b}{2}$$

分散

$$\frac{(b - a + 1)^2 - 1}{12}$$

使用例

サイコロのように事象が等しい確率で起こる際に使用される。

ベルヌーイ分布

NumPyroでの実装

list_dist = [
    dist.Bernoulli(probs=0.2),
]
plot_dist(list_dist, kde=False)

Probability density function（PDF）

$$Bernoulli(x | a, b) = \left\{ \begin{array}{ll} 1 - p & (x = 0) \\ p & (x = 1) \end{array} \right.$$

support

$$0 \space or \space 1$$

parameters

$$0 \le p \le 1$$

平均

$$p$$

分散

$$p(1 - p)$$

使用例
コイン投げで表か裏か、実験が成功か失敗かなど２値で表すことができる場合に使用されます。また0か0以外かの２値を表すベルヌーイ分布と他の分布を組み合わせることでゼロ過剰モデルやハードルモデルを定義できたりします。

その他
パラメータpは0~1の範囲の実数なので、何からの値に対してロジスティック関数を適用し0~1の範囲にした結果をpとして使用する場合が多いです。そのため、NumPyroでは引数の値に対してロジスティック関数を適用してくれるdist.BernoulliLogitsという関数も用意されています。
また、二項分布の特殊なケースで、１回の試行が行われた二項分布とみなすことができす。

二項分布

NumPyroでの実装

list_dist = [
    dist.Binomial(total_count=10, probs=0.2)
]
plot_dist(list_dist, kde=False)

Probability density function（PDF）

$$Binomial(x | n, p) = \frac{n!}{x!(n-x)!}p^x(1-p)^{n-x}$$

support

$$0 \sim nまでの整数$$

parameters

$$n：正の整数 \\ p：0 \sim 1の範囲の実数$$

平均

$$np$$

分散

$$np(1-p)$$

使用例
10回コイントスを行った中で表だった回数、10回実験を行った中で上手くいった回数を表す際に使用される。

その他
二項分布は全ての試行回数で同じ確率pを持つベルヌーイ分布を想定しているが、異なる確率を仮定する場合はPoisson binomial distributionやBeta-binomial distributionが使用できる。二項分布はPoisson binomial distributionのpが全て同じ特殊ケースとみなすことができる。
また、nが十分大きい場合は正規分布$Normal(np, np(1-p))$として近似することができる。
nが大きく、pが十分小さい場合はポアソン分布$Poisson(np)$として近似することができる。

カテゴリカル分布

NumPyroでの実装

list_dist = [
    dist.Categorical(probs=jnp.ones(6)*1/6)
]
plot_dist(list_dist, kde=False)

Probability density function（PDF）

$$Categorical(x | \pi) = \prod_{k=1}^K \pi_k^{s_k} \\ s_k \in {0, 1} \space and \space \sum s_k = 1$$

support

$$1 \sim Kの整数$$

parameters

$$K：正の整数 \\ \pi_i：0 \le \pi_i \le 1. \sum \pi_i = 1. 各カテゴリの確率$$

平均

$$s_kの平均値：\pi_k$$

分散

$$s_kの分散：\pi_k(1-\pi_k)$$

使用例
ベルヌーイ分布の多変数版。カテゴリやインデックスを表す際に使用される。例えば、K=6, $\pi_k=1/6$にすれば6面サイコロの出る目を表すことになる。

その他
パラメータ$\pi$は0~1の範囲の実数のK次元のベクトルなので、K次元のベクトルに対してsoftmax関数を適用し各要素が0~1の範囲で合計値が1になるK次元ベクトルにした結果をpとして使用する場合が多いです。
パラメータ$\pi$の事前分布としてはベータ分布を多次元化したディリクレ分布が使用されることが多い。

多項分布

NumPyroでの実装

ここでは0番目のカテゴリの回数の分布を表示した。

samples = dist.Multinomial(total_count=6, probs=jnp.ones(6)*1/6).sample(random.PRNGKey(0), (1000,))
sns.histplot(samples[:, 0], kde=False, stat="density")

Probability density function（PDF）

$$Multinomial(\vec{x} | n, \vec{p}) = \frac{n!}{\prod_{k=1}^{K}p_k!}\prod_{k=1}^{K}p_k^{x_k} \\$$

support

$$x_i：0 \sim nの整数で、\sum_i^{K} x_i = nを満たす$$

parameters

$$n：正の整数 \\ p_i：0 \le p_i \le 1, \sum p_i=1を満たす実数$$

平均

$$x_iの平均：np_i$$

分散

$$x_iの分散：np_i(1 - p_i)$$

使用例
サイコロを10回投げた時の各出目が何回ずつ出るかを表す。

その他
パラメータ$\pi$の事前分布としてはベータ分布を多次元化したディリクレ分布が使用されることが多い。
二項分布の時の内容と同じように過分散に対応した分布としてDirichlet-multinomial distributionがあります。この分布の確率を固定したものが多項分布です。

ポアソン分布

NumPyroでの実装

list_dist = [
    dist.Poisson(rate=1),
    dist.Poisson(rate=4),
    dist.Poisson(rate=10),
]
plot_dist(list_dist, kde=False)

Probability density function（PDF）

$$Poisson(x | \lambda) = \frac{1}{x!}\lambda^x\exp(-\lambda)$$

support

$$0以上の整数値$$

parameters

$$\lambda：正の実数値$$

平均

$$\lambda$$

分散

$$\lambda$$

使用例
ある一定期間で機械の故障などのイベントが発生する回数を表す際に使用される。

その他
パラメータ$\beta$の指数分布に従う時間間隔のイベントがある際に、その時間間隔の間で起こるイベントの回数はパラメータ$\beta$のPoisson分布に従う。そのため、指数分布の時と同様に、ある一定期間内に起こるイベントの発生割合が一定ではない場合は、通常のポアソン分布では上手く近似できないかもしれない。この場合は、混合ポアソン分布が使用できる可能性がある。また、必ず１回以上イベントが起こる場合は、zero-truncated Poisson distribution.を使用できる。イベントがゼロ回の数がポアソン分布で予測される数より多い場合は、Zero-inflated Poisson distributionを使用してモデル化することができる。また、過分散と過小分散のポアソン分布としてパラメータを追加したConway–Maxwell–Poisson distributionがある。
$\lambda$が十分大きい時は、$Normal(\lambda, \sqrt{\lambda})$で近似できる。また、Nが十分大きく、pが小さい場合は、$Binomial(n, p)$は$Poisson(np)$に近似できる。

事前分布

以下のStan関連のリンクに非常によくまとまっていました。

https://github.com/stan-dev/stan/wiki/Prior-Choice-Recommendations
https://mc-stan.org/docs/stan-users-guide/regression-priors.html
https://mc-stan.org/rstanarm/reference/priors.html

最後に

以上で「各分布に関してまとめ」は終わりです。情報の羅列なので適宜参考にしてください。次回は「ガウス過程」です。