TensorFlowにおける確率分布の基本的な使い方

TensorFlowに限った話ではありませんが、最近は多くの深層学習ライブラリで確率分布関係の実装が充実してきているようです。生成系のモデルを作るときに使えるらしいので、学んだことをメモしておきます。

# Distributionクラス
TensorFlowのPython APIでは、確率分布はすべて`Distribution`クラスを基底とするクラスとして定義されています。僕の理解の範囲では、使うことがありそうな`Distribution`クラスの属性とメソッドはこのあたりです。詳細は[tf.distributions.Distribution](https://www.tensorflow.org/api_docs/python/tf/distributions/Distribution)のドキュメントをご参照ください。

## 属性

|名前|説明|
|:--|:--|
|batch_shape| バッチの形、つまり個々のサンプルの並びの形 |
|event_shape| 個々のサンプルの形 |
|parameters| 各確率モデルが持つパラメータ |

## メソッド

|名前|説明|
|:--|:--|
|prob| 確率変数が入力値と等しくなる確率(密度)を計算する |
|log_prob| `prob`の対数を取ったものを計算する |
|entropy|確率分布のエントロピーを計算する|
|cross_entropy|確率分布間の交差エントロピーを計算する|
|kl_divergence|確率分布間のKL距離を計算する|
|mean|確率変数が取る平均値を計算する|
|variance|確率変数の各次元の分散を計算する|
|covariance|確率変数の分散・共分散行列を計算する|

# 使用例
具体例として、独立した2次元の確率変数上の正規分布を`MultivariateNormalDiag`を使って定義し、そこからサンプリングを行いました。以下のコードは簡潔さのために、TensorFlowのEager Execution機能を利用しています。

```python
import numpy as np
import tensorflow as tf
import tensorflow.contrib.distributions as tfd
import tensorflow.contrib.eager as tfe
import matplotlib.pyplot as plt

tfe.enable_eager_execution()

# モデル1を定義
# 平均: [0, 0]
# 分散・共分散行列:
#   [[2, 0],
#    [0, 1]]
# の正規分布
model1 = tfd.MultivariateNormalDiag(
    loc = np.zeros(2, dtype=np.float32),
    scale_diag = np.array([2,1], dtype=np.float32)
)
# モデル1から1000個の点をサンプリング
sampled = model1.sample(1000).numpy() # shape=(1000, 2)

# プロット
plt.figure(figsize=(5,5))
plt.scatter(x=sampled[:,0], y=sampled[:,1])
plt.xlim([-5.0, 5.0])
plt.ylim([-5.0, 5.0])
```
> 出力:
> ![mul_gaussian_plot.png](https://qiita-image-store.s3.amazonaws.com/0/238956/bbe15b42-ae37-ed7b-d04c-8a53d0f06b9e.png)

## KL距離の計算
同じ種類の確率分布の間であれば、KL距離等の計算が簡単にできます。モデルをもうひとつ定義して計算してみます。

```python
# モデル2を定義
# 平均: [3, 3]
# 分散・共分散行列:
#   [[1, 0],
#    [0, 1]]
# の正規分布
model2 = tfd.MultivariateNormalDiag(
    loc = np.array([3,3], dtype=np.float32),
    scale_diag = np.array([1,1], dtype=np.float32)
)
print(model1.kl_divergence(other=model2))
```
> 出力: tf.Tensor(9.806852, shape=(), dtype=float32)

これで、例えば変分自己符号化器(VAE)の損失関数の項が簡単に計算できたりします[(gist)](https://gist.github.com/kozamb/9edac660e882462d591c7e549328218c)。

# Independentクラスによるbatch_shapeとevent_shapeの指定
例として、10x28x28の大きさの確率変数をもった確率分布があったとします。この確率分布は、28x28のサイズの画像を10枚生成するためのモデルかも知れませんし、そうではなく本当に10x28x28の大きさの3次元配列を生成するためのモデルかもしれません。形式的な違いですが、この違いを`Distribution`クラスの属性である`batch_shape`と`event_shape`を使って表すことができます。

例として、次のようにモデルを定義してみます。

```python
model = tfd.Bernoulli(np.zeros((10,28,28), dtype=np.float32))
```
デフォルトではこの`model`の`batch_shape`と`event_shape`は次のようになっています。

|変数|値|
|:--|:-:|
|model.batch_shape|[10,28,28]|
|model.event_shape|[ ]|

これは、ゼロ次元、すなわちスカラー値である確率変数を[10x28x28]個生成するモデルであることを表しています。これを[28x28]の大きさの確率変数を10個生成するモデルとして扱いたい場合には、`Independent`クラスを利用して変形させます。

```python
model = tfd.Independent(model, reinterpreted_batch_ndims=2)
```
引数の`reinterpreted_batch_ndims`で、元の`batch_shape`の後ろからいくつを`event_shape`に移すかを指定しています。これを指定しない場合は、(`batch_shape`の次元)-1を指定したのと同じ扱いになるようです。バッチサイズは1次元であることがほとんどでしょうから、実際に指定する場面は稀でしょう。上の処理によって、`model`の中身は次のようになりました。

|変数|値|
|:--|:-:|
|model.batch_shape|[10]|
|model.event_shape|[28,28]|

意図した形になったことがわかります。

# Mixtureクラスによる混合分布の作成
ガウス混合分布(GMM)等、複数の分布を重み付けして足し合わせたいことがあります。そのような時は`Mixture`クラスを利用します。先ほど定義した`model1`と`model2`を同じ重みで足しあわせてサンプリングしてみました。

```python
# model1とmodel2を同じ重み0.5で足し合わせる
gm = tfd.Mixture(
    cat=tfd.Categorical([.5, .5]),
    components=[model1, model2])

# プロット
sampled = gm.sample(1000).numpy()
plt.figure(figsize=(5,5))
plt.scatter(x=sampled[:,0], y=sampled[:,1])
plt.xlim([-5.0, 5.0])
plt.ylim([-5.0, 5.0])
```
> 出力:
> ![gaussian_mixture.png](https://qiita-image-store.s3.amazonaws.com/0/238956/dfb8840f-c219-947c-c05d-0d7ca1a89f7c.png)

出力を見るに、混合分布からうまくサンプリングできているようです。注意が必要なのは、`MultivariateNormalDiag`や`Bernoulli`といったピュアな確率分布で使えていた`entropy`や`kl_divergence`といったメソッドが、`Mixture`にした時点で使えなくなるという点です。latent spaceに複雑な分布を想定したモデルの実装等はまだ難しいかもしれません。

# まとめ
TensorFlowにおける確率分布の基本的な使い方を調べました。

# 参考
* [Module: tf.contrib.distributions | TensorFlow]( https://www.tensorflow.org/api_docs/python/tf/contrib/distributions)
* [[1711.10604] TensorFlow Distributions](https://arxiv.org/abs/1711.10604)