PyStanで混合ベルヌーイ分布を実装

目的

0と1を要素とする長さNのベクトルが作るn個のデータがあるとき、その集合をクラスタリングするモデルとして混合ベルヌーイ分布がある。ここではそのベイズ推測をPythonによるStanのインターフェースであるPyStanで行う方法をまとめる。

混合ベルヌーイ分布とは, $\theta$ と $\mu$ をパラメータとする $x_i \in \{0, 1\}^N$ の分布であり, 次式で定義されるものである.

$\begin{eqnarray} \Pr(x_i | \boldsymbol{\theta}, \boldsymbol{\mu}) = \sum_{r=0}^{k - 1} \theta_r \prod_{j=0}^{N - 1} \mu_{rj}^{x_{ij}} ( 1 - \mu_{rj})^{1 - x_{ij}} \end{eqnarray}$

今回は, PyStanで混合ベルヌーイ分布のパラメータのサンプリングを行ってみた.

サンプリングを行うにあたって, パラメータ $\theta$ , $\mu$ の事前分布を

$\begin{eqnarray} \theta &\sim& {\rm Dirichlet}(\alpha, N) \\ \mu &\sim& {\rm Beta}(\beta_1, \beta_2) \end{eqnarray}$

とし, 入力するデータを, 既知の混合ベルヌーイ分布から生成された120個のサンプル $\{x_i\}$
(各 $x_i$ は $N=80$ 次元のベクトル)とする.

(ただし, $\theta = \{0.25, 0.25, 0.25, 0.25 \}$ ,
各 $r = 0, 1, 2, 3$ , $j = 0, 1, ..., 79$ について $20 * r \leq j < 20 * (r + 1)$ のとき $\mu_{rj} = 0.9$ , それ以外は $\mu_{rj} = 0.1$ とする)

また, サンプリングされた $\theta$ と $\mu$ を用いて, 各サンプル $x_i$ が $r$ に所属する事後確率 $y_{ir}$ を計算し, どのクラスタに所属しているかを推定する.

$y_{ir}$ は次のように表される.(ここで, $g_i$ は $x_i$ の所属するクラスタを表す )

$\begin{eqnarray} y_{ir} = \Pr(g_i = r | \{ x_i \}, \boldsymbol{\theta}, \boldsymbol{\mu}) = \frac{\theta_r \prod_j \mu_{rj}^{x_{ij}} ( 1 - \mu_{rj})^{1 - x_{ij}}}{\sum_s \theta_s \prod_j \mu_{sj}^{x_{ij}} ( 1 - \mu_{sj})^{1 - x_{ij}}} \end{eqnarray}$

データの生成と入力, 結果の出力を行うPythonのコードは次のように実装する.

import pystan
import numpy as np
import matplotlib.pyplot as plt
import collections


def make_cluster(result, n):
    cluster = collections.defaultdict(list)

    for i in range(n):
        cluster[np.argmax(result[i])].append(i)

    return cluster


def draw_matrix(X, n, cluster, k, title):
    plt.figure(figsize=(7,7))
    divide_pos = -0.5
    min_lim = -0.5
    max_lim = n - 0.5
    X_new = [X[i] for r in range(k) for i in cluster[r]]

    for r in range(k):
        if r == k - 1:
            break
        divide_pos += len(cluster[r])
        plt.hlines([divide_pos], min_lim, max_lim, linestyles="dashed")

    plt.imshow(X_new, cmap="Greys", interpolation="none")
    plt.tick_params(labelbottom='off')
    plt.tick_params(labelleft='off')
    plt.savefig(title)
    plt.show()


if __name__ == '__main__':
    n = 120  # サンプル数
    N = 80  # サンプル一つあたりの次元
    k = 4  # クラスタ数

    # テストデータのthetaの値
    real_theta = [1. / k for i in range(k)]

    # テストデータのmuの値
    real_mu = [[0.9 if (N / k) * i <= j < (N / k) * (i + 1)
                else 0.1 for j in range(N)] for i in range(k)]

    # テストデータの生成
    X = np.zeros((n, N))

    for i in range(n):
        c = int(np.random.choice(k, 1, p=real_theta))
        for j in range(N):
            mu = real_mu[c][j]
            X[i][j] = np.random.choice(2, 1, p=[1.0 - mu, mu])

    X = X.astype(int)

    # Stanのコードに渡すデータ
    stan_data = {'n': n, 'N': N, 'k': k, 'X': X}

    # Stanのコードで定義されたモデルを生成
    sm = pystan.StanModel(file='mixture_bernoulli.stan')

    # サンプリング
    fit = sm.sampling(data=stan_data, chains=4)

    fit.plot()
    samples = fit.extract()

    # 集めたサンプルからクラスタ構造を推定
    cluster = make_cluster(samples['y'].mean(axis=0), n)

    # 結果の表示
    draw_matrix(X, n, cluster, k, 'y_sample')

Pythonのコードに読み込ませるStanのコードは次の通り.
事前分布のハイパーパラメータは次のように設定した.

$\begin{eqnarray} \alpha &=& \{ 1.0 \} \\ \beta_1 &=& 1.0 \\ \beta_2 &=& 3.0 \end{eqnarray}$

data {
    int<lower=0> n;
    int<lower=0> N;
    int<lower=0> k;
    int<lower=0, upper=1> X[n, N];
}

parameters {
    simplex[k] theta;
    real<lower=0, upper=1> mu[k, N];
}

transformed parameters {
    vector[k] ps[n];

    for(i in 1:n){
        for(r in 1:k){
            ps[i, r] = log(theta[r]);
            for(j in 1:N){
                ps[i, r] += bernoulli_lpmf(X[i, j] | mu[r, j]);
            }
        }
    }
}

model {
    theta ~ dirichlet(rep_vector(1.0, k));
    for(r in 1:k){
        for(j in 1:N){
            mu[r, j] ~ beta(1.0, 3.0);
        }
    }
    for(i in 1:n){
        target += log_sum_exp(ps[i]);
    }

}

generated quantities {
    simplex[k] y[n];

    for(i in 1:n){
        y[i] = softmax(ps[i]);
    }
}

推測されたパラメータは次のようになった.
f:id:gakushukun1:20190504151136p:plain
パラメータから推測されたクラスタ分割から, サンプル $\{x_i\}$ の各ベクトル(0=白, 1=黒)を縦に並べなおすと次のようになる. おおむねあらかじめ設定したサンプル生成元の分布と一致することがわかる.
f:id:gakushukun1:20190504155058p:plain