隠れマルコフ連鎖モデルの変分推論(Rで実装)

今読んでいるベイズ推論のテキストベイズ推論による機械学習入門の内容より、隠れマルコフモデルの変分推論についてRで実装したので記事化してみます。

1.隠れマルコフモデルにおける変分推論

モデルの確認

$K$ 個の状態 ${1, 2, ..., K}$ の間を移りあうマルコフ連鎖を考えます。この状態が潜在変数で表されているモデルを隠れマルコフモデルといいます。

いくつか記号を用意します。状態を表す潜在変数 $s_{n, k}$ を時点 $n$ における状態が $k$ の場合1, それ以外の場合は0であると定めます。また、 $\underline{s_n} = (s _ {n, 1}, ..., s _ {n, K})^ {\sf T}, \quad {\sf S} = (\underline{s_1},..., \underline{s_N})$ とまとめて書きます。

2つのハイパーパラメーター

$\pi = (\pi_1, ..., \pi_K) , \quad (\sum _ {i = 1} ^ K \pi_k = 1)$ : $\underline{s_1}$ の分布

${\sf A} = (A _ {i, j}) _ {1 \leq I, j \leq K}, \quad (\sum _ {l = 1} ^ K A _ {l, k}) = 1$ : 遷移行列

を与えると、

$( \underline{s_n}) _ {1 \leq n \leq N}$ はマルコフ連鎖に従うので、同時確率 $p({\sf S} | \pi, {\sf A})$ が次の式で与えられます。

$p({\sf S} | \pi, {\sf A}) = p( \underline{s _ 1} | \pi ) \prod _ {n = 2} ^ N p( \underline{s _ n} | \underline{s _ {n - 1}}, {\sf A})$

また、 $\underline{s_n}$ は直接観測されない代わりに以下の式で与えられる確率変数 $x_n$ が観測データとして与えられるとします。

$p(x_n | \underline{s_n}, \underline{\theta}) = \prod _ {k = 1} ^ K {\rm Poisson}(x_n | \theta_k) ^ {s _ {n, k}}$

ただし、 $\underline{\theta} = (\theta _ 1, ..., \theta _K)$ はハイパーパラメーターとします。

以上の設定のもと、観測データ ${\sf X} = (x_1, ..., x_N)$ に対する同時分布が次の式で与えられることになります。

$p({\sf X}, {\sf S}, \underline{\theta} , \pi, {\sf A})$

$\quad = p({\sf X} | {\sf S}, \underline{\theta} ) p({\sf S} | \pi, {\sf A}) p(\underline{\theta})p(\pi)p({\sf A})$

$\quad = p(\underline{\theta})p(\pi)p({\sf A}) p(x_1 | \underline{s _ 1}, \theta _ 1 ) p( \underline{s _ 1} | \pi _ 1) \prod _ {n = 2} ^ N p(x_n |\underline{s _ n} , \underline{\theta}) p( \underline{s _ n} | \underline{s _ {n - 1}}, {\sf A})$

完全分解変分推論

隠れマルコフモデルの事後分布 $p({\sf S}, \underline{\theta} , \pi, {\sf A} | {\sf X})$ を近似的に求める方法として、変分推論を適用することを考えます。

完全分解型変分推論( ${\sf completely \ factorized \ variational \ inference}$ )では事後分布を次のような分布の積によって近似を行います。

$p({\sf S}, \underline{\theta} , \pi, {\sf A} | {\sf X}) \approx q( \underline{\theta}, {\sf A}, \pi ) \prod _ {n = 1} ^ N q( \underline{s _ n})$

テキストでは次の事前分布に対する仮定のもと、完全分解型変分推論に対する更新アルゴリズムの導出が載っていますので気になった方はご参照ください。

$\pi \sim {\rm Dirichlet} (\underline{\alpha}), \quad \underline{\alpha} = (\alpha _ 1, ..., \alpha _ K)$

${\sf A} \sim \prod _ {k = 1} {\rm Dirichlet} (\underline{ \beta _ k }), \quad \underline{ \beta _ k } = (\beta _ {1, k}, ..., \beta _ {K, k})$

$\theta _ k \sim {\rm Gamma} (a, b)$

2.Rによる実装


library(tidyverse)
library(glue)
#'@param N
#'@param K 
#'@param X array([N])
#'@param prior list(alpha[K], beta[K, K], theta list(a[K], b[K]))
#'@param init list(alpha[K], beta[K, K], theta list(a[K], b[K]))
#'@param iterMax
CFVI <- function(N, K, X, prior, init, iterMax) {

    softmax <- function(s) exp(s) / sum(exp(s))
    entropy.Gamma <- function(a, b) - digamma(a) + log(b)
    entropy.Dirichlet <- function(alpha) - digamma(alpha) + digamma(sum(alpha))
    
    hyp.params = list(
        S = array(rep(0, K * N), dim = c(K, N)),
        pi = init$alpha,
        beta = init$beta,
        theta = list(a = rep(init$theta$a, K), b = rep(init$theta$b, K)))
    
    for (iter in 1:iterMax) {
        #update of hyperparams of S
        hyp.params$S <- c(1:N) %>% 
            map_dfc(function(n) {
                entropy.beta <- hyp.params$beta %>% 
                    as.data.frame %>%
                    map_dfc(~ {entropy.Dirichlet(.)}) %>% 
                    as.matrix
                entropy.theta <- entropy.Gamma(hyp.params$theta$a, hyp.params$theta$b)
                exp.theta <- hyp.params$theta$a / hyp.params$theta$b
                
                if (n == 1) {
                    return(c(- X[1] * entropy.theta - exp.theta - 
                               entropy.Dirichlet(hyp.params$pi) -
                               hyp.params$S[,2] %*% entropy.beta))
                } else if (n == N) {
                    return(c(- X[N] * entropy.theta - exp.theta - 
                                hyp.params$S[,N - 1] %*% entropy.beta))
                } else {
                    return(c(- X[n] * entropy.theta - exp.theta -
                               hyp.params$S[,n - 1] %*% t(entropy.beta) - 
                               hyp.params$S[,n + 1] %*% entropy.beta))
                }
            } %>% softmax) %>% as.matrix
        
        #update of hyperparams of alpha
        hyp.params$pi <- hyp.params$S[,1] + prior$alpha
        
        #update of hyperparams of beta
        hyp.params$beta <- hyp.params$S[,2:N] %*% t(hyp.params$S[,1:(N-1)]) + prior$beta
        
        #update of hyperparams of theta
        hyp.params$theta <- list(
            a = c(hyp.params$S %*% X) + prior$theta$a,
            b = rowSums(hyp.params$S) + prior$theta$b)
    }
    return(hyp.params)
}

#test###########################################################################
test <- function(){
    #'generate sample from hidden markov model with 2 - states = {1, 2}
    test.N <- 100L
    test.K <- 2L
    #'transition matrix
    test.A <- array(c(0.2, 0.8, 0.4, 0.6), dim = c(2, 2))
    
    nextState <- function(current, A) {
        sample(c(1, 2), size = 1, replace = TRUE, prob = A[,current])
    }
    test.state <- purrr::reduce(1:test.N, ~ {append(.x, nextState(tail(.x, 1), test.A))})
    
    #test data
    test.theta <- c(2, 10)
    test.X <- purrr::map_dbl(test.state, ~ {rpois(1, test.theta[.x])})
    
    #目視チェック
    plot <- ggplot(data.frame(x = c(1:test.N), val = test.X), aes(x = x, y = val)) + 
        geom_point() + 
        labs(title='plot of test data')
    print(plot)
    
    #事前分布のパラメーターを与える
    test.prior <- list(alpha = c(1, 1), 
                       beta = array(rep(0.5, 4), dim = c(2, 2)), 
                       theta = list(a = 1, b = 1))
    test.init <- list(alpha = c(0.1, 1), 
                      beta = array(rep(0.5, 4), dim = c(2, 2)), 
                      theta = list(a = 1, b = 1))
    
    estim <- CFVI(test.N, test.K, test.X, test.prior, test.init, 20L)
    
    #'状態を推定し可視化
    estim.state <- ifelse(estim$S[1,] < 0.5, 1, 2)
    estim.plot <- ggplot(data.frame(x = c(1:test.N), val = test.X, state = as.factor(estim.state)), 
                         aes(x = x, y = val, colour = state)) + 
        geom_point() + 
        labs(title='estimation of states')
    print(estim.plot)
    
    #'正解率([推定された状態 = 真の状態]であるデータの個数 / 全データ数)
    print(glue('正解率 : {sum(ifelse(estim.state == test.state, 1, 0)) / test.N}'))
}

test()

推定結果による判別の正解率( ${\sf accuracy \ rate}$ )を