Refactoring HMC API #436

xukai92 · 2018-04-16T23:58:23Z

Addressing #431

Still in working.

xukai92 · 2018-04-16T23:58:47Z

2a4dfb3 solves #434 (comment)

xukai92 · 2018-04-17T09:31:35Z

@yebai I've refactored core HMC code. Basically I added three functions _leapfrog, _find_H and _sample_momentum and make the original corresponding functions calling them by wrapping Turing.jl internal things. Can you have a look at current change? and let's discuss what's the best way to write the unit tests, e.g. do we simply write another non-Turing.jl based HMC and test them? or something else.

yebai · 2018-04-17T19:56:16Z

The code looks good to me. For unit tests, perhaps we can write 2-3 models without constrained parameters (e.g. gdemo; Bayesian logistic regression, stochastic volatility from the NUTS paper),

implement them in both Turing and plain Julia (using AD to compute gradients for the plain Julia version of the models)
run Turing's HMC sampler on both the Turing version and Julia version

This would allow us to check whether there are errors in the compiler, transformation or gradients. After this step is done, we can focus on debugging the HMC sampler itself by running it on the plain Julia version of models.

xukai92 · 2018-04-26T16:25:29Z

Just a small update:

I'm working on this PR today. It turns out that rather than only having Turing-free HMC core functions, we should actually have Turing-free complete HMC step first. Then handling with vi, spl and model outside this level of abstraction (at the moment I think it's hard to do that on the whole HMC level). I'm working on this now.

xukai92 · 2018-04-28T16:26:01Z

Simple Gaussian (with prior on mean only)
Bayesian logistic regression
~~stochastic volatility~~

…into refactor-hmc-api

xukai92 · 2018-04-30T21:25:49Z

@yebai So now we have a simple Gaussian and a simple Bayesian linear regression. I'm going to implement a stochastic volatility today or tomorrow. But for the existing two without constrained variables, how should we really debug them?

xukai92 · 2018-04-30T21:26:31Z

We actually want to implement a NUTS in the same way like _hmc_step() do we?

yebai · 2018-05-08T17:07:56Z

But for the existing two without constrained variables, how should we really debug them?

@xukai92 It's fine to first test HMC and NUTS on models without constrained variables, just to verify dual averaging and leap-frog integrator is correct. We can change the prior to truncated Gaussian when the initial tests are passed.

yebai · 2018-05-08T17:08:46Z

We actually want to implement a NUTS in the same way like _hmc_step() do we?

Yes, we would like to be able to unit test NUTS using the same set of models as HMC.

xukai92 · 2018-05-22T22:50:37Z

Agreed work plan

implement Turing.jl free gdemo
implement Turing.jl free LDA

The followings are based on these two models. "Check" means doing unit tests.

check gradient
check leaf-frog
~~check dual-averaging~~
- Didn't use Turing-free DA in the end
NUTS Turing.jl free implementation
check NUTS w/out DA
check NUTS w/ DA
~~check pre-cond. adaptation~~
- Didn't use Turing-free pre-cond. adapt. in the end
- Online update of pre-cond. already has its unit test
check initialization

xukai92 · 2018-06-21T19:49:06Z

Thanks @willtebbutt and @wesselb - address them soon.

xukai92 · 2018-06-25T21:13:14Z

Previous comments are resolved.

yebai · 2018-06-27T08:57:58Z

@willtebbutt could you take a look before I merge this PR? Thanks!

yebai · 2018-06-27T08:58:23Z

Previous comments are resolved.

Thanks, Kai!

willtebbutt · 2018-06-27T10:09:31Z

Thanks for addressing our concerns @xukai92, just one remaining point. Are the @gen_local_grad_func (and related) macros really necessary? Would closures not suffice? i.e.

function gen_grad_func(vi, spl, model)
    return function (θ::AbstractVector)
        if ADBACKEND == :forward_diff
            vi[spl] = θ
            grad = gradient(vi, model, spl)
        elseif ADBACKEND == :reverse_diff
            grad = gradient_r(θ, vi, model, spl)
        else
            error("An appropriate error.")
        end
        return getlogp(vi), grad
    end
end

It's really a matter of taste, but my impression is that a good general rule is to only use metaprogramming where it's really necessary, essentially just because it's tricky to read.

xukai92 · 2018-06-28T00:45:54Z

I've replaced macros with closures. Thanks for pointing this - I agree we'd better use closures here.

xukai92 · 2018-06-28T23:26:59Z

@willtebbutt Any other comments?

willtebbutt · 2018-06-29T13:29:35Z

Apologies for the delay. Will review the technical details of the NUTS implementation tomorrow morning and, assuming that everything looks fine, I'm happy for this to be merged. (Wessel and I had to refresh our memories regarding NUTS, hence the delay)

willtebbutt

Overall looks good. We couldn't find any obvious technical errors in the NUTS implementation, other than concerns about the capping of j to 5 (see comment for details). There are a few stylistic things that it would be good to have addressed.

willtebbutt · 2018-06-30T15:26:42Z

src/samplers/nuts.jl

+      - ϵ           : leapfrog step size
+      - H0          : initial H
+      - lj_func     : function for log-joint
+      - grad_func   : function for the gradient of log-joint


This documentation should be outside of the function. See here for examples: https://docs.julialang.org/en/stable/manual/documentation/

In this particular case, we need something like:

""" build_tree(θ::T, r::Vector, logu::Float64, v::Int, j::Int, ϵ::Float64, H0::Float64, lj_func::Function, grad_func::Function, stds::Vector) where {T<:Union{Vector,SubArray}} Recursively build balanced tree. Ref: Algorithm 6 on http://www.stat.columbia.edu/~gelman/research/published/nuts.pdf # Arguments: - θ: an argument - r: another argument """

to be consistent with the standard Julia conventions. If we then call ?_build_tree we will get the appropriate information displayed correctly.

I know the format. That should be a mistake during copying and pasting. Will fix it

willtebbutt · 2018-06-30T15:27:24Z

src/samplers/nuts.jl

+Ref: Algorithm 6 on http://www.stat.columbia.edu/~gelman/research/published/nuts.pdf
+"""
+function _build_tree(θ::T, r::Vector, logu::Float64, v::Int, j::Int, ϵ::Float64, H0::Float64,
+                    lj_func::Function, grad_func::Function, stds::Vector; Δ_max=1000) where {T<:Union{Vector,SubArray}}


Is there a particular reason that we're using Union{Vector, SubArray} rather than just AbstractArray here?

willtebbutt · 2018-06-30T15:40:54Z

src/samplers/nuts.jl

+Recursively build balanced tree.
+
+Ref: Algorithm 6 on http://www.stat.columbia.edu/~gelman/research/published/nuts.pdf
+"""


Please add documentation for stds and Δ_max

willtebbutt · 2018-06-30T15:44:57Z

src/samplers/nuts.jl

+
+Ref: Algorithm 6 on http://www.stat.columbia.edu/~gelman/research/published/nuts.pdf
+"""
+function _build_tree(θ::T, r::Vector, logu::Float64, v::Int, j::Int, ϵ::Float64, H0::Float64,


Lots of arguments' types appear to be overly constrained. Is this necessary? For example, it's not clear to me that logu can't be constrained to be an AbstractFloat as opposed to Float64 -- similarly for ϵ andH0. Conversely, the type of Δ_max is not constrained at all. Should this be constrained to be an AbstractFloat also?

The same comment applies to r: could it be an AbstractVector instead of just a Vector. Same for stds.

Make sense! Thanks for pointing out this!

willtebbutt · 2018-06-30T15:59:22Z

src/samplers/nuts.jl

+    end
+  end
+
+function _nuts_step(θ, ϵ, lj_func, grad_func, stds)


It would be good to document this function properly, despite that fact it's not exposed publicly.

willtebbutt · 2018-06-30T16:04:34Z

src/samplers/nuts.jl

+  θm = θ; θp = θ; rm = r0; rp = r0; j = 0; θ_new = θ; n = 1; s = 1
+  local da_stat
+
+  while s == 1 && j <= 5


Why is j arbitrarily capped at 5? Surely this could potentially re-introduce the random-walk behaviour that we are trying to avoid?

In practise we have to set such a condition for j otherwise the sampler can go super slow if the step-size happens to be very small. Stan uses 10 (maximum 2^10 evaluations) as default and 5 (maximum 2^5 evaluations) was originally chosen by Hong because we were bit slow back then.

Let me make this number an optional argument in the function. We probably need another interface change to support a user specified maximum j.

willtebbutt · 2018-06-30T16:05:10Z

src/samplers/nuts.jl

+    v = rand([-1, 1])
+
+    if v == -1
+


Inconsistent addition of blank lines here. Maybe remove these?

yebai · 2018-07-03T14:35:55Z

@willtebbutt @xukai92 @wesselb many thanks for the hard work!

* move gmm overwrite out core source code * refactor find_H * refactor sample momentum * refactor lf step * hmc step abstraction v1.0 done * make hmcda using _hmc_step * add a note * add bayes lr * add things to runtests * add sv turing * add Turing-free nuts * bug free Turing-free nuts * update reference * add grad check unit test * add gdemo * add gdemo nuts * restructure hmc_core tests * NUTS works * Change rand seed for test * fix adapation condition bug * add test REQUIRE * Remove all benchmarks * change test nuts file name * rearrange hmc codes * add unit test for leapfrog * clean nuts test * Remove obsolete dependence on deps.jl * fix typo * add new lines to the end of files * rename file * add new lines to the end of files * use macros to gen functions with the same pattern * resolve indentation * add explict return * resolve indentation * more stable mh accept * Remove unrelated notebook * Unify the use of mh_accept * replace macro by closure to gen local funcs * improve doc

move gmm overwrite out core source code

2a4dfb3

xukai92 added 3 commits April 17, 2018 01:51

refactor find_H

841a22b

refactor sample momentum

0784177

refactor lf step

6df29f5

Merge branch 'master' into refactor-hmc-api

956102d

xukai92 added 3 commits April 27, 2018 14:28

hmc step abstraction v1.0 done

5d21541

make hmcda using _hmc_step

ee5659e

add a note

48b4e2f

xukai92 added 3 commits April 30, 2018 16:48

Merge branch 'refactor-hmc-api' of https://github.com/yebai/Turing.jl …

02b2706

…into refactor-hmc-api

add bayes lr

4b3f0bb

add things to runtests

cc46ac7

xukai92 added 2 commits May 2, 2018 21:30

add sv turing

3ec6eba

Merge branch 'master' into refactor-hmc-api

32bc26a

xukai92 added 7 commits May 24, 2018 00:13

add Turing-free nuts

6d9c5e5

bug free Turing-free nuts

6ddf92c

update reference

aa3d564

add grad check unit test

66c4d53

add gdemo

8d15cea

add gdemo nuts

fecd1c9

restructure hmc_core tests

8a3fbad

xukai92 added 11 commits June 21, 2018 20:49

fix typo

58be06b

add new lines to the end of files

e533fa5

rename file

7c8c7be

add new lines to the end of files

02623bc

use macros to gen functions with the same pattern

dbc71d2

resolve indentation

84fd884

add explict return

5eda1ff

resolve indentation

b31047e

more stable mh accept

26e72fb

Remove unrelated notebook

2073f5b

Unify the use of mh_accept

5c66f27

xukai92 mentioned this pull request Jun 25, 2018

Update Turing inference SciML/DiffEqBayes.jl#48

Merged

replace macro by closure to gen local funcs

0c2bfc2

willtebbutt requested changes Jun 30, 2018

View reviewed changes

improve doc

b31447f

yebai merged commit b67e9e1 into master Jul 3, 2018

This was referenced Jul 3, 2018

Refactoring HMC/NUTS related functions for better unit testing #431

Closed

Make sure HMCDA is not buggy #289

Closed

yebai deleted the refactor-hmc-api branch July 9, 2018 12:56

xukai92 mentioned this pull request Nov 22, 2018

Pre-cond seems unstable/buggy #284

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactoring HMC API #436

Refactoring HMC API #436

xukai92 commented Apr 16, 2018

xukai92 commented Apr 16, 2018

xukai92 commented Apr 17, 2018

yebai commented Apr 17, 2018

xukai92 commented Apr 26, 2018

xukai92 commented Apr 28, 2018 •

edited

Loading

xukai92 commented Apr 30, 2018

xukai92 commented Apr 30, 2018

yebai commented May 8, 2018

yebai commented May 8, 2018

xukai92 commented May 22, 2018 •

edited

Loading

xukai92 commented Jun 21, 2018

xukai92 commented Jun 25, 2018

yebai commented Jun 27, 2018

yebai commented Jun 27, 2018

willtebbutt commented Jun 27, 2018 •

edited

Loading

xukai92 commented Jun 28, 2018

xukai92 commented Jun 28, 2018

willtebbutt commented Jun 29, 2018 •

edited

Loading

willtebbutt left a comment

willtebbutt Jun 30, 2018

xukai92 Jun 30, 2018

willtebbutt Jun 30, 2018

willtebbutt Jun 30, 2018

willtebbutt Jun 30, 2018

xukai92 Jun 30, 2018

willtebbutt Jun 30, 2018

willtebbutt Jun 30, 2018

xukai92 Jun 30, 2018

willtebbutt Jun 30, 2018

yebai commented Jul 3, 2018

Refactoring HMC API #436

Refactoring HMC API #436

Conversation

xukai92 commented Apr 16, 2018

xukai92 commented Apr 16, 2018

xukai92 commented Apr 17, 2018

yebai commented Apr 17, 2018

xukai92 commented Apr 26, 2018

xukai92 commented Apr 28, 2018 • edited Loading

xukai92 commented Apr 30, 2018

xukai92 commented Apr 30, 2018

yebai commented May 8, 2018

yebai commented May 8, 2018

xukai92 commented May 22, 2018 • edited Loading

xukai92 commented Jun 21, 2018

xukai92 commented Jun 25, 2018

yebai commented Jun 27, 2018

yebai commented Jun 27, 2018

willtebbutt commented Jun 27, 2018 • edited Loading

xukai92 commented Jun 28, 2018

xukai92 commented Jun 28, 2018

willtebbutt commented Jun 29, 2018 • edited Loading

willtebbutt left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yebai commented Jul 3, 2018

xukai92 commented Apr 28, 2018 •

edited

Loading

xukai92 commented May 22, 2018 •

edited

Loading

willtebbutt commented Jun 27, 2018 •

edited

Loading

willtebbutt commented Jun 29, 2018 •

edited

Loading