How to calculate Effective sample size per gradient evaluation?

numpyro
1

I wish to calculate the effective sample size per gradient evaluation. Is there some way to do this in Numpyro?

2

you can get ess by print summary and number of grad steps by using extra_fields.

1 Like
3

Thank you, @fehiepsi. Is it also possible to know the number of gradient evaluations in SVI?

4

Also, what do I need to pass in the extra_fields to get the number of gradient evaluations? I could not find it in the documentation.

5

you can search forum for extra fields + num_steps I think. The number of grad in SVI is the number of steps (unless you use specific objectives).

6

Hi, I found no relevant posts on the forum related to gradient evaluations (num_steps). Could you please refer me to some relevant post?

7

You can see this post Get number of steps, sizes and acc. prob over all iterations

1 Like
8

Thanks for providing the relevant post. Upon using it, the output was an array of numbers with dimensions equivalent to the number of samples. Is the total number of gradient evaluations the sum of all the elements in the array? Also, does it take into account the gradients computed during warmup?

9

Yes, it is the sum. To collect warmup samples and info, you can use the warmup method.

1 Like
10

Thank you. But, the num_steps array corresponding to the warmup comes out to an array full of zero. Is it correct? The code that I used is shown below:

nuts_kernel = NUTS(funnel_model)
mcmc = MCMC(nuts_kernel, num_warmup=1000, num_samples=1000)
mcmc.warmup(rng_key, extra_fields=('num_steps',))
warmup_steps = mcmc.get_extra_fields()['num_steps']
11

I’m not sure. Sounds like a bug to me. Does your sampling look normal? You can also try using HMC with adapt_step_size=False to make sure to have a constant number of leapfrog steps per mcmc step - if num_steps is 0 then it would be a bug. Please open a github issue in that case. Thanks

12

Yes, it still returns a zero matrix. I’ll create an issue regarding this in the repository.

13

Does varying the number of particles in ELBO increases the number of gradient evaluations, or is it still only related to the number of steps?

14

for Trace_ELBO: # of gradients = num_particles x num_steps

1 Like