はじめに
TensorFlow Probabiliryのコードの書き方を数式と対応付けて記載してまとめておきます。
確率分布まわりの基本操作
正規分布を使って色々見ていきます。
確率分布
$$ p(x \mid \mu, \sigma) = \mathcal N(x∣\mu ,\sigma) $$
p = tfd.Normal(loc=mu, scale=sigma)
確率分布からのサンプリング
上記で宣言した確率分布からデータ $x$ を1つだけサンプリングする場合 $$ x \sim p(x\mid \mu, \sigma) $$
x = p.sample()
宣言した確率分布から独立にデータを100個サンプリングする場合
$$ \begin{align} x_i &\sim p(x_i \mid \mu, \sigma) \\ X &= \{x_1, \cdots, x_{100} \} \end{align} $$
X = p.sample(100)
対数尤度の計算
既にサンプリングされたデータ$X$ について、これが確率分布 $p$ から生ずる確率を計算したものを尤度と言う。 下記はサンプル数がN個の場合の尤度の計算である。
$$ \begin{align} {\rm likelihood}(X, \mu, \sigma) & = \prod_{i=1}^{N} p(x_i \mid \mu, \sigma) \\ & = \prod_{i=1}^{N} \mathcal N(x_i \mid∣\mu ,\sigma) \\ & =\prod_{i=1}^{N} \frac{1}{\sqrt{2\pi \sigma^2}} \exp \left(−\frac{(x_i - \mu)^2}{2\sigma^2}\right) \end{align} $$
この計算は掛け算が入っており少々ややこしい。なので通常は下記のように対数を取った「対数尤度」を用いる。
$$
\begin{align}
{\rm log\ likelihood}(X, \mu, \sigma)
& = \log ({\rm likelihood}(X, \mu, \sigma)) \\
& = \log \left(\prod_{i=1}^{N} \mathcal N(x_i \mid \mu ,\sigma)\right) \\
& = \sum_{i=1}^{N} \log(\mathcal N(x_i \mid \mu, \sigma)) \\
& =-\frac{N}{2}\log(2\pi\sigma^2) - \frac{1}{2\sigma^2}\sum_{i=1}^N (x_i-\mu)^2
\end{align}
$$
対数を取ると「掛け算は足し算」になる。そのお陰で、個々のデータの尤度（発生確率）の対数を取って、最後に和を取ればデータ全体の対数尤度の計算が可能になる。また「割り算は引き算」、「累乗は掛け算」になる。この性質を使えば最後の式が簡単に出てくる。対数関数の単調増加性によって、尤度が増加するならば対数尤度も増加するので勾配法などでも計算が簡単な対数尤度を使う。
likelihood = tf.reduce_sum(p.log_prob(X))
最尤推定
仮定するモデルとその対数尤度関数
$X = \{x_1,\cdots, x_N \}$とデータが手元にある場合、モデルを宣言（仮定）するところから推定をはじめる。どのような仮定とするかについては、データの可視化を行ったり事前知識を用いて決める。今回はその工程を一気に飛ばしてデータ $X$ が正規分布していると仮定する。
この場合に求めたいのは データをよく表すことのできる正規分布の形であり、正規分布の形を決めるのは平均と標準偏差を表すパラメータ $\mu, \sigma$ である。最尤推定とは仮定したモデルのパラメータ（今回の場合 $\mu, \sigma$）を尤度関数が最大になるように決定することである。既に述べたように、実際には対数尤度を用いる。正規分布の対数尤度は下記であった。この関数は $\mu, \sigma$ の関数であることを意識すればよい。$X$ については既に手に入っているデータであり、数式から見たら定数にすぎない。
$$
\begin{align}
{\rm log\ likelihood}(\mu, \sigma)
& = \sum_{i=1}^{N} \log(\mathcal N(x_i \mid \mu, \sigma)) \\
& =-\frac{N}{2}\log(2\pi\sigma^2) - \frac{1}{2\sigma^2}\sum_{i=1}^N (x_i-\mu)^2
\end{align}
$$
この式を最大化するのが最尤推定であったのでTensorFlow Probabilityでは、上記の関数の符号を入れ替えたものを「損失関数」だとみなして学習を実施すれば良い。わざわざ式を書き下したが、コード上では2つ目の式だけを意識すればよく、
def log_likelihood_fn(X, mu, sigma): p = tfd.Normal(loc=mu, scale=sigma) log_likelihood = p.log_prob(X) return tf.reduce_sum(log_likelihood) # init variable mu = tf.Variable(0.) sigma = tf.Variable(1.) # loss_fn is function of mu and sigma loss_fn = lambda mu, sigma : - log_likelihood_fn(X=X_train, mu=mu, sigma=sigma)
という関数を書けば良い。ここまでで「モデルを仮定する」ことは終了している。仮に平均 $\mu$ だけを求めたいのであれば、loss_fnをmuだけの関数にしてしまえばいい（log_likelihood_fn()にsigma=tf.constant(1.0)などを与えてしまえばいい）。兎にも角にも最尤推定では尤度関数を最大化する。尤度関数では「どんな確率分布を仮定しているのか」を明記する。そして尤度関数の対数を取った対数尤度関数を作り、更にその負を損失関数として最小化を試みる。損失関数を作る時点で「何を定数とみなし、何をパラメータとみなす」かを決めてやれば、いよいよ準備完了というわけだ。
学習ループ
適当な最適化手法を決める。勾配方向の計算と大きさをどう調整するかがポイントであるが、今回はよく使われるAdamという手法を使うことにする。正規分布の最尤推定は凸最適化であるし、収束の速度を気にするほど難しい問題でもない。なので別になんでも良い（そもそも解析的に解ける）。
optimizer = tf.train.AdamOptimizer()
value_and_grads = tf.contrib.eager.implicit_value_and_gradients(loss)
for i in range(10000):
loss, grads = value_and_grads(mu, sigma)
optimizer.apply_gradients(grads)
if i % 1000 == 0:
print("likelihood ", - loss.numpy())
下記は、青が学習データのヒストグラム、緑が学習後のパラメータを使った正規分布からのサンプリングである。いい具合に平均と分散を調整してくれているのが分かる。
