zhusuan.variational

ELBO

class ELBO(generator, variational, estimator='sgvb', transform=None, transform_var=[], auxillary_var=[])

Bases: torch.nn.modules.module.Module

The class that represents the evidence lower bound (ELBO) objective for variational inference. It can be constructed like a Jittor’s Module by passing 2 BayesianNet instances. For example, the generator network and the variational inference network in VAE. The model can calculate the ELBO’s value with observations passed.

See also

For more details and examples, please refer to Variational Autoencoders and Bayesian Neural Networks

Parameters

• generator – A :class’~zhusuan.framework.BayesianNet` instance or a log joint probability function. For the latter, it must accepts a dictionary argument of (string, Tensor) pairs, which are mappings from all node names in the model to their observed values. The function should return a Tensor, representing the log joint likelihood of the model.

• variational – A BayesianNet instance that defines the variational family.

• estimator – gradient estimate method, including sgvb and reinforce

• transform – A RevNet instance that transform Specified variables,

returns the transformed variable and the log_det_J i.e log-determinant of transition Jacobian matrix :param transform_var: a list of names of variable to be transformed, all tensor that correspond to these names will be placed into tuple by order and feed to the transform network :param auxillary_var: auxillary variable name list that need to be passed to transform network

forward(observed, reduce_mean=True, **kwargs)

observe nodes, transform latent variables, return evidence lower bound :return: evidence lower bound

log_joint(nodes)

The default log joint probability function. It works by summing over all the conditional log probabilities of stochastic nodes evaluated at their current values (samples or observations).

Returns

A Var.

reinforce(logpxz, logqz, reduce_mean=True, baseline=None, variance_reduction=True, decay=0.8)

Implements the score function gradient estimator for the ELBO, with optional variance reduction using moving mean estimate or “baseline”. Also known as “REINFORCE” (Williams, 1992), “NVIL” (Mnih, 2014), and “likelihood-ratio estimator” (Glynn, 1990).

It works for all types of latent StochasticTensor s.

Note

To use the reinforce() estimator, the is_reparameterized property of each reparameterizable latent StochasticTensor must be set False.

Parameters

• logpxz – log joint of generator

• logqz – log joint of variational

• reduce_mean – whether reduce to a scalar by mean operation

• baseline – A Tensor that can broadcast to match the shape returned by log_joint. A trainable estimation for the scale of the elbo value, which is typically dependent on observed values, e.g., a neural network with observed values as inputs. This will be additional.

• variance_reduction – Bool. Whether to use variance reduction. By default will subtract the learning signal with a moving mean estimation of it. Users can pass an additional customized baseline using the baseline argument, in that way the returned will be a tuple of costs, the former for the gradient estimator, the latter for adapting the baseline.

• decay – Float. The moving average decay for variance normalization.

Returns

A Tensor. The surrogate cost for optimizers to minimize.

sgvb(logpxz, logqz, reduce_mean=True, log_det=None)

Implements the stochastic gradient variational bayes (SGVB) gradient estimator for the objective, also known as “reparameterization trick” or “path derivative estimator”. It was first used for importance weighted objectives in (Burda, 2015), where it’s named “IWAE”.

It only works for latent StochasticTensor s that can be reparameterized (Kingma, 2013). For example, Normal and Concrete.

Note

To use the sgvb() estimator, the is_reparameterized property of each latent StochasticTensor must be True (which is the default setting when they are constructed).

Returns

A Tensor. The surrogate cost for optimizers to minimize.

training: bool

ImportanceWeightedObjective

class ImportanceWeightedObjective(generator, variational, axis=None, estimator='sgvb')

Bases: torch.nn.modules.module.Module

The class that represents the importance weighted objective for variational inference (Burda, 2015)

As a variational objective, ImportanceWeightedObjective provides two gradient estimators for the variational (proposal) parameters:

• sgvb(): The Stochastic Gradient Variational Bayes (SGVB) estimator, also known as “the reparameterization trick”, or “path derivative estimator”.

• vimco(): The multi-sample score function estimator with variance reduction, also known as “VIMCO”.

Parameters

• generator – generator part of importance weighted objective

• variational – variational part of importance weighted objective

• axis – The sample dimension(s) to reduce when computing the outer expectation in the objective.

• estimator – the estimator, a str in either ‘sgvb’ or ‘vimco’

forward(observed, reduce_mean=True)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

log_joint(nodes)

sgvb(logpxz, logqz, reduce_mean=True)

Implements the stochastic gradient variational bayes (SGVB) gradient estimator for the objective, also known as “reparameterization trick” or “path derivative estimator”. It was first used for importance weighted objectives in (Burda, 2015), where it’s named “IWAE”.

It only works for latent StochasticTensor s that can be reparameterized (Kingma, 2013). For example, Normal and Concrete.

Note

To use the sgvb() estimator, the is_reparameterized property of each latent StochasticTensor must be True (which is the default setting when they are constructed).

Returns

A Tensor. The surrogate cost for optimizers to minimize.

training: bool

vimco(logpxz, logqz, reduce_mean=True)

Implements the multi-sample score function gradient estimator for the objective, also known as “VIMCO”, which is named by authors of the original paper (Minh, 2016).

It works for all kinds of latent StochasticTensor s.

Note

To use the vimco() estimator, the is_reparameterized property of each reparameterizable latent StochasticTensor must be set False.

Returns

A Tensor. The surrogate cost for optimizers to minimize.