Memory requirements of ADVI in Stan : Out of Memory

[Vivek Kulkarni]

Hi,
I am trying to use ADVI on a variant of LDA where the number of documents is around 100K, with a vocabulary size of 100K. However using variational inference on my model with this data, takes up all the RAM on my machine and just terminates by running out of memory. I have a generously proportioned server with 384GB RAM and yet ADVI takes up all of it. What are the scaling limits of using inference with ADVI in Stan? My understanding is that ADVI estimates the ELBO from a few samples and so might scale well both in terms of time and memory. I can post the stan code for my model if needed.

Best,
Vivek.

[Dustin Tran]

Hi Vivek, Stan requires loading all data in RAM so that’s one bottleneck. Do you know how much memory that takes by itself? In ADVI, there are a few inefficient copies of the parameter vector that occur, which in this case is pretty massive as ADVI will treat all parameters including document-specific/word-specific parameters as global. You can try some hacks that do data subsampling, c.f.,  https://github.com/stan-dev/stan/blob/adsvi/how_to_ADSVI.md

Dustin

--
You received this message because you are subscribed to the Google Groups "Stan users mailing list" group.
To unsubscribe from this group and stop receiving emails from it, send an email to stan-users+...@googlegroups.com.
To post to this group, send email to stan-...@googlegroups.com.
For more options, visit https://groups.google.com/d/optout.

[Krzysztof Sakrejda]

Posting Stan code always helps to figure out memory issues.  It helps to see the dimensions of everything---data/transformed data/parameters/transformed parameters/data/generated quantities.  Krzysztof

[Vivek Kulkarni]

Thanks a lot. I will try the subsampling trick you mentioned.  Also posting the Stan code if that gives any insights
P = 5 # Number of super topics
K = 25 # Number of topics
V = 100000 # Number of words
R = 1 # Number of roles
U = 90000  # Number of users
N =  264166617

data {
    int<lower=1> U; // num users
    int<lower=2> P; // num super topics
    int<lower=2> K; // num topics
    int<lower=2> V; // vocabulary size
    int<lower=1> N; // total word instances
    int<lower=1> R; // total number of roles 
    int<lower=1,upper=V> words[N]; // word n
    int<lower=1,upper=U> userids[N]; // User ID for word n
    int<lower=1,upper=R> roleids[N]; //  Role ID for word n
    vector<lower=0>[P] alpha; // super-topics prior
    vector<lower=0>[V] beta; // word prior
    vector<lower=0>[K] gamma; // topic prior
}
parameters {
    simplex[P] theta[U]; // super-topic dist for user u
    simplex[K] psi[P, R]; // topic dist for supertopic p
    simplex[V] phi[K]; // word dist for topic k
}
model {
    for (u in 1:U)
        theta[u] ~ dirichlet(alpha); // prior
    for (p in 1:P) {
        for (r in 1:R) {
            psi[p,r] ~ dirichlet(gamma); // prior
        }
    }
    for (k in 1:K)
        phi[k] ~ dirichlet(beta); // prior

    for (n in 1:N) {
        matrix[P,K] probs;
        for (p in 1:P) {
            for (z in 1:K) {
                probs[p,z] <- log(theta[userids[n], p]) + log(psi[p, roleids[n], z]) + log(phi[z,words[n]]); 
            }
        }
        increment_log_prob(log_sum_exp(probs)); // likelihood
    }
}

[terrance savitsky]

Hi Dustin,

> In ADVI, there are a few inefficient copies of the parameter vector that occur, which in this case is pretty massive as ADVI will treat all parameters including document-specific/word-specific parameters as global

I have a similar RAM management issue when estimating a GP model with ADVI due to high dimensional covariance matrices.   Will declaring the covariance matrices as local (in local variable blocks) see them treated as local under ADVI estimation or is your point that ADVI treats all parameters as global, no matter that they are declared as local (in local variable blocks)?  In other words, is the only solution path to subsample the data?

> You can try some hacks that do data subsampling, c.f.,  https://github.com/stan-dev/stan/blob/adsvi/how_to_ADSVI.md

Any idea how long until subsampling is formally launched in rstan??

Thank you,

Terrance

--

Thank you, Terrance Savitsky

[Bob Carpenter]

It's not ADVI per se, it's calculating gradients. Every variable
whose value is a function involving a parameter goes onto the autodiff
stack (which manages where the chain rule gets passed during gradient
calculations).

No concrete targets for data streaming. The streaming from Alp's
papers was all done with modified versions of Stan rather than
the release versions.

- Bob

[terrance savitsky]

hi Bob,  Okay, then.  You had explained this point to me and I also pulled the paper on the arxiv.  I get that point.  I read too much into Dustin's statement to mean that even a local variable is never let go from RAM and passed back R, which should not be the case based on your follow-up point.  A local variable function of a parameter will take longer to be released from RAM at the c++ layer, but it will be released.

[Bob Carpenter]

The actual local variable acts just like a local variable. It
goes off the stack when the local block is done executing. But
in C++ terms, creating a variable has the side effect of adding
things to global memory that will not be freed until after the
derivative is calculated.

- Bob