Efficient data updates for EM algorithm

[Paul vdb]

Hello,

I've been using RTMB for fitting mixture models using the EM algorithm. Things are working but I think it is inefficient. The main complexity is that for each maximization step, I need to pass the updated expectations from the previous estimates. The full worked example is attached. The core objective function is pasted below:

## Objective Function to maximize:
objectiveFunction <- function(pars){
  getAll(pars)
  n <- length(y)
  p <- c(exp(logitp), 1)/(sum(exp(logitp)) + 1)
  logp <- log(p)
  sigma <- exp(logsigma)

  empty <- rep(mu[1], 0)
  ## Updates based on expectation step
  getDat = function() {
    .GlobalEnv$post_probs_vec
  }  
  data.vec <- DataEval(getDat, empty)
  postp <- matrix(data.vec, nrow = n, ncol = K, byrow = FALSE)
 
  obj <- 0
  for( k in 1:K ){
    obj <- obj - sum(postp[,k]*(logp[k] + dnorm(y, mu[k], sigma[k], log = TRUE)))
  }
  obj  
}

Maybe there is a trick to do the expectation within the same objective function to avoid passing new data back and forth?

Other question, what is the most efficient way to pass the data to the functions? Based on another thread I am using:

dataenv <- local({y <- y; K <- K; environment()})
environment(objectiveFunction) <- dataenv

Thanks so much for any advice.

Paul

[Kasper Kristensen]

> Maybe there is a trick to do the expectation within the same objective function to avoid passing new data back and forth?

Yes, that's possible. You can use nested tapes to easily pass 'postp' from one context to another. In addition, you can add a solver to the tape, that performs the M-step.

See attached simpleEM2.R

> what is the most efficient way to pass the data to the functions?

I think the best way is to use closures: https://kaskr.r-universe.dev/RTMB/doc/RTMB-introduction.html#using-an-explicit-data-argument

[Paul vdb]

That's neat. I wouldn't have thought of that with my head stuck in MakeADFun. The nested taping makes sense. Thanks for the suggestions!