RでEMアルゴリズムによる混合ベルヌーイモデル最尤推定

PRML 9.3.3に記載の通り、EMアルゴリズムによって混合ベルヌーイモデルの最尤推定が行われる過程と、対数尤度関数の収束の様子を示します。

10次元2値サンプルを、3種のベルヌーイ分布に従って、150サンプルずつ順に生成し、それらのサンプルに対する対数尤度を最大化します。EMステップ毎に、各サンプルについて、その混合要素のうち、最大の負担率γ(z_nk)の値を、その混合要素を示す色で記します。また、各混合要素について、ベルヌーイ分布のパラメータμkの値を示します。

また、MNISTの手書き数字文字データから生成した784(=28x28)次元2値データについて、数字2,3,4を各150サンプルずつ利用し、それらのサンプルに対する対数尤度を最大化します。

frame()
set.seed(0)
par(mfrow=c(8, 4))
par(mar=c(2.3, 2.5, 1, 0.1))
par(mgp=c(1.3, .5, 0))
MNIST <- F  # to load the MNIST handwritten digits
K <- 3
N <- 450

rbern <- function(mu) {
    ifelse(runif(length(mu)) < mu, 1, 0)
}
logdbern <- function(x, mu) {
    sum(log(mu ^ x * (1 - mu) ^ (1 - x)))
}
logsumexp <- function (x) {
  m <- max(x)
  m + log(sum(exp(x - m)))
}

if (MNIST) {
    # the following code is taken from https://gist.github.com/brendano/39760
    load_mnist <- function() {
      load_image_file <- function(filename) {
        ret = list()
        f = file(filename,'rb')
        readBin(f,'integer',n=1,size=4,endian='big')
        ret$n = readBin(f,'integer',n=1,size=4,endian='big')
        nrow = readBin(f,'integer',n=1,size=4,endian='big')
        ncol = readBin(f,'integer',n=1,size=4,endian='big')
        x = readBin(f,'integer',n=ret$n*nrow*ncol,size=1,signed=F)
        ret$x = matrix(x, ncol=nrow*ncol, byrow=T)
        close(f)
        ret
      }
      load_label_file <- function(filename) {
        f = file(filename,'rb')
        readBin(f,'integer',n=1,size=4,endian='big')
        n = readBin(f,'integer',n=1,size=4,endian='big')
        y = readBin(f,'integer',n=n,size=1,signed=F)
        close(f)
        y
      }
      train <<- load_image_file('mnist/train-images-idx3-ubyte')
      train$y <<- load_label_file('mnist/train-labels-idx1-ubyte')
    }
    # the data is available at http://yann.lecun.com/exdb/mnist/
    # extract the files in the "mnist" directory
    setwd("C:/Users/Public/Documents")
    load_mnist()
    x <- floor(rbind(
        train$x[train$y == 2, ][1:(N / 3), ],  
        train$x[train$y == 3, ][1:(N / 3), ],  
        train$x[train$y == 4, ][1:(N / 3), ]) / 128)
    D <- ncol(x)
} else {
    PHI <- 0.9
    PLO <- 0.1
    muorg <- matrix(c(
        PHI, PHI, PHI, PHI, PHI, PLO, PLO, PLO, PLO, PLO, 
        PLO, PLO, PLO, PLO, PLO, PHI, PHI, PHI, PHI, PHI, 
        PHI, PHI, PHI, PLO, PLO, PLO, PLO, PHI, PHI, PHI
        ), 3, byrow=T)
    z <- rep(1:3, each=N/3)
    x <- rbern(muorg[z, ])
    D <- 10
}
image(1:N, 1:D, x, xlab="n", ylab="i", breaks=seq(0, 1, 0.1), 
    col=hsv(0, seq(0.1, 1, 0.1), 1))
title("sample")

mu <- matrix(runif(D * K), K, byrow=T)
pz <- rep(1 / K, K)
gamma <- matrix(NA, nrow=N, ncol=K)
likelihood <- numeric()

iteration <- 0
repeat {
    cat("mu\n");print(mu)
    cat("pi\n");print(pz)

    if (!is.na(gamma[1, 1])) {
        plot(apply(gamma, 1, max), col=hsv(apply(gamma, 1, which.max) / K, 1, 1), 
            ylim=c(0, 1.05), pch=20, xlab="n", ylab=expression(gamma(z_nk)))
        title(paste0("gamma#", iteration))
    }

    if (MNIST) {
        image(1:28, 1:(28*K), matrix(as.vector(t(mu + rep(0:(K-1), D) + 1.0E-8)), nrow=28)[, (28*K):1], 
            axes=F, xlab="i", ylab="k", 
            breaks=seq(0, K, 0.1), 
            col=outer(seq(0.1, 1, 0.1), (1:K)/K, function(x1, x2) hsv(x2, x1, 1)))
        axis(1)
        axis(2, at=0:(K-1) * 28 + 14, labels=1:K)
    } else {
        image(1:K, 1:D, mu + rep(0:(K-1), D) + 1.0E-8, 
            axes=F, xlab="k", ylab="i",
            breaks=seq(0, K, 0.1), 
            col=outer(seq(0.1, 1, 0.1), (1:K)/K, function(x1, x2) hsv(x2, x1, 1)))
        axis(1, at=1:K)
        axis(2)
    }
    title(paste0("mu#", iteration))

    # E step
    for (n in 1:N) {
        pzx <- sapply(
            1:K, 
            function(k) log(pz[k]) + logdbern(x[n, ], mu[k, ])
            )
        pzx <- pzx - max(pzx)
        gamma[n, ] <- exp(pzx) / sum(exp(pzx))
    }

    # M step
    nk <- colSums(gamma)
    for (k in 1:K) {
        mu[k, ] <- colSums(x * gamma[, k]) / nk[k]
        pz[k] <- nk[k] / N
    }

    # likelihood
    likelihood <- c(likelihood, sum(sapply(1:N, function(n)
            logsumexp(sapply(
                1:K, 
                function(k)
                    log(pz[k]) + logdbern(x[n, ], mu[k, ])
                ))
            )))

    if (length(likelihood) > 1 
        && likelihood[length(likelihood)] - likelihood[length(likelihood) - 1] < 1.0E-2) {
        break
    }
    iteration <- iteration + 1
}

plot(likelihood, type="l", xlab="iteration", ylab="ln p(X)")
title("ln p(X)")