Rao-Blackwellized HMC

[Michael Betancourt]

One of the immediate generalizations of algorithms like the No-U-Turn
sampler is to average over the entirety of each trajectory, effectively
Rao-Blackwellizing the algorithm. I had hesitated implementing this
for a long time because intuitively I didn’t think there could be much
gained if we’re already generating nearly-independent samples. With
more motivation from the integration time paper that I put out a few
months ago I finally got around to playing with it, and it looks like my
intuitions were largely correct.

Given the marginal improvements I wasn’t planning on putting together
a full paper or anything.

[Andrew Gelman]

This is worth publishing as is (just adding a as a research note. I am of course happy to post on blog (and I guess you can put on Arxiv too) but I think the point is important enough that it should be cited.
A

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

[Marcus Brubaker]

How are you measuring computational cost in plot (d)?  Is that wall-clock time or gradient evaluations or ?  The reason I ask is two fold.  First the time spent computing the gradient for a Gaussian is going to be negligible because of the simple form of the posterior and heavy optimization that's been done for a multivariate normal.  The second is that Stan's I/O (rendering floating point values into plain text) is pretty inefficient and dramatically increasing the number of samples that need to be kept is going to bog things down.  (How much will also depend heavily on whether you're on a SSD, HDD or a networked disk.)  I think showing both (c) and (d) relative to number of gradient evaluations provides a more relevant measure of computational effort than either number of transitions or wall clock time.

Tangentially, Radford proposed this a while back as part of [1] (See section 6 and reference [5] therein) and found some significant benefits to basic HMC.  Also interesting, was that you don't need to use the same weights which are used in the Markov transition kernel for selecting states from the trajectory.  This could, for instance, allow you to store only subsets of the simulated trajectory and reduce the I/O cost.

Cheers,

Marcus

[1] Neal, R. M. (1994). An improved acceptance procedure for the hybrid Monte Carlo algorithm. Journal of Computational Physics, 111:194–203.

[Bob Carpenter]

I'm also curious how you're doing it, because I wouldn't think
the combination overhead would be too onerous in either time or
memory.

If you're only saving the averages per iteration, there's
the issue of how we deal with the interfaces and whether it'll
be possible to compute proper posterior intervals.

I think there's a lot of possiblity for online algorithms.
For instance, if you're just doing Rao-Blackwell and you want
to normalize, that is you get a stream of

w[n], p[n]

and you want to compute a weighted average

sum(w .* p) / sum(w)

then you can do it all with two accumulators on the log scale:

log(sum(w .* p) / sum(w))

= log(sum(w .* p)) - log(sum(w))

= log_sum_exp(log(w) .+ log(p)) - log_sum_exp(log(w))

But to compute p% quantiles, we'd need to save p% of the draws.
Or perhaps we could get a noisier estimate with subsampling (does
the inverse of Rao-Blackwell have a name?).

- Bob

[Michael Betancourt]

> How are you measuring computational cost in plot (d)? Is that wall-clock time or gradient evaluations or ?

Wall clock time for running Stan, not for the additional post processing to compute expectations and the like.

> The reason I ask is two fold. First the time spent computing the gradient for a Gaussian is going to be negligible because of the simple form of the posterior and heavy optimization that's been done for a multivariate normal. The second is that Stan's I/O (rendering floating point values into plain text) is pretty inefficient and dramatically increasing the number of samples that need to be kept is going to bog things down. (How much will also depend heavily on whether you're on a SSD, HDD or a networked disk.) I think showing both (c) and (d) relative to number of gradient evaluations provides a more relevant measure of computational effort than either number of transitions or wall clock time.

None of this is going to matter in practice. Even if all of the I/O overhead could be avoided, the improvements
are marginal.

> Tangentially, Radford proposed this a while back as part of [1] (See section 6 and reference [5] therein) and found some significant benefits to basic HMC. Also interesting, was that you don't need to use the same weights which are used in the Markov transition kernel for selecting states from the trajectory. This could, for instance, allow you to store only subsets of the simulated trajectory and reduce the I/O cost.

The improvements Radford saw were only in comparison to an HMC algorithm that takes the final point of
the trajectory. Uniformly sampling from the trajectory provides the significant improvement, and then
averaging over the entire trajectory yields only a marginal improvement.

You can use different weights but you still have to ensure the correct marginal probabilities or else you’ll
bias the expectations. In particular, any subsampling would have to include random decisions as to which
states to keep which reduces the efficacy of the Rao-Blackwellization.

Because I used the marginal probabilities exactly, this is the best that you’re going to do.

[Michael Betancourt]

In the note I linked to the branch with the implementation.
so you can see all the details directly.

But online won’t work because expectations don’t commute
with nonlinear transformations. In particular, the trajectories
are in the unconstrained space so we have to constrains the
states (and run generated quantities!) _before_ averaging
the states together, and the code isn’t factored in any way
to do that.

Then there’s the additional problem that you’d have to do
this separately for _every expectation that you want to compute_.
That includes means and variances, but not quantiles as
quantile estimators really aren’t amenable to Rao-Blackwellization.
Anyways, you then end up having to carry around a bunch of
different aggregators in memory, which starts defeating the purpose.

[Bob Carpenter]

Thanks for clarifying. I'd have thought it would have reduced
MCMC std err for expectations considerably.

And yes, that was my intuition as to why it helps so much with
basic HMC. But I'm usually nowhere near independent draws per
iteration with most of the models I'm fitting (many parameters
still run at a factor of 10 less).

When I was suggesting subsampling, I just meant that we could run
regular old NUTS to draw a small sample from the actual posterior
to compute quantiles at the same time the basic expectation calcs
get Rao-Blackwellized.

But obviously a no-go if the basic MCMC std errors for expectations
aren't reduced with Rao-Blackwellization. Because it'd induce a doc
and user education headache the likes of which would probably
dwarf those of the ~.

- Bob

[Bob Carpenter]

> On Apr 27, 2016, at 1:20 PM, Michael Betancourt <betan...@gmail.com> wrote:
>
> In the note I linked to the branch with the implementation.
> so you can see all the details directly.
>
> But online won’t work because expectations don’t commute
> with nonlinear transformations. In particular, the trajectories
> are in the unconstrained space so we have to constrains the
> states (and run generated quantities!) _before_ averaging
> the states together, and the code isn’t factored in any way
> to do that.

Good point. Damn those nonlinearities.

> Then there’s the additional problem that you’d have to do
> this separately for _every expectation that you want to compute_.
> That includes means and variances, but not quantiles as
> quantile estimators really aren’t amenable to Rao-Blackwellization.

Right --- that's why I was suggesting subsampling or just doing
efficient quantile algorithms.

> Anyways, you then end up having to carry around a bunch of
> different aggregators in memory, which starts defeating the purpose.

You don't need to carry these in memory. Just transform each point and
aggregate the results as you go.

Too bad we couldn't lose the draws altogether. We could do R-hat online,
but I don't think we could calculate n_eff and thus MCMC std err online.

- Bob

[Michael Betancourt]

On Apr 27, 2016, at 6:28 PM, Bob Carpenter <ca...@alias-i.com> wrote:

Anyways, you then end up having to carry around a bunch of
different aggregators in memory, which starts defeating the purpose.

You don't need to carry these in memory.  Just transform each point and
aggregate the results as you go.  

No, I mean you’d need a separate aggregator for the mean, variance,

etc.  All of these methods are focuses on computing just one expectation,

which conflicts with our board summary-statistic approach.

Too bad we couldn't lose the draws altogether.  We could do R-hat online,
but I don't think we could calculate n_eff and thus MCMC std err online.

Not easily, but it could potentially be done.   Just not using the slick

FFT trick.

[Aki Vehtari]

How this relates to http://arxiv.org/abs/1511.06925 which claims that using intermediate steps "can lead to substantial gains in computational efficiency"?

Aki

[Andrew Gelman]

Oooh, controversy!

[Bob Carpenter]

No controversy here, move along (unless you want details).

Specifically, Nishimura and Dunson are simply wrong
in their characterization of "current algorithms" when they
say in the abstract of their arXiv paper:

The utility of such [HMC] approaches is highlighted in the
software package STAN ...

Current algorithms use only the final step, and wastefully
discard all the intermediate steps.

That's *not* what NUTS does in Stan. Michael already explained
why the Rao-Blackwellization is more useful for standard HMC
than for NUTS when he commented on the last entry in this thread
that said Radford already showed Rao-Blackwellization helped way
back in 1994 or something like that.

It's sad how badly mischaracterized other people's approaches
are in papers.

- Bob

[Andrew Gelman]

I can send an email to Dunson on this.
A

[Michael Betancourt]

Apparently I have a bad attitude because I point
out that papers like these are terribly misguided…
Good thing academia is mature and not petty at
all!

In any case, the Nishimura-Dunson paper is
just another scheme for sampling from a trajectory.
In the integration time paper I note that you can
use a slice sampler (as in the original NUTS
implementation) or simple multinomial sampling,
and the Metropolis procedure they do in the paper
would be a third option. But it will yield near identical
results for either sampling a single point or doing
a Rao-Blackwellization.

[Aki Vehtari]

Thanks for confirming this (I thought this would be the explanation, but wasn't completely sure).

Aki