Efficient sampling for Gaussian linear regression with arbitrary priors という論文を読んだときのアルゴリズムに関するメモです。
楕円スライスサンプリングについて
以下のような分解が可能な分布 p(\bm \Delta) を考える。
\begin{aligned}
p(\bm \Delta) \propto \mathcal N(\bm \Delta \mid \bm 0, \bm V) L(\bm \Delta).
\end{aligned}
分布 p(\bm \Delta) に従うサンプルを生成するアルゴリズムとして、文献では楕円スライスサンプリング という方法が提案されている。楕円スライスサンプリングは以下の原理を使う。
!
基本原理
\bm v_0 \sim \mathcal N(\bm 0, \bm V), \bm v_1 \sim \mathcal N(\bm 0, \bm V) で、\bm v_0 と \bm v_1 が独立のとき、 \forall \theta \in [0, 2\pi) に対して
\begin{aligned}
\bm \Delta = \sin(\theta) \bm v_0 + \cos(\theta) \bm v_1
\end{aligned}
とすれば、
\begin{aligned}
\bm \Delta \sim \mathcal N(\bm 0, \bm V).
\end{aligned}
略証
\bm \Delta はGauss乱数の線形結合だから \bm \Delta 自体もGauss乱数である。その平均ベクトルは
\begin{aligned}
\mathbb E[\bm \Delta] = \sin(\theta) \mathbb E[\bm v_0] + \cos(\theta) \mathbb E[\bm v_1] = \bm 0
\end{aligned}
共分散行列は
\begin{aligned}
\operatorname{Cov}[\bm \Delta]
&= \mathbb E[\bm \Delta \bm \Delta^{\top}] \\
&= \mathbb E[(\sin(\theta) \bm v_0 + \cos(\theta) \bm v_1)(\sin(\theta) \bm v_0 + \cos(\theta) \bm v_1)^{\top}] \\
&= \sin^2(\theta) \mathbb E[\bm v_0 \bm v_0^{\top}] + \cos^2(\theta) \mathbb E[\bm v_1 \bm v_1^{\top}] \\
&= \sin^2(\theta) \bm V + \cos^2(\theta) \bm V \\
&= \bm V.
\end{aligned}
もとの確率分布
\begin{aligned}
p(\bm \Delta) \propto \mathcal N(\bm \Delta \mid \bm 0, \bm V) L(\bm \Delta)
\end{aligned}
の代わりに、基本原理の3つの確率変数 \bm v_0, \bm v_1, \theta の同時分布
\begin{aligned}
p(\bm v_0, \bm v_1, \theta)
\propto
\mathcal N(\bm v_0 \mid \bm 0, \bm V)
\mathcal N(\bm v_1 \mid \bm 0, \bm V)
\pi(\theta)
L(\sin(\theta) \bm v_0 + \cos(\theta) \bm v_1)
\end{aligned}
からサンプルを生成できるならば、p(\bm \Delta) に従うサンプルを
\begin{aligned}
\bm \Delta = \sin(\theta) \bm v_0 + \cos(\theta) \bm v_1
\end{aligned}
と構成することができる。手続きは以下の通り。
ところが p(\bm v_0, \bm v_1, \theta) からのサンプリングが難しすぎるので、この手続きは不可能である。代わりに \bm \Delta も確率変数とみなして、4変数の同時分布
\begin{aligned}
p(\bm v_0, \bm v_1, \bm \Delta, \theta)
\propto p(\bm v_0, \bm v_1 \bm \Delta \mid \theta) \pi (\theta) L(\sin(\theta) \bm v_0 + \cos(\theta) \bm v_1)
\end{aligned}
に従うサンプルを生成する。
3つの確率変数 \bm \Delta, \bm v_0, \bm v_1 はいずれも同一の平均ベクトルと共分散行列をもつGauss分布に従うが、\bm \Delta は \bm v_0, \bm v_1 それぞれと相関を持ち、\bm v_0 と \bm v_1 は相関を持たない。このことに注意して、条件付き同時分布 p(\bm v_0, \bm v_1, \bm \Delta \mid \theta) を記述すると、
\begin{aligned}
p(\bm v_0, \bm v_1, \bm \Delta \mid \theta)
&=
\mathcal N\!\left(
\begin{bmatrix}
\bm v_0 \\
\bm v_1 \\
\bm \Delta
\end{bmatrix}
~ \middle| ~
\bm 0,
\begin{bmatrix}
\bm V & \bm O & \bm V \sin \theta \\
\bm O & \bm V & \bm V \cos \theta \\
\bm V \sin \theta & \bm V \cos \theta & \bm V
\end{bmatrix}
\right)
\end{aligned}
が得られる。また \pi(\theta) については、[0, 2\pi) 上の一様分布を仮定する。目標となる同時分布は以下のように書き換えられる。
\begin{aligned}
p(\bm v_0, \bm v_1, \bm \Delta, \theta)
&\propto
p(\bm v_0, \bm v_1 \bm \Delta \mid \theta)
\pi(\theta)
L(\sin(\theta) \bm v_0 + \cos(\theta) \bm v_1)
\\
&\propto
\mathcal N\!\left(
\begin{bmatrix}
\bm v_0 \\
\bm v_1 \\
\bm \Delta
\end{bmatrix}
~ \middle| ~
\bm 0,
\begin{bmatrix}
\bm V & \bm O & \bm V \sin \theta \\
\bm O & \bm V & \bm V \cos \theta \\
\bm V \sin \theta & \bm V \cos \theta & \bm V
\end{bmatrix}
\right)
\\ & \qquad \times
L(\sin(\theta) \bm v_0 + \cos(\theta) \bm v_1).
\end{aligned}
アルゴリズムは以下の2ブロックのサンプリングで構成される。
!
楕円スライスサンプリング
\bm v_0, \bm v_1 \mid \bm \Delta, \theta の生成 (楕円回転)
\bm v \sim \mathcal N(\bm 0, \bm V) を生成する。
\bm v_0 \gets \sin(\theta) \bm \Delta + \cos(\theta) \bm v とする。
\bm v_1 \gets \cos(\theta) \bm \Delta - \sin(\theta) \bm v とする。
\bm \Delta, \theta \mid \bm v_0, \bm v_1 の生成 (スライスサンプリング)
a \gets 0, b \gets 2\pi とする。
\ell \sim \operatorname{Uniform}(0, L(\sin(\theta) \bm v_0 + \cos(\theta) \bm v_1)) を生成する。
\theta^\ast \sim \operatorname{Uniform}(a, b) を生成する。
もしも L(\sin(\theta^\ast) \bm v_0 + \cos(\theta^\ast) \bm v_1) \lt \ell ならば
もしも \theta^\ast < \theta ならば
a \gets \theta^\ast として、ステップ 2.-3. に戻る。
そうでなければ
b \gets \theta^\ast として、ステップ 2.-3. に戻る。
\theta \gets \theta^\ast とする。
\bm \Delta \gets \sin(\theta) \bm v_0 + \cos(\theta) \bm v_1 とする。
1.-2. を繰り返す。
線形回帰への適用
線形回帰モデル
\begin{aligned}
\bm y &= \bm X \bm \beta + \bm \varepsilon, \\
\bm \varepsilon &\sim \mathcal N(\bm 0, \sigma^2 \bm I_n), \\
\bm X &\in \mathbb R^{n \times p},
\end{aligned}
に対して \bm \beta の事前分布として \pi(\bm \beta, \lambda) を用いる。事後分布は
\begin{aligned}
p(\bm \beta \mid \bm X, \bm y, \sigma)
\propto
\mathcal N(\bm y \mid \bm X \bm \beta, \sigma^2 \bm I_n)
\pi(\bm \beta, \lambda).
\end{aligned}
これをBayes Ridge回帰の要領でちょっと変形すると
\begin{aligned}
p(\bm \beta \mid \bm X, \bm y, \sigma)
&\propto
\mathcal N(\bm y \mid \bm X \bm \beta, \sigma^2 \bm I_n)
\mathcal N(\bm \beta \mid \bm 0, \sigma^2 c \bm I_p)
\times
\frac{\pi(\bm \beta, \lambda)}{\mathcal N(\bm \beta \mid \bm 0, \sigma^2 c \bm I_p)}
\\
&\propto
\mathcal N\!\left( \bm \beta ~\middle|~ \hat{\bm \beta}, \sigma^2 \left( \bm X^\top \bm X + \frac{1}{c} \bm I_p \right)^{-1} \right)
\times
\frac{\pi(\bm \beta, \lambda)}{\mathcal N(\bm \beta \mid \bm 0, \sigma^2 c \bm I_p)}
\\
&=
\mathcal N\!\left( \bm \beta - \hat{\bm \beta} ~\middle|~ \bm 0, \sigma^2 \left( \bm X^\top \bm X + \frac{1}{c} \bm I_p \right)^{-1} \right)
\times
\frac{\pi(\bm \beta, \lambda)}{\mathcal N(\bm \beta \mid \bm 0, \sigma^2 c \bm I_p)}.
\end{aligned}
ただし \hat{\bm \beta} は次式で与えられるRidge推定量である。
\begin{aligned}
\hat{\bm \beta} = \left( \bm X^\top \bm X + \frac{1}{c} \bm I_p \right)^{-1} \bm X^\top \bm y.
\end{aligned}
この分布の形と p(\bm \Delta) \propto \mathcal N(\bm \Delta \mid \bm 0, \bm V) L(\bm \Delta) の形を見比べると、
\bm \Delta = \bm \beta - \hat{\bm \beta}
\bm V = \sigma^2 (\bm X^\top \bm X + c^{-1} \bm I_p)^{-1}
L(\bm \Delta) = \dfrac{\pi(\bm \beta, \lambda)}{\mathcal N(\bm \beta \mid \bm 0, \sigma^2 c \bm I_p)}
という対応で、楕円スライスサンプリングを用いて \bm \beta - \hat{\bm \beta} のサンプリングを行うことができる。
Gibbsサンプリングの導入
全部の \bm \beta に対して楕円スライスサンプリングを行うと計算量が大きくなるので、\bm \beta を K ~ (\le p) 個のブロックに分割する。K 個のブロックのうち k 番目のブロックを \bm \beta_k 、それ以外を \bm \beta_{\setminus k} とし、さらに \hat{\bm \beta} についても同様に \hat{\bm \beta}_k, \hat{\bm \beta}_{\setminus k} とする。そして \pi(\bm \beta, \lambda) が \pi(\bm \beta_k, \lambda) \pi(\bm \beta_{\setminus k}, \lambda) へと分解可能であるとする。このとき
\begin{aligned}
p(\bm \beta_k, \bm \beta_{\setminus k} \mid \bm X, \bm y, \sigma)
&\propto
\mathcal N\!\left(
\begin{bmatrix}
\bm \beta_k \\
\bm \beta_{\setminus k}
\end{bmatrix}
~\middle|~
\begin{bmatrix}
\hat{\bm \beta}_k \\
\hat{\bm \beta}_{\setminus k}
\end{bmatrix},
\sigma^2
\begin{bmatrix}
\bm \Sigma_{kk} & \bm \Sigma_{k, \setminus k} \\
\bm \Sigma_{\setminus k, k} & \bm \Sigma_{\setminus k, \setminus k}
\end{bmatrix}
\right)
\\ &\qquad
\times
\frac{\pi(\bm \beta_k) \pi(\bm \beta_{\setminus k})}{
\mathcal N(\bm \beta_k \mid \bm 0, \sigma^2 c \bm I)
\mathcal N(\bm \beta_{\setminus k} \mid \bm 0, \sigma^2 c \bm I)
},
\end{aligned}
\\
\begin{aligned}
\begin{bmatrix}
\bm \Sigma_{kk} & \bm \Sigma_{k, \setminus k} \\
\bm \Sigma_{\setminus k, k} & \bm \Sigma_{\setminus k, \setminus k}
\end{bmatrix}
&=
\left(
\bm X^\top \bm X + \frac{1}{c} \bm I_p
\right)^{-1},
\end{aligned}
\\
\begin{aligned}
\begin{bmatrix}
\hat{\bm \beta}_k \\
\hat{\bm \beta}_{\setminus k}
\end{bmatrix}
&=
\left(
\bm X^\top \bm X + \frac{1}{c} \bm I_p
\right)^{-1}
\bm X^\top \bm y.
\end{aligned}
\\
これを使うと多変量正規分布の条件付き確率の性質から、
\begin{aligned}
p(\bm \beta_k \mid \bm \beta_{\setminus k}, \bm X, \bm y, \sigma)
&\propto
\mathcal N\!\left(
\bm \beta_k
~\middle|~
\tilde{\bm \beta}_k,
\sigma^2 \tilde{\bm \Sigma}_k
\right)
\times
\frac{\pi(\bm \beta_k, \lambda)}{\mathcal N\!\left(
\bm \beta_k
~\middle|~
\bm 0, \sigma^2 c \bm I
\right)},
\\
&=
\mathcal N\!\left(
\bm \beta_k - \tilde{\bm \beta}_k
~\middle|~
\bm 0,
\sigma^2 \tilde{\bm \Sigma}_k
\right)
\times
\frac{\pi(\bm \beta_k, \lambda)}{\mathcal N\!\left(
\bm \beta_k
~\middle|~
\bm 0, \sigma^2 c \bm I
\right)},
\end{aligned}
\\
\begin{aligned}
\tilde{\bm \beta}_k
&=
\hat{\bm \beta}_k + \bm \Sigma_{k, \setminus k} \bm \Sigma_{\setminus k, \setminus k}^{-1} (\bm \beta_{\setminus k} - \hat{\bm \beta}_{\setminus k}), \\
\tilde{\bm \Sigma}_k
&=
\bm \Sigma_{kk} - \bm \Sigma_{k, \setminus k} \bm \Sigma_{\setminus k, \setminus k}^{-1} \bm \Sigma_{\setminus k, k}
\end{aligned}
が構成できる。こうやって構成された \bm \beta_k に関して、改めて楕円スライスサンプリングを行うことで、効率的にサンプリングを行うことができる。
この方法では、各ブロックの構築に必要な行列
\begin{aligned}
& \bm \Sigma_{k, \setminus k} \bm \Sigma_{\setminus k, \setminus k}^{-1}, \\
& \tilde{\bm \Sigma_k} = \bm \Sigma_{kk} - \bm \Sigma_{k, \setminus k} \bm \Sigma_{\setminus k, \setminus k}^{-1} \bm \Sigma_{\setminus k, k}
\end{aligned}
および、各ブロックにおいて \bm v_k を生成するときに必要な Cholesky 分解
\begin{aligned}
\tilde{\bm \Sigma}_k = \bm L_k \bm L_k^\top
\end{aligned}
をあらかじめ計算しておくことができ、各ブロックのサンプリングに必要な計算量を大幅に削減できる。しかもこれらの事前計算は並列化できるので、大規模なデータセットに対しても高速にサンプリングできる。論文著者らは K=p 、すなわち回帰係数 \bm \beta をスカラーにまで分解するのが最も効率的だと報告している。
Discussion