Objectives¶

Variational objectives: ELBO, IWAEBound, RenyiBound, VRIWAEBound, ChiVI, RWS, and DReGsBound, paired with the gradient-estimator strategies in Estimators.

objectives ¶

Variational objectives.

An Objective is a torch.nn.Module-callable that, given (model, guide, x, observations), returns a scalar loss whose backward() produces a gradient on the model and guide parameters. The most common objective is ELBO; tighter multi-sample bounds (IWAEBound, RenyiBound, VRIWAEBound, DReGsBound) trade compute for bound-tightness. ChiVI gives an upper bound on the marginal likelihood via the chi-squared divergence, and RWS implements reweighted wake-sleep for inference nets over discrete latents.

Every objective accepts a quivers.inference.estimators.GradientEstimator strategy that decides how the per-particle log-density tensors are turned into a scalar loss whose gradient is the chosen estimator. The default is quivers.inference.estimators.Reparameterized.

Per-particle log-densities are stacked along a leading torch axis of shape (K, batch) before the estimator turns them into a scalar loss. The particle draws themselves run in a Python for loop over num_particles, since the model's runtime path and the underlying torch.distributions calls are not vectorized over a Monte Carlo dimension; the per-step cost scales linearly in K.

References¶

Standard ELBO: Kingma and Welling 2013.
IWAE: Burda, Grosse and Salakhutdinov 2016.
Rényi divergence VI: Li and Turner 2016.
VR-IWAE: Daudel, Douc and Roueff 2023.
CUBO / chi-square VI: Dieng, Tran, Ranganath, Paisley and Blei 2017.
Reweighted wake-sleep: Bornschein and Bengio 2015.
DReG surrogate: Tucker, Lawson, Gu and Maddison 2019.

Objective ¶

Objective(estimator: GradientEstimator | None = None)

Bases: Module, ABC

Base class for variational objectives.

Subclasses implement forward to return a scalar loss (negated objective). The estimator attribute is the gradient-estimation strategy applied to the per-particle log-densities.

Source code in src/quivers/inference/objectives.py

def __init__(self, estimator: GradientEstimator | None = None) -> None:
    super().__init__()
    self.estimator = estimator if estimator is not None else Reparameterized()

ELBO ¶

ELBO(num_particles: int = 1, estimator: GradientEstimator | None = None)

Bases: Objective

Evidence lower bound objective.

.. math::

\mathcal{L}_{\mathrm{ELBO}}
    = \mathbb{E}_{q_\phi(z)} \bigl[ \log p(z, y) - \log q_\phi(z) \bigr].

Returns the negated ELBO so Objective.forward can be plugged into a minimizer. num_particles averages independent Monte-Carlo estimates; num_particles == 1 is the standard reparameterization-trick ELBO.

PARAMETER	DESCRIPTION
`num_particles`	Number of independent guide samples per step. Default `1`. TYPE: `int` DEFAULT: `1`
`estimator`	Gradient-estimator strategy. Default `Reparameterized`. TYPE: `GradientEstimator` DEFAULT: `None`

Source code in src/quivers/inference/objectives.py

def __init__(
    self,
    num_particles: int = 1,
    estimator: GradientEstimator | None = None,
) -> None:
    super().__init__(estimator=estimator)
    if num_particles < 1:
        raise ValueError(f"ELBO: num_particles must be >= 1, got {num_particles}")
    self.num_particles = num_particles

IWAEBound ¶

IWAEBound(num_particles: int = 8, estimator: GradientEstimator | None = None)

Bases: Objective

Importance-weighted bound (Burda-Grosse-Salakhutdinov 2016).

.. math::

\mathcal{L}_{\mathrm{IWAE}}
    = \mathbb{E}\Bigl[\log \frac{1}{K} \sum_{k=1}^{K}
        \frac{p(z_k, y)}{q_\phi(z_k)}\Bigr],

a tighter lower bound on :math:\log p(y) than the ELBO. Approaches the marginal likelihood as :math:K \to \infty.

The default estimator is DoublyReparameterized because the naive reparameterized gradient's signal-to-noise ratio for the inference network collapses as :math:K grows (Tucker-Lawson-Gu-Maddison 2019).

Source code in src/quivers/inference/objectives.py

def __init__(
    self,
    num_particles: int = 8,
    estimator: GradientEstimator | None = None,
) -> None:
    if estimator is None:
        estimator = DoublyReparameterized()
    super().__init__(estimator=estimator)
    if num_particles < 1:
        raise ValueError(
            f"IWAEBound: num_particles must be >= 1, got {num_particles}"
        )
    self.num_particles = num_particles

RenyiBound ¶

RenyiBound(alpha: float = 0.5, num_particles: int = 8, estimator: GradientEstimator | None = None)

Bases: Objective

Rényi α-divergence variational bound (Li-Turner 2016).

.. math::

\mathcal{L}_\alpha = \frac{1}{1 - \alpha}
    \log \mathbb{E}_q\Bigl[ \bigl(p(z, y) / q_\phi(z)\bigr)^{1-\alpha}\Bigr].

Recovers the ELBO at :math:\alpha = 1 (in the limit) and the IWAE bound at :math:\alpha = 0. The interesting regime is :math:\alpha < 0, which gives an upper bound on :math:\log p(y) and so a tighter posterior-mode estimate when the variational family is too narrow.

PARAMETER	DESCRIPTION
`alpha`	Divergence order. `alpha != 1`; values close to 1 may be numerically unstable. TYPE: `float` DEFAULT: `0.5`
`num_particles`	Number of guide samples per step. TYPE: `int` DEFAULT: `8`

Source code in src/quivers/inference/objectives.py

def __init__(
    self,
    alpha: float = 0.5,
    num_particles: int = 8,
    estimator: GradientEstimator | None = None,
) -> None:
    super().__init__(estimator=estimator)
    if alpha == 1.0:
        raise ValueError(
            "RenyiBound: alpha == 1.0 recovers the ELBO in the "
            "limit but is numerically singular here. Use the "
            "ELBO objective instead."
        )
    if num_particles < 1:
        raise ValueError(
            f"RenyiBound: num_particles must be >= 1, got {num_particles}"
        )
    self.alpha = alpha
    self.num_particles = num_particles

VRIWAEBound ¶

VRIWAEBound(alpha: float = 0.0, num_particles: int = 8, estimator: GradientEstimator | None = None)

Bases: Objective

Variational Rényi-IWAE bound (Daudel-Douc-Roueff 2023).

Unifies ELBO, IWAEBound, and RenyiBound into a single bound parameterized by alpha and num_particles:

.. math::

\mathcal{L}_{\mathrm{VR\text{-}IWAE}}
    = \frac{1}{1 - \alpha} \,\log\,
      \frac{1}{K} \sum_{k=1}^{K} \Bigl(\frac{p}{q}\Bigr)^{1-\alpha}.

Special cases:

alpha = 0, K > 1 → IWAE bound.
alpha = 0, K = 1 → ELBO.
alpha != 0, K = 1 → Rényi α-VI.

For intermediate alpha the bound interpolates between "cheap, biased" (high α) and "expensive, tight" (low α).

Source code in src/quivers/inference/objectives.py

def __init__(
    self,
    alpha: float = 0.0,
    num_particles: int = 8,
    estimator: GradientEstimator | None = None,
) -> None:
    super().__init__(estimator=estimator)
    if alpha == 1.0:
        raise ValueError("VRIWAEBound: alpha == 1.0 is singular. Use ELBO instead.")
    if num_particles < 1:
        raise ValueError(
            f"VRIWAEBound: num_particles must be >= 1, got {num_particles}"
        )
    self.alpha = alpha
    self.num_particles = num_particles

ChiVI ¶

ChiVI(n: float = 2.0, num_particles: int = 8, estimator: GradientEstimator | None = None)

Bases: Objective

Chi-squared variational upper bound (Dieng et al. 2017).

Minimises the chi-squared divergence :math:\chi^2(p \| q_\phi) between the true posterior and the guide via the CUBO surrogate

.. math::

\mathcal{L}_{n\text{-CUBO}}(\phi)
    = \tfrac{1}{n} \log \mathbb{E}_{q_\phi}
      \Bigl[\bigl(p(z, y) / q_\phi(z)\bigr)^{n}\Bigr]
    \geq \log p(y).

In contrast to the ELBO (a lower bound), this is an upper bound on :math:\log p(y), so minimising it drives the guide to over-cover the true posterior. Useful for posterior calibration and for sandwich estimates when paired with ELBO.

Reference: Dieng, Tran, Ranganath, Paisley and Blei 2017.

PARAMETER	DESCRIPTION
`n`	Divergence order (`n = 2` for standard chi-squared). Must be `> 0`. TYPE: `float` DEFAULT: `2.0`
`num_particles`	Number of guide samples per step. TYPE: `int` DEFAULT: `8`
`estimator`	Gradient-estimator strategy. Default `Reparameterized`. TYPE: `GradientEstimator` DEFAULT: `None`

Source code in src/quivers/inference/objectives.py

def __init__(
    self,
    n: float = 2.0,
    num_particles: int = 8,
    estimator: GradientEstimator | None = None,
) -> None:
    super().__init__(estimator=estimator)
    if n <= 0.0:
        raise ValueError(f"ChiVI: n must be > 0, got {n}")
    if num_particles < 1:
        raise ValueError(f"ChiVI: num_particles must be >= 1, got {num_particles}")
    self.n = float(n)
    self.num_particles = num_particles

RWS ¶

RWS(num_particles: int = 8, estimator: GradientEstimator | None = None)

Bases: Objective

Reweighted wake-sleep (Bornschein and Bengio 2015).

Combines a wake-phase gradient on the model (importance-weighted log-likelihood) with a sleep-phase gradient on the guide (KL from the self-normalised importance-weighted posterior). RWS handles discrete latents where the reparameterization trick does not apply and where score-function estimators are high-variance.

The loss returned combines both phases:

.. math::

\mathcal{L}_{\mathrm{RWS}}
    = -\mathbb{E}_{q}\bigl[\tilde w_k \log p(z_k, y)\bigr]
      -\mathbb{E}_{q}\bigl[\tilde w_k \log q_\phi(z_k)\bigr],

where :math:\tilde w_k are self-normalised importance weights :math:w_k / \sum_j w_j with :math:w_k = p(z_k, y) / q(z_k), detached from the gradient path so both phases are unbiased.

Reference: Bornschein and Bengio 2015.

PARAMETER	DESCRIPTION
`num_particles`	Number of guide samples per step. TYPE: `int` DEFAULT: `8`
`estimator`	Gradient-estimator strategy. Default `Reparameterized`; the RWS surrogate is estimator-agnostic because the weights are detached. TYPE: `GradientEstimator` DEFAULT: `None`

Source code in src/quivers/inference/objectives.py

def __init__(
    self,
    num_particles: int = 8,
    estimator: GradientEstimator | None = None,
) -> None:
    super().__init__(estimator=estimator)
    if num_particles < 1:
        raise ValueError(f"RWS: num_particles must be >= 1, got {num_particles}")
    self.num_particles = num_particles

DReGsBound ¶

DReGsBound(num_particles: int = 8)

Bases: Objective

Doubly reparameterised IWAE bound (Tucker et al. 2019).

The bound is the IWAE bound; the estimator is the DReG surrogate that removes the score-function term whose signal- to-noise ratio collapses as :math:K grows in the naive reparameterised IWAE gradient. Equivalent to IWAEBound(num_particles=K, estimator=DoublyReparameterized()) but exposed as a first-class objective for callers that want to switch bound + estimator together.

Reference: Tucker, Lawson, Gu and Maddison 2019.

Note: DReGsBound here is a bound + estimator pair; the scalar gradient-estimator strategy also called "DReG" lives at quivers.inference.estimators.DoublyReparameterized and can be attached to any objective through the estimator argument.

PARAMETER	DESCRIPTION
`num_particles`	Number of guide samples per step. DReG's variance- reduction benefit is largest for `K >= 8`. TYPE: `int` DEFAULT: `8`

Source code in src/quivers/inference/objectives.py

def __init__(self, num_particles: int = 8) -> None:
    super().__init__(estimator=DoublyReparameterized())
    if num_particles < 1:
        raise ValueError(
            f"DReGsBound: num_particles must be >= 1, got {num_particles}"
        )
    self.num_particles = num_particles