【R】WAICを用いた予測モデル評価を混合Poisson分布で試す

こちらの本の6, 7章の汎化損失・WAICの内容をまとめてみました。

作者:浜田宏,石田淳,清水裕士
出版社/メーカー: 朝倉書店
発売日: 2019/12/01
メディア: 単行本（ソフトカバー）

理論と数値計算(R + RStan)がバランスよく書かれているのでおすすめです。

予測モデルの評価と汎化損失

真の分布 $q(X)$ が未知であり、真の分布からの独立な $n$ 個のサンプルの実現値 $\underline{x} = (x_1, \ldots, x_n)$ が与えられた状況を考えます。この実現値を用いて、予測モデル $p^\ast (X)$ を作成したとき予測パフォーマンスのよさを測る指標としてつぎの汎化損失があります。

汎化損失
$G_n := {\mathbb E}_{q(X)}[ - \log p^\ast (X)] = - \int q(X)\log p^\ast (X) dX$

真の分布が未知の場合、汎化損失の計算は直接計算が行えないことに注意が必要です。

以下、予測モデルとしてベイズモデルを使う状況を考えます。

モデル(尤度) : $p(\underline{x} | \theta)$
事前分布 : $\phi(\underline{x})$

このモデルの予測分布(predictive distribution)は、つぎの期待値で定義されます。

$p^\ast (X) := {\mathbb E}_{p(\theta |\underline{x})}[p(X |\theta ] = \int p(X | \theta) p(\theta |\underline{x}) d\theta$

特に、ベイズ予測モデルについての汎化損失はつぎの式で表せることになります。

$G_n = - \int q(X) \log \Bigl( \int p(X | \theta) p(\theta |\underline{x}) d\theta \Bigr) dX$

WAIC

ここではベイズ予測モデルの汎化損失について、真の分布が未知の場合にも計算が可能な推定量となるWAICを導入します。

WAIC
$\underline{x} = (x_1, \ldots, x_n)$ を真の分布のサンプルの実現値とする。
経験損失(training loss) $T_n$ と汎関数分散(general function variance)を以下で定義する
$T_n := - \frac{1}{n} \sum_{i = 1}^n {\mathbb E}_{p(\theta | \underline{x})} [\log p(x_i|\theta)]$
$V_n = \sum_{i = 1}^n \{ {\mathbb E}_{p(\theta | \underline{x})} [ (\log p(x_i|\theta))^2 ] - {\mathbb E}_{p(\theta | \underline{x})} [ \log p(x_i|\theta) ]^2 \}.$
このとき、WAIC(Watanabe-Akaike Information Criterion)は次で定義される
${\sf WAIC} := T_n + \frac{V_n}{n}$

定義からわかるようにWAICはサンプルの実現値の出方によってばらつきがあります。しかし、つぎの結果によりサンプルの実現値に関する平均をとったときには、WAICは近似的に汎化損失の推定量を与えることが知られています。

$q(\underline{X}) \sim \prod_i^n q(X_i)$ を真の分布のサンプルとする。このとき次が成り立つ。
${\mathbb E}_{q(\underline{X})}[{\sf G}_n] = {\mathbb E}_{q(\underline{X})}[ {\sf WAIC}_n ] + o(1/n)$
(ただし、 $o(-)$ はランダウの記号を表す。)

混合Poissonモデルで数値実験

以下混合Poisson分布モデルを題材に、WAICを用いた予測モデル評価をRを使って行います。

真の分布として、2つのPoisson分布 ${\sf Poisson}(\lambda_i), i = 1, 2$ を比率 $\mu, ( \mu \in (0, 1))$ で混合して得られる分布 $q(X)$ を考えます。このとき確率質量関数は

$q(X) = \mu {\sf Poisson}(X | \lambda_1) + (1 - \mu) {\sf Poisson}(X | \lambda_2)$

となります。以下 $\mu = 0.5, \lambda_1=2, \lambda_2=8$ としてシミュレーションを行います。

まず、真の分布から適当な数のサンプルを作成します。


# 真の分布のパラメーター
target.param <- list(
    mu = 0.5,
    lambda = c(2, 8)
)

# 真の分布からsampler
MixPois.Sampler <- function(param) {
    return(function(n) {
        choices <- rbinom(n, 1, param$mu)
        return(mapply(function(choice, x, y) {
            return(ifelse(choice == 1, x, y))},
            choices, rpois(n, param$lambda[1]), rpois(n, param$lambda[2]))
        )
    })
}

# サンプリング
N <- 200L
target.Sampler <- MixPois.Sampler(target.param)
target.sample <- target.Sampler(N)
hist(target.sample)

f:id:kimigayoseishou:20200101122444j:plain — 真の分布からのサンプル(N = 200)

このサンプルに対して、混合Poisson分布を用いたモデリングを行います。

$N = 200$ をサンプル数, $\underline{x} = (x_1, \ldots, x_N)$ サンプルとして

$K$ > 1個のクラスターをもつ混合Poisson分布の尤度は

$p(\underline{x} | \underline{\lambda}, \underline{\mu}) = \prod_{n =1}^N \sum_{k = 1}^K \mu_k {\sf Poisson}(x_n | \lambda_k)$

で定義されます。ただし、 $\underline{\lambda}, \underline{\mu}$ は以下のようなハイパーパラメーターとします。

$\lambda_k \sim {\sf Gamma}(a, b),\quad a, b \in (0, \infty) , k = 1, \ldots, K$
$\underline{\mu} \sim {\sf Dirichlet}(\underline{r}), \quad \underline{r}=(r_1, \ldots, r_K) \in {\mathbb R}_{+}^K$

以下、この予測モデルについてクラスター数 $K$ を変化させたときの予測精度の変化をWAICで評価してみます。

WAICを求める上で事後分布 $p(\underline{\lambda}, \underline{\mu} |\underline{x})$ が必要ですが、この分布は解析的に解くことが難しいためRStanを使ったモンテカルロ法で処理します。

$K$ が2以上の場合

以下stanコードになります。


// Reference
// http://rpubs.com/kaz_yos/fmm2
// https://mc-stan.org/docs/2_21/stan-users-guide/soft-k-means.html
data {
  int<lower=1> K;//number of Cluster 
  int<lower=1> N;//data size
  int<lower=0> x[N];//sample 
  real<lower=0> a;
  real<lower=0> b;
  vector[K] r;//Dirichlet parameter
}

parameters {
  positive_ordered[K] lambda;
  simplex[K] mu;
}

transformed parameters {
  real lp[N, K];
  for(n in 1:N) {
    for(k in 1:K) {
      lp[n, k] = log(mu[k]) + poisson_lpmf(x[n]|lambda[k]); 
    }
  }
}

model {
  //prior 
  mu ~ dirichlet(r);
  
  for(k in 1:K) {
    lambda[k] ~ gamma(a, b);
  }
  
  //model
  for(n in 1:N) {
    target+= log_sum_exp(lp[n]);
  }
}

generated quantities {
  real lp_sample[N, K];
  real log_lik[N];
  for(n in 1:N) {
    for(k in 1:K) {
      lp_sample[n, k] = log(mu[k]) + poisson_lpmf(x[n]|lambda[k]); 
    }
    log_lik[n] = log_sum_exp(lp_sample[n]);
  }
}

このコードをRからキックして、事後分布からの乱数が得られます。

以下は $K = 2$ の場合のRコードになります。


.model = rstan::stan_model(file = './mixture_of_poisson.stan')
model.data <- list(
    K = 2L,# クラスター数
    N = N,# サンプルサイズ
    x = target.sample,# サンプル
    a = 1, b = 1,
    r = c(1, 1) # K-vector 
)

# 各chain 試行回数2000 & warmup = 1000で4回サンプリング
iter <- 4000L
.fit = rstan::sampling(.model, data = model.data, iter=iter, chains=4)

# Rhat < 1.1に収まっているか収束確認
summary <- rstan::summary(.fit)$summary
summary[summary[,'Rhat'] >= 1.1,] %>% View

# 収束を目視で確認
.pars <- c(str_glue('lambda[{c(1:model.data$K)}]'),
           str_glue('mu[{c(1:model.data$K)}]'))
rstan::stan_hist(.fit, pars = .pars)

WAICは発生させたサンプルからつぎのように計算できます。


waic_mcmc <- function(log_lik) {
    train_loss <- mean( - log(colMeans(exp(log_lik))))
    gen_fun_var.mean <- mean(apply(log_lik, MARGIN = 2, var))
    waic <- train_loss + gen_fun_var.mean
    return(waic)
}
model.sample <- rstan::extract(.fit)
model.waic_mcmc <- waic_mcmc(model.sample$log_lik)

$K=1$ の場合

この場合はモデルは単にPoissonモデルになりますので、stanコードは以下のようになります。サンプリング & WAICのシミュレーションは $K$ > 1の場合と同様なので省略します。


// Reference
data {
  int<lower=1> N;//data size
  int<lower=0> x[N];//sample 
  real<lower=0> a;
  real<lower=0> b;
}

parameters {
  real<lower=0> lambda;
}

model {
  //prior 
  lambda ~ gamma(a, b);
  
  //model
  target+= poisson_lpmf(x | lambda);
}

generated quantities {
  real log_lik[N];
  for(n in 1:N) {
    log_lik[n] = poisson_lpmf(x[n]|lambda);
  }
}

予測分布のシミュレーション

$\{\theta_s \}_{s = 1, ... S}$ を事後分布 $p(\theta | \underline{x})$ からのサンプルとしたとき、モンテカルロ法により予測分布はつぎのように近似的に求められます

$p(X | \underline{x}) = \int p(X | \theta ) p(\theta | D) d\theta \approx \frac{1}{S} \sum_{s = 1}^S p(X | \theta_s).$

今回の数値実験ではR側で事後分布の計算を行うことにします。


model.pred_pmf <- function(x, sample) {
    return(x %>% sapply(function(.x) {
        if (is.null(model.sample$mu)) {
            # K = 1の場合
            model.sample$mu <- rep(1, length(x))
        }
        pred_sample <- model.sample$mu %*% t(dpois(.x, model.sample$lambda))
        return(mean(pred_sample))
    }))
}
# 事後分布p(X | \underline{x})をX = 0, ..., 20についてシミュレーション
pred_pmf <- model.pred_pmf(c(0:20L), model.sample)

$K = 1~4$ の各モデルについて予測分布をシミュレーションした結果がつぎの図になります。図を見ると、 $K=2, 3, 4$ のモデルでは真の分布のサンプルに見られた2つのピークをうまく表現できていのに対し、 $K=1$ のモデルではピークをうまく表現できていないことがわかります。

f:id:kimigayoseishou:20200102201445p:plain — 各モデルの予測分布と真の分布からのサンプルの度数の比較

WAICのシミュレーション結果

$K=1, 2, 3, 4$ でWAICのシミュレーションした結果が以下の表になります。

$K$	WAIC
1	2.98
2	2.60
3	2.60
4	2.60

この表から以下のような観察が可能です。

$K=1$ のときのモデル(Poisson分布)よりも、(真の分布のクラスター数である) $K=2$ のモデルの方がWAICが小さく、予測のパフォーマンスが高い。
$K= 2, 3, 4$ のモデル間ではWAICの変化が見られず、クラスター数を増やすことによる予測パフォーマンスの改善は( $K$ < 5 の範囲では)見られない。