Memory spike with RTMB

[Jan-Ole Koslik]

Hi,

I'm experiencing something similar to here - I didn't find a similar post focussing on RTMB.

I am trying to run a hidden Markov model with 2 random effects per time point and 10K observations, i.e. 20K random effects.

The likelihood function basically looks like

for (track in 1:ntracks) {

   nll <- nll - track_llk(..., ~ 120 random effects for this track)

}

The inner part necessarily introduces dependence (due to the forward algorithm), leading to a banded Hessian with bandwidth ~ 120.

Doing the rough calculation, that means I should have ~ 120 * 20K = 2.4M non-zero entries in the Hessian. 

After the tape has been optimized and (I guess) the sparsity pattern of the Hessian has been matched, it takes up ~ 4 Gb of memory, but when calling MakeADFun() or in the next step, calling obj$fn() / obj$env$spHess() (I'm not totally sure about that), memory spikes to ~ 17 Gb which is a problem, particularly because I want to scale up the problem and the peak memory usage seems to grow quadratically in the number of random effects.

Can I do anything in my code to optimize for memory usage? For example would replacing the outer loop by sapply() help? Or is the size just inherent to the problem?

Thanks in advance :)

[Kasper Kristensen]

The fact that you go from 17Gb to 4Gb after tape optimization, indicates that the sparse hessian tape generated a lot of redundant operations temporarily. So, at least in theory, it should be possible to eliminate the spike by representing the tapes differently. Hard to give further advice without knowing exactly what's inside the track_nll().

FWIW I recently came across a model with similar symptoms that turned out to be caused by a fairly large number of matrix-vector multiplies. It was solved in this PR: https://github.com/spacetime-ecologist/spacetime-ecologists-RTMB/pull/3

[Jan-Ole Koslik]

Hmm, track_nll() looks something like 

track_nll <- function(delta, Gamma, allprobs_track, k) {
  tracklen <- nrow(allprobs_track)
  track_l <- 0
  foo <- delta * allprobs_track[1, , drop = FALSE]
  phi <- foo / sum(foo)
  for(t in 1:k){
    foo <- (phi %*% Gamma) * allprobs_track[t, , drop = FALSE]
    phi <- foo / sum(foo)
  }
  for(t in (k+1):tracklen){
    foo <- (phi %*% Gamma) * allprobs[t, , drop = FALSE]
    sumfoo <- sum(foo)
    track_l <- track_l + log(sumfoo)
    phi <- foo / sumfoo
  }
  -track_l
}

I did now actually compare doing 

for(track in 1:ntracks) {

   nll <- nll + track_nll(...)

}

to 

nll <- sum(sapply(1:ntracks, function(track) track_nll(...)))

and in the smaller example I tried, memory usage for the first one goes to 6 GB while only going to 2.5 for the second one.

[Jan-Ole Koslik]

I don't think it's about matrix multiplication here, the matrices are 2x2 :D

[Jan-Ole Koslik]

Actually, the sapply() part is wrong, same memory requirements. I was mislead here.