Improve advice / diagnostics when VI gives terrible results

[ricardoV94]

I have never used VI before, so I don't know if it's expected to fail so easily:

import pymc as pm

with pm.Model():
    x = pm.Normal("x")
    # increasing the distance between prior and posterior makes no difference for VI
    # e.g., sigma=1e-3 or observed=100
    y = pm.Normal("y", x, sigma=0.5, observed=10)  
    
    nuts = pm.sample()
    vi = pm.fit()
    
print((
    nuts.posterior["x"].mean().values, 
    vi.sample().posterior["x"].mean().values
))  # (array(8.01670159), array(3.11480314))

[fonnesbeck]

Interesting. When I do loosen the sigma, the results to end up being closer for me. For example, with sigma=10 I get:

(array(0.09860912), array(0.11399919))

which is not a great estimate from either method!

[ricardoV94]

which is not a great estimate from either method!

The example is conditioning on a single observation. With a looser likelihood it's normal (no pun intended) for the posterior to get closer to the prior.

Here is the example I tried first, which I thought might have been too unfair for VI:

import pymc as pm

with pm.Model():
    x = pm.Normal("x")
    y = pm.Normal("y", x, sigma=1e-5, observed=100)

    nuts = pm.sample()
    vi = pm.fit()

print((
    nuts.posterior["x"].mean().values,
    vi.sample().posterior["x"].mean().values
))  # (array(100.00000015), array(3.12762018))

[fonnesbeck]

Ah, you're right. I was confusing the likelihood and the prior.

Yes, this is odd behavior. FWIW, I get pretty reasonable results on my bigger models (at least as far as point estimates go).

[ricardoV94]

V3 gives the same results...

[ColCarroll]

Just so we're on the same page, I think the correct posterior here is

x | y ~ N(obs / (1 + sigma^2), sigma^2 / (1 + sigma^2))

so the most recent example with obs = 100, sigma = 1e-5 should be

N(99.99999999,  0.9999e-11)

and the first example with obs=10, sigma=0.5 should be

N(8, 0.2)

[fonnesbeck]

VI should not have a problem approximating this well. Any thoughts, @ferrine ?

[ferrine]

The learning rate seems to be too small for this model. This code helps to investigate

with model:
    advi = pm.ADVI()
    tracker = pm.callbacks.Tracker(
        mean=advi.approx.mean.eval,  # callable that returns mean
        std=advi.approx.std.eval,  # callable that returns std
    )
    approx = advi.fit(callbacks=[tracker])
fig = plt.figure(figsize=(16, 9))
mu_ax = fig.add_subplot(221)
std_ax = fig.add_subplot(222)
hist_ax = fig.add_subplot(212)
mu_ax.plot(tracker["mean"])
mu_ax.set_title("Mean track")
std_ax.plot(tracker["std"])
std_ax.set_title("Std track")
hist_ax.plot(advi.hist)
hist_ax.set_title("Negative ELBO track");

When I increase learning rate to be 1e-2:

approx = advi.fit(callbacks=[tracker], obj_optimizer=pm.adagrad_window(learning_rate=1e-2))

I get a different picture