機械学習
Julia

トピックモデルをJuliaで実装した

都内の某IT企業にてデータサイエンティストとして働いている @yusaku-i です.
メイン言語を から に移行して半年ほど経ちます.
ローカル環境で粛々と開発してきた機械学習手法たちを整理してGitHubにパッケージとして公開しています.

今回はトピックモデルをスクラッチ実装したので紹介したいと思います.
体系的に学びたい方は「トピックモデル(機械学習プロフェッショナルシリーズ)」や「トピックモデルによる統計的潜在意味解析 (自然言語処理シリーズ)」を参考にしてください.

トピックモデルとは

各文書には潜在トピックがあると仮定されます(例えば,「スポーツ」に関する記事や「経済」に関する記事といった感じ).
トピックモデル(latent topic model)とは,これらの文書が生成される過程を確率的にモデリングすることで潜在トピックを抽出するための手法です.
具体的には,トピックモデルを文書データに適用することで「文書 $D$ はスポーツに関する文書」であったり「単語 $W$ はスポーツに関する単語」といった情報を知ることができます.
またトピックモデルは文書データのみならず,画像や購買履歴,ソーシャルネットなどの離散データにも適用可能な汎用モデルとなっています.

LDAとは

LDA(latent Dirichlet allocation; 潜在的ディリクレ配分法)はトピックモデルの代表的な手法です.

lda.png

LDAでは各文書が潜在トピック $z$ の混合割合を表すトピック分布 ${\boldsymbol \Theta}
= \{ {\boldsymbol \theta} _d \}$ をもち,潜在トピックは単語分布 ${\boldsymbol \Phi} = \{ {\boldsymbol \phi}_k \} _{k=1}^{K}$ をもちます.
ここで,$K$ はトピック数を意味し,LDAの学習では事前に決める必要があります.

LDAの生成過程は以下のようになります.

\begin{align}
\{ {\boldsymbol \phi}_k \}_{k=1}^{K} &\sim {\rm Dirichlet}({\boldsymbol \beta}) \\
\{ {\boldsymbol \theta}_d \}_{d=1}^{D} &\sim {\rm Dirichlet}({\boldsymbol \alpha}) \\
\{ \{ z_{dn} \}_{n=1}^{N_d} \}_{d=1}^{D} &\sim {\rm Categorical}({\boldsymbol \theta}_d) \\
\{ \{ w_{dn} \}_{n=1}^{N_d} \}_{d=1}^{D} &\sim {\rm Categorical}({\boldsymbol \phi}_{z_{dn}})
\end{align}

トピック分布 ${\boldsymbol \Theta}$ と単語分布 ${\boldsymbol \Phi}$ はそれぞれ共役事前分布であるディリクレ分布を事前分布とします.
ちなみにディリクレ分布を仮定しない手法はPLSA(probabilistic latent semantic analysis)と呼ばれます(PLSAもトピックモデルのひとつです).

LDAの同時確率は次のようになります.
$$p({\bf W}, {\bf Z}, {\boldsymbol \Theta}, {\bf \Phi}\ \vert\ {\boldsymbol \alpha}, {\boldsymbol \beta}) = p({\bf Z}\ \vert\ {\boldsymbol \Theta}) \cdot p({\boldsymbol \Theta}\ \vert\ {\boldsymbol \alpha}) \cdot p({\bf W}\ \vert\ {\bf Z}, {\bf \Phi}) \cdot p({\bf \Phi}\ \vert\ {\boldsymbol \beta})$$

第1項は
$$p({\bf Z}\ \vert\ {\boldsymbol \Theta}) = \prod_{d=1}^{D} \prod_{k=1}^{K} \prod_{n=1}^{N_d} \theta_{dk}^{\delta_{(z_{dn}=k)}} = \prod_{d=1}^{D} \prod_{k=1}^{K} \theta_{dk}^{N_{dk}}$$

第2項は
$$p({\boldsymbol \Theta}\ \vert\ {\boldsymbol \alpha}) = \prod_{d=1}^{D} \frac{\Gamma(\sum_{k=1}^{K} \alpha_k)}{\prod_{k=1}^{K} \Gamma(\alpha_k)} \prod_{k=1}^{K} \theta_{dk}^{\alpha_k - 1}$$

第3項は
$$p({\bf W}\ \vert\ {\bf Z}, {\bf \Phi}) = \prod_{d=1}^{D} \prod_{k=1}^{K} \prod_{n=1}^{N_d} \left[ \phi_{kw_{dn}}^{N_{dw_{dn}}} \right]^{\delta_{(z_{dn}=k)}} = \prod_{k=1}^{K} \prod_{v=1}^{V} \phi_{kv}^{N_{kv}}$$

第4項は
$$p({\bf \Phi}\ \vert\ {\boldsymbol \beta}) = \prod_{k=1}^{K} \frac{\Gamma(\sum_{v=1}^{V} \beta_v)}{\prod_{v=1}^{V} \Gamma(\beta_v)} \prod_{v=1}^{V} \phi_{kv}^{\beta_v - 1}$$

となります.
なお,$N_{dk} = \sum_n \delta_{(z_{dn}=k)}$,$N_{kv} = \sum_d \sum_n N_{dw_{dn}} \delta_{(z_{dn}=k, w_{dn}=v)}$ と計算されます.

モデル推論

潜在トピック ${\bf Z}$ の推論方法として collapsed Gibbs sampling(CGS)を紹介します.
collapsed とはモデルパラメータ $\{{\boldsymbol \Theta}, {\boldsymbol \Phi}\}$ を積分消去(周辺化)することを意味します.

周辺同時確率は
$$p({\bf W}, {\bf Z}\ \vert\ {\boldsymbol \alpha}, {\boldsymbol \beta}) = p({\bf Z}\ \vert\ {\boldsymbol \alpha}) \cdot p({\bf W}\ \vert\ {\bf Z}, {\boldsymbol \beta})$$
であり,第1項は
$$p({\bf Z}\ \vert\ {\boldsymbol \alpha}) = \int p({\bf Z}\ \vert\ {\boldsymbol \Theta}) \cdot p({\boldsymbol \Theta}\ \vert\ {\boldsymbol \alpha}) d{\boldsymbol \Theta} = \prod_{d=1}^{D} \frac{\Gamma(\sum_{k=1}^{K} \alpha_k)}{\prod_{k=1}^{K} \Gamma(\alpha_k)} \frac{\prod_{k=1}^{K} \Gamma(N_{dk} + \alpha_k)}{\Gamma(N_d + \sum_{k=1}^{K} \alpha_k)}$$
第2項は
$$p({\bf W}\ \vert\ {\bf Z}, {\boldsymbol \beta}) = \int p({\bf W}\ \vert\ {\bf Z}, {\bf \Phi}) \cdot p({\bf \Phi}\ \vert\ {\boldsymbol \beta}) d{\bf \Phi} = \prod_{k=1}^{K} \frac{\Gamma(\sum_{v=1}^{V} \beta_v)}{\prod_{v=1}^{V} \Gamma(\beta_v)} \frac{\prod_{v=1}^{V} \Gamma(N_{kv} + \beta_v)}{\Gamma(N_k + \sum_{v=1}^{V} \beta_v)}$$
となります.

CGSではトピック $\{z_{dn}\}$ を繰り返しサンプリングすることで事後分布を推論します.

$z_{dn}$ のサンプリング式は次のようになります.
$$p(z_{dn}=k\ \vert\ {\bf W}, {\bf Z}_{\backslash dn}, {\boldsymbol \alpha}, {\boldsymbol \beta}) \propto p(z _{dn}=k\ \vert\ {\bf Z} _{\backslash dn}, {\boldsymbol \alpha}) \cdot p(w _{dn}\ \vert\ {\bf W} _{\backslash dn}, z _{dn}=k, {\bf Z} _{\backslash dn}, {\boldsymbol \beta})$$

第1項は
$$p(z_ {dn}=k\ \vert\ {\bf Z} _{\backslash dn}, {\boldsymbol \alpha}) = \frac{N _{dk\backslash dn} + \alpha _k}{N _d - 1 + \sum _{k=1}^{K} \alpha _k}$$

第2項は
$$p(w _{dn} \vert {\bf W} _{\backslash dn}, z _{dn}=k, {\bf Z} _{\backslash dn}, {\boldsymbol \beta}) = \frac{N _{kw _{dn}\backslash dn} + \beta _{w _{dn}}}{N _{k\backslash dn} + \sum _{v=1}^V \beta _v}$$

なお,$N_{dk\backslash dn} = \sum_{n'\neq n} \delta_{(z_{dn'}=k)}$,$N_{kv\backslash dn} = \sum_d \sum_{n'\neq n} N_{dw_{dn'}} \delta_{(z_{dn'}=k, w_{dn'}=v)}$ と計算されます.

モデル実装

LDAのスクラッチ実装は以下になります.

using StatsBase

const α = 0.5
const β = 0.1

# estimate model parameters with collapsed Gibbs sampling
function cgs(X_::Array, max_iter::Int, n_words::Int, n_components::Int)
    z_dn = [zeros(Int, N_d) for (N_d, X) in X_]
    N_dk = zeros(Int, length(X_), n_components)
    N_kv = zeros(Int, n_components, n_words)
    N_k  = zeros(Int, n_components)

    srand(0)
    @inbounds for iter in 1:max_iter
        @inbounds for (d, (N_d, X)) in enumerate(X_)
            n = 0
            @inbounds for (v, N_dv) in X
                @inbounds for _ in 1:N_dv
                    n += 1

                    # remove w_dv’s statistics
                    if z_dn[d][n] != 0
                        k = z_dn[d][n]
                        N_dk[d, k] -= 1
                        N_kv[k, v] -= 1
                        N_k[k]     -= 1
                    end

                    log_p_k = zeros(n_components)
                    @inbounds for k in 1:n_components
                        log_p_k[k]  = log(N_dk[d, k] + α)
                        log_p_k[k] += log(N_kv[k, v] + β)
                        log_p_k[k] -= log(N_k[k] + β * n_words)
                    end

                    # sample z_d after normalizing
                    p_k = exp.(log_p_k - maximum(log_p_k))
                    z_dn[d][n] = sample(1:n_components, WeightVec(p_k / sum(p_k)))

                    # add x_d’s statistics
                    k = z_dn[d][n]
                    N_dk[d, k] += 1
                    N_kv[k, v] += 1
                    N_k[k]     += 1
                end
            end
        end
    end
end

HDPとは

LDAではトピック数 $K$ を事前に決める必要があります.
これに対してノンパラベイズ拡張を行うことで適切なトピック数を推定することができるようになります.

しかし,LDAでは文書ごとにトピック分布 ${\boldsymbol \theta}$ をもつと仮定しているため,文書固有のトピック分布に無限次元の離散分布を用いても異なる文書間でトピックが共有されません.
そこで,LDAをノンパラベイズ拡張するための手法としてHDP(hierarchical Dirichlet process; 階層ディリクレ過程)が提案されています.

HDPでは文書ごとの固有分布 ${\boldsymbol \Theta}
= \{ {\boldsymbol \theta} _d \}$ と文書共有分布 ${\boldsymbol \psi}$ が導入されます.
各文書の単語は文書固有の潜在クラスタ $\{t _{dn}\}$ が割り当てられ,固有分布 ${\boldsymbol \theta}$ はそれら潜在クラスタの混合割合として表されます.
そして各クラスタには潜在トピック $\{z _{d\ell}\}$ が割り当てられ,共有分布 ${\boldsymbol \psi}$ は潜在トピックの混合割合として表されます.

HDP-LDAの生成過程は以下のようになります.

\begin{align}
{\boldsymbol \psi} &\sim {\rm SBP}(\gamma) \\
\{ {\boldsymbol \phi}_k \}_{k=1}^{\infty} &\sim {\rm Dirichlet}({\boldsymbol \beta}) \\
\{ {\boldsymbol \theta}_d \}_{d=1}^{D} &\sim {\rm SBP}(\alpha) \\
\{ \{ t_{dn} \}_{n=1}^{N_d} \}_{d=1}^{D} &\sim {\rm Categorical}({\boldsymbol \theta}_d) \\
\{ \{ z_{d\ell} \}_{\ell=1}^{\infty} \}_{d=1}^{D} &\sim {\rm Categorical}({\boldsymbol \psi}) \\
\{ \{ w_{dn} \}_{n=1}^{N_d} \}_{d=1}^{D} &\sim {\rm Categorical}({\boldsymbol \phi}_{z_{dt_{dn}}})
\end{align}

SBP(stick-breaking process; 棒折り過程)は無限次元の離散分布から生成されることを意味します.

モデル推論

LDAと同様にHDP-LDAでの collapsed Gibbs sampling(CGS)を紹介します.

まず,各単語のクラスタ $t_{dn}$ を以下でサンプリングします.

p(t _{dn}=\ell\ \vert\ {\bf W}, {\bf T} _{\backslash dn}, {\bf Z}, \alpha, \gamma, {\boldsymbol \beta}) = \begin{cases}
    \frac{N _{d\ell\backslash dn}}{N _d - 1 + \alpha} \cdot \frac{N _{z _{d\ell}w _{dn}\backslash dn} + \beta _{w _{dn}}}{N _{z _{d\ell}\backslash dn} + \sum _{v=1}^{V} \beta _v} & \text{if $\ell$ is an existing component} \\
    \frac{\alpha}{N _d - 1 + \alpha} \cdot \frac{M _{k\backslash dn}}{M _{\backslash dn} + \gamma} \cdot \frac{N _{kw _{dn}\backslash dn} + \beta _{w _{dn}}}{N _{k\backslash dn} + \sum _{v=1}^{V} \beta _v} & \text{if $\ell$ is a new component and $k$ is an existing component} \\
    \frac{\alpha}{N _d - 1 + \alpha} \cdot \frac{\gamma}{M _{\backslash dn} + \gamma} \cdot \frac{\beta _{w _{dn}}}{\sum _{v=1}^{V} \beta _v} & \text{if $\ell$ and $k$ are new components}    
  \end{cases}

なお,$N_{d\ell\backslash dn} = \sum_{n'\neq n} \delta_{(t_{dn'}=\ell)}$,$M_{k\backslash dn} = \sum_d \sum_{\ell\neq t_{dn}} \delta_{(z_{d\ell}=k)}$ と計算されます.

次に各クラスタのトピック $z_{d\ell}$ を以下でサンプリングします.

p(z_{d\ell}=k\ \vert\ {\bf W}, {\bf T}, {\bf Z}_{\backslash d\ell}, \gamma, {\boldsymbol \beta}) = \begin{cases}
    \frac{M_{k\backslash d\ell}}{M - 1 + \gamma} \cdot \frac{\Gamma(N_{k\backslash d\ell} + \sum_{v=1}^{V} \beta_v)}{\Gamma(N_{k\backslash d\ell} + N_{d\ell} + \sum_{v=1}^{V} \beta_v)} \prod_{v=1}^{V} \frac{\Gamma(N_{kv\backslash d\ell} + N_{d\ell v} + \beta_v)}{\Gamma(N_{kv\backslash d\ell} + \beta_v)} & \text{if $k$ is an existing component} \\
    \frac{\gamma}{M - 1 + \gamma} \cdot \frac{\Gamma(\sum_{v=1}^{V} \beta_v)}{\Gamma(N_{d\ell} + \sum_{v=1}^{V} \beta_v)} \prod_{v=1}^{V} \frac{\Gamma(N_{d\ell v} + \beta_v)}{\Gamma(\beta_v)} & \text{if $k$ is a new component}
  \end{cases}

モデル実装

HDP-LDAのスクラッチ実装は以下になります.

const γ = 0.5

function cgs(X_::Array, max_iter::Int, n_words::Int)
    t_dn  = [zeros(Int, N_d) for (N_d, X) in X_]
    z_dℓ  = [zeros(Int, 0) for _ in X_]
    N_dℓv = [zeros(Int, 0, n_words) for _ in X_]
    N_dℓ  = [zeros(Int, 0) for _ in X_]
    N_kv  = zeros(Int, 0, n_words)
    N_k   = zeros(Int, 0)
    M_k   = zeros(Int, 0)
    M     = 0
    n_components = 0
    n_tables     = zeros(Int, length(X_))

    srand(0)
    @inbounds for iter in 1:max_iter
        @inbounds for (d, (N_d, X)) in enumerate(X_)

            # sampling t_dn
            n = 0
            @inbounds for (v, N_dv) in X
                @inbounds for _ in 1:N_dv
                    n += 1

                    # remove w_dv’s statistics
                    if t_dn[d][n] != 0
                         = t_dn[d][n]
                        k = z_dℓ[d][]
                        N_dℓv[d][, v] -= 1
                        N_dℓ[d][]     -= 1
                        N_kv[k, v]     -= 1
                        N_k[k]         -= 1

                        # if any component is empty, remove it and decrease n_tables
                        if N_dℓ[d][] == 0
                            n_tables[d] -= 1
                            M_k[k]      -= 1
                            M           -= 1
                            deleteat!(z_dℓ[d], )
                            deleteat!(N_dℓ[d], )
                            N_dℓv[d] = N_dℓv[d][setdiff(1:end, ), :]
                            @inbounds for (tmp_n, tmp_t) in enumerate(t_dn[d])
                                if tmp_t > 
                                    t_dn[d][tmp_n] -= 1
                                end
                            end

                            # if any component is empty, remove it and decrease n_components
                            if M_k[k] == 0
                                n_components -= 1
                                deleteat!(M_k, k)
                                deleteat!(N_k, k)
                                N_kv = N_kv[setdiff(1:end, k), :]

                                @inbounds for tmp_d in 1:length(X_)
                                    @inbounds for (tmp_ℓ, tmp_z) in enumerate(z_dℓ[tmp_d])
                                        if tmp_z > k
                                            z_dℓ[tmp_d][tmp_ℓ] -= 1
                                        end
                                    end
                                end
                            end
                        end
                    end

                    L       = n_tables[d] + n_components + 1
                    log_p_ℓ = zeros(L)

                    # existing components
                    @inbounds for  in 1:n_tables[d]
                        log_p_ℓ[]  = log(N_dℓ[d][])
                        log_p_ℓ[] += log(N_kv[z_dℓ[d][], v] + β)
                        log_p_ℓ[] -= log(N_k[z_dℓ[d][]] + β * n_words)
                    end

                    # new components
                    @inbounds for k in 1:n_components
                         = n_tables[d] + k
                        log_p_ℓ[]  = log(α)
                        log_p_ℓ[] += log(M_k[k])
                        log_p_ℓ[] -= log(M + γ)
                        log_p_ℓ[] += log(N_kv[k, v] + β)
                        log_p_ℓ[] -= log(N_k[k] + β * n_words)
                    end

                    log_p_ℓ[L]  = log(α)
                    log_p_ℓ[L] += log(γ)
                    log_p_ℓ[L] -= log(M + γ)
                    log_p_ℓ[L] += log(β)
                    log_p_ℓ[L] -= log(β * n_words)

                    # sample t_dn after normalizing
                    p_ℓ = exp.(log_p_ℓ - maximum(log_p_ℓ))
                     = sample(1:L, WeightVec(p_ℓ / sum(p_ℓ)))

                    if  <= n_tables[d]
                        t_dn[d][n] = 
                    else
                        t_dn[d][n] = n_tables[d] + 1
                        push!(z_dℓ[d], 0)
                        push!(N_dℓ[d], 0)
                        N_dℓv[d] = cat(1, N_dℓv[d], zeros(Int, 1, n_words))

                        if  - n_tables[d] <= n_components
                            z_dℓ[d][t_dn[d][n]] =  - n_tables[d]
                        else
                            z_dℓ[d][t_dn[d][n]] = n_components + 1
                            n_components += 1
                            push!(M_k, 0)
                            push!(N_k, 0)
                            N_kv = cat(1, N_kv, zeros(Int, 1, n_words))
                        end

                        n_tables[d] += 1
                        M_k[z_dℓ[d][t_dn[d][n]]] += 1
                        M                        += 1
                    end

                    # add w_dv’s statistics
                     = t_dn[d][n]
                    k = z_dℓ[d][]
                    N_dℓv[d][, v] += 1
                    N_dℓ[d][]     += 1
                    N_kv[k, v]     += 1
                    N_k[k]         += 1
                end
            end

            # sampling z_dℓ
            @inbounds for  in 1:n_tables[d]

                # remove w_dℓ’s statistics
                if z_dℓ[d][] != 0
                    k = z_dℓ[d][]
                    M_k[k] -= 1
                    @inbounds for (v, N_dv) in X
                        N_kv[k, v] -= N_dℓv[d][, v]
                        N_k[k]     -= N_dℓv[d][, v]
                    end

                    # if any component is empty, remove it and decrease n_components
                    if M_k[k] == 0
                        n_components -= 1
                        deleteat!(M_k, k)
                        deleteat!(N_k, k)
                        N_kv = N_kv[setdiff(1:end, k), :]

                        @inbounds for tmp_d in 1:length(X_)
                            @inbounds for (tmp_ℓ, tmp_z) in enumerate(z_dℓ[tmp_d])
                                if tmp_z > k
                                    z_dℓ[tmp_d][tmp_ℓ] -= 1
                                end
                            end
                        end
                    end
                end

                K = n_components + 1
                log_p_k = zeros(K)

                # existing components
                @inbounds for k in 1:n_components
                    log_p_k[k]  = log(M_k[k])
                    log_p_k[k] += lgamma(N_k[k] + β * n_words)
                    log_p_k[k] -= lgamma(N_k[k] + N_dℓ[d][] + β * n_words)
                    @inbounds for (v, N_dv) in X
                        log_p_k[k] += lgamma(N_kv[k, v] + N_dℓv[d][, v] + β)
                        log_p_k[k] -= lgamma(N_kv[k, v] + β)
                    end
                end

                # a new component
                log_p_k[K]  = log(γ)
                log_p_k[K] += lgamma(β * n_words)
                log_p_k[K] -= lgamma(N_dℓ[d][] + β * n_words)
                @inbounds for (v, N_dv) in X
                    log_p_k[K] += lgamma(N_dℓv[d][, v] + β)
                    log_p_k[K] -= lgamma(β)
                end

                # sample z_dℓ after normalizing
                p_k = exp.(log_p_k - maximum(log_p_k))
                z_dℓ[d][] = sample(1:K, WeightVec(p_k / sum(p_k)))

                if z_dℓ[d][] == K
                    n_components += 1
                    push!(M_k, 0)
                    push!(N_k, 0)
                    N_kv = cat(1, N_kv, zeros(Int, 1, n_words))
               end

                # add w_dℓ’s statistics
                k = z_dℓ[d][]
                M_k[k] += 1
                @inbounds for (v, N_dv) in X
                    N_kv[k, v] += N_dℓv[d][, v]
                    N_k[k]     += N_dℓv[d][, v]
                end
            end
        end
    end

    N_dk = zeros(Int, length(X_), n_components)
    @inbounds for d in 1:length(X_)
        @inbounds for (, k) in enumerate(z_dℓ[d])
            N_dk[d, k] += N_dℓ[d][]
        end
    end
end

まとめ

トピックモデルについて紹介しました.
スクラッチ実装したLDAとHDP-LDAは https://github.com/yusaku-i/Bayesian_mixtures.jl/tree/master/src/LDA で公開しています.
簡易データセットでの動作検証はしてますが,まともな精度検証はできていないので,ぜひ動作検証ついでに触ってみてください.

最後にJulia×データ分析に関する執筆依頼お待ちしております.