正規分布のハミルトニアンモンテカルロ法(HMC法)によるベイズ推定をRで書いてみた


今この本を読み進めています。 5章の終わりまではRでの実装が出来ましたので、5章の最後のハミルトニアンモンテカルロ法を用いた正規分布の母数(平均 { \mu }、分散 { \sigma ^2 } )の推定プログラムを紹介します。 母数が多次元(2次元)の場合の推定です。

平均170、分散49の正規分布から取り出した100個のサンプルが手に入っているとします。

set.seed(77)
x <- rnorm(100, mean = 170, sd = 7)
x
g <- ggplot(data_frame(x), aes(x = x))
g <- g + geom_histogram()
plot(g)
# [1] 166.1525 177.6374 174.4784 177.2980 171.1879 177.9646 163.2061 169.0772 171.0236
# [10] 180.0890 149.4103 168.3001 169.0159 169.7714 171.9587 174.1310 177.1701 184.7513
# [19] 171.0824 176.3915 168.2204 180.6360 182.4689 163.8481 159.2955 170.9514 165.0377
# [28] 160.1330 182.8166 179.0314 153.4630 166.1432 167.8674 164.7503 171.0061 166.1540
# [37] 171.1195 169.3863 170.5679 176.2938 170.0218 166.2814 165.0303 167.9607 176.1927
# [46] 168.9197 163.3198 174.6708 172.7194 173.8441 160.8651 176.2086 186.3518 173.5221
# [55] 154.1242 155.7792 169.0335 163.3324 181.0436 161.2650 172.7907 170.1869 162.9896
# [64] 170.0623 168.8515 172.9581 167.0328 166.7133 171.2051 182.9289 158.5322 167.6728
# [73] 173.1356 171.9069 168.2195 186.4382 168.5089 170.3545 154.2780 176.4378 175.2604
# [82] 169.5340 166.4753 168.8772 169.3487 173.3175 170.9141 175.0213 169.9780 154.7639
# [91] 176.1373 175.8453 169.6634 146.1331 159.4070 171.7665 166.6858 176.8403 187.0456
# [100] 169.4234

f:id:zoetaka38:20160529110713p:plain

細かい説明は本に任せますが、ざっくりと以下の数式があればハミルトンモンテカルロ法アルゴリズムが組めるようになります。

  1. 対数事後分布のカーネルの式
  2. 対数事後分布のカーベルをそれぞれの母数で微分した式
  3. ハミルトニアンの計算式
  4. リープフロッグ法の計算式

まず、1.から。

正規分布の式は

[tex:{\displaystyle f(x) = \frac{1}{\sqrt{2\pi\sigma^{2}}} \exp\left(-\frac{(x-\mu)^2}{2\sigma^2} \right)}]

なので、対数事後分布のカーネル

[tex:{\displaystyle \log f(x|\mu, \sigma^2) = \frac{-n}{2}\log\sigma^2 +\frac{-1}{2\sigma^2}\sum_{i=1}^n (x_i-\mu)^2}]

そして、2.は、上記の対数事後分布のカーネルを母数である[tex:{ \mu }]と[tex:{ \sigma }]それぞれで微分した式となる。

[tex:{\displaystyle \frac{d}{d\mu} \log f(x|\mu, \sigma^2) = \frac{1}{\sigma^2}\sum_{i=1}^n (x_i-\mu)}]
[tex:{\displaystyle \frac{d}{d\sigma^2} \log f(x|\mu, \sigma^2) = \frac{-n}{2\sigma^2} +\frac{1}{2\sigma^4}\sum_{i=1}^n (x_i-\mu)^2}]

ハミルトンの計算式は以下の通り。([tex:{\boldsymbol{\theta}}]には、[tex:{ \mu }]と[tex:{ \sigma }]が保持されます。最後の合計の[tex:{ d }]は[tex:{\boldsymbol{\theta}}]もしくは、[tex:{\boldsymbol{p}}]の次元数を示します。

[tex:{\displaystyle H( \boldsymbol{\theta},\boldsymbol{p}) = h(\boldsymbol{\theta}) + \frac{1}{2}\sum_{i=1}^d \boldsymbol{p}^2}]

そして、4.のリープフロッグ法の計算式は以下の2つ。

[tex:{\displaystyle \frac{d}{d\sigma^2} \log f(x|\mu, \sigma^2) = \frac{-n}{2\sigma^2} +\frac{1}{2\sigma^4}\sum_{i=1}^n (x_i-\mu)^2}]

コード

ここからがコードです。 まずは、HMC法で必要な式を関数にするところ。

# 事後分布のカーネル
# 5.59
h <- function(mu, sigma){
  tmp <- function(x) (x - mu)**2
  return(length(x)* 0.5 *log(sigma) + 0.5 * sum(sapply(x, tmp)) / (sigma))
}

# 事後分布
# 5.60
# hをmuで微分した式
dh_dmu <- function(mu, sigma){
  tmp <- function(x) (x - mu)
  return(-1 * sum(sapply(x, tmp)) / sigma)
}
# 5.61
# hをsigmaで微分した式
dh_dsigma <- function(mu, sigma){
  tmp <- function(x) (x - mu)**2
  return(length(x) / (2*sigma) - 0.5 / (sigma**2)* sum(sapply(x, tmp)))
}

# 5.36
# ハミルトニアンの計算式
hamiltonian <- function(p, mu, sigma){
  return(h(mu, sigma) + 0.5*sum(p**2))
}

# リープフロッグ
# 5.32
# 運動量を1/2単位時間分更新
leapfrog_nexthalf_p <- function(p, d, mu, sigma, eps = 0.01){
  return(p - 0.5 * eps * d(mu,sigma))
}

# リープフロッグ
# 5.33
# 位置を1単位時間分更新
leapfrog_next_param <- function(p, param, eps = 0.01){
  return(param + eps * p)
}

上記で定義した関数を用いて、シミュレーションの実践を行います。 本に記載のある通り、L=100,eps=0.01,T=100000で行います。

# HMC simulation
scale_p = 1

# initial param
mu <- 160
sigma <- 55
p = rnorm(n=2, mean=0, sd=scale_p)
eps = 0.01
L = 100
T = 100000
prev_hamiltonian <- hamiltonian(p, mu, sigma)
sim_result <- data_frame(num = 0, p[1], p[2], mu, sigma, hamiltonian = prev_hamiltonian, logi=TRUE)

ti<-proc.time() # 実行時間計測

for(t in 1:T){
  prev_p <- p
  prev_mu <- mu
  prev_sigma <- sigma
  prev_hamiltonian <- hamiltonian(p, mu, sigma)
  
  for(i in 1:L){
    tmp_mu <- mu
    tmp_sigma <- sigma
    p[1] <- leapfrog_nexthalf_p(p[1], dh_dmu, tmp_mu, tmp_sigma, eps)
    mu <- leapfrog_next_param(p[1], mu, eps)
    p[1] <- leapfrog_nexthalf_p(p[1], dh_dmu, mu, tmp_sigma, eps)
    
    p[2] <- leapfrog_nexthalf_p(p[2], dh_dsigma, tmp_mu, tmp_sigma, eps)
    sigma <- leapfrog_next_param(p[2], sigma, eps)
    p[2] <- leapfrog_nexthalf_p(p[2], dh_dsigma, tmp_mu, sigma, eps)
  }
  H <- hamiltonian(p, mu, sigma) # 提案された新ハミルトニアン
  r <- exp(prev_hamiltonian - H) # ハミルトニアンの計算上の誤差を修正
  if(r > 1){
    sim_result <- rbind(sim_result, c(t, p, mu, sigma, hamiltonian(p,mu, sigma), TRUE))
  }else if(r > 0 & runif(1) < r){
    sim_result <- rbind(sim_result, c(t, p, mu, sigma, hamiltonian(p,mu, sigma), TRUE))
  }else{
    sim_result <- rbind(sim_result, c(t, p, mu, sigma, hamiltonian(p,mu, sigma), FALSE))
    mu <- prev_mu
    sigma <- prev_sigma
  }
  p <- rnorm(2, mean = 0,sd = scale_p)
}

# バーンイン期間を1000とし、シミュレーション結果が受容されたもののみをresultに格納する
sim_result2 <- sim_result[-(1:1000), ]
result <- sim_result2[sim_result2$logi==1, ]
ti <- proc.time()-ti
ti/60
    user   system  elapsed 
229.0732   0.3102 229.8342 

最尤推定量での計算(標本平均、標本分散)

mean(x)
sd(x)**2
> mean(x)
[1] 170.0342
> sd(x)**2
[1] 60.08921

事後分布のシミュレーションの結果

mean(result$mu)
mean(result$sigma)
> mean(result$mu)
[1] 170.038
> mean(result$sigma)
[1] 62.43348

トレースプロット({ \mu })

g <- ggplot(
  result,
  aes(
    x = num,
    y = mu
  )
)
g <- g + geom_line(
  colour = "magenta",
  linetype = 1,
  size = 0.5
)
plot(g)

f:id:zoetaka38:20160625192044p:plain

収束していると思われますね。

トレースプロット({ \sigma })

g <- ggplot(
  result,
  aes(
    x = num,
    y = sigma
  )
)
g <- g + geom_line(
  colour = "magenta",
  linetype = 1,
  size = 0.5
)
plot(g)

f:id:zoetaka38:20160625192223p:plain

こちらもうまく収束しているのではないかと。

最後に注釈

うまく収束したように思えますが、収束の尺度であるR_hatなどは計算していないので、実際のところはうまくいったかは不明です。 また、今回はうまくステップ数Lやステップの幅epsを調整できたのですが、実際はNUTS法と呼ばれる自動調整アルゴリズムを実装するのが普通みたいです。 NUTS法を使っていないせいか、本の6章にある正規分布を使った推定では、{ \sigma }がうまく収束してくれませんでした。 結論、ある程度理解できたらおとなしくStan使っとけということですね(笑)Rで頑張ったらどうしても処理は遅くなりますし。