bonnevie/rebar

A Python implementation of the gradient REBAR estimator.

rebar's Issues

Critical vectorization issue

Assuming the loss and logp tensors passed to RELAX are meaned over the number of samples leads to problems when

approx_gap = (loss - weight*conditional_control)
[...]
score_grad = approx * score

is encountered, as the product in score_grad (meant to be element-wise over samples) becomes a product of two means, expanding into a quadratic number of terms, inducing significant bias that ironically gets more severe with the number of samples used.

Clarifying LSE in categorical_backward function

I didn't get why you add logsumexp to gumbel random variables here.

rebar/relaxflow/reparam.py

Line 237 in 0ff4715

topgumbels = gumbel + tf.reduce_logsumexp(alpha, axis=-1, keep_dims=True)

The paper does not do so. Does it have any advantage over the original version?

Equation for binary_backward is not the same in REBAR

The function binary_backward defined here does not use the same equation as what is written in the REBAR paper (Appendix C).

nan issue and a typo?

Hi,

Thanks for sharing the code. Just wondering the reason of using killnan() function? It seems you have eps for all tf.log, so I am not sure where nan can arise..

Also, in this line, shouldn't self.b be self.zb?

Explain truncated gumbel sampling

I didn't know why truncated gumbel can sampled in this way
return -tf.log(EPSILON + tf.exp(-gumbel) + tf.exp(-truncation))
can u describe it?