WIP: Implement opvi

[ferrine]

This PR suggests new approach to variational inference via OPVI framework

Variational inference is an umbrella term for algorithms which cast Bayesian inference as optimization. Classically, variational inference uses the Kullback-Leibler divergence to define the optimization. Though this divergence has been widely used, the resultant posterior approximation can suffer from undesirable statistical properties. To address this, we reexamine variational inference from its roots as an optimization problem. We use operators, or functions of functions, to design variational objectives. As one example, we design a variational objective with a Langevin-Stein operator. We develop a black box algorithm, operator variational inference (OPVI), for optimizing any operator objective. Importantly, operators enable us to make explicit the statistical and computational tradeoffs for variational inference. We can characterize different properties of variational objectives, such as objectives that admit data subsampling---allowing inference to scale to massive data---as well as objectives that admit variational programs---a rich class of posterior approximations that does not require a tractable density. We illustrate the benefits of OPVI on a mixture model and a generative model of images.

Convenient interface is still a question, I'm open to suggestions

[ferrine]

There is no support for aevb yet, it is going to be implemented soon

[springcoil]

Why is this not ready yet?

[ferrine]

I've copied this file from #1600, interface is different so tests will fail. I'll refactor them later when interface will be convenient

[springcoil]

Ok I get you :)

[springcoil]

Maybe rng needs expanded here in the docstring.

[springcoil]

This is good :)

[ferrine]

In the paper authors suggest the way how to deal with discrete vars. The approach is described in the appendix. See https://arxiv.org/abs/1610.09033

[ferrine]

I think that it can be possibly implemented in the future

[springcoil]

I'll need to review the papers to really get on top of this, which will be difficult. Is anyone else au fait with this? Maybe @taku-y or @ColCarroll or @twiecki

[fonnesbeck]

Thanks for working on this!

[twiecki]

I think we need blank lines in-between.

[taku-y]

It would be better to use logdet(S), which is more numerically stable log(det(S)) especially for high dimensional variables. See #1012.

[ferrine]

That PR in Theano is not merged yet. I can copy and paste it into theanof or our math maybe

[twiecki]

#1777 is merged now.

[taku-y]

It would be better to use def instead of lambda, PEP8.

[taku-y]

updates() looks good to me.

[ferrine]

Rebased

[springcoil]

This looks good. Are there changes that need made to Theano for this to work?

[twiecki]

I suppose this is still not functional?

[ferrine]

Yes, this does not work yet. I have great vacations and will resume to work in 14 days:)
Welcome for feedback)

[ferrine]

@springcoil no changes for Theano are required. The only serious problem I know is with langevine stein operator. I need theoretical help with it. Here I'm trying to resolve the it.

[twiecki]

@ferrine I think we should forge ahead here without langevin stein.

[ferrine]

@twiecki I've thought about it. Good idea.

[twiecki]

@ferrine #1777 is merged.

[ferrine]

I've found problems with minibatch training. Can't get why that happens.

[ferrine]

Rebased

[ferrine]

Finally

Usage

# assuming neural_network, sigmoid_output, input_var are defined above
with neural_network:
    inference = ADVI()
    # BTW inference.approx is already created
    approx = inference.fit(10000)
    new_input = tt.matrix()
    proba = approx.apply_replacements(
        sigmoid_output,
        deterministic=False, # It's default, just for clarity
        more_replacements={input_var:new_input}
    )
    predict_proba = theano.function([new_input], proba)

Features

• Unified interface for variational methods - scope for future work
• Simple and clear interface
• Custom callbacks while training (add predefined?)
• Make replacements for prediction with one line of code
• Fast sampling from posterior into NDArray trace
• FullRankADVI - catch mutual correlations in posterior for a computational cost
• No minibatch pain - use generators in observed nodes and specify total_size and use regular way for inference
• Object oriented approach - easy access to the parts of implementation (useful for experiments)

Benchmarks

I've checked if this works with http://twiecki.github.io/blog/2016/06/01/bayesian-deep-learning/
and got good results:

• OPVI ADVI performs 6-8% slower than previous implementation. I don't think that's big deal
• FullRankADVI worked for that example too. Assuming ADVI speed is 1, FullRank alternative is 0.12 of this. Tracking ELBO matters a lot here: without we have 0.12, 0.05 otherwise.

TODO:

• Add documentation
• Maybe add some shortcuts

Further ideas

• Construct Approximation from given trace (sample from trace object directly, applying replacements as usual)
• Implement Normalizing Flows
• Implement Langevin Stein Operator
• Implement Approximation transfer (pretrain MeanField and initialize FullRank with it)
• Implement Structured Variational Inference (https://arxiv.org/abs/1404.4114)
• Replace old advi with new one?? (in future PR ofc)
• Support pickling?
• Other Ideas @twiecki, @taku-y, @springcoil?

[ferrine]

What about merge?

[taku-y]

It might be better to update example notebooks with the new interface, while commenting out the old one.

[taku-y]

c is scaled with the number of elements in x: c -> c * x.ravel().shape[0] ?

[ferrine]

There are elemwise operations, so that's ok here

[taku-y]

It seems conflicting with the code: order, flat_view, view -> flat_view, order, view

[fonnesbeck]

This looks really good. My only recommendation is a tweak to the API to make it more consistent with the rest of the PyMC fitting methods. At the moment, it is identical to the scikit-learn interface, where you instantiate a class and call its fit method. Can we eliminate the separation between the two and have something along the lines of:

with neural_network:
    approx = advi(10000)

this makes it closer to what we do for MCMC. Even better might be something like:

with neural_network:
    approx = approximate(10000, method='ADVI')

which mirrors sample (a verb) with approximate (also a verb).

I think a consistent interface is important, so I'd like to hear what others think. If we like a scikit-learn-like API better, then I'd suggest making changes across the board to accomodate it.

[ferrine]

@fonnesbeck Yes, it is good idea. I think that I'll still leave OOP interface for advanced users (for training model sequentially) while adding a simple shortcut approximate

[fonnesbeck]

Also, we probably want to modify one or more of the ADVI examples to use this interface.

[ferrine]

Rebased.

[twiecki]

Congrats @ferrine, this is a huge one.

[fonnesbeck]

Yes, well done! I will be putting this interface to use right away.

[fonnesbeck]

I don't see a fit function here -- just the method.

Edit: never mind, I see it now.