ベイズ線形回帰の事前分布を色々比較してみた

はじめに

ベイズモデルは柔軟性が高く、伝統的な線形モデルや一般化線形モデルの他に、筆者がこの記事のように、

様々なモデル・分布を自由自在に組み合わせることによって、生データに隠された構造を可視化できる。

ただ、ベイズモデルを利用する際に、パラメータの事前分布を指定する必要があり、しかも事前分布のチョイス次第でモデルの挙動が大きく変わることがあるため、初心者にとっては少々難易度が高いと思われる。

本記事は単純な線形回帰モデルを事例に、事前分布を正規分布、ラプラス分布、ガンマラッソ（Taddy 2017）に変えることでパラメータ（係数）の挙動がどうなるかを実際に可視化して確認する。

モデル説明

ここではまずモデルの説明から入る。

一つ目は切片（β）に正規分布を設定するモデル：

\sigma \sim Gamma(1,1) \\
\rho \sim Gamma(1, 1) \\
\beta \sim Normal(0, 1/\rho^2)\\
Y \sim Normal(alpha + X \beta,  1/\sigma^2)

これはベイズ版のリッジ回帰と見なしても良いと思われる（たとえばTheodoridis 2020 p. 598を参照）。

二つ目は切片にラプラス分布を設定するモデル：

\sigma \sim Gamma(1,1) \\
\rho \sim Gamma(1, 1) \\
\beta \sim Laplace(0, 1/\rho^2)\\
Y \sim Normal(alpha + X \beta,  1/\sigma^2)

これはベイズ版のL1正則化回帰（ラッソ）と見なしても良いと思われる（たとえばTheodoridis 2020 p. 666を参照）。

最後の三番目はTaddy（2017）が提案したガンマラッソ（gamma lasso）モデルのベイズ的な解釈として紹介したモデル：

\sigma \sim Gamma(1,1) \\
\rho_{k} \sim Gamma(1, 1) \\
\beta_{k} \sim Laplace(0, 1/\rho_{k}^2)\\
Y \sim Normal(alpha + X \beta,  1/\sigma^2)

このモデルは二番目のベイズ版のL1正則化回帰と似ているが、パラメータごとに異なる「分散」（dispersion）を設定するところが特徴である。

詳細はTaddy（2017）の4.1のところ（p. 530）を参照してください。Taddy（2017）で紹介されたガンマラッソの頻度論版（？）はRのgamlrパッケージから利用できる：

実装

Stanのコード

ここでは、大量の独立変数を入れるため、観測値がP種類のグループの中の一つのグループに属し、そのグループに属している観測値は皆同じパラメータを持つ状況を想定する。

この場合ダミー変数をP個もしくはP-1個作成して巨大な行列を作る方法もあるが、本記事では直接indexをStanに渡す方法でモデルを実装する。

まず一つ目のベイズ版のリッジ回帰は下記のようで、これをnormal_prior.stanとして保存する。

data {
  int<lower=1> N; //測値
  int<lower=1> P; //説明変数の数
  array[N] int<lower=1,upper=P> X; //観測値の説明変数。観測値が属するグループのようなイメージ
  array[N] real Y; //結果変数の値
}
parameters {
  real<lower=0> rho;
  real<lower=0> sigma;
  real alpha;
  vector[P] beta;
}
model{
  rho ~ gamma(1, 1);
  sigma ~ gamma(1, 1);
  beta ~ normal(0, 1/rho^2);
  Y ~ normal(alpha + beta[X], 1/sigma^2);
}

二つ目のベイズ版のL1正則化回帰は下記のようで、これをlaplace_prior.stanとして保存する。

data {
  int<lower=1> N; //測値
  int<lower=1> P; //説明変数の数
  array[N] int<lower=1,upper=P> X; //観測値の説明変数。観測値が属するグループのようなイメージ
  array[N] real Y; //結果変数の値
}
parameters {
  real<lower=0> rho;
  real<lower=0> sigma;
  real alpha;
  vector[P] beta;
}
model{
  rho ~ gamma(1, 1);
  sigma ~ gamma(1, 1);
  beta ~ double_exponential(0, 1/rho^2); //これがLaplace分布
  Y ~ normal(alpha + beta[X], 1/sigma^2);
}

三つ目のベイズ版のガンマラッソ（gamma lasso）モデルは下記のようで、これをgamma_lasso.stanとして保存する。

data {
  int<lower=1> N; //測値
  int<lower=1> P; //説明変数の数
  array[N] int<lower=1,upper=P> X; //観測値の説明変数。観測値が属するグループのようなイメージ
  array[N] real Y; //結果変数の値
}
parameters {
  vector<lower=0>[P] rho;
  real<lower=0> sigma;
  real alpha;
  vector[P] beta;
}
model{
  rho ~ gamma(1, 1);
  sigma ~ gamma(1, 1);
  beta ~ double_exponential(0, 1.0./(rho.^2)); //ブロードキャスト処理をする際は1をreal型の1.0にしないとエラーが出る
  Y ~ normal(alpha + beta[X], 1/sigma^2);
}

モデル推定

まずは必要なパッケージを入れる。

library(tidyverse)
library(patchwork)
library(cmdstanr)

次に、データのシミュレーションを行う。具体的には、観測値の数と独立変数（グループ数）は5000で、独立変数のパラメータの最初の250個の値は20、251個目から500個目は-20で、残りは全て0という状況を仮定してデータを生成する。

set.seed(12345)
N <- 5000
P <- 5000
P_nonzero_1 <- 250
P_nonzero_2 <- 250

beta <- c(
  rep(20, P_nonzero_1),
  rep(-20, P_nonzero_2),
  rep(0, P - (P_nonzero_1 + P_nonzero_2))
)

#観測値が属するグループをサンプリング
#1からPまでのラベルをN個復元抽出する
X <- sample(1:P, size = N, replace = TRUE)

Y <- beta[X] + rnorm(N)

次にモデルの推定を行う

m_normal_init <- cmdstan_model("normal_prior.stan")
m_laplace_init <- cmdstan_model("laplace_prior.stan")
m_gamma_lasso_init <- cmdstan_model("gamma_lasso.stan")

data_list <- list(
  N = N,
  P = P,
  X = X,
  Y = Y
)

m_normal_estimate <- m_normal_init$variational(data = data_list)
m_normal_summary <- m_normal_estimate$summary()


m_laplace_estimate <- m_laplace_init$variational(data = data_list)
m_laplace_summary <- m_laplace_estimate$summary()

m_gamma_lasso_estimate <- m_gamma_lasso_init$variational(data = data_list)
m_gamma_lasso_summary <- m_gamma_lasso_estimate$summary()

可視化

ここでは本記事の主な内容で、実際に異なる事前分布の元で推定されたパラメータの挙動を可視化して、それぞれの事前分布が意味する前提設定とそれが分析結果に与える影響を説明する。

g_normal <- m_normal_summary[which(str_detect(m_normal_summary$variable, "beta")),] %>%
  ggplot() + 
  geom_point(aes(x = 1:P, y = mean), size = 0.1) + 
  geom_ribbon(aes(x = 1:P, ymin = q5, ymax = q95), fill = alpha("blue", 0.3)) + 
  geom_vline(xintercept = (P_nonzero_1 + P_nonzero_2), linetype = "dashed", color = "red") + 
  xlab("係数") + 
  ylab("係数の推定値") + 
  ggtitle("事前分布：正規分布") + 
  ylim(-30,30) + 
  theme_gray (base_family = "HiraKakuPro-W3")

g_laplace <- m_laplace_summary[which(str_detect(m_laplace_summary$variable, "beta")),] %>%
  ggplot() + 
  geom_point(aes(x = 1:P, y = mean), size = 0.1) + 
  geom_ribbon(aes(x = 1:P, ymin = q5, ymax = q95), fill = alpha("blue", 0.3)) + 
  geom_vline(xintercept = (P_nonzero_1 + P_nonzero_2), linetype = "dashed", color = "red") + 
  xlab("係数") + 
  ylab("係数の推定値") + 
  ggtitle("事前分布：ラプラス分布") + 
  ylim(-30,30) + 
  theme_gray (base_family = "HiraKakuPro-W3")

g_gamma_lasso <- m_gamma_lasso_summary[which(str_detect(m_gamma_lasso_summary$variable, "beta")),] %>%
  ggplot() + 
  geom_point(aes(x = 1:P, y = mean), size = 0.1) + 
  geom_ribbon(aes(x = 1:P, ymin = q5, ymax = q95), fill = alpha("blue", 0.3)) + 
  geom_vline(xintercept = (P_nonzero_1 + P_nonzero_2), linetype = "dashed", color = "red") + 
  xlab("係数") + 
  ylab("係数の推定値") + 
  ggtitle("事前分布：gamma lasso分布") + 
  ylim(-30,30) + 
  theme_gray (base_family = "HiraKakuPro-W3")

g_normal | g_laplace | g_gamma_lasso

まず、一つ目のリッジ回帰（正規分布が事前分布）の場合、1個目から250個目のパラメータも、251個目から500個目のパラメータも正解の値（それぞれ20と-20）からかけ離れたゼロ近辺の値になっている。これはTheodoridis（2020）のpp. 666 - 667の説明によると、正規分布は裾が軽いため、モデルは事前分布の平均値付近（ゼロ付近）の値を好み、ゼロから離れた値に対するペナルティが強いからである。

二つ目のベイズ版のL1正則化回帰（ラプラス分布が事前分布）の場合、事前分布のラプラス分布は裾が重いため、平均値付近（ゼロ付近）以外の値も受け入れる傾向があり、結果として1個目から250個目のパラメータも、251個目から500個目のパラメータも正解の値（それぞれ20と-20）近辺の値になっている。

最後のガンマラッソ（gamma lasso）モデルは二つ目のベイズ版のL1正則化回帰と同じくラプラス分布を事前分布としているため、全体的な挙動は似ているが、よくよく確認したら501個目以降の正解がゼロのパラメータの推定値がベイズ版のL1正則化回帰よりゼロ付近に密集している。これはガンマラッソがパラメータごとに異なる「分散」（dispersion）を設定する効果と理解しても良いと思われる。

参考文献

Taddy, Matt. "One-step estimator paths for concave regularization." Journal of Computational and Graphical Statistics 26.3 (2017): 525-536.

Theodoridis, Sergios. (2020). Machine learning: a Bayesian and optimization perspective 2nd Edition. Academic press.