超幾何分布をマルコフ基底を使って計算する

こちらはリブセンス Advent Calendar 2023 の 7 日目の記事です。社内 LT 会で喋ったことをより詳しく解説しています。

参考文献

本記事は共立出版の代数的統計モデルを参考に書いています。フィッシャーの正確確率検定とか超幾何分布とか十分統計量とかは既知のものとして書かれているので、はじめは読むのに手こずりました。しかし話題はとても面白いですし、他に類書があまり無いので、良い本だと思います。

分割表の統計モデル

分割表とは

分割表とは、こんな感じのよくある集計表のことです (値に意味はありません)。各セルは非負整数で、一般には $N_0 \times N_1 \times \cdots$ のように多次元まで考えることができます。

	シナリオ X	シナリオ Y	シナリオ Z
介入 A	24	81	35
介入 B	41	44	55
介入 C	32	51	57

よくある A/B テストの仮定

簡単のために $2 \times 2$ の分割表を考えましょう。 $2 \times 2$ の分割表は A/B テストで目にすることが多いんじゃないでしょうか。

	クリックした	クリックしなかった
A パターン	400	600
B パターン	500	500

普通 A/B テストは、各パターンごとに収集するサンプルサイズを予め決めた上で行います (実際にはサンプルサイズではなくテスト期間を決めておくやり方もあると思いますが)。つまりサンプルサイズを所与として、そのうちどれほどのサンプルでクリックが発生したかを調べていることになります。

例えば上の表では A パターンのサンプルサイズは 1 行目の和の 1000 であり、そのうちいくつのサンプルでクリックされたかに興味があります。B パターンのサンプルサイズも 2 行目の和となり同様です。

したがって通常の A/B テストは、各行和を固定した上での分割表の統計的性質を調べていることになります。スローガン的に書くと以下のようになります。

条件付き確率 $p(\text{分割表} \mid \text{各行和})$ はなんぼのもんじゃい

各行和ごとに固定しているため、各行 (各パターン) は独立です。また、各行は二項分布に従うものとして考えます。実際にはラプラス定理を使って正規近似を行って確率の計算するのが普通でしょう。

フィッシャーの正確確率検定

たまに「サンプルサイズが小さいときは正確な確率が計算できるフィッシャーの正確確率検定を使いましょう」という主張を見かけることがあります。フィッシャーの正確確率検定で考えているのは実は次です。

条件付き確率 $p(\text{分割表} \mid \text{各行和と各列和})$ はなんぼのもんじゃい

行和だけではなく列和も固定して考えるのがフィッシャーの正確確率検定です。サンプルサイズが小さいかどうか、確率が正確に計算できるかどうか、というよりはそもそも統計的仮定が異なるのです。

A/B テストで列和を所与としていいものなのか、ということも別途議論が必要なポイントではありますが、本記事では数学的な側面について議論していきます。

さて、行和と列和を固定すると、分割表は超幾何分布に従います。記号の導入とともに定式化しておきましょう。

記号

$I \times J$ の 2 元分割表を $\bm{x} = \{x_{ij}\}$ と書く。
$i$ 行の和を $x_{i+}$ と書き、 $j$ 列の和を $x_{+j}$ と書く。
すべてのセルの総和を $n$ とする。

用語

行和と列和を固定した分割表全体をファイバーと呼ぶ。

分割表の分布

ファイバー上の分割表は以下のような超幾何分布に従う。

\begin{aligned} p(\bm{x} \mid x_{1+}, \ldots, x_{I+}, x_{+1}, \ldots, x_{+J}) &= \frac{ \begin{pmatrix} n \\ x_{11}, \ldots, x_{IJ} \end{pmatrix} }{ \begin{pmatrix} n \\ x_{1+}, \ldots, x_{I+} \end{pmatrix} \begin{pmatrix} n \\ x_{+1}, \ldots, x_{+J} \end{pmatrix} } \\[3em] &= \frac{ \prod_{i=1}^I x_{i+}! \ \prod_{j=1}^J x_{+j}! }{ n! \ \prod_{i, j} x_{ij}! } \end{aligned}

具体的な分割表を例にどんな分布になるのかを見ておきましょう。

行和が 7 と 11、列和が 8 と 10 であるような分割表は全部で 8 つありますが、それらの出現確率はこちらの通りです。 $[3,4,5,6]$ が最頻値 (最頻表？) ですね。

\begin{aligned} &\begin{bmatrix} 0 & 7 \\ 8 & 3 \end{bmatrix} \qquad 0.004 \\ &\begin{bmatrix} 1 & 6 \\ 7 & 4 \end{bmatrix} \qquad 0.053 \\ &\begin{bmatrix} 2 & 5 \\ 6 & 5 \end{bmatrix} \qquad 0.222 \\ &\begin{bmatrix} 3 & 4 \\ 5 & 6 \end{bmatrix} \qquad 0.370 \\ &\begin{bmatrix} 4 & 3 \\ 4 & 7 \end{bmatrix} \qquad 0.264 \\ &\begin{bmatrix} 5 & 2 \\ 3 & 8 \end{bmatrix} \qquad 0.079 \\ &\begin{bmatrix} 6 & 1 \\ 2 & 9 \end{bmatrix} \qquad 0.009 \\ &\begin{bmatrix} 7 & 0 \\ 1 & 10 \end{bmatrix} \qquad 0.000 \\ \end{aligned}

超幾何分布を定常分布とするマルコフ連鎖

上の例では $2 \times 2$ の分割表で、行和・列和ともに小さかかったため、8 つの確率を計算すれば分布全体を理解することができました。しかし一般には取りうる分割表の数が爆発的に増えますし、たくさんの階乗を計算しなくてはならないため、個別に確率値を計算するのは現実的ではありません。

そこで、超幾何分布を定常分布とするファイバー上のマルコフ連鎖を生成することで解決しようというのが本記事の主旨です。

マルコフ基底

マルコフ連鎖を生成するためには、ある分割表から同一ファイバー内の別の分割表へと遷移するステップを構成する必要があります。同一ファイバー内で遷移するということは、行和も列和も変化しないということです。

以下のように行和と列和が $0$ であるような整数 (非負整数ではない) からなる表を足し引きすれば、同一ファイバー内で動き回ることができます。

\begin{bmatrix} 1 & 2 & 3 \\ 4 & 5 & 6 \\ 7 & 8 & 9 \\ \end{bmatrix} +\begin{bmatrix} 1 & -1 & 0 \\ -1 & 1 & 0 \\ 0 & 0 & 0 \\ \end{bmatrix} =\begin{bmatrix} 2 & 1 & 3 \\ 3 & 6 & 6 \\ 7 & 8 & 9 \\ \end{bmatrix}

行和と列和が $0$ であるような表の加減を繰り返し、ファイバー内の任意の分割表間を遷移することができれば (到達可能性)、マルコフ連鎖が構成できそうな感じがしますね。

ではマルコフ連鎖の素となるマルコフ基底を導入しましょう。

mod

$\mathcal{B}$ を、行和と列和が $0$ であるような表の集合とする。 $\mathcal{B}$ の要素を足し引きすることによって、ファイバー $\mathcal{F}$ 内の分割表 $x$ から $y$ へ移れるとき、 $x \sim y \ (\mathrm{mod} \ \mathcal{B})$ と書く。

\begin{aligned} x \sim y \ (\mathrm{mod} \ \mathcal{B}) \stackrel{\text{def}}{\iff} &\exists L, z_i \in \mathcal{B},\ \epsilon = \pm 1,\ i = 1,\ldots,L, \\ &\mathrm{s.t.}\quad y = x + \sum_{i = 1}^L \epsilon_i z_i,\ x + \sum_{i + 1}^l \epsilon_i z_i \in \mathcal{F},\ l = 1,\ldots,L \end{aligned}

マルコフ基底

$\mathcal{B}$ がマルコフ基底であるとは、以下を満たすことである。

\forall \mathcal{F},\ \forall x, y \in \mathcal{F},\ x \sim y \ (\mathrm{mod}\ \mathcal{B})

マルコフ基底の要素をいくつか足し引きすれば、ファイバー内を自由に動き回れるということです。しかも、ファイバーごとにマルコフ基底を定めているわけではなく、全てのファイバーに対してマルコフ基底がひとつ定義されていることにも注意です。

マルコフ基底の例

行和と列和が $0$ であるような表全体は自明なマルコフ基底ですね。任意の 2 つの分割表間を 1 ステップで遷移することができます。

また、分解可能モデル (たとえば 2 元分割表の独立モデルなど) と呼ばれるモデルにおいては、以下のような一箇所の $2 \times 2$ の部分だけがプラマイ $1$ になったような表全体もマルコフ基底になります。このような表を基本移動と呼びます。基本移動を足し引きするというのは、すぐ隣の分割表に遷移するようなイメージですね。

\begin{bmatrix} 0 & 0 & 0 & 0 & 0 \\ 0 & 1 & -1 & 0 & 0 \\ 0 & -1 & 1 & 0 & 0 \\ 0 & 0 & 0 & 0 & 0 \\ \end{bmatrix}

アルゴリズム

天下り的ではありますが、そのアルゴリズムを書き下しましょう。とは言っても、よく見るとメトロポリス・ヘイスティングス法を超幾何分布に具体化したものなので、とくに難しいことはしていないことが分かるでしょう。というわけでこのアルゴリズムで得られるサンプルは、漸近的に超幾何分布に従います。

設定

観測した分割表: $x^o$
マルコフ基底: $\mathcal{B}$
超幾何分布: $f$
ファイバー: $\mathcal{F}$

アルゴリズム

ステップ 1: $x = x^o$

ステップ 2: $z \in \mathcal{B}$ をランダムに選ぶ。符号 $\epsilon \in \{-1, 1\}$ を等確率 $1/2$ で選ぶ。

ステップ 3: もし $x + \epsilon z \notin \mathcal{F}$ ならば、 $x_{\mathrm{next}} = x$ と置きステップ 2 に進む。もし $x + \epsilon z \in \mathcal{F}$ ならば、 $u$ を区間 $[0, 1]$ の一様乱数とする。

ステップ 4: もし $u \leq \frac{f(x + \epsilon z)}{f(x)}$ ならば $x_{\mathrm{next}} = x + \epsilon z$ と置いてステップ 2 に進む。もし $u > \frac{f(x + \epsilon z)}{f(x)}$ ならば $x_{\mathrm{next}} = x$ と置いてステップ 2 に進む。

これを予め決めたステップ数だけ繰り返します。適宜 burn in を設けたり、サンプルを配列に格納したりは用途に応じて実装すれば良いでしょう。

上のアルゴリズムを見て、あれ、超幾何分布の計算コストが高いから MCMC を使ってるのにアルゴリズム内で超幾何分布使ってるじゃんと思われた方もいるでしょう。一般化されたアルゴリズムにおいては数式上超幾何分布が現れますが、基本移動だけを考えるともっと簡単になります。そもそも計算対象の確率比を考えることで、基準化定数を知らなくていい、みたいなところもメトロポリス・ヘイスティングス法の強みでしたよね。

さて、超幾何分布は以下のような数式で表されるのでした。

\begin{aligned} f(\bm{x}) = \frac{ \prod_{i=1}^I x_{i+}! \ \prod_{j=1}^J x_{+j}! }{ n! \ \prod_{i, j} x_{ij}! } \end{aligned}

$x$ に、例えば以下の基本移動を足した分割表を $y$ としましょう。

\begin{bmatrix} 1 & -1 & 0 & \cdots \\ -1 & 1 & 0 & \cdots \\ 0 & 0 & 0 & \cdots \\ \vdots & \vdots & \vdots & \ddots \\ \end{bmatrix}

$x$ と $y$ は行和、列和、総和が等しいことを踏まえると、2 つの超幾何分布の値の比はかなり簡単になることがわかります。

\begin{aligned} \frac{f(y)}{f(x)} &= \frac{ \prod_{i=1}^I y_{i+}! \ \prod_{j=1}^J y_{+j}! }{ n! \ \prod_{i, j} y_{ij}! } \Bigg/ \frac{ \prod_{i=1}^I x_{i+}! \ \prod_{j=1}^J x_{+j}! }{ n! \ \prod_{i, j} x_{ij}! } \\ &= \frac{\prod_{i, j} x_{ij}!}{\prod_{i, j} y_{ij}!} \\[1em] &= \frac{x_{11}! x_{12}! x_{21}! x_{22}!}{y_{11}! y_{12}! y_{21}! y_{22}!} \\[1em] &= \frac{x_{11}! x_{12}! x_{21}! x_{22}!}{(x_{11} + 1)! (x_{12} - 1)! (x_{21} - 1)! (x_{22} + 1)!} \\[1em] &= \frac{x_{12} x_{21}}{(x_{11} + 1)(x_{22} + 1)} \end{aligned}

実装

$2 \times 2$ の分割表に限定した実装は logicoffee/fisher-exact-test-mcmc に置いてあります。特筆すべきこともなく、上で説明したものをそのまま書き起こした感じですね。

README にも書いていますが、こんな感じでサンプリングの様子を観察することができます。

$ python3 main.py visualize 3 4 5 6 --sample 10

1
┏━━━┳━━━┓
┃  3  ┃  4  ┃
┣━━━╋━━━┫
┃  5  ┃  6  ┃
┗━━━┻━━━┛


2
┏━━━┳━━━┓
┃  2  ┃  5  ┃
┣━━━╋━━━┫
┃  6  ┃  5  ┃
┗━━━┻━━━┛


3
┏━━━┳━━━┓
┃  3  ┃  4  ┃
┣━━━╋━━━┫
┃  5  ┃  6  ┃
┗━━━┻━━━┛

...

まとめ

普段 Z 検定やカイ二乗検定くらいでしか触れない分割表ですが、結構奥が深いものですね。それでは本記事の要点をまとめておきましょう。

行和と列和を固定すると、分割表は超幾何分布に従う。
超幾何分布を定常分布とするマルコフ連鎖をマルコフ基底から構成できる。
マルコフ基底によっては各ステップの計算がかなりシンプルになる。