Problem with KLqp inference for Poisson distribution



I’m a newbie to Edward and Tensorflow so I decided to implement something very simple for a start. For my first exercise I am implementing bayesian inference for a Poisson distribution. This should be easy, right?

Unfortunately I ran into problems using KLqp inference for this problem. I can’t get it to converge to a correct posterior distribution no matter how I tweak the parameters (n_iter, n_samples).

# Generate sample data
N = 100
lamb_true = 1.0
data = np.random.poisson(lam=lamb_true, size=N).astype('float32')

# Setup generative model
alpha_0 = 1.0
beta_0 = 1.0
lamb = Gamma(concentration=tf.constant(alpha_0), rate=tf.constant(beta_0))
x = Poisson(lamb, sample_shape=N)

# Inference using KLqp: NOT WORKING CORRECTLY 
qlamb_alpha = tf.Variable(1.0) 
qlamb_beta = tf.Variable(1.0)
qlamb = Gamma(concentration=qlamb_alpha, rate=qlamb_beta)
inference = ed.KLqp({lamb: qlamb}, data={x: data})
inference.initialize(n_iter=20000, n_samples=5)
for t in xrange(inference.n_iter):

Full code with results is here

In attempt to figure out the problem I have implemented three other ways to estimate posterior distribution and all these methods give very consistent results, but KLqp inference is completely off.

First, the exact Gamma posterior is available for Gamma prior:

qlamb_exact = Gamma(concentration=tf.constant([alpha_0 + np.sum(data)]), rate=tf.constant(beta_0 + len(data)))

Second, I implemented Gibbs sampling with Gamma prior (using ed.Gibbs and the same generative model)

T = 500
x = Poisson(lamb, sample_shape=N)
qlamb2 = Empirical(tf.Variable(tf.zeros([T])))

inference2 = ed.Gibbs({lamb: qlamb2}, data={x: data})

Third is numerical maximization using analytic formula for ELBO using Gamma prior and Gamma posterior (see notebook for details). This one proves that minimizing KL-divergence is a valid method for this problem and the optimization problem is well posed.

def elbo(alpha, beta, alpha_0, beta_0, data):
   log_beta = np.log(beta)
   digamma_alpha = digamma(alpha)
   kl = (alpha - alpha_0) * digamma_alpha
   kl = kl - gammaln(alpha) + gammaln(alpha_0)
   kl = kl + alpha_0 * (log_beta - np.log(beta_0))
   kl = kl + alpha * (beta_0 - beta) / beta
   return - kl - len(data) * alpha / beta + (digamma_alpha - log_beta) * np.sum(data)

opt = minimize(
    fun=lambda x: -elbo(x[0], x[1], alpha_0, beta_0, data),
    bounds=[(1e-4, None), (1.e-4, None)],
    x0=(1.0, 1.0)

Please see full code with results and plots in the notebook.
I would appreciate any help to figure out what am I doing wrong with KLqp inference! Thank you!