Probability: How do I port my Probabilistic MF Edward code to TFP code?

[ankur-manikandan]

I have coded a Probabilistic Matrix Factorization model in Edward. I am trying to port it over to TFP, but I am not sure how to define the log-likelihood and KL divergence terms. Here is the code in Edward -

# MODEL
U = Normal(
    loc=0.0,
    scale=1.0,
    sample_shape=[n_latent_dims, batch_size])
V = Normal(
    loc=0.0,
    scale=1.0,
    sample_shape=[n_latent_dims, n_features])
R = Bernoulli(logits=tf.matmul(tf.transpose(U), V))
R_ph = tf.placeholder(tf.int32, [None, n_features])

# INFERENCE
qU = Normal(
    loc=tf.get_variable("qU/loc",
                        [n_latent_dims, batch_size]),
    scale=tf.nn.softplus(
        tf.get_variable("qU/scale",
                        [n_latent_dims, batch_size])))
qV = Normal(
    loc=tf.get_variable("qV/loc",
                        [n_latent_dims, n_features]),
    scale=tf.nn.softplus(
        tf.get_variable("qV/scale",
                        [n_latent_dims, n_features])))
qR = Bernoulli(logits=tf.matmul(tf.transpose(qU), qV))
qR_avg = Bernoulli(
    logits=tf.reduce_mean(qR.parameters['logits'], 0))
log_likli = tf.reduce_mean(qR_avg.log_prob(R_ph), 1)
inference = ed.KLqp(
    {
        U: qU,
        V: qV
    }, data={self.R: self.R_ph})
inference_init = inference.initialize()
init = tf.global_variables_initializer()

[junpenglao]

See the PCA example https://github.com/tensorflow/probability/blob/master/tensorflow_probability/examples/jupyter_notebooks/Probabilistic_PCA.ipynb

You can also try with the new JointDistributionSeqential class, here is an end-to-end example (in TF2.0):

batch_size = 200
n_latent_dims = 10
n_features = 50
X = (np.random.rand(10000, 50) >= .5).astype(int)

# MODEL
MF = tfd.JointDistributionSequential([
    # U
    tfd.Sample(tfd.Normal(0., 1.), [batch_size, n_latent_dims]),
    # V
    tfd.Sample(tfd.Normal(0., 1.), [n_latent_dims, n_features]),
    # R
    lambda V, U: tfd.Independent(
        tfd.Bernoulli(logits=tf.matmul(U, V)),
        reinterpreted_batch_ndims=2
    )
], validate_args=True)

# Free parameters for Meanfield Approximation
qU_loc = tf.Variable(tf.random.normal([batch_size, n_latent_dims]), name='qU_loc')
qU_rho = tf.Variable(tf.random.normal([batch_size, n_latent_dims]), name='qU_rho')
qV_loc = tf.Variable(tf.random.normal([n_latent_dims, n_features]), name='qV_loc')
qV_rho = tf.Variable(tf.random.normal([n_latent_dims, n_features]), name='qV_rho')

qMF_func = lambda qU_loc, qU_rho, qV_loc, qV_rho: tfd.JointDistributionSequential([
    # qU
    tfd.Independent(tfd.Normal(qU_loc, tf.nn.softplus(qU_rho)), 2),
    # qV
    tfd.Independent(tfd.Normal(qV_loc, tf.nn.softplus(qV_rho)), 2)],
    validate_args=True)

free_param = [qU_loc, qU_rho, qV_loc, qV_rho]

@tf.function
def train(observed):

  def MF_log_prob(q_samples):
    return MF.log_prob(list(q_samples) + [observed])

  def elbo_loss():
    qMF = qMF_func(qU_loc, qU_rho, qV_loc, qV_rho)
    return tfp.vi.monte_carlo_csiszar_f_divergence(
      f=tfp.vi.kl_reverse,  # same as: Evidence Lower BOund
      p_log_prob=MF_log_prob,
      q=qMF,
      num_draws=5)

  opt.minimize(elbo_loss, var_list=free_param)
  return elbo_loss()

opt = tf.optimizers.Adam(.5)
niter = 100

dataset = tf.data.Dataset.from_tensor_slices(X)
dataset = dataset.batch(batch_size, drop_remainder=True)
for step, x in enumerate(dataset):
  for iter_ in range(niter):
    loss = train(x)
  print("step:{:>4}  loss:{:.3f}".format(step, loss))

Maybe @jvdillon and other TFP devs know a way to write it more concise in Keras.

There is also another example of how to do ADVI with JointDistribution* here https://colab.research.google.com/github/tensorflow/probability/blob/master/tensorflow_probability/examples/jupyter_notebooks/Modeling_with_JointDistribution.ipynb

[ankur-manikandan]

@junpenglao This is great! Thank you so much!! I have a follow-up question - Say I have a test dataset and now I want to retrieve the log likelihood of the test dataset, how do I do that? I was able to that using the Probabilistic PCA example and get the values. But, if I was to leverage the TF2.0 code you wrote above how would I go about it? I'm sorry if my question is trivial. I am new to this and still figuring it out.

[junpenglao]

Glad you find it helpful!

Say I have a test dataset and now I want to retrieve the log likelihood of the test dataset, how do I do that?

I assume that you meant the log-likelihood of the test dataset _conditioned on_ the fitted parameters. I would do:

MF.log_prob([qU_loc, qV_loc, test_dataset])

which conditioned on the (approximated) posterior mean of U and V.

[ankur-manikandan]

Great, that works! Thanks a lot @junpenglao! Really appreciate your help with this :)

[junpenglao]

You are welcome! I will close this for now but feel free to keep discussing here if you have any follow up question.

[jan-glx]

Thanks for the the TF2.0 example, @junpenglao! To make it run with current beta/nightly () I did need to add an asterix in MF_log_prob and was advised by a deprecation warning to replace montecarlo_csiszar_f_divergence with monte_carlo_variational_loss. These changes result in:

import tensorflow as tf
from tensorflow_probability import distributions as tfd
import tensorflow_probability as tfp
import numpy as np

n_epochs = 3
train_data_size = 10000
batch_size = 200
n_latent_dims = 10
n_features = 50

V = tfd.Sample(tfd.Normal(0., 1.), [n_latent_dims, n_features])


# MODEL
def model(n_observations):
    return tfd.JointDistributionSequential([
        # U
        tfd.Sample(tfd.Normal(0., 1.), [n_observations, n_latent_dims]),
        # V
        V,
        # R
        lambda V, U: tfd.Independent(
            tfd.Bernoulli(logits=tf.matmul(U, V)),
            reinterpreted_batch_ndims=2
        )
    ], validate_args=True)


U_data, V_data, X_data = model(train_data_size).sample()
MF = model(batch_size)


# Free parameters for Meanfield Approximation
qU_loc = tf.Variable(tf.random.normal([batch_size, n_latent_dims]), name='qU_loc')
qU_rho = tf.Variable(tf.random.normal([batch_size, n_latent_dims]), name='qU_rho')
qV_loc = tf.Variable(tf.random.normal([n_latent_dims, n_features]), name='qV_loc')
qV_rho = tf.Variable(tf.random.normal([n_latent_dims, n_features]), name='qV_rho')

qMF_func = lambda qU_loc, qU_rho, qV_loc, qV_rho: tfd.JointDistributionSequential([
    # qU
    tfd.Independent(tfd.Normal(qU_loc, tf.nn.softplus(qU_rho)), 2),
    # qV
    tfd.Independent(tfd.Normal(qV_loc, tf.nn.softplus(qV_rho)), 2)],
    validate_args=True)

free_param = [qU_loc, qU_rho, qV_loc, qV_rho]


@tf.function
def train(observed):
    def MF_log_prob(*q_samples):
        return MF.log_prob(list(q_samples) + [observed])

    def elbo_loss():
        qMF = qMF_func(qU_loc, qU_rho, qV_loc, qV_rho)
        return tfp.vi.monte_carlo_variational_loss(
            target_log_prob_fn=MF_log_prob,
            surrogate_posterior=qMF,
            sample_size=5)

    opt.minimize(elbo_loss, var_list=free_param)
    return elbo_loss()


opt = tf.optimizers.Adam(.5)
niter = 100

dataset = tf.data.Dataset.from_tensor_slices(X_data)
dataset = dataset.repeat(n_epochs)
dataset = dataset.shuffle(train_data_size*2)
dataset = dataset.batch(batch_size, drop_remainder=True)
for step, x in enumerate(dataset):
    loss = train(x)
    print("step:{:>4}  loss:{:.3f}".format(step, loss))

Furthermore, sampling here uses the model instead of random data such that actually something useful can be learned.