LIVESENSE Data Analytics Blog

リブセンスのデータ分析、機械学習、分析基盤に関する取り組みをご紹介するブログです。

コレスキー分解を利用した相関係数のベイズ推定

 こんにちは、リブセンスでデータサイエンティストをしている北原です。今回は、多変量正規分布の分散共分散行列を扱うときに有用であることが知られているコレスキー分解を取り上げます。

 多変量正規分布を使ったモデリングをしたいことはよくありますが、複雑な分布であるため計算時間が長くなりやすかったり不安定になりやすかったりします。コレスキー分解を利用することで、この問題が緩和されます。今回は、コレスキー分解を利用した具体的な例として相関係数の推定を扱います。コードはRとStanです。

相関係数

 まず、基本の確認のため、簡単に相関係数について説明します。 

 相関係数は二変量の線形な関係性を定量的に示す指標です。実際には相関係数と呼ばれるものはいろいろありますが、ここで扱うのは最も基本的なピアソンの積率相関係数です。ピアソンの積率相関係数では、偏差に正規分布を仮定します。

 具体的には、2変量正規分布

$$ \begin{align} f(\mathbf{x})&={\frac {1}{2\pi \sqrt {|{\boldsymbol {\Sigma }}|}}} \exp \left(-{\frac {1}{2}}({\mathbf {x} }-{\boldsymbol {\mu }})^{\mathrm {T} }{\boldsymbol {\Sigma }}^{-1}({\mathbf {x} }-{\boldsymbol {\mu }})\right) \\ {\boldsymbol {\mu }}&={\begin{pmatrix}\mu _{1}\\\mu_{2}\end{pmatrix}}\\ {\boldsymbol {\Sigma }}&={\begin{pmatrix}\sigma _{1}^{2}&\rho \sigma _{1}\sigma _{2}\\\rho \sigma _{1}\sigma _{2}&\sigma _{2}^{2}\end{pmatrix}} \end{align} $$

における、分散共分散行列$\mathbf{\Sigma}$の要素$\rho$が相関係数、$\sigma_i$が分散です。

 Rでサンプルデータを作成して、実際に相関係数を計算してみましょう。多変量正規分布からのサンプリングにはmvtnormパッケージのrmvnorm()関数を使います。相関係数の計算にはcorrelationパッケージのcorrelation()関数が便利なので、これを使います。なお、ここで作成されたサンプルデータは後で紹介するベイズ推定でも流用します。

library(tidyverse)

library(mvtnorm)
library(correlation)

# サンプルデータ生成
set.seed(1)
N <- 300
mu <- c(2, 5)       # 平均
sig <- c(0.2, 0.3)  # 標準偏差
rho <- 0.7          # 相関係数
Sig <-              # 分散共分散行列
  matrix(
    c(sig[1]^2,              sig[1] * sig[2] * rho,
      sig[1] * sig[2] * rho, sig[2]^2),
    nrow = 2)

dat <- 
  rmvnorm(n = N, mean = mu, sigma = Sig) %>% 
  as_tibble(.name_repair = "minimal") %>% 
  setNames(c("x1", "x2"))

# 相関係数を算出
dat %>% 
  dplyr::select(x1, x2) %>% 
  correlation()

次のような出力結果が得られます。

# Correlation Matrix (pearson-method)

Parameter1 | Parameter2 |    r |       95% CI | t(298) |         p
------------------------------------------------------------------
x1         |         x2 | 0.68 | [0.61, 0.73] |  15.79 | < .001***

p-value adjustment method: Holm (1979)
Observations: 300

相関係数のベイズ推定

 まず、コレスキー分解を使わないで相関係数をベイズ推定してみましょう。なお、動作確認しているcmdstanのバージョンは2.31.0です。

 前節の多変量正規分布をStanでそのまま表すと次のようになります。

// cor.stan
data{
  int<lower=0> N;
  array[N] vector[2] x;
}
parameters {
  vector[2] mu;
  array[2] real<lower=0> sig;  // 標準偏差
  real<lower=-1, upper=1> rho; // 相関係数
}
transformed parameters {
  cov_matrix[2] Sig = [[sig[1]^2, rho * sig[1] * sig[2]], [rho * sig[1] * sig[2], sig[2]^2]];
}
model {
  x ~ multi_normal(mu, Sig);
  sig ~ cauchy(0, 5);
}

Stanを実行するRコードは以下の通りです。

library(cmdstanr)
library(doParallel)
registerDoParallel(detectCores())
options(mc.cores = parallel::detectCores())

mod_cor <- cmdstan_model("cor.stan")
fit <- mod_cor$sample(
  data = list(N = nrow(dat), 
              x = dat %>% dplyr::select(x1, x2))
    )

fit$summary() %>% 
  mutate_at(vars(-variable), ~ round(., digits = 2)) %>% 
  kable() %>% 
  print()

実行結果は次のようになります。意図通り推定されていることがわかります。

variable mean median sd mad q5 q95 rhat ess_bulk ess_tail
lp__ 637.42 637.76 1.55 1.36 634.47 639.30 1 2007.37 2717.26
mu[1] 2.00 2.00 0.01 0.01 1.98 2.02 1 3091.49 2876.91
mu[2] 5.01 5.01 0.02 0.02 4.98 5.04 1 3139.37 2498.59
sig[1] 0.20 0.20 0.01 0.01 0.19 0.21 1 2870.80 3007.41
sig[2] 0.29 0.29 0.01 0.01 0.27 0.31 1 3057.67 3008.42
rho 0.67 0.67 0.03 0.03 0.62 0.72 1 2950.49 3097.38
Sig[1,1] 0.04 0.04 0.00 0.00 0.04 0.05 1 2870.85 3007.41
Sig[2,1] 0.04 0.04 0.00 0.00 0.03 0.05 1 2457.67 2297.42
Sig[1,2] 0.04 0.04 0.00 0.00 0.03 0.05 1 2457.67 2297.42
Sig[2,2] 0.09 0.09 0.01 0.01 0.08 0.10 1 3057.65 3008.42

コレスキー分解

 コレスキー分解は、対称正定値行列$\mathbf{A}$を正の対角要素をもつ下三角行列$\mathbf{L}$を使って$\mathbf{A} = \mathbf{L} \mathbf{L}^T$と分解することを指します。下三角行列であれば効率的な計算が可能です。分散共分散行列や相関行列は対称正定値行列なので、基本的にはコレスキー分解ができます。

 例えば、2変量正規分布の分散共分散行列$\mathbf{\Sigma }$は

$$ \begin{align} {\mathbf {\Sigma}} &={\begin{pmatrix}\sigma _{1}^{2}&\rho \sigma _{1}\sigma _{2}\\\rho \sigma _{1}\sigma _{2}&\sigma _{2}^{2}\end{pmatrix}} \nonumber \\ &={\begin{pmatrix}\sigma _{1} & 0 \\\ \rho \sigma _{2}& \sigma _{2} \sqrt{1 - \rho^2} \end{pmatrix}} {\begin{pmatrix}\sigma _{1} & \rho \sigma _{2} \\\ 0& \sigma _{2} \sqrt{1 - \rho^2} \end{pmatrix}} \nonumber \\ &= \mathbf{L} \mathbf{L}^T \end{align} $$

と下三角行列$\mathbf{L}$にコレスキー分解されます。Rでは関数chol()、Stanでは関数cholesky_decompose()を使えばコレスキー分解が簡単にできるので、必ずしも手計算する必要はありません。

 分散共分散行列$\mathbf{\Sigma }$は、分散$\mathbf{\sigma}$の対角行列$\mathbf{D}_{\sigma}$と相関行列$\mathbf{R}$のコレスキー分解で

$$ \begin{align} \mathbf {\Sigma} &= \mathbf{D}_{\sigma} \mathbf{R} \mathbf{D}_{\sigma} \nonumber \\ &= \mathbf{D}_{\sigma} \mathbf{L}_{\rho} \mathbf{L}_{\rho}^T \mathbf{D}_{\sigma} \end{align} $$

と表すこともできます。2変量正規分布の場合、

$$ \begin{align}{\mathbf {\Sigma}} &= {\begin{pmatrix} \sigma_1 & 0 \\ 0 & \sigma_2 \end{pmatrix}} {\begin{pmatrix} 1 & \rho \\ \rho & 1 \end{pmatrix}} {\begin{pmatrix} \sigma_1 & 0 \\ 0 & \sigma_2 \end{pmatrix}} \nonumber \\ &= {\begin{pmatrix} \sigma_1 & 0 \\ 0 & \sigma_2 \end{pmatrix}} {\begin{pmatrix} 1 & 0 \\ \rho & \sqrt{1 - \rho^2} \end{pmatrix}} {\begin{pmatrix} 1 & \rho \\ 0 & \sqrt{1 - \rho^2} \end{pmatrix}} {\begin{pmatrix} \sigma_1 & 0 \\ 0 & \sigma_2 \end{pmatrix}} \\ \end{align} $$

となるので、これを利用して推定を行うこともできます。

コレスキー分解を利用した相関係数のベイズ推定

 では、コレスキー分解を利用して相関係数を推定してみましょう。

 Stanにはコレスキー分解に対応した多変量正規分布のサンプリングステートメントmulti_normal_cholesky()があるので、これを利用することでコレスキー分解を利用した効率的なサンプリングが可能になります。multi_normal()で分散共分散行列を引数としていたところに、コレスキー分解した分散共分散行列を指定すれば同じ結果が得られます。Stanコードは以下のようになります。変更箇所は一行のみです。

// cor_cholesky.stan
data{
  int<lower=0> N;
  array[N] vector[2] x;
}
parameters {
  vector[2] mu;
  array[2] real<lower=0> sig;  // 標準偏差
  real<lower=-1, upper=1> rho; // 相関係数
}
transformed parameters {
  cov_matrix[2] Sig = [[sig[1]^2, rho * sig[1] * sig[2]], [rho * sig[1] * sig[2], sig[2]^2]];
}
model {
  x ~ multi_normal_cholesky(mu, cholesky_decompose(Sig)); // コレスキー分解を利用
  sig ~ cauchy(0, 5);
}

実行結果は以下のようになります。コレスキー分解を利用しない時とほぼ同じ結果であることがわかります。

variable mean median sd mad q5 q95 rhat ess_bulk ess_tail
lp__ 637.37 637.72 1.57 1.36 634.37 639.27 1 2094.49 2518.01
mu[1] 2.00 2.00 0.01 0.01 1.98 2.02 1 3318.35 2930.70
mu[2] 5.01 5.01 0.02 0.02 4.98 5.04 1 3428.26 3280.05
sig[1] 0.20 0.20 0.01 0.01 0.19 0.21 1 2718.84 2540.19
sig[2] 0.29 0.29 0.01 0.01 0.28 0.32 1 2880.14 2839.61
rho 0.67 0.67 0.03 0.03 0.62 0.72 1 2962.34 2721.85
Sig[1,1] 0.04 0.04 0.00 0.00 0.04 0.05 1 2718.96 2540.19
Sig[2,1] 0.04 0.04 0.00 0.00 0.03 0.05 1 2314.92 2305.13
Sig[1,2] 0.04 0.04 0.00 0.00 0.03 0.05 1 2314.92 2305.13
Sig[2,2] 0.09 0.09 0.01 0.01 0.08 0.10 1 2880.09 2839.61

 行列の要素を明示せずにモデリングしたいこともあるので、次に、分散の対角行列と相関行列のコレスキー分解を利用して推定を行ってみましょう。Stanコードは以下のようになります。

// cor_cholesky2.stan
data{
  int<lower=0> N;
  array[N] vector[2] x;
}
parameters {
  vector[2] mu;
  corr_matrix[2] R;       // 相関行列
  vector<lower=0>[2] sig; // 標準偏差
}
model {
  x ~ multi_normal_cholesky(mu, diag_matrix(sig) * cholesky_decompose(R));
  sig ~ cauchy(0, 5);
}

実行結果は以下のようになります。相関行列の要素から相関係数がわかります。こちらでも、コレスキー分解を利用しない時とほぼ同じ結果であることが確認できます。

variable mean median sd mad q5 q95 rhat ess_bulk ess_tail
lp__ 638.01 638.35 1.61 1.44 634.93 639.97 1 1853.54 2576.67
mu[1] 2.00 2.00 0.01 0.01 1.98 2.02 1 3019.93 2926.09
mu[2] 5.01 5.01 0.02 0.02 4.98 5.04 1 2858.86 3019.28
R[1,1] 1.00 1.00 0.00 0.00 1.00 1.00 NA NA NA
R[2,1] 0.67 0.67 0.03 0.03 0.62 0.72 1 2557.81 2616.92
R[1,2] 0.67 0.67 0.03 0.03 0.62 0.72 1 2557.81 2616.92
R[2,2] 1.00 1.00 0.00 0.00 1.00 1.00 NA NA NA
sig[1] 0.20 0.20 0.01 0.01 0.19 0.21 1 2739.05 2872.83
sig[2] 0.29 0.29 0.01 0.01 0.28 0.32 1 2569.37 2731.76

 私の計算機環境ではコレスキー分解を利用することで2倍近く推定速度が速くなりました。簡単に推定時間を短くできるので知っておくと便利なテクニックです。

まとめ

 今回はStanで多変量正規分布を使ったモデリングを行うときに有用なコレスキー分解を紹介しました。コレスキー分解を利用した例として相関係数の推定を取り上げました。Stanにはコレスキー分解に対応した関数が用意されているので、容易に利用することができ、推定速度を向上させることができることを示しました。今回のようなテクニックは、知っている人にとっては常識的でほとんど説明がされないため忘れてしまいがちですが、覚えておくと便利です。