def monte_carlo_csiszar_f_divergence( f, p_log_prob, q, num_draws, use_reparametrization=None, seed=None, name=None): """Monte-Carlo approximation of the Csiszar f-Divergence. A Csiszar-function is a member of, ```none F = { f:R_+ to R : f convex }. ``` The Csiszar f-Divergence for Csiszar-function f is given by: ```none D_f[p(X), q(X)] := E_{q(X)}[ f( p(X) / q(X) ) ] ~= m**-1 sum_j^m f( p(x_j) / q(x_j) ), where x_j ~iid q(X) ``` Tricks: Reparameterization and Score-Gradient When q is "reparameterized", i.e., a diffeomorphic transformation of a parameterless distribution (e.g., `Normal(Y; m, s) <=> Y = sX + m, X ~ Normal(0,1)`), we can swap gradient and expectation, i.e., `grad[Avg{ s_i : i=1...n }] = Avg{ grad[s_i] : i=1...n }` where `S_n=Avg{s_i}` and `s_i = f(x_i), x_i ~iid q(X)`. However, if q is not reparameterized, TensorFlow's gradient will be incorrect since the chain-rule stops at samples of unreparameterized distributions. In this circumstance using the Score-Gradient trick results in an unbiased gradient, i.e., ```none grad[ E_q[f(X)] ] = grad[ int dx q(x) f(x) ] = int dx grad[ q(x) f(x) ] = int dx [ q'(x) f(x) + q(x) f'(x) ] = int dx q(x) [q'(x) / q(x) f(x) + f'(x) ] = int dx q(x) grad[ f(x) q(x) / stop_grad[q(x)] ] = E_q[ grad[ f(x) q(x) / stop_grad[q(x)] ] ] ``` Unless `q.reparameterization_type != tfd.FULLY_REPARAMETERIZED` it is usually preferable to set `use_reparametrization = True`. Example Application: The Csiszar f-Divergence is a useful framework for variational inference. I.e., observe that, ```none f(p(x)) = f( E_{q(Z | x)}[ p(x, Z) / q(Z | x) ] ) <= E_{q(Z | x)}[ f( p(x, Z) / q(Z | x) ) ] := D_f[p(x, Z), q(Z | x)] ``` The inequality follows from the fact that the "perspective" of `f`, i.e., `(s, t) |-> t f(s / t))`, is convex in `(s, t)` when `s/t in domain(f)` and `t` is a real. Since the above framework includes the popular Evidence Lower BOund (ELBO) as a special case, i.e., `f(u) = -log(u)`, we call this framework "Evidence Divergence Bound Optimization" (EDBO). Args: f: Python `callable` representing a Csiszar-function in log-space, i.e., takes `p_log_prob(q_samples) - q.log_prob(q_samples)`. p_log_prob: Python `callable` taking (a batch of) samples from `q` and returning the natural-log of the probability under distribution `p`. (In variational inference `p` is the joint distribution.) q: `tf.Distribution`-like instance; must implement: `reparameterization_type`, `sample(n, seed)`, and `log_prob(x)`. (In variational inference `q` is the approximate posterior distribution.) num_draws: Integer scalar number of draws used to approximate the f-Divergence expectation. use_reparametrization: Python `bool`. When `None` (the default), automatically set to: `q.reparameterization_type == tfd.FULLY_REPARAMETERIZED`. When `True` uses the standard Monte-Carlo average. When `False` uses the score-gradient trick. (See above for details.) When `False`, consider using `csiszar_vimco`. seed: Python `int` seed for `q.sample`. name: Python `str` name prefixed to Ops created by this function. Returns: monte_carlo_csiszar_f_divergence: `float`-like `Tensor` Monte Carlo approximation of the Csiszar f-Divergence. Raises: ValueError: if `q` is not a reparameterized distribution and `use_reparametrization = True`. A distribution `q` is said to be "reparameterized" when its samples are generated by transforming the samples of another distribution which does not depend on the parameterization of `q`. This property ensures the gradient (with respect to parameters) is valid. TypeError: if `p_log_prob` is not a Python `callable`. """ with tf.name_scope(name, "monte_carlo_csiszar_f_divergence", [num_draws]): if use_reparametrization is None: use_reparametrization = (q.reparameterization_type == tfd.FULLY_REPARAMETERIZED) elif (use_reparametrization and q.reparameterization_type != tfd.FULLY_REPARAMETERIZED): # TODO(jvdillon): Consider only raising an exception if the gradient is # requested. raise ValueError( "Distribution `q` must be reparameterized, i.e., a diffeomorphic " "transformation of a parameterless distribution. (Otherwise this " "function has a biased gradient.)") if not callable(p_log_prob): raise TypeError("`p_log_prob` must be a Python `callable` function.") return monte_carlo.expectation( f=lambda q_samples: f(p_log_prob(q_samples) - q.log_prob(q_samples)), samples=q.sample(num_draws, seed=seed), log_prob=q.log_prob, # Only used if use_reparametrization=False. use_reparametrization=use_reparametrization)
def monte_carlo_variational_loss(target_log_prob_fn, surrogate_posterior, sample_size=1, importance_sample_size=1, discrepancy_fn=kl_reverse, use_reparameterization=None, seed=None, name=None): """Monte-Carlo approximation of an f-Divergence variational loss. Variational losses measure the divergence between an unnormalized target distribution `p` (provided via `target_log_prob_fn`) and a surrogate distribution `q` (provided as `surrogate_posterior`). When the target distribution is an unnormalized posterior from conditioning a model on data, minimizing the loss with respect to the parameters of `surrogate_posterior` performs approximate posterior inference. This function defines losses of the form `E_q[discrepancy_fn(log(u))]`, where `u = p(z) / q(z)` in the (default) case where `importance_sample_size == 1`, and `u = mean([p(z[k]) / q(z[k]) for k in range(importance_sample_size)]))` more generally. These losses are sometimes known as f-divergences [1, 2]. The default behavior (`discrepancy_fn == tfp.vi.kl_reverse`, where `tfp.vi.kl_reverse = lambda logu: -logu`, and `importance_sample_size == 1`) computes an unbiased estimate of the standard evidence lower bound (ELBO) [3]. The bound may be tightened by setting `importance_sample_size > 1` [4], and the variance of the estimate reduced by setting `sample_size > 1`. Other discrepancies of interest available under `tfp.vi` include the forward `KL[p||q]`, total variation distance, Amari alpha-divergences, and [more]( https://en.wikipedia.org/wiki/F-divergence). Args: target_log_prob_fn: Python callable that takes a set of `Tensor` arguments and returns a `Tensor` log-density. Given `q_sample = surrogate_posterior.sample(sample_size)`, this will be called as `target_log_prob_fn(*q_sample)` if `q_sample` is a list or a tuple, `target_log_prob_fn(**q_sample)` if `q_sample` is a dictionary, or `target_log_prob_fn(q_sample)` if `q_sample` is a `Tensor`. It should support batched evaluation, i.e., should return a result of shape `[sample_size]`. surrogate_posterior: A `tfp.distributions.Distribution` instance defining a variational posterior (could be a `tfd.JointDistribution`). Crucially, the distribution's `log_prob` and (if reparameterizeable) `sample` methods must directly invoke all ops that generate gradients to the underlying variables. One way to ensure this is to use `tfp.util.TransformedVariable` and/or `tfp.util.DeferredTensor` to represent any parameters defined as transformations of unconstrained variables, so that the transformations execute at runtime instead of at distribution creation. sample_size: Integer scalar number of Monte Carlo samples used to approximate the variational divergence. Larger values may stabilize the optimization, but at higher cost per step in time and memory. Default value: `1`. importance_sample_size: Python `int` number of terms used to define an importance-weighted divergence. If `importance_sample_size > 1`, then the `surrogate_posterior` is optimized to function as an importance-sampling proposal distribution. In this case it often makes sense to use importance sampling to approximate posterior expectations (see `tfp.vi.fit_surrogate_posterior` for an example). Default value: `1`. discrepancy_fn: Python `callable` representing a Csiszar `f` function in in log-space. That is, `discrepancy_fn(log(u)) = f(u)`, where `f` is convex in `u`. Default value: `tfp.vi.kl_reverse`. use_reparameterization: Python `bool`. When `None` (the default), automatically set to: `surrogate_posterior.reparameterization_type == tfd.FULLY_REPARAMETERIZED`. When `True` uses the standard Monte-Carlo average. When `False` uses the score-gradient trick. (See above for details.) When `False`, consider using `csiszar_vimco`. seed: PRNG seed for `surrogate_posterior.sample`; see `tfp.random.sanitize_seed` for details. name: Python `str` name prefixed to Ops created by this function. Returns: monte_carlo_variational_loss: `float`-like `Tensor` Monte Carlo approximation of the Csiszar f-Divergence. Raises: ValueError: if `surrogate_posterior` is not a reparameterized distribution and `use_reparameterization = True`. A distribution is said to be "reparameterized" when its samples are generated by transforming the samples of another distribution that does not depend on the first distribution's parameters. This property ensures the gradient with respect to parameters is valid. TypeError: if `target_log_prob_fn` is not a Python `callable`. #### Csiszar f-divergences A Csiszar function `f` is a convex function from `R^+` (the positive reals) to `R`. The Csiszar f-Divergence is given by: ```none D_f[p(X), q(X)] := E_{q(X)}[ f( p(X) / q(X) ) ] ~= m**-1 sum_j^m f( p(x_j) / q(x_j) ), where x_j ~iid q(X) ``` For example, `f = lambda u: -log(u)` recovers `KL[q||p]`, while `f = lambda u: u * log(u)` recovers the forward `KL[p||q]`. These and other functions are available in `tfp.vi`. #### Tricks: Reparameterization and Score-Gradient When q is "reparameterized", i.e., a diffeomorphic transformation of a parameterless distribution (e.g., `Normal(Y; m, s) <=> Y = sX + m, X ~ Normal(0,1)`), we can swap gradient and expectation, i.e., `grad[Avg{ s_i : i=1...n }] = Avg{ grad[s_i] : i=1...n }` where `S_n=Avg{s_i}` and `s_i = f(x_i), x_i ~iid q(X)`. However, if q is not reparameterized, TensorFlow's gradient will be incorrect since the chain-rule stops at samples of unreparameterized distributions. In this circumstance using the Score-Gradient trick results in an unbiased gradient, i.e., ```none grad[ E_q[f(X)] ] = grad[ int dx q(x) f(x) ] = int dx grad[ q(x) f(x) ] = int dx [ q'(x) f(x) + q(x) f'(x) ] = int dx q(x) [q'(x) / q(x) f(x) + f'(x) ] = int dx q(x) grad[ f(x) q(x) / stop_grad[q(x)] ] = E_q[ grad[ f(x) q(x) / stop_grad[q(x)] ] ] ``` Unless `q.reparameterization_type != tfd.FULLY_REPARAMETERIZED` it is usually preferable to set `use_reparameterization = True`. #### Example Application: The Csiszar f-Divergence is a useful framework for variational inference. I.e., observe that, ```none f(p(x)) = f( E_{q(Z | x)}[ p(x, Z) / q(Z | x) ] ) <= E_{q(Z | x)}[ f( p(x, Z) / q(Z | x) ) ] := D_f[p(x, Z), q(Z | x)] ``` The inequality follows from the fact that the "perspective" of `f`, i.e., `(s, t) |-> t f(s / t))`, is convex in `(s, t)` when `s/t in domain(f)` and `t` is a real. Since the above framework includes the popular Evidence Lower BOund (ELBO) as a special case, i.e., `f(u) = -log(u)`, we call this framework "Evidence Divergence Bound Optimization" (EDBO). #### References: [1]: https://en.wikipedia.org/wiki/F-divergence [2]: Ali, Syed Mumtaz, and Samuel D. Silvey. "A general class of coefficients of divergence of one distribution from another." Journal of the Royal Statistical Society: Series B (Methodological) 28.1 (1966): 131-142. [3]: Christopher M. Bishop. Pattern Recognition and Machine Learning. Springer, 2006. [4] Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance Weighted Autoencoders. In _International Conference on Learning Representations_, 2016. https://arxiv.org/abs/1509.00519 """ with tf.name_scope(name or 'monte_carlo_variational_loss'): reparameterization_types = tf.nest.flatten( surrogate_posterior.reparameterization_type) if use_reparameterization is None: use_reparameterization = all( reparameterization_type == FULLY_REPARAMETERIZED for reparameterization_type in reparameterization_types) elif (use_reparameterization and any( reparameterization_type != FULLY_REPARAMETERIZED for reparameterization_type in reparameterization_types)): # TODO(jvdillon): Consider only raising an exception if the gradient is # requested. raise ValueError( 'Distribution `surrogate_posterior` must be reparameterized, i.e.,' 'a diffeomorphic transformation of a parameterless distribution. ' '(Otherwise this function has a biased gradient.)') if not callable(target_log_prob_fn): raise TypeError('`target_log_prob_fn` must be a Python `callable`' 'function.') if use_reparameterization: # Attempt to avoid bijector inverses by computing the surrogate log prob # during the forward sampling pass. q_samples, q_lp = surrogate_posterior.experimental_sample_and_log_prob( [sample_size * importance_sample_size], seed=seed) else: # Score fn objective requires explicit gradients of `log_prob`. q_samples = surrogate_posterior.sample( [sample_size * importance_sample_size], seed=seed) q_lp = None return monte_carlo.expectation( f=_make_importance_weighted_divergence_fn( target_log_prob_fn, surrogate_posterior=surrogate_posterior, discrepancy_fn=discrepancy_fn, precomputed_surrogate_log_prob=q_lp, importance_sample_size=importance_sample_size), samples=q_samples, # Log-prob is only used if use_reparameterization=False. log_prob=surrogate_posterior.log_prob, use_reparameterization=use_reparameterization)
def monte_carlo_csiszar_f_divergence(f, p_log_prob, q, num_draws, use_reparametrization=None, seed=None, name=None): """Monte-Carlo approximation of the Csiszar f-Divergence. A Csiszar-function is a member of, ```none F = { f:R_+ to R : f convex }. ``` The Csiszar f-Divergence for Csiszar-function f is given by: ```none D_f[p(X), q(X)] := E_{q(X)}[ f( p(X) / q(X) ) ] ~= m**-1 sum_j^m f( p(x_j) / q(x_j) ), where x_j ~iid q(X) ``` Tricks: Reparameterization and Score-Gradient When q is "reparameterized", i.e., a diffeomorphic transformation of a parameterless distribution (e.g., `Normal(Y; m, s) <=> Y = sX + m, X ~ Normal(0,1)`), we can swap gradient and expectation, i.e., `grad[Avg{ s_i : i=1...n }] = Avg{ grad[s_i] : i=1...n }` where `S_n=Avg{s_i}` and `s_i = f(x_i), x_i ~iid q(X)`. However, if q is not reparameterized, TensorFlow's gradient will be incorrect since the chain-rule stops at samples of unreparameterized distributions. In this circumstance using the Score-Gradient trick results in an unbiased gradient, i.e., ```none grad[ E_q[f(X)] ] = grad[ int dx q(x) f(x) ] = int dx grad[ q(x) f(x) ] = int dx [ q'(x) f(x) + q(x) f'(x) ] = int dx q(x) [q'(x) / q(x) f(x) + f'(x) ] = int dx q(x) grad[ f(x) q(x) / stop_grad[q(x)] ] = E_q[ grad[ f(x) q(x) / stop_grad[q(x)] ] ] ``` Unless `q.reparameterization_type != tfd.FULLY_REPARAMETERIZED` it is usually preferable to set `use_reparametrization = True`. Example Application: The Csiszar f-Divergence is a useful framework for variational inference. I.e., observe that, ```none f(p(x)) = f( E_{q(Z | x)}[ p(x, Z) / q(Z | x) ] ) <= E_{q(Z | x)}[ f( p(x, Z) / q(Z | x) ) ] := D_f[p(x, Z), q(Z | x)] ``` The inequality follows from the fact that the "perspective" of `f`, i.e., `(s, t) |-> t f(s / t))`, is convex in `(s, t)` when `s/t in domain(f)` and `t` is a real. Since the above framework includes the popular Evidence Lower BOund (ELBO) as a special case, i.e., `f(u) = -log(u)`, we call this framework "Evidence Divergence Bound Optimization" (EDBO). Args: f: Python `callable` representing a Csiszar-function in log-space, i.e., takes `p_log_prob(q_samples) - q.log_prob(q_samples)`. p_log_prob: Python `callable` taking (a batch of) samples from `q` and returning the natural-log of the probability under distribution `p`. (In variational inference `p` is the joint distribution.) q: `tf.Distribution`-like instance; must implement: `reparameterization_type`, `sample(n, seed)`, and `log_prob(x)`. (In variational inference `q` is the approximate posterior distribution.) num_draws: Integer scalar number of draws used to approximate the f-Divergence expectation. use_reparametrization: Python `bool`. When `None` (the default), automatically set to: `q.reparameterization_type == tfd.FULLY_REPARAMETERIZED`. When `True` uses the standard Monte-Carlo average. When `False` uses the score-gradient trick. (See above for details.) When `False`, consider using `csiszar_vimco`. seed: Python `int` seed for `q.sample`. name: Python `str` name prefixed to Ops created by this function. Returns: monte_carlo_csiszar_f_divergence: `float`-like `Tensor` Monte Carlo approximation of the Csiszar f-Divergence. Raises: ValueError: if `q` is not a reparameterized distribution and `use_reparametrization = True`. A distribution `q` is said to be "reparameterized" when its samples are generated by transforming the samples of another distribution which does not depend on the parameterization of `q`. This property ensures the gradient (with respect to parameters) is valid. TypeError: if `p_log_prob` is not a Python `callable`. """ with tf.name_scope(name, "monte_carlo_csiszar_f_divergence", [num_draws]): if use_reparametrization is None: use_reparametrization = ( q.reparameterization_type == tfd.FULLY_REPARAMETERIZED) elif (use_reparametrization and q.reparameterization_type != tfd.FULLY_REPARAMETERIZED): # TODO(jvdillon): Consider only raising an exception if the gradient is # requested. raise ValueError( "Distribution `q` must be reparameterized, i.e., a diffeomorphic " "transformation of a parameterless distribution. (Otherwise this " "function has a biased gradient.)") if not callable(p_log_prob): raise TypeError( "`p_log_prob` must be a Python `callable` function.") return monte_carlo.expectation( f=lambda q_samples: f( p_log_prob(q_samples) - q.log_prob(q_samples)), samples=q.sample(num_draws, seed=seed), log_prob=q.log_prob, # Only used if use_reparametrization=False. use_reparametrization=use_reparametrization)
def monte_carlo_variational_loss(target_log_prob_fn, surrogate_posterior, sample_size=1, importance_sample_size=1, discrepancy_fn=kl_reverse, use_reparameterization=None, gradient_estimator=None, stopped_surrogate_posterior=None, seed=None, name=None): """Monte-Carlo approximation of an f-Divergence variational loss. Variational losses measure the divergence between an unnormalized target distribution `p` (provided via `target_log_prob_fn`) and a surrogate distribution `q` (provided as `surrogate_posterior`). When the target distribution is an unnormalized posterior from conditioning a model on data, minimizing the loss with respect to the parameters of `surrogate_posterior` performs approximate posterior inference. This function defines losses of the form `E_q[discrepancy_fn(log(u))]`, where `u = p(z) / q(z)` in the (default) case where `importance_sample_size == 1`, and `u = mean([p(z[k]) / q(z[k]) for k in range(importance_sample_size)]))` more generally. These losses are sometimes known as f-divergences [1, 2]. The default behavior (`discrepancy_fn == tfp.vi.kl_reverse`, where `tfp.vi.kl_reverse = lambda logu: -logu`, and `importance_sample_size == 1`) computes an unbiased estimate of the standard evidence lower bound (ELBO) [3]. The bound may be tightened by setting `importance_sample_size > 1` [4], and the variance of the estimate reduced by setting `sample_size > 1`. Other discrepancies of interest available under `tfp.vi` include the forward `KL[p||q]`, total variation distance, Amari alpha-divergences, and [more]( https://en.wikipedia.org/wiki/F-divergence). Args: target_log_prob_fn: Python callable that takes a set of `Tensor` arguments and returns a `Tensor` log-density. Given `q_sample = surrogate_posterior.sample(sample_size)`, this will be called as `target_log_prob_fn(*q_sample)` if `q_sample` is a list or a tuple, `target_log_prob_fn(**q_sample)` if `q_sample` is a dictionary, or `target_log_prob_fn(q_sample)` if `q_sample` is a `Tensor`. It should support batched evaluation, i.e., should return a result of shape `[sample_size]`. surrogate_posterior: A `tfp.distributions.Distribution` instance defining a variational posterior (could be a `tfd.JointDistribution`). If using `tf.Variable` parameters, the distribution's `log_prob` and (if reparameterizeable) `sample` methods must directly invoke all ops that generate gradients to the underlying variables. One way to ensure this is to use `tfp.util.TransformedVariable` and/or `tfp.util.DeferredTensor` to represent any parameters defined as transformations of unconstrained variables, so that the transformations execute at runtime instead of at distribution creation. sample_size: Integer scalar number of Monte Carlo samples used to approximate the variational divergence. Larger values may stabilize the optimization, but at higher cost per step in time and memory. Default value: `1`. importance_sample_size: Python `int` number of terms used to define an importance-weighted divergence. If `importance_sample_size > 1`, then the `surrogate_posterior` is optimized to function as an importance-sampling proposal distribution. In this case it often makes sense to use importance sampling to approximate posterior expectations (see `tfp.vi.fit_surrogate_posterior` for an example). Default value: `1`. discrepancy_fn: Python `callable` representing a Csiszar `f` function in in log-space. That is, `discrepancy_fn(log(u)) = f(u)`, where `f` is convex in `u`. Default value: `tfp.vi.kl_reverse`. use_reparameterization: Deprecated; use `gradient_estimator` instead. gradient_estimator: Optional element from `tfp.vi.GradientEstimators` specifying the stochastic gradient estimator to associate with the variational loss. If `None`, a default estimator (either score-function or reparameterization) is chosen based on `surrogate_posterior.reparameterization_type`. Default value: `None`. stopped_surrogate_posterior: Optional copy of `surrogate_posterior` with stopped gradients to the parameters, e.g., `tfd.Normal(loc=tf.stop_gradient(loc), scale=tf.stop_gradient(scale))`. Required if and only if `gradient_estimator == tfp.vi.GradientEstimators.DOUBLY_REPARAMETERIZED`. Default value: `None`. seed: PRNG seed for `surrogate_posterior.sample`; see `tfp.random.sanitize_seed` for details. name: Python `str` name prefixed to Ops created by this function. Returns: monte_carlo_variational_loss: `float`-like `Tensor` Monte Carlo approximation of the Csiszar f-Divergence. Raises: ValueError: if `surrogate_posterior` is not a reparameterized distribution and `use_reparameterization = True`. A distribution is said to be "reparameterized" when its samples are generated by transforming the samples of another distribution that does not depend on the first distribution's parameters. This property ensures the gradient with respect to parameters is valid. TypeError: if `target_log_prob_fn` is not a Python `callable`. #### Csiszar f-divergences A Csiszar function `f` is a convex function from `R^+` (the positive reals) to `R`. The Csiszar f-Divergence is given by: ```none D_f[p(X), q(X)] := E_{q(X)}[ f( p(X) / q(X) ) ] ~= m**-1 sum_j^m f( p(x_j) / q(x_j) ), where x_j ~iid q(X) ``` For example, `f = lambda u: -log(u)` recovers `KL[q||p]`, while `f = lambda u: u * log(u)` recovers the forward `KL[p||q]`. These and other functions are available in `tfp.vi`. #### Example Application: The Csiszar f-Divergence is a useful framework for variational inference. I.e., observe that, ```none f(p(x)) = f( E_{q(Z | x)}[ p(x, Z) / q(Z | x) ] ) <= E_{q(Z | x)}[ f( p(x, Z) / q(Z | x) ) ] := D_f[p(x, Z), q(Z | x)] ``` The inequality follows from the fact that the "perspective" of `f`, i.e., `(s, t) |-> t f(s / t))`, is convex in `(s, t)` when `s/t in domain(f)` and `t` is a real. Since the above framework includes the popular Evidence Lower BOund (ELBO) as a special case, i.e., `f(u) = -log(u)`, we call this framework "Evidence Divergence Bound Optimization" (EDBO). #### References: [1]: https://en.wikipedia.org/wiki/F-divergence [2]: Ali, Syed Mumtaz, and Samuel D. Silvey. "A general class of coefficients of divergence of one distribution from another." Journal of the Royal Statistical Society: Series B (Methodological) 28.1 (1966): 131-142. [3]: Christopher M. Bishop. Pattern Recognition and Machine Learning. Springer, 2006. [4] Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance Weighted Autoencoders. In _International Conference on Learning Representations_, 2016. https://arxiv.org/abs/1509.00519 """ with tf.name_scope(name or 'monte_carlo_variational_loss'): if not callable(target_log_prob_fn): raise TypeError('`target_log_prob_fn` must be a Python `callable`' 'function.') reparameterization_types = tf.nest.flatten( surrogate_posterior.reparameterization_type) if gradient_estimator is None: gradient_estimator = _choose_gradient_estimator( use_reparameterization=use_reparameterization, reparameterization_types=reparameterization_types) if gradient_estimator == GradientEstimators.VIMCO: return csiszar_vimco(f=discrepancy_fn, p_log_prob=target_log_prob_fn, q=surrogate_posterior, num_draws=importance_sample_size, num_batch_draws=sample_size, seed=seed) if gradient_estimator == GradientEstimators.SCORE_FUNCTION: if tf.get_static_value(importance_sample_size) != 1: # TODO(b/213378570): Support score function gradients for # importance-weighted bounds. raise ValueError( 'Score-function gradients are not supported for ' 'losses with `importance_sample_size != 1`.') # Score fn objective requires explicit gradients of `log_prob`. q_samples = surrogate_posterior.sample( [sample_size * importance_sample_size], seed=seed) q_lp = None else: if any(reparameterization_type != FULLY_REPARAMETERIZED for reparameterization_type in reparameterization_types): warnings.warn( 'Reparameterization gradients requested, but ' '`surrogate_posterior.reparameterization_type` is not fully ' 'reparameterized (saw: {}). Gradient estimates may be ' 'biased.'.format( surrogate_posterior.reparameterization_type)) # Attempt to avoid bijector inverses by computing the surrogate log prob # during the forward sampling pass. q_samples, q_lp = surrogate_posterior.experimental_sample_and_log_prob( [sample_size * importance_sample_size], seed=seed) return monte_carlo.expectation( f=_make_importance_weighted_divergence_fn( target_log_prob_fn, surrogate_posterior=surrogate_posterior, discrepancy_fn=discrepancy_fn, precomputed_surrogate_log_prob=q_lp, importance_sample_size=importance_sample_size, gradient_estimator=gradient_estimator, stopped_surrogate_posterior=(stopped_surrogate_posterior)), samples=q_samples, # Log-prob is only used if `gradient_estimator == SCORE_FUNCTION`. log_prob=surrogate_posterior.log_prob, use_reparameterization=(gradient_estimator != GradientEstimators.SCORE_FUNCTION))