統計コンサルの議事メモ

統計や機械学習の話題を中心に、思うがままに

WAICを計算してみる

ある統計モデルの予測精度を表す指標としてAICAkaike Information Criterion)というものがよく使われますが、stanを用いてモデリングした場合にはAICを計算することができません。しかし2009年に渡辺澄夫先生が発表したWAICWidely Applicable Information Criterion、広く使える情報量規準)はAICをより一般化したものとして定義され、MCMCによる事後分布から計算することができます。

現在のところ{rstan}では直接WAICを計算することができませんが、{loo}というパッケージを用いることでWAICを得ることができます。この記事では、stanの結果からWAICを計算により求め、それがlooの結果と一致することを確認します。

WAICの定義

まず始めにWAICの定義を確認しましょう*1。WAICは以下のように定義されます:

WAIC = -2(lppd - p_{WAIC})

ここで

\displaystyle lppd = \sum_{i = 1}^{N}\log{Pr(y_{i})}

であり、また

\displaystyle p_{WAIC} = \sum_{i = 1}^{N}V(y_{i})

です。Pr(y_{i})およびV(y_{i})は、それぞれ観測値iの事後分布から得られる対数尤度の平均および対数尤度の分散を表してます。なおlppdはlog-pointwise-predictive-densityです。

シミュレーション

簡単な数値例を発生させ、そのWAICを計算してみましょう。パラメータの推定には当然ながら{rstan}を利用します。

# シミュレーションデータの発生
set.seed(123)
N <- 100 # サンプルサイズ
b <- 1.2 # 回帰係数
X <- rnorm(N, 0, 1) # 説明変数
E <- rnorm(N, 0, 2) # 誤差項
Y <- b * X + E
D <- data.frame(Y, X) # データフレーム

回帰係数1.2、誤差標準偏差2として発生させたデータに対し、{rstan}を用いて線形回帰モデルを当てはめてみましょう。いつものように{rstan}をロードします。

library(rstan)
rstan_options(auto_write = TRUE)
options(mc.cores = parallel::detectCores())

stanに渡すデータを以下のように定義します。今回の例では単純な線形モデルを用いているので、渡すパラメータもほとんどありません。

dat_stan <- list(
   N = N,
   X = D$X,
   Y = D$Y
)

定義したdat_stanをstanに渡し、フィッティングします。

fit_01 <- stan(file = './WAIC.stan', 
               data = dat_stan, 
               iter = 3000,
               chains = 4,
               seed = 1234)

ここでWAIC.stanの内容は以下の通りです。簡単です。

data {
   int<lower=1> N;
   vector[N] X;
   vector[N] Y;
}

parameters {
   real b0;
   real b1;
   real<lower=0> sigma;
}

model {
   Y ~ normal(b0 + b1 * X, sigma);
}

フィッティングが終わると、以下のような結果が得られます:

> summary(fit_01)$summary
             mean     se_mean        sd         2.5%          25%
b0      -0.201859 0.002533264 0.1962258   -0.5791061   -0.3360031
b1       1.092886 0.002755139 0.2134122    0.6655672    0.9504853
sigma    1.964687 0.001920081 0.1421800    1.7080842    1.8637401
lp__  -116.169685 0.021008195 1.2014656 -119.2035323 -116.7161937
               50%           75%        97.5%    n_eff     Rhat
b0      -0.2018323   -0.06931008    0.1877817 6000.000 1.000041
b1       1.0919604    1.23311123    1.5082127 6000.000 1.000551
sigma    1.9570918    2.05945046    2.2612886 5483.250 1.000177
lp__  -115.8524219 -115.28245822 -114.7994731 3270.734 1.000839

色々と書いてあって情報が多いので、必要なところを抜き出しましょう。真の回帰係数および誤差標準偏差はそれぞれ1.2と2でした。

> summary(fit_01)$summary[c("b1", "sigma"), c("mean", "50%")]
          mean      50%
b1    1.092886 1.091960
sigma 1.964687 1.957092

回帰係数が少し低いようですが、まぁ良しとしましょう。

WAICの計算

それではこの結果を用いて、WAICを計算してみましょう。まずはlppdからですが、lppdは前述の通り以下によって計算されます:

\displaystyle lppd = \sum_{i = 1}^{N}\log{Pr(y_{i})}

まずはstanの結果からposterior sampleを得ます。stanを実行するときにwarmupを指定していないためsampleサイズ(ite)の半分が切り捨てられた結果、残る6,000サンプル(1,500 * 4 chain)がposterior sampleとして保存されています。

post_samples <- extract(fit_01)
> str(post_samples)
List of 4
 $ b0   : num [1:6000(1d)] -0.146 -0.13 -0.195 -0.535 -0.162 ...
  ..- attr(*, "dimnames")=List of 1
  .. ..$ iterations: NULL
 $ b1   : num [1:6000(1d)] 1.036 0.928 1.004 1.309 1.005 ...
  ..- attr(*, "dimnames")=List of 1
  .. ..$ iterations: NULL
 $ sigma: num [1:6000(1d)] 2.2 2.01 1.64 2.03 1.86 ...
  ..- attr(*, "dimnames")=List of 1
  .. ..$ iterations: NULL
 $ lp__ : num [1:6000(1d)] -116 -115 -118 -117 -115 ...
  ..- attr(*, "dimnames")=List of 1
  .. ..$ iterations: NULL

この6,000サンプルのそれぞれをパラメータとして与えた時の各 y_{i}の対数尤度を求めると:

Des <- cbind(1, X) # 計画行列(Design Matrix)
B   <- cbind(post_samples$b0, post_samples$b1) # パラメータ(Beta)
tmp <- matrix(NA, length(post_samples$b0), N) # 6000行、 100列の行列
for (i in 1:N) {
   tmp[, i] <- dnorm(Y[i], mean = B %*% Des[i, ], 
                     sd = post_samples$sigma)
}

となります。ここからlppdは:

lppd <- sum(log(colMeans(tmp)))

で得られます。
同じように p_{waic}を求めてみましょう。 p_{waic}の定義は以下のようでした:

\displaystyle p_{WAIC} = \sum_{i = 1}^{N}V(y_{i})

ここから、

pwaic <- sum(apply(tmp, 2, var))

として求められます。
最後に、これらの値を用いてWAICを求めます:

WAIC <- -2 * (lppd - pwaic)

これでようやくWAICを得ることがことができましたので、これらの値を出力してみると:

> cat("WAIC:", WAIC, "\n",
+     "lppd:", lppd, "\n",
+     "pwaic:", pwaic, "\n")
WAIC: 420.8022 
lppd: -207.2305 
pwaic: 3.170572 

となります。ちなみに同じモデルについて、lmを用いてフィッティングした時のAICを求めると:

fit_02 <- lm(Y ~ X, D)
> AIC(fit_02)
[1] 420.4523

であり、非常に近い値となります。

パッケージ{loo}を使ってWAICを求める。

では、今度は上記の行程で求めた数値をパッケージで得た値と比較してみましょう。WAICの計算には{loo}を用います。

library(loo)

looによりWAICを計算させるためには、stanのスクリプトにgenerated quantitiesブロックを追加しておく必要があります。以下をmodelブロックの後に追加します。

generated quantities {
   vector[N] log_lik;
   
   for (n in 1:N)
      log_lik[n] = normal_lpdf(Y[n] | b0 + b1 * X[n], sigma);
}

normal_lpdfというのはRで言うところのdnormですね。
このスクリプトをWAIC_02.stanとして保存し、stanを実行します。

fit_03 <- stan(file = './WAIC_02.stan', 
               data = dat_stan, 
               iter = 3000,
               chains = 4,
               seed = 1234)

結果は以下のようになります。前の結果と同じであることを確認してください。

> summary(fit_03)$summary[c("b1", "sigma"), c("mean", "50%")]
          mean      50%
b1    1.092886 1.091960
sigma 1.964687 1.957092

では、今度は{loo}のextract_log_likを使ってposterior sampleを抽出し、loo::waicによりWAICを計算します:

tmp02   <- extract_log_lik(fit_03)
WAIC_02 <- waic(tmp02)

> cat("WAIC by loo:", WAIC_02$waic, "\n",
+     "WAIC by manual:", WAIC, "\n")
WAIC by loo: 420.8022 
WAIC by manual: 420.8022 

ちゃんと同じ値が得られていますね!

終わりに

分析内容は以上となります。この記事ではWAICの定義を確認した上で、その計算のための方法を示し、既存のパッケージによる数値と一致することを確かめました。WAICは階層モデルのような複雑なものであっても適用可能で、stanを使っていく上では抑えておきたい指標なので今後も勉強を続けていきます。

*1:Statistical RethinkingのChapter 6から