Example Page

In this tutorial, I’ll share my top 10 tips for getting started with Academic:

Tip 1

Tip 2

Data Distribution

$$ q(\mathbf{x}) $$

x_train_dataset = make_dataset().shuffle(buffer_size=shuffle_buffer_size) \
                                .batch(batch_size, drop_remainder=True) \
                                .dataset.repeat(num_epochs)
x_train_iterator = x_train_dataset.make_one_shot_iterator()
x_data_sample = x_train_iterator.get_next()

Prior Distribution

$$ p(\mathbf{z}) = \mathcal{N}(\mathbf{0}, \mathbf{I}) $$

prior = tfp.distributions.Independent(
    tfp.distributions.MultivariateNormalDiag(scale_diag=tf.ones(latent_shape)),
    reinterpreted_batch_ndims=2)

Conditional Distributions

Likelihood

$$ p_{\theta}(\mathbf{x} \mid \mathbf{z}) = \mathcal{N}(\mathbf{x} \mid \mu_{\theta}(\mathbf{z}), \mathrm{diag}(\sigma_{\theta}^2(\mathbf{z}))) $$

Variational Posterior

$$ q_{\phi}(\mathbf{z} \mid \mathbf{x}) = \mathcal{N}(\mathbf{z} \mid m_{\phi}(\mathbf{z}), \mathrm{diag}(s_{\phi}^2(\mathbf{x}))) $$

def make_gaussian(fn):
    """
    `p( . |u) = N( . | mu(u), sigma(u)^2 I )`

    where `mu(u), sigma(u) = fn(u)`
    """
    def gaussian(u):

        loc, scale = fn(u)

        dist = tfp.distributions.MultivariateNormalDiag(loc=loc, 
            scale_diag=tfp.trainable_distributions.softplus_and_shift(scale))

        return tfp.distributions.Independent(dist, reinterpreted_batch_ndims=2)

    return gaussian
inference_network = build_residual_network(observed_shape, latent_chns,
                                           filters=filters,
                                           num_blocks=num_blocks,
                                           return_scale=True)

generative_network = build_residual_network(latent_shape, observed_chns,
                                            filters=filters,
                                            num_blocks=num_blocks,
                                            return_scale=True)

posterior = make_gaussian(inference_network)
likelihood = make_gaussian(generative_network)

$$ \mathbf{z}^{(1)}, \dotsc, \mathbf{z}^{(M)} \sim q_{\phi}(\mathbf{z} \mid \mathbf{x}) $$

z_posterior_sample = posterior(x_data_sample).sample(mc_sample_size)

$$ \log p_{\theta}(\mathbf{x}^{(j)} \mid \mathbf{z}^{(i)}) $$

ell_local = likelihood(z_posterior_sample).log_prob(x_data_sample)

$$\log q_{\phi}(\mathbf{z}^{(i)} \mid \mathbf{x}^{(j)}) - \log p(\mathbf{z}^{(i)})$$

kl_local = (posterior(x_data_sample).log_prob(z_posterior_sample) -
            prior.log_prob(z_posterior_sample))

$$ \log \sum_{i=1}^M e^{\log p_{\theta}(\mathbf{x}^{(j)} \mid \mathbf{z}^{(i)}) + \log p(\mathbf{z}^{(i)}) - \log q_{\phi}(\mathbf{z}^{(i)} \mid \mathbf{x}^{(j)})} - \log M = \log \frac{1}{M} \sum_{i=1}^M \frac{p_{\theta}(\mathbf{x}^{(j)}, \mathbf{z}^{(i)})}{q_{\phi}(\mathbf{z}^{(i)} \mid \mathbf{x}^{(j)})} $$

elbo_local = (tf.reduce_logsumexp(ell_local - kl_local, axis=0) -
              tf.log(tf.to_float(mc_sample_size)))

$$ \begin{align} \mathcal{L}_M(\theta, \phi) & := \frac{1}{N} \sum_{j=1}^N \log \frac{1}{M} \sum_{i=1}^M \frac{p_{\theta}(\mathbf{x}^{(j)}, \mathbf{z}^{(i)})}{q_{\phi}(\mathbf{z}^{(i)} \mid \mathbf{x}^{(j)})} \newline & \approx \mathbb{E}_{q(\mathbf{x})} \left [ \log \frac{1}{M} \sum_{i=1}^M \frac{p_{\theta}(\mathbf{x}, \mathbf{z}^{(i)})}{q_{\phi}(\mathbf{z}^{(i)} \mid \mathbf{x})} \right ] \end{align} $$

elbo = tf.reduce_mean(elbo_local)