google-deepmind / annealed_flow_transport Goto Github PK

1. Ineffective option hmc_steps_per_iter has no effect. In markov_kernel.py, there should be a loop and run hmc_steps_per_iter times hmc_wrapped()
2. optimization: the normalization/log_solftmax() operation in reweight_no_flow() does not seem to be necessary. The unnormalized weights would be normalized in the end anyway in get_log_normalizer_increment_no_flow()
3. optimization: using Eq. (15) in Del Moral, Douce, and Jasra (2006) is more optimal, as many normalization/log_solftmax() and expectation/ logsumexp() operations can be spared in between resampling events. In all the models we ran, only a few (<10) resampling events occurred.

Not sure how much juice one can squeeze out by removing these redundant operations though.