StatModeling Memorandum

StatModeling Memorandum

StanとRとPythonでベイズ統計モデリングします. たまに書評.

PythonのSymPyで変分ベイズの例題を理解する

この記事の続きです。

ここではPRMLの10.1.3項の一変数ガウス分布の例題(WikipediaVariational_Bayesian_methodsのA basic exampleと同じ)をSymPyで解きます。すなわちデータが

  Y_n \sim Normal(\mu, \tau^{-1})   n = 1,..,N

に従い*1 \mu \tauが、

  \mu \sim Normal(\mu_0, (\lambda_0 \tau)^{-1})

  \tau \sim Gamma(a_0, b_0)

に従うという状況です。ここでデータ Y_n n=1,...,N)が得られたとして事後分布 p(\mu, \tau | \boldsymbol{Y})を変分ベイズで求めます。

まずはじめに、上記の確率モデルから同時分布 p(\boldsymbol{Y}, \mu, \tau)を書き下しておきます。

  p(\boldsymbol{Y}, \mu, \tau) = p(\boldsymbol{Y} | \mu, \tau) p(\mu | \tau) p(\tau)

なので、

  p(\boldsymbol{Y}, \mu, \tau) = \prod_{n=1}^N Normal(Y_n | \mu, \tau^{-1}) \cdot Normal(\mu | \mu_0, (\lambda_0 \tau)^{-1}) \cdot Gamma(\tau | a_0, b_0)

となります。

この問題は単純なので事後分布は厳密に求まるのですが、ここでは変分ベイズで解きます。すなわち、事後分布 p(\mu, \tau | \boldsymbol{Y}) q(\mu, \tau)で近似します。さらに q(\mu, \tau) = q(\mu) q(\tau)と因子分解可能と仮定します。そして、前の記事の最後の2つの式を使って、 q(\mu) q(\tau)が収束するまで繰り返し交互に更新して求めるのでした。以下ではこれをSymPyでやります。

from sympy import *
from sympy.stats import *
init_printing(use_unicode=True)

from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity = 'all'
  • 3行目: 僕は基本的にはJupyter Notebookで実行しています。この行を追加することで、数式がMathJaxで綺麗に表示されます。
  • 5~6行目: セルの途中で出力しても数式が綺麗に表示されるようにしています。こちらの記事を参考にしました。
y, mu, mu0 = symbols('y mu mu0', real=True)
Y_vec = symbols('Y1:4', real=True)
tau, lambda0, a0, b0 = symbols('tau lambda0 a0 b0', positive=True)
  • 1行目: SymPyで使う変数はsymbols関数で作成しておく必要があります。real=Trueと指定することで、実数と仮定することができます。何も指定しなければ複素数になります。このように仮定を入れておかないと、のちの式変形や積分がうまくいかない場合があります。
  • 2行目: このように変数のリストを作成することもできます。
  • 3行目: positive=Trueと指定することで、正の実数だと仮定することができます。

なお、SymPyでは要素数やデータ数をNとするような一般の場合の式変形は基本的に難しいです。しかし具体的な値に決めれば実行できます。そこで、ここでは2行目でデータ数をY1,Y2,Y33個として先に進めます。あとで値を色々変えて試すとNの場合の見当がつくので、そこから一般の場合を証明することもできます。

p_y = density(Normal('', mu, 1/sqrt(tau)))(y)
p_mu = density(Normal('', mu0, 1/sqrt(lambda0*tau)))(mu)
p_tau = density(Gamma('', a0, 1/b0))(tau)

sympy.statsには確率分布の密度関数の式がありますので、それを使っています。ここではデータ1つあたりのyの分布とmutauの事前分布を定義しています。

SymPyの正規分布Normal(平均, 標準偏差)なので、精度であるtau1/sqrt(tau)を代入しています。また、PRMLWikipediaのガンマ分布はGamma(shape, rate)である一方*2、SymPyのガンマ分布はGamma(shape, scale)なので、1/b0を代入しています。

integrate(p_mu, (mu, -oo, oo))
simplify(integrate(p_tau, (tau, 0, oo)))

試しにmuの分布を -\inftyから \inftyまで積分してみましょう。期待通り1が返ります。tauの分布でも同様に積分すると1にならずに整理されていない式が返ってきますが、simplify関数で整理すると1になります。

同時分布の対数( log\ p)の準備

前の記事の最後の2つの式でやっていることを日本語で書くと以下です。

  • 同時分布の対数 log\ p(\boldsymbol{Y}, \mu, \tau) q(\tau)を掛けて \tau積分して、 \muの分布 q(\mu)を求める。
  • 同時分布の対数 log\ p(\boldsymbol{Y}, \mu, \tau) q(\mu)を掛けて \mu積分して、 \tauの分布 q(\tau)を求める。

そこでまず同時分布の対数を準備します。

log_p = sum([log(p_y.subs(y, x)) for x in Y_vec]) + log(p_mu) + log(p_tau)
log_p = simplify(log_p)
log_p

  - \frac{Y_1^{2} \tau}{2} + Y_1 \mu \tau - \frac{Y_2^{2} \tau}{2} + Y_2 \mu \tau - \frac{Y_3^{2} \tau}{2} + Y_3 \mu \tau + a_0 \log{\left (b_0 \right )} + a_0 \log{\left (\tau \right )} - b_0 \tau \\ - \frac{\lambda_0 \tau}{2} \mu^{2} + \lambda_0 \mu \mu_0 \tau - \frac{\lambda_0 \tau}{2} \mu_0^{2} - \frac{3 \tau}{2} \mu^{2} + \log{\left (a_0 \right )} + \frac{1}{2} \log{\left (\lambda_0 \right )} + \log{\left (\tau \right )} - \log{\left (\Gamma{\left(a_0 + 1 \right)} \right )} \\ - 2 \log{\left (\pi \right )} - 2 \log{\left (2 \right )}

  • 1行目: expr.subs(y, x)は式expryxを代入します。

次に積分にすすみます。

 \muを含まない項に \muを含まない分布 q(\tau)を掛けて \tau積分したところで、やはり \muに関係がない定数になります。定数は最後に規格化して求めればよいので、途中の計算はなるべく簡単になるように余計な項を取り除きます。これがSymPyで計算をうまくさせるポイントになります。

log_p_for_mu = integrate(diff(log_p, mu), mu)
log_p_for_mu = collect(log_p_for_mu, mu)
log_p_for_mu

  \mu^{2} \left(- \frac{\lambda_0 \tau}{2} - \frac{3 \tau}{2}\right) + \mu \left(Y_1 \tau + Y_2 \tau + Y_3 \tau + \lambda_0 \mu_0 \tau\right)

  • 1行目: log_pmu微分してmu積分することで、muを含まない項を取り除いています。
  • 2行目: collect関数はmuの関数として式をみたときに共通部分をくくります。
log_p_for_tau = integrate(diff(log_p, tau), tau)
log_p_for_tau = collect(log_p_for_tau, tau)
log_p_for_tau

  \tau \left(- \frac{Y_1^{2}}{2} + Y_1 \mu - \frac{Y_2^{2}}{2} + Y_2 \mu - \frac{Y_3^{2}}{2} + Y_3 \mu - b_0 - \frac{\lambda_0 \mu^{2}}{2} + \lambda_0 \mu \mu_0 - \frac{\lambda_0 \mu_0^{2}}{2} - \frac{3 \mu^{2}}{2}\right) + \left(a_0 + 1\right) \log{\left (\tau \right )}

 \mu積分して q(\tau)を求める方も同様なのでそうしておきます。

できるところまで解析的に求める

SymPyの練習のため、事前分布から積分を1回実行して q(\mu) q(\tau)を求めるところをやってみます。

log_q1_mu = integrate(log_p_for_mu * p_tau, (tau, 0, oo))
log_q1_mu
log_q1_mu = simplify(log_q1_mu)
log_q1_mu
log_q1_mu = collect(expand(log_q1_mu), mu)
log_q1_mu

  \mu^{2} \left(- \frac{a_0 \lambda_0}{2 b_0} - \frac{3 a_0}{2 b_0}\right) + \mu \left(\frac{Y_1 a_0}{b_0} + \frac{Y_2 a_0}{b_0} + \frac{Y_3 a_0}{b_0} + \frac{a_0 \mu_0}{b_0} \lambda_0\right)

  • 5行目: 3行目でsimplifyしていますが、解析者が意図しない形になることはよくあります。ここでは、expandしてcollectすることでmu多項式にしています。

log_q1_muの式は \muの二次関数のマイナスなので、このすぐあとのq1_mu正規分布になることが分かります。共役事前分布を使っているからです。規格化定数をもとめて規格化しましょう。

q1_mu = exp(log_q1_mu)
const = simplify(integrate(q1_mu, (mu, -oo, oo)))
const
q1_mu = 1/const * exp(log_q1_mu)
q1_mu

constが規格化定数になります。以下の部分です。

  \frac{\sqrt{2} \sqrt{\pi} \sqrt{b_0}}{\sqrt{a_0} \sqrt{\lambda_0 + 3}} e^{\frac{a_0 \left(Y_1 + Y_2 + Y_3 + \lambda_0 \mu_0\right)^{2}}{2 b_0 \left(\lambda_0 + 3\right)}}

q1_muは規格化された分布の q(\mu)です。以下になります。

  \frac{\sqrt{2} \sqrt{a_0} \sqrt{\lambda_0 + 3}}{2 \sqrt{\pi} \sqrt{b_0}} e^{- \frac{a_0 \left(Y_1 + Y_2 + Y_3 + \lambda_0 \mu_0\right)^{2}}{2 b_0 \left(\lambda_0 + 3\right)}} e^{\mu^{2} \left(- \frac{a_0 \lambda_0}{2 b_0} - \frac{3 a_0}{2 b_0}\right) + \mu \left(\frac{Y_1 a_0}{b_0} + \frac{Y_2 a_0}{b_0} + \frac{Y_3 a_0}{b_0} + \frac{a_0 \mu_0}{b_0} \lambda_0\right)}

同じようにq1_tauを求めます。変分ベイズの手順としては、上で求めたばかりの q(\mu)を掛けて \mu積分します。しかしSymPyではその計算は重くて実行できないので、ここでは \muの事前分布p_muを使ってq1_tauを求めてみます。

log_q1_tau = integrate(log_p_for_tau * p_mu, (mu, -oo, oo))
log_q1_tau
log_q1_tau = integrate(diff(log_q1_tau, tau), tau)
log_q1_tau
log_q1_tau = collect(log_q1_tau, tau)
log_q1_tau
  • 3行目: あとで規格化定数を求めればよいので定数項は取り除いておきます。

このすぐあとのq1_tauはガンマ分布になることが分かります。これも共役事前分布を使っているからです。

q1_tau = logcombine(exp(log_q1_tau))
q1_tau
# const = integrate(q1_tau, (tau, 0, oo))
# const
# q1_tau = 1/const * q1_tau
# q1_tau
  • 1行目: logcombine関数を使うことで exp(log(x)) xにします。simplify関数だとこの変形をやってくれないことがあります。
  • 3行目: これで素直に積分できればよいのですが、残念ながらできません。

q1_tauは以下です。

  \tau^{a_0 + 1} e^{\tau \left(- \frac{Y_1^{2}}{2} + Y_1 \mu_0 - \frac{Y_2^{2}}{2} + Y_2 \mu_0 - \frac{Y_3^{2}}{2} + Y_3 \mu_0 - b_0 - \frac{3 \mu_0^{2}}{2}\right)}

この expの肩にのっている \tauの係数が負だとSymPyが分からないから積分できないようです。ちなみにこのあたりはMathematicaの方が圧倒的に賢くて、例えば以下の入力できちんと積分できます。

Integrate[tau^(a+1)*Exp[tau * (-1/2*x^2 + x*mu - 1/2* y^2 + y*mu - b - mu^2)], {tau, 0, Infinity}, Assumptions -> {b > 0, a > 0, Element[x, Reals], Element[y, Reals], Element[mu, Reals]} ]

これをうまく積分させるには、 \tauの係数が負であることを確認してから変数で置き換えて実行します。

まず \tauの係数が負であることを確認します。

coef = collect(log_q1_tau, tau).coeff(tau)
coef
sol = solve(diff(coef, Y_vec[0]), Y_vec[0])[0]
sol #=> mu0
replacements = [(var, sol) for var in Y_vec]
coef.subs(replacements) #=> -b0
  • 1行目:  \tauの係数coefを取得しています。
  • 3行目: coefの最大値が負であることを示せばOKです。まずはY_vec[0]についてcoefが最大になる値を探します。それは微分して0(&2階微分が負)になる点を求めればOKです。Y_vec[0]と他のY_vec[*]は区別がある形ではないので、Y_vec[*]についても同じ点でcoefが最大となります。
  • 5~6行目: それをまとめて代入しています。最大値は-b0と分かるので、 \tauの係数は負であることがわかります。

次に変数で置き換えて積分します。

xi = symbols('xi', positive=True)
const = simplify(integrate(tau**(a0+1)*exp(-xi*tau), (tau, 0, oo)))
const = const.subs(xi, -coef)
const
q1_tau = 1/const * q1_tau
q1_tau

constが規格化定数になります。以下の部分です。

  \left(\frac{Y_1^{2}}{2} - Y_1 \mu_0 + \frac{Y_2^{2}}{2} - Y_2 \mu_0 + \frac{Y_3^{2}}{2} - Y_3 \mu_0 + b_0 + \frac{3 \mu_0^{2}}{2}\right)^{- a_0 - 2} \Gamma{\left(a_0 + 2 \right)}

q1_tauは正規化された分布の q(\tau)です。以下になります。

  \frac{\tau^{a_0 + 1}}{\Gamma{\left(a_0 + 2 \right)}} \left(\frac{Y_1^{2}}{2} - Y_1 \mu_0 + \frac{Y_2^{2}}{2} - Y_2 \mu_0 + \frac{Y_3^{2}}{2} - Y_3 \mu_0 + b_0 + \frac{3 \mu_0^{2}}{2}\right)^{a_0 + 2} e^{\tau \left(- \frac{Y_1^{2}}{2} + Y_1 \mu_0 - \frac{Y_2^{2}}{2} + Y_2 \mu_0 - \frac{Y_3^{2}}{2} + Y_3 \mu_0 - b_0 - \frac{3 \mu_0^{2}}{2}\right)}

このように解析解を求めることはコンセプトの理解に役立ちます。一方で、積分を繰り返して事後分布 q(\mu, \tau)が収束するか確認するようなことは数値的に求めた方が分かりやすいです。

数値的に求める

仮に得られたデータY_vec1.1,1.0,1.3とします。また、事前分布はa0 = 1, b0 = 1, mu0 = 0, lambda0 = 1とします。

replacements = [(a0, 1), (b0, 1), (mu0, 0), (lambda0, 1)]
data_vec = [1.1, 1.0, 1.3]
replacements.extend([(var, val) for var, val in zip(Y_vec, data_vec)])
log_p_for_mu_subs = log_p_for_mu.subs(replacements)
log_p_for_tau_subs = log_p_for_tau.subs(replacements)
[log_p_for_mu_subs, log_p_for_tau_subs]

  \left [ - 2 \mu^{2} \tau + 3.4 \mu \tau, \quad \tau \left(- 2 \mu^{2} + 3.4 \mu - 2.95\right) + 2 \log{\left (\tau \right )}\right ]

  • 1行目: 事前分布の分の代入を作っています。
  • 2~3行目: データの分の代入を追加しています。

 \tauの初期分布をp_tauとして、 q(\mu)を求める→ q(\tau)を求める→ q(\mu)を求める→...と7回ほど繰り返してみます。

q_tau = N(p_tau.subs(replacements))
q_tau

for i in range(7):
    log_q_mu = N(integrate(log_p_for_mu_subs * q_tau, (tau, 0, oo)))
    const = N(integrate(exp(log_q_mu), (mu, -oo, oo)))
    q_mu = 1/const * exp(log_q_mu)

    log_q_tau = N(integrate(log_p_for_tau_subs * q_mu, (mu, -oo, oo)))
    const = N(integrate(exp(log_q_tau), (tau, 0, oo)))
    q_tau = 1/const * exp(log_q_tau)

    [q_mu, q_tau]

  \left [ 0.188098154753774 e^{- 2.0 \mu^{2} + 3.4 \mu}, \quad 4.03007506250001 \tau^{2.0} e^{- 2.005 \tau}\right ]   \left [ 0.112320150163227 e^{- 2.99251870324189 \mu^{2} + 5.08728179551122 \mu}, \quad 3.11052191637731 \tau^{2.0} e^{- 1.83916666666667 \tau}\right ]   \left [ 0.0965024138432034 e^{- 3.26234707748074 \mu^{2} + 5.54599003171727 \mu}, \quad 2.97238456804457 \tau^{2.0} e^{- 1.81152777777778 \tau}\right ]   \left [ 0.0938011432750369 e^{- 3.31212144445296 \mu^{2} + 5.63060645557004 \mu}, \quad 2.94976700750461 \tau^{2.0} e^{- 1.8069212962963 \tau}\right ]   \left [ 0.0933494031271016 e^{- 3.32056521349236 \mu^{2} + 5.64496086293701 \mu}, \quad 2.94600860516469 \tau^{2.0} e^{- 1.80615354938272 \tau}\right ]   \left [ 0.0932740721319143 e^{- 3.32197669575248 \mu^{2} + 5.64736038277921 \mu}, \quad 2.94538251532282 \tau^{2.0} e^{- 1.80602559156379 \tau}\right ]   \left [ 0.0932615158350226 e^{- 3.32221205946743 \mu^{2} + 5.64776050109463 \mu}, \quad 2.94527817564072 \tau^{2.0} e^{- 1.80600426526063 \tau}\right ]

  • 1行目: N関数は数値による近似を求める関数です。

7回ほどの繰り返しのあとでほぼ収束していそうなことがわかります。

最後に求めた事後分布(の近似) q(\mu, \tau) = q(\mu) q(\tau)を可視化してみましょう。SymPyにもsympy.plottingsympy.plotting.plotが存在するのですが、ちょっと凝った図を書こうとするとすぐ厳しくなってしまいます。そこで、得られた事後分布をlambdify関数で関数化し、NumPyとMatplotlibで描くのが拡張性が高くてオススメです。

from sympy.utilities.lambdify import lambdify
import numpy as np
import matplotlib.pyplot as plt

delta = 0.05
x = np.arange(-1.0, 3.0, delta)
y = np.arange(0.0, 6.0, delta)
X, Y = np.meshgrid(x, y)
func = lambdify((mu, tau), q_mu * q_tau, 'numpy')
Z = func(X, Y)

plt.figure()
CS = plt.contour(X, Y, Z)
plt.clabel(CS, inline=1, fontsize=10)

f:id:StatModeling:20201106161401p:plain

まとめ

  • SymPyはデータサイエンスや機械学習の書籍や論文を読み進める上で、非常に有用な補助ツールです。
  • 現状では細かいところでMathematicaにまだ負けていると思います。プロにはMathematicaがオススメ。オープンソース重視の人やPython好きな人にはSymPyがオススメ。
  • 式変形には「一般的な場合のようにコンセプトが重要で深く理解しなければならない式変形」と「SymPyなどの数式処理ソフトで追えれば十分であるような式変形」があると個人的に思っています。専門書や技術書を執筆する場合は、その二つを区別すると読者にとって親切かなぁと思いました。

Enjoy!

謝辞

北大電子研の佐藤勝彦氏に感謝します。僕が院生の頃に輪読していた ニコリス プリゴジーヌ『散逸構造』の例題をMathematicaで10分ぐらいで一般解を求めるという衝撃のデモを見せてもらい、その後もたまにMathematicaを教えてもらい、数式処理を学ぶきっかけをもらいました。

*1:いつもはStanとの相性を考えて Normal(平均, 標準偏差)で書いてますが、この記事では Normal(平均, 分散)で書きます。

*2:Stanもね。