SoftCVI: Contrastive variational inference with self-generated soft labels

## Abstract

Abstract Estimating a distribution given access to its unnormalized density is pivotal in Bayesian inference, where the posterior is generally known only up to an unknown normalizing constant. Variational inference and Markov chain Monte Carlo methods are the predominant tools for this task; however, both are often challenging to apply reliably, particularly when the posterior has complex geometry. Here, we introduce Soft Contrastive Variational Inference (SoftCVI) , which allows a family of variational objectives to be derived through a contrastive estimation framework. The approach parameterizes a classifier in terms of a variational distribution, reframing the inference task as a contrastive estimation problem aiming to identify a single true posterior sample among a set of samples. Despite this framing, we do not require positive or negative samples, but rather learn by sampling the variational distribution and computing ground truth soft classification labels from the unnormalized posterior itself. The objectives have zero variance gradient when the variational approximation is exact, without the need for specialized gradient estimators. We empirically investigate the performance on a variety of Bayesian inference tasks, using both simple (e.g. normal) and expressive (normalizing flow) variational distributions. We find that SoftCVI can be used to form objectives which are stable to train and mass-covering, frequently outperforming inference with other variational approaches.

## 1 Introduction

Consider a probabilistic model with a set of latent variables $\bm{\theta}$, for which given a set of observations ${\bm{x}_{\text{obs}}}$ we wish to infer a posterior $p(\bm{\theta}|{\bm{x}_{\text{obs}}})$. Unless the model takes a particularly convenient form, the posterior cannot be directly computed. However, typically $p(\bm{\theta},{\bm{x}_{\text{obs}}})$ is available, related to the posterior by $p(\bm{\theta}|{\bm{x}_{\text{obs}}})=p(\bm{\theta},{\bm{x}_{\text{obs}}})/p({
\bm{x}_{\text{obs}}})$, where $p({\bm{x}_{\text{obs}}})$ is an intractable normalizing constant $p({\bm{x}_{\text{obs}}})=\int p(\bm{\theta},{\bm{x}_{\text{obs}}})d\bm{\theta}$. In these cases, computational methods such as Markov chain Monte Carlo (MCMC) Metropolis et al. (); Hastings () or variational inference Jordan et al. (); Kingma & Welling () are required to perform inference.

Variational inference approaches the inference task as an optimization problem by defining a variational distribution $q_{\bm{\phi}}(\bm{\theta})$ and minimizing its divergence to the true posterior $p(\bm{\theta}|{\bm{x}_{\text{obs}}})$. The performance and reliability of variational inference is dependent on numerous factors, including the choice of divergence, variational distribution and parameter initialization. Whilst choosing a divergence that favors mass-covering posterior estimates may facilitate reliable inference, in many cases these divergences introduce bias or are less stable to train, leading to worse performance Dhaka et al. (); Naesseth et al. (). Alongside difficulties in assessing performance Yao et al. (), these issues hinder practical applications of variational inference, particularly in applications such as scientific research, where accurate uncertainty quantification is crucial.

In contrast to variational inference, contrastive learning is generally used to perform inference when the likelihood function is only available through an intractable integral, such as for fitting energy-based models or for performing simulation-based inference Gutmann et al. (). Learning is achieved by contrasting positive samples with negative samples, where the negative samples are often generated by augmenting positive samples or by drawing samples from a carefully chosen (possibly highly structured) noise distribution. A common theme is to leverage the invariance of classification-based objectives to unknown normalizing constants to enable learning where likelihood-based methods are infeasible.

Contribution
We show that contrastive estimation can be used to derive a family of variational objectives, terming the approach Soft Contrastive Variational Inference (SoftCVI). The task of fitting the posterior approximation is reframed as a classification task, aiming to identify a single true posterior sample among a set of samples. Instead of using explicitly positive and negative samples, we show that for arbitrary samples from a proposal distribution, we can generate ground truth soft classification labels using the unnormalized posterior density itself. The samples and corresponding labels are used for fitting a classifier parameterized in terms of the variational distribution, such that the optimal classifier recovers the true posterior. We find SoftCVI enables derivation of stable and mass-covering objectives, and demonstrate the performance across a series of experiments with both simple (normal) and flexible (normalizing flow) variational distributions. Compared to alternative variational objectives, we find SoftCVI generally yields better calibrated posteriors with a lower forward Kullback-Leibler (KL) divergence to the true posterior. We provide a pair of Python packages, [softcvi](https://github.com/danielward27/softcvi) and [softcvi_validation](https://github.com/danielward27/softcvi_validation), which provide the implementation, and the code for reproducing the results of this paper, respectively.

## 2 SoftCVI

In order to fit a variational distribution with SoftCVI, we must define a proposal distribution $\pi(\bm{\theta})$, a negative distribution $p^{-}(\bm{\theta})$, and the variational distribution itself, $q_{\bm{\phi}}(\bm{\theta})$. At each optimization step, three steps are performed which allow fitting the variational distribution:

- 1.
Sample latent variables $\{\bm{\theta}_{k}\}_{k=1}^{K}\sim\pi(\bm{\theta})$ from the proposal distribution $\pi(\bm{\theta})$.
- 2.
Generate corresponding ground truth soft labels $\bm{y}\in(0,1)^{K}$ for the task of classifying between positive and negative samples, presumed to be from $p(\bm{\theta}|{\bm{x}_{\text{obs}}})$ and $p^{-}(\bm{\theta})$, respectively.
- 3.
Use $\{\bm{\theta}_{k}\}_{k=1}^{K}$ along with the soft labels $\bm{y}$ to optimize a classifier parameterized in terms of the variational distribution $q_{\bm{\phi}}(\bm{\theta})$, such that the optimal classifier recovers the true posterior.

The choice of proposal distribution in step one will influence the region in which learning is focused. Throughout the experiments here, we take the intuitive and convenient choice of using the variational distribution itself as the proposal distribution $\pi(\bm{\theta})=q_{\bm{\phi}}(\bm{\theta})$, which over the course of training directs learning towards regions with reasonable posterior mass. The remainder of this section is structured as follows: section 2.1 details the assignment of ground truth labels to arbitrary proposal samples; section 2.2 discusses the parameterization and optimization of the classifier; finally, section 2.3 considers the choice of negative distribution.

### 2.1 Generating ground truth soft labels

Consider for simplicity the case that we have access to positive (true posterior) and negative samples. Specifically, consider the scenario in which we have $K$ samples $\{\bm{\theta}_{k}\}_{k=1}^{K}$, consisting of exactly one positive sample from the true posterior $p(\bm{\theta}|{\bm{x}_{\text{obs}}})$, and $K-1$ negative samples from a negative distribution $p^{-}(\bm{\theta})$. The optimal classifier can be parameterized in terms of the ratio between the positive and negative densities $p(\bm{\theta}|{\bm{x}_{\text{obs}}})/p^{-}(\bm{\theta})$ (e.g Gutmann & Hyvärinen, ; Oord et al., ; Thomas et al., ). Letting $k^{*}$ denote the index of the true sample, and assuming each element $\bm{\theta}_{k}$ is equally likely to be the true sample a priori, the optimal classifier is

$$ $\displaystyle y_{k}$ $\displaystyle=P(k^{*}=k|\{\bm{\theta}_{k}\}_{k=1}^{K},{\bm{x}_{\text{obs}}})= \frac{p(\bm{\theta}_{k}|{\bm{x}_{\text{obs}}})/p^{-}(\bm{\theta}_{k})}{\sum_{k ^{\prime}=1}^{K}p(\bm{\theta}_{k^{\prime}}|{\bm{x}_{\text{obs}}})/p^{-}(\bm{ \theta}_{k^{\prime}})},$ (1) $$

where $y_{k}=P(k^{*}=k|\{\bm{\theta}_{k}\}_{k=1}^{K},{\bm{x}_{\text{obs}}})$ is the probability on the interval $(0,1)$ that $\bm{\theta}_{k}$ is the positive sample $\bm{\theta}_{k^{*}}$ among all the considered samples, and $\sum_{k=1}^{K}y_{k}=1$. The above computation is invariant to multiplicative constants to the ratio $p(\bm{\theta}|{\bm{x}_{\text{obs}}})/p^{-}(\bm{\theta})$, meaning that access only to unnormalized versions of $p(\bm{\theta}|{\bm{x}_{\text{obs}}})$ and $p^{-}(\bm{\theta})$ is sufficient for computing the labels. Since typically $p(\bm{\theta},{\bm{x}_{\text{obs}}})$ is known and proportional to $p(\bm{\theta}|{\bm{x}_{\text{obs}}})$, we can compute the labels as

$$ $\displaystyle y_{k}$ $\displaystyle=\frac{p(\bm{\theta}_{k},{\bm{x}_{\text{obs}}})/p^{-}(\bm{\theta} _{k})}{\sum_{k^{\prime}=1}^{K}p(\bm{\theta}_{k^{\prime}},{\bm{x}_{\text{obs}}} )/p^{-}(\bm{\theta}_{k^{\prime}})},$ (2) $\displaystyle\bm{y}$ $\displaystyle=\operatorname{softmax}(\bm{z}),\quad\text{ where }z_{k}=\log \frac{p(\bm{\theta}_{k},{\bm{x}_{\text{obs}}})}{p^{-}(\bm{\theta}_{k})},$ $$

and $p^{-}(\bm{\theta})$ may be unnormalized. Intuitively, the ratio $p(\bm{\theta}_{k},{\bm{x}_{\text{obs}}})/p^{-}(\bm{\theta}_{k})$ encodes the relative likelihood of a sample $\bm{\theta}_{k}$ being the true sample, and to compute the labels the ratios are normalized over the set of considered samples.

While contrastive estimation methods generally use explicitly positive and negative examples with hard labels $\bm{y}\in\{0,1\}^{K}$, here we lack access to positive samples from $p(\bm{\theta}|{\bm{x}_{\text{obs}}})$, and we may not have access to samples from $p^{-}(\bm{\theta})$. Nonetheless, despite the framing of the classification task, we can instead generate samples from a proposal distribution, $\{\bm{\theta}_{k}\}_{k=1}^{K}\sim\pi(\bm{\theta})$, and by assigning labels to these samples it is still possible to train a classifier which at optimality recovers the true posterior, as we will discuss in the subsequent section.

### 2.2 Parameterization and optimization of the classifier

Analogously to the computation of the labels in eq. 1, we can parameterize the classifier in terms of the ratio between the variational and negative distribution

$$ $\displaystyle\hat{y}_{k}$ $\displaystyle=\frac{q_{\bm{\phi}}(\bm{\theta}_{k})/p^{-}(\bm{\theta}_{k})}{ \sum_{k^{\prime}=1}^{K}q_{\bm{\phi}}(\bm{\theta}_{k^{\prime}})/p^{-}(\bm{ \theta}_{k^{\prime}})},$ (3) $\displaystyle\hat{\bm{y}}$ $\displaystyle=\operatorname{Softmax}(\hat{\bm{z}}),\quad\text{where }\hat{z}_{ k}=\log\frac{q_{\bm{\phi}}(\bm{\theta}_{k})}{p^{-}(\bm{\theta}_{k})}.$ $$

Given a set of samples $\{\bm{\theta}_{k}\}_{k=1}^{K}\sim\pi(\bm{\theta})$ and corresponding labels $\bm{y}$, an optimization step can be taken to minimize the softmax cross-entropy loss function (i.e. the negative multinomial log-likelihood)

$$ $\displaystyle\mathcal{L}_{\text{SoftCVI}}(\bm{\phi};\{\bm{\theta}_{k}\}_{k=1}^ {K},\bm{y})=-\sum_{k=1}^{K}y_{k}\log\left(\hat{y}_{k}\right)=-\sum_{k=1}^{K}y_ {k}\log\left(\frac{q_{\bm{\phi}}(\bm{\theta}_{k})/p^{-}(\bm{\theta}_{k})}{\sum _{k^{\prime}=1}^{K}q_{\bm{\phi}}(\bm{\theta}_{k^{\prime}})/p^{-}(\bm{\theta}_{ k^{\prime}})}\right).$ (4) $$

Let $\Theta$ represent the support of the proposal distribution, with $p(\bm{\theta}|{\bm{x}_{\text{obs}}})$ and $q_{\bm{\phi}}(\bm{\theta})$ supported on the same set, and $p^{-}(\bm{\theta})$ supported on a superset of $\Theta$. Under these conditions, the optimal classifier parameterized as in eq. 3 recovers the true posterior, $p(\bm{\theta}|{\bm{x}_{\text{obs}}})=q_{\bm{\phi}}(\bm{\theta})$. This result follows from the fact the optimal classifier must learn the density ratio between the positive and negative distributions up to a constant across $\Theta$

$$ $\displaystyle p(\bm{\theta}|{\bm{x}_{\text{obs}}})/p^{-}(\bm{\theta})$ $\displaystyle=c\cdot q_{\bm{\phi}}(\bm{\theta})/p^{-}(\bm{\theta}),$ $\displaystyle p(\bm{\theta}|{\bm{x}_{\text{obs}}})$ $\displaystyle=c\cdot q_{\bm{\phi}}(\bm{\theta}).$ Due to the shared support, integrating both sides over the proposal support $\Theta$ gives $\displaystyle\int_{\Theta}p(\bm{\theta}|{\bm{x}_{\text{obs}}})$ $\displaystyle=c\cdot\int_{\Theta}q_{\bm{\phi}}(\bm{\theta}),$ (5) $\displaystyle c$ $\displaystyle=1.$ Thus the optimal classifier recovers the true posterior $\displaystyle p(\bm{\theta}|{\bm{x}_{\text{obs}}})$ $\displaystyle=q_{\bm{\phi}}(\bm{\theta}).$ (6) $$

Importantly, this suggests if $q_{\bm{\phi}}(\bm{\theta})$ or $p(\bm{\theta}|{\bm{x}_{\text{obs}}})$ have non-zero regions outside the support of $\pi(\bm{\theta})$, then $p(\bm{\theta}|{\bm{x}_{\text{obs}}})=c\cdot q_{\bm{\phi}}(\bm{\theta})$ can be satisfied within $\Theta$ without recovering the true posterior, as the integrals in eq. 5 will not both evaluate to one. This is a consequence of the incentive to learn the density ratio (up to a constant) being localized to the sampled region. In practice, ensuring $\pi(\bm{\theta})$, $q_{\bm{\phi}}(\bm{\theta})$ and $p(\bm{\theta}|{\bm{x}_{\text{obs}}})$ are supported on the same set $\Theta$ can be achieved by defining a variational distribution supported on the same set as the posterior $p(\bm{\theta}|{\bm{x}_{\text{obs}}})$, alongside using the variational distribution as the proposal distribution $\pi(\bm{\theta})=q_{\bm{\phi}}(\bm{\theta})$.

When the optimal classifier is achieved, we show in section A.1 that the gradient is zero for finite $K$, and consequently has zero variance. This is a generally desirable property that is not present in many variational objectives, although specialized gradient estimators have been developed which in some cases can address this issue Roeder et al. (); Tucker et al. (). An algorithm outlining the overall approach of SoftCVI is shown in algorithm 1.

Figure: Algorithm 1 SoftCVI

### 2.3 Choice of the negative distribution

The choice of negative distribution is a well-known challenge in contrastive learning, with the optimal negative distribution often differing significantly from the positive distribution Chehab et al. (). In contrast to standard applications of contrastive estimation, in SoftCVI, the negative distribution is never sampled, meaning the impact of the choice is limited to its influence on the objective properties. To better understand this, we can rewrite the softmax cross-entropy objective from eq. 4 by separating out the log term

$$ $\displaystyle\mathcal{L}_{\text{SoftCVI}}(\bm{\phi};\{\bm{\theta}_{k}\}_{k=1}^ {K},\bm{y})$ $\displaystyle=-\sum_{k=1}^{K}y_{k}\log\left(\frac{q_{\bm{\phi}}(\bm{\theta}_{k })}{p^{-}(\bm{\theta}_{k})}\right)+\sum_{k=1}^{K}y_{k}\log\sum_{k^{\prime}=1}^ {K}\frac{q_{\bm{\phi}}(\bm{\theta}_{k^{\prime}})}{p^{-}(\bm{\theta}_{k^{\prime }})},$ $\displaystyle=-\sum_{k=1}^{K}y_{k}\log q_{\bm{\phi}}(\bm{\theta}_{k})+\log\sum _{k=1}^{K}\frac{q_{\bm{\phi}}(\bm{\theta}_{k})}{p^{-}(\bm{\theta}_{k})}+\text{ const},$ (7) $$

where in the last line we remove the denominator from the first term into a constant as it is independent of $\bm{\phi}$, and use $\sum_{k=1}^{K}y_{k}=1$. The first term encourages placing posterior mass on the samples likely to be from the true posterior, and the second term penalizes the sum of the ratios. In addition to appearing in the denominator of the second term, the negative distribution choice also influences the labels $y_{k}$ through eq. 2.

We focus on setting $p^{-}(\bm{\theta})=\pi(\bm{\theta})^{\alpha}$, where $\alpha\in[0,1]$ is a tempering hyperparameter which interpolates between directly using the proposal distribution as the negative distribution when $\alpha=1$ and using an improper flat negative distribution when $\alpha=0$. Lower values of $\alpha$ tend to favor more mass-covering solutions by reducing the relative penalization of the ratios for samples in lower density regions in $q_{\bm{\phi}}(\bm{\theta})$. However, choosing a too small $\alpha$ causes a large discrepancy between the positive and negative distributions, risking the classification problem becoming too easy. Specifically, a too small $\alpha$ can lead to problematically high variances of the log ratios $z_{k}=\log[p(\bm{\theta}_{k},{\bm{x}_{\text{obs}}})/p^{-}(\bm{\theta}_{k})]$ and $\hat{z}_{k}=\log[q_{\bm{\phi}}(\bm{\theta}_{k})/p^{-}(\bm{\theta}_{k})]$. This can result in a large proportion of samples contributing negligibly to the loss, both being assigned labels near zero, and being easily classified as negative by poor posterior models. In terms of eq. 7, the high variance of $\hat{z}_{k}$ leads to a failure to sufficiently punish the ratios in lower density regions of $q_{\bm{\phi}}(\bm{\theta})$, as $\log\sum_{k=1}^{K}\exp(\hat{z}_{k})$ is dominated by the few largest ratios. In practice, we empirically show this leads to “leakage” of mass into regions of negligible posterior density (fig. 2), and a weak signal-to-noise ratio (SNR) (see section A.3).

As described in the introduction to section 2, we choose the proposal distribution to equal the variational distribution, meaning the choice above is equivalent to using $p^{-}(\bm{\theta})=q_{\bm{\phi}}(\bm{\theta})^{\alpha}$. However, when computing the objective gradient, we treat $p^{-}(\bm{\theta})$ as independent of $\bm{\phi}$, which practically can be achieved by applying the stop-gradient operator present in automatic differentiation packages.(^1^11Without preventing gradient flow through $p^{-}(\bm{\theta})$, the choice $p^{-}(\bm{\theta})=\pi(\bm{\theta})=q_{\bm{\phi}}(\bm{\theta})$ leads to zero gradients as the parameterization of the classifier becomes $q_{\bm{\phi}}(\bm{\theta})/q_{\bm{\phi}}(\bm{\theta})$.)

Another possible choice of negative distribution is $p^{-}(\bm{\theta})=p(\bm{\theta},{\bm{x}_{\text{obs}}})\propto p(\bm{\theta}|{
\bm{x}_{\text{obs}}})$. However, this introduces the ratio $q_{\bm{\phi}}(\bm{\theta})/p(\bm{\theta}|{\bm{x}_{\text{obs}}})$ which heavily punishes $q_{\bm{\phi}}(\bm{\theta})\gg p(\bm{\theta}|{\bm{x}_{\text{obs}}})$, favoring mode-seeking solutions. For this reason, we leave further investigation of this choice until section A.4.

## 3 Related work

### 3.1 Variational inference

Given access to an unnormalized posterior density, $p(\bm{\theta},{\bm{x}_{\text{obs}}})$, variational inference allows optimizing a variational distribution $q_{\bm{\phi}}(\bm{\theta})$ to approximate the posterior. For certain model classes there exist closed form solutions which can be exploited during optimization (e.g. Parisi, ; Ghahramani & Beal, ; Blei et al., ). However, the lack of broad applicability of these methods hinders the ability of practitioners to freely alter the model and variational family. As such, we instead focus on methods which place minimal restrictions on the model form and variational family, often termed black box variational inference Ranganath et al. ().

#### 3.1.1 Evidence lower bound

The most commonly used variational objective is to minimize the negative Evidence Lower Bound (ELBO) or equivalently, minimizing the reverse KL divergence between the posterior approximation and the true posterior

$$ $\displaystyle D_{KL}[q_{\bm{\phi}}(\bm{\theta})||p(\bm{\theta}|{\bm{x}_{\text{ obs}}})]$ $\displaystyle=\mathbb{E}_{q_{\bm{\phi}}(\bm{\theta})}\left[\log q_{\bm{\phi}}( \bm{\theta})-\log p(\bm{\theta}|{\bm{x}_{\text{obs}}})\right],$ $\displaystyle=\mathbb{E}_{q_{\bm{\phi}}(\bm{\theta})}[\log q_{\bm{\phi}}(\bm{ \theta})-\log p(\bm{\theta},{\bm{x}_{\text{obs}}})]+\text{const}.$ (8) $$

As there is no general closed-form solution for the divergence, a Monte Carlo approximation is often used Kingma & Welling ()

$$ $\displaystyle\mathcal{L}_{\text{ELBO}}(\bm{\phi})$ $\displaystyle=\frac{1}{K}\sum_{k=1}^{K}\left[\log q_{\bm{\phi}}(\bm{\theta}_{k })-\log p(\bm{\theta}_{k},{\bm{x}_{\text{obs}}})\right],\text{ where }\{\bm{ \theta}_{k}\}_{k=1}^{K}\sim q_{\bm{\phi}}(\bm{\theta}).$ (9) $$

Although straightforward to apply, the reverse KL divergence heavily punishes placing mass in regions with little mass in the true posterior leading to mode-seeking behavior and lighter tails than $p(\bm{\theta}|{\bm{x}_{\text{obs}}})$. In addition to underestimating uncertainty, the light tails may also lead to the approximation performing poorly in downstream tasks, such as when acting as a proposal distribution in importance sampling Chatterjee & Diaconis (); Gelman & Meng (); Yao et al. (); Müller et al. () or for reparameterizing MCMC Hoffman et al. ().

#### 3.1.2 Self-normalized importance sampling forward KL divergence

In order to address the limitations of the ELBO, numerous alternative objectives have been proposed that encourage more mass-covering behavior, such as the importance weighted ELBO Burda et al. (), the Rényi divergence, Li & Turner (), and methods targeting the forward KL divergence Jerfel et al. (); Naesseth et al. (). In this section, we will focus on the Self-Normalized Importance Sampling Forward Kullback-Leibler (SNIS-fKL) divergence estimator introduced by Jerfel et al. (). Specifically, a standard importance weighted estimate of the forward KL divergence is given by

$$ $\displaystyle D_{KL}[p(\bm{\theta}|{\bm{x}_{\text{obs}}})||q_{\bm{\phi}}(\bm{ \theta})]$ $\displaystyle=\mathbb{E}_{\pi(\bm{\theta})}\left[w(\bm{\theta})\log\frac{p(\bm {\theta},{\bm{x}_{\text{obs}}})}{q_{\bm{\phi}}(\bm{\theta})}\right]+\text{ const},$ (10) $$

where $w(\bm{\theta})=p(\bm{\theta}|{\bm{x}_{\text{obs}}})/\pi(\bm{\theta})$ is the importance weights. Computing a set of self-normalized weights
$\tilde{w}(\bm{\theta}_{k})=\frac{p(\bm{\theta}_{k},{\bm{x}_{\text{obs}}})/\pi(
\bm{\theta}_{k})}{\sum_{k^{\prime}=1}^{K}p(\bm{\theta}_{k^{\prime}},{\bm{x}_{
\text{obs}}})/\pi(\bm{\theta}_{k^{\prime}})}$
and using these alongside a Monte Carlo approximation of eq. 10, yields the objective

$$ $\displaystyle\mathcal{L}_{\text{SNIS-fKL}}(\bm{\phi};\{\bm{\theta}_{k}\}_{k=1} ^{K})=\sum_{k=1}^{K}\tilde{w}(\bm{\theta}_{k})\log\frac{p(\bm{\theta}_{k},{\bm {x}_{\text{obs}}})}{q_{\bm{\phi}}(\bm{\theta}_{k})},$ (11) $$

which introduces bias in the approximation of the forward KL vanishing with order $\mathcal{O}(1/K)$ Agapiou et al. (). Generally, the proposal is chosen to equal the variational distribution $\pi(\bm{\theta})=q_{\bm{\phi}}(\bm{\theta})$, and the proposal parameters are held constant under differentiation Glöckler et al. (); Dhaka et al. ().

### 3.2 Comparison of SoftCVI and SNIS-fKL

In this section, we demonstrate that in the special case of choosing $p^{-}(\bm{\theta})=\pi(\bm{\theta})=q_{\bm{\phi}}(\bm{\theta})$, both SoftCVI and SNIS-fKL produce gradients that are equal in expectation, but the SoftCVI objective naturally includes a control variate which ensures the variance of the gradient decreases to zero as the variational distribution approaches the true posterior. We use this result to suggest an alternative, lower-variance gradient estimator for the SNIS-fKL objective, exactly equivalent to optimizing the SoftCVI objective with $p^{-}(\bm{\theta})=\pi(\bm{\theta})=q_{\bm{\phi}}(\bm{\theta})$.

Noticing that from eq. 11, the numerator term does not depend on $\bm{\phi}$, we can equivalently write the SNIS-fKL objective as

$$ $\displaystyle\mathcal{L}_{\text{SNIS-fKL}}(\bm{\phi};\{\bm{\theta}_{k}\}_{k=1} ^{K})=-\sum_{k=1}^{K}\tilde{w}(\bm{\theta}_{k})\log q_{\bm{\phi}}(\bm{\theta}_ {k})+\text{const}.$ (12) $$

In the special case under consideration, the self-normalized weights $\tilde{w}(\bm{\theta})$ from the SNIS-fKL objective, and the ground truth labels $\bm{y}$ from SoftCVI are computed identically. This allows rewriting the SoftCVI objective from eq. 7 as the sum of the SNIS-fKL objective and the normalization term

$$ $\displaystyle\mathcal{L}_{\text{SoftCVI}}(\bm{\phi};\{\bm{\theta}_{k}\}_{k=1}^ {K})=\mathcal{L}_{\text{SNIS-fKL}}(\bm{\phi};\{\bm{\theta}_{k}\}_{k=1}^{K})+ \log\sum_{k=1}^{K}\frac{q_{\bm{\phi}}(\bm{\theta}_{k})}{p^{-}(\bm{\theta}_{k}) }+\text{const},$ (13) $$

where we show in section A.2 that when $p^{-}(\bm{\theta})=q_{\bm{\phi}}(\bm{\theta})$, the gradient of the normalization term is

$$ $\displaystyle\nabla_{\bm{\phi}}\log\sum_{k=1}^{K}\frac{q_{\bm{\phi}}(\bm{ \theta}_{k})}{p^{-}(\bm{\theta}_{k})}=\frac{1}{K}\sum_{k=1}^{K}\nabla_{\bm{ \phi}}\log q_{\bm{\phi}}(\bm{\theta}_{k}).$ (14) $$

Since $\mathbb{E}_{\pi(\bm{\theta})}[\nabla_{\bm{\phi}}\log q_{\bm{\phi}}(\bm{\theta}
)]=\mathbb{E}_{q_{\bm{\phi}}(\bm{\theta})}[\nabla_{\bm{\phi}}\log q_{\bm{\phi}
}(\bm{\theta})]=\bm{0}$, the inclusion or omission of the normalization term does not change the gradient in expectation over sets of $\{\bm{\theta}_{k}\}_{k=1}^{K}$, but it significantly influences the variance. As shown in section A.1, the gradient variance for SoftCVI decreases to zero as the variational distribution approaches the posterior. In contrast, this implies the variance of the SNIS-fKL objective approaches the variance of $\frac{1}{K}\sum_{k=1}^{K}\nabla_{\bm{\phi}}\log q_{\bm{\phi}}(\bm{\theta}_{k})$, which is positive, scaling inversely with $K$.

In addition to providing a novel perspective of the SNIS-fKL objective as training an (unnormalized) classifier, this result also implies a straightforward modification to form a lower-variance SNIS-fKL gradient estimator

$$ $\displaystyle\mathcal{L}_{\text{SNIS-fKL-LV}}(\bm{\phi};\{\bm{\theta}_{k}\}_{k =1}^{K})=\sum_{k=1}^{K}\tilde{w}(\bm{\theta}_{k})\log\frac{p(\bm{\theta}_{k},{ \bm{x}_{\text{obs}}})}{q_{\bm{\phi}}(\bm{\theta}_{k})}+\frac{1}{K}\sum_{k=1}^{ K}\log q_{\bm{\phi}}(\bm{\theta}_{k}),$ (15) $$

which is exactly equivalent to optimizing the SoftCVI objective with $p^{-}(\bm{\theta})=\pi(\bm{\theta})=q_{\bm{\phi}}(\bm{\theta})$. This approach of lowering the variance of a gradient estimator by utilizing $\nabla_{\bm{\phi}}\log q_{\bm{\phi}}(\bm{\theta})$ as a control variate has also been applied to other variational objectives, such as the sticking the landing estimator of the ELBO Roeder et al. (); Tucker et al. (). There is no guarantee that the gradient variance will be lower in all instances. However, the variance provably reduces to zero as the variational distribution approaches the true posterior, which allows a reasonable SNR to be maintained even when near convergence, which is likely beneficial for performance and stable convergence. We support this claim with the empirical results in section 5, in addition to investigating the signal-to-noise ratio of the objectives in section A.3.

### 3.3 Contrastive learning

Contrastive learning most commonly allows learning of distributions through the comparison of true samples to a set of negative (noise or augmented) samples Gutmann & Hyvärinen (). Contrastive learning has been applied to numerous domains, including representation learning Oord et al. (); He et al. (); Chen et al. (), image generation Aneja et al. () and natural language processing Mnih & Kavukcuoglu (). Generally, and notably in contrast to to the current work, applications of contrastive learning focus on problems where the likelihood is unavailable, such as for fitting energy-based models Gutmann & Hyvärinen (); Rhodes et al. (); Gao et al. (); Gutmann et al. (); Rhodes & Gutmann () or performing simulation-based inference Hermans et al. (); Thomas et al. (); Greenberg et al. (); Miller et al. (); Gutmann et al. (); Durkan et al. ().

Probably the most widely used objective function in contrastive learning is the InfoNCE loss, proposed by Oord et al. (). Given a set of samples $\{\bm{\theta}_{k}\}_{k=1}^{K}$, containing a single true sample with index $k^{*}$, and $K-1$ negative samples, the loss can be computed as

$$ $\displaystyle\mathcal{L}_{\text{InfoNCE}}(\bm{\phi};\{\bm{\theta}_{k}\}_{k=1}^ {K})=-\log\frac{f_{\bm{\phi}}(\bm{\theta}_{k^{*}})}{\sum_{k=1}^{K}f_{\bm{\phi} }(\bm{\theta}_{k})}.$ (16) $$

where $f_{\bm{\phi}}(\bm{\theta})$ approximates the ratio between the positive and negative distributions (up to a constant). This is commonly also presented with an additional sum (or expectation) over different sets of $\{\bm{\theta}_{k}\}_{k=1}^{K}$. The InfoNCE loss can be derived from the softmax cross-entropy loss, eq. 4, by inputting a one-hot encoded vector of labels $\bm{y}$ with $y_{k^{*}}=1$ and all other elements 0 0. This results in only the $k^{*}$-th summation term being non-zero, recovering the InfoNCE loss.

In contrastive methods for simulation-based inference, learning is achieved through comparing positive parameter-output pairs from $p(\bm{x}|\bm{\theta})\pi(\bm{\theta})$, to negative (mismatched) pairs drawn marginally from $p(\bm{x})\pi(\bm{\theta})$. Similar to the current work, the classifier is often parameterized in terms of a normalized approximation to the posterior, such that after training the posterior approximation is directly available for sampling and density evaluation. In these methods, the classifier is parameterized using the ratio between the approximate posterior and the prior $q_{\bm{\phi}}(\bm{\theta}|\bm{x})/p(\bm{\theta})$, and optimized using the InfoNCE objective Durkan et al. (); Greenberg et al. ().

Recently, there has been investigation of the use of soft (or ranked) labels in contrastive learning, frequently using cross-entropy-like loss functions. In contrast to the current work, these methods focus on improving performance in the standard context of contrastive learning, where evaluation of the likelihood is infeasible. As such, the soft labels cannot be computed exactly, and instead are generated using other methods, such as label smoothing Hugger & Uhlmann () or by utilizing a similarity metric such as cosine similarity between embeddings of positive and negative samples Park et al. (); Feng & Patras (); Hoffmann et al. ().

## 4 Experiments

We evaluate the performance of SoftCVI on a variety of Bayesian inference tasks for which reference (ground truth) posterior samples, $\{\bm{\theta}^{*}_{i}\}_{i=1}^{N_{\text{ref}}}$, are available. We use $p^{-}(\bm{\theta})=\pi(\bm{\theta})^{\alpha}$ and focus results on two choices, $\alpha=0.75$ and $\alpha=1$. We compare performance to variational inference either with the ELBO or the often more mass-covering SNIS-fKL divergence Jerfel et al. (). For each task, we perform 50 independent runs with different random seeds. Where possible (i.e. an analytical posterior is available), we use a different observation for every run. For tasks where the reference posterior is created through sampling methods we rely on reference posteriors provided by PosteriorDB Magnusson et al. () or SBI-Benchmark Lueckmann et al. () and for each run we sample from the available observations if multiple options are present. We run all considered methods for 50,000 optimization steps using the Adam optimizer Kingma & Ba (), and use $K=8$ samples from $q_{\bm{\phi}}(\bm{\theta})$ during computation of the objectives.

Software. Our implementation and experiments made wide use of the python packages JAX Bradbury et al. (), equinox Kidger & Garcia (), numpyro Phan et al. (), flowjax Ward () and optax DeepMind et al. ().

### 4.1 Metrics

Coverage probabilities.
Posterior coverage probabilities have been widely used to assess the reliability of posteriors
(e.g. Cook et al., ; Talts et al., ; Prangle et al., ; Hermans et al., ; Ward et al., ; Cannon et al., ). Given a nominal frequency $\gamma\in[0,1]$, the metric assesses the frequency at which true posterior samples fall within the $100\gamma\%$ highest posterior density region (credible region) of the approximate posterior. For a given $\gamma$, if the actual frequency exceeds $\gamma$, the posterior is conservative for that coverage probability; if it is lower, then the posterior is overconfident. A posterior is said to be well calibrated if the actual frequency matches $\gamma$ for any choice of $\gamma$. A well calibrated or somewhat conservative posterior is needed for drawing reliable scientific conclusions, and as such is an important property to investigate. We estimate the actual coverage frequency for a posterior estimate as

$$ $\displaystyle\frac{1}{N_{\text{ref}}}\sum_{i=1}^{N_{\text{ref}}}\mathds{1} \left\{{\bm{\theta}_{i}^{*}\in\operatorname{HDR}_{q_{\bm{\phi}}(\bm{\theta}|{ \bm{x}_{\text{obs}}})}(\gamma)}\right\}$ (17) $$

where $\mathds{1}$ is the indicator function, and $\operatorname{HDR}$ represents the highest posterior density region, inferred using the density quantile approach of Hyndman ().

Log probability of $\bm{\theta}^{*}$. Looking at the probability of either reference posterior samples or the ground truth parameters in a posterior approximation is a common metric for assessing performance (e.g. Lueckmann et al., ; Papamakarios & Murray, ; Greenberg et al., ). We compute this independently for each posterior approximation as follows

$$ $\displaystyle\frac{1}{N_{\text{ref}}}\sum_{i=1}^{N_{\text{ref}}}\log q_{\bm{ \phi}}(\bm{\theta}^{*}_{i}).$ (18) $$

This metric also approximates the negative forward KL divergence between the true and approximate posterior up to a constant, which is a quantity known to control error in importance sampling Chatterjee & Diaconis ()

$$ $\displaystyle-D_{\text{KL}}[p(\bm{\theta}|{\bm{x}_{\text{obs}}})||q_{\bm{\phi} }(\bm{\theta})]$ $\displaystyle=-\mathbb{E}_{p(\bm{\theta}|{\bm{x}_{\text{obs}}})}[\log p(\bm{ \theta}|{\bm{x}_{\text{obs}}})-\log q_{\bm{\phi}}(\bm{\theta})],$ $\displaystyle\approx\frac{1}{N_{\text{ref}}}\sum_{i=1}^{N_{\text{ref}}}\log q_ {\bm{\phi}}(\bm{\theta}_{i}^{*})+\text{const},$ (19) $$

Posterior mean accuracy.
A posterior that is overconfident, but unbiased would perform poorly based on the aforementioned metrics, but could form good point estimates for the latent variables. To assess this, we choose to measure the accuracy using the negative $L^{2}$-norm of the standardized difference in posterior means

$$ $\displaystyle-\left\lVert\frac{\text{mean}(\bm{\theta}^{*})-\text{mean}(\bm{ \theta})}{\text{std}(\bm{\theta}^{*})}\right\rVert_{2},$ (20) $$

where $\text{mean}(\bm{\theta}^{*})$ and $\text{mean}(\bm{\theta})$ are the mean vectors of the reference and posterior approximation samples respectively, and $\text{std}(\bm{\theta}^{*})$ is the vector of standard deviations of the reference samples.

### 4.2 Tasks

Where possible, we make use of reparameterizations (such as non-centering) to reduce the dependencies between the latent variables, and ensure variables are reasonably scaled, which is generally considered to improve performance Papaspiliopoulos et al. (); Gorinova et al. (); Betancourt & Girolami ().

Eight schools. A classic hierarchical Bayesian inference problem, where the aim is to infer the treatment effects of a coaching program applied to eight schools, which are assumed to be exchangeable Rubin (); Gelman et al. (). The parameter set is $\bm{\theta}=\{\mu,\tau,\bm{m}\}$, where $\mu$ is the average treatment effect across the schools, $\tau$ is the standard deviation of the treatment effects across the schools and $\bm{m}$ is the eight-dimensional vector of treatment effects for each school. For the posterior approximation $q_{\bm{\phi}}(\bm{\theta})$, we use a normal distribution for $\mu$, a folded Student’s t distribution for $\tau$ (where folding is equivalent to taking an absolute value transform), and a Student’s t distribution for $\bm{m}$. We use the reference posterior samples available from PosteriorDB Magnusson et al. ().

Figure: Figure 1: The posterior performance metrics (described in detail in section 4.1). a) The nominal coverage frequency against the difference between the nominal and actual coverage frequency, averaged across the runs. Well-calibrated methods follow the black dotted line ($y=0$), whereas conservative methods fall above, and overconfident methods below. b) The average probability of the reference posterior samples in the approximate posterior. c) The accuracy of the approximate posteriors mean, calculated as the negative l2-norm between the mean of the standardized reference and approximate posterior samples.
Refer to caption: x1.png

Linear regression. A Bayesian linear regression model with standard normal priors on the parameters $\bm{\theta}=\{\bm{\beta},\mu\}$, where $\bm{\beta}\in\mathbb{R}^{50}$ is the regression coefficients, and $\mu\in\mathbb{R}$ is the bias parameter. The covariates $\bm{X}\in\mathbb{R}^{(200\times 50)}$ are generated from a standard normal distribution, with targets sampled according to $\bm{y}\sim N(\bm{X}\bm{\beta}+\mu,1)$. The posterior approximation $q_{\bm{\phi}}(\bm{\theta})$ is implemented as a fully factorized normal distribution. This task yields an analytical form for the posterior.

SLCP. The Simple Likelihood Complex Posterior (SLCP) task introduced in Papamakarios et al. (), parameterizes a multivariate normal distribution using a 5-dimensional vector $\bm{\theta}$. Due to squaring in the parameterization, the posterior contains four symmetric modes. A uniform prior $\{\bm{\theta}_{i}\}_{i=1}^{5}\sim\text{Uniform}(-3,3)$ is used, which introduces sharp cut-offs in the posterior density. For the posterior approximation $q_{\bm{\phi}}(\bm{\theta})$, we use a masked autoregressive flow Papamakarios et al. (); Kingma et al. (); Germain et al. (), with a rational quadratic spline transformer Durkan et al. (). We use the observations and reference posteriors available from the SBI-Benchmark package Lueckmann et al. ().

GARCH(1,1). A Generalized Autoregressive Conditional heteroscedasticity (GARCH) model Bollerslev (). GARCH models are used for modeling the changing volatility of time series data by accounting for time-varying and autocorrelated variance in the error terms. The observation consists of a 200-dimensional time series $\bm{x}$, where each element $x_{t}$ is drawn from a normal distribution with mean $\mu$ and time-varying variance $\sigma_{t}^{2}$. The variance $\sigma_{t}^{2}$ is defined recursively by the update $\sigma^{2}_{t}=\alpha_{0}+\alpha_{1}(x_{t-1}-\mu)^{2}+\beta_{1}\sigma^{2}_{t-1}$.
To parameterize $q_{\bm{\phi}}(\bm{\theta})$, we use a normal distribution for $\mu$ and a log normal distribution for $\alpha_{0}$. For $\alpha_{1}$ and $\beta_{1}$ we use uniform distributions constrained to the prior support, transformed with a rational quadratic spline. To allow modeling of posterior dependencies, the distribution over $\beta_{1}$ is parameterized as a function of $\alpha_{0}$ and $\alpha_{1}$ using a neural network. We use the reference posterior samples available from PosteriorDB Magnusson et al. ().

## 5 Results

Across all tasks and metrics considered, the SoftCVI derived objectives performed competitively with the ELBO and SNIS-fKL objectives (fig. 1). Overall, SoftCVI using a negative distribution $\pi(\bm{\theta})^{\alpha}$ with $\alpha=0.75$ outperformed the other methods, giving rise to better calibrated posteriors and tending to place more mass on average on the reference posterior samples, indicating a lower forward KL divergence to the true posterior.

A key comparison, is between SoftCVI with $\alpha=1$ and the SNIS-fKL objective, which give gradient approximators that are equal in expectation (see section 3.2). With the exception of the eight schools task, where both methods performed similarly, SoftCVI with $\alpha=1$ tended to place more mass on the reference samples and yielded better calibrated posteriors. These results highlight the benefit of the reduced gradient variance provided by SoftCVI when the variational distribution is sufficiently close to the true posterior.
Further, the performance discrepancy becomes more pronounced for a smaller choice of $K$ (see fig. 6 in the appendix). This can be explained due to the variance of the SNIS-fKL objective gradient scaling inversely with $K$, meaning the control variate naturally included by the SoftCVI objective becomes more crucial.

On the SLCP task, which yields complex posterior geometry with four symmetric modes, only the SoftCVI objectives performed well. The ELBO objective often resulted in poor distribution of the mass across the modes, sometimes missing modes entirely (fig. 2; see also fig. 7 in the appendix). In contrast, SNIS-fKL, though less mode-seeking, frequently approximated the individual modes poorly. This observation aligns with previous work suggesting the unreliability of existing mass-covering objectives Dhaka et al. ().
To illustrate the impact of the choice of negative distribution, fig. 2 also shows a posterior trained with a flat negative distribution by setting $\alpha=0$. The flat negative distribution does not sufficiently penalize placing mass in $q_{\bm{\phi}}(\bm{\theta})$ in regions of negligible posterior density, leading to “leakage” of mass (see section 2.3).

Figure: Figure 2: A 2-dimensional posterior marginal for a single run of the SLCP task, with the reference posterior samples shown in green.
Refer to caption: x2.png

## 6 Conclusion

In this work, we introduced SoftCVI, a novel framework for deriving variational objectives motivated through contrastive learning. Our experiments across various Bayesian inference tasks indicate that SoftCVI often outperforms other variational objectives, producing posteriors approximations with better coverage properties and a lower forward KL divergence to the true posterior. The performance difference is particularly notable for tasks with complex posterior geometries which require flexible density estimators. SoftCVI bridges between variational inference and contrastive estimation, which we hope will open up new avenues for theoretical and experimental research. Theoretically, it would be interesting to explore if other variational objectives can be reframed using SoftCVI. Additionally, both experimental and theoretical work is needed to further guide the choice of negative distribution. Finally, it could be valuable to investigate how advances from classification and contrastive learning, such as label smoothing Müller et al. () and temperature scaling Wang & Liu (), could be adapted to SoftCVI to enhance training stability and further control posterior calibration.

## Appendix A Appendix

### A.1 Zero variance gradient at optimum with finite K

When the optimal classifier is reached, we can show that the gradient is zero (and hence has zero variance), even for finite $K$. To avoid clashes with the parameters $\bm{\phi}$, we will use superscripts for indices. We have the objective

$$ $\displaystyle\mathcal{L}(\bm{\phi};\{\bm{\theta}^{k}\}_{k=1}^{K},\bm{y})$ $\displaystyle=-\sum_{k=1}^{K}y^{k}\log\left(\frac{q_{\bm{\phi}}(\bm{\theta}^{k })/p^{-}(\bm{\theta}^{k})}{\sum_{k^{\prime}=1}^{K}q_{\bm{\phi}}(\bm{\theta}^{k ^{\prime}})/p^{-}(\bm{\theta}^{k^{\prime}})}\right).$ Letting $\hat{y}_{\bm{\phi}}^{k}$ replace the term inside the log $\displaystyle=-\sum_{k=1}^{K}y^{k}\log\hat{y}_{\bm{\phi}}^{k},$ (21) $\displaystyle\nabla\mathcal{L}(\bm{\phi};\{\bm{\theta}^{k}\}_{k=1}^{K},\bm{y})$ $\displaystyle=-\sum_{k=1}^{K}y^{k}\nabla\log\hat{y}_{\bm{\phi}}^{k},$ (22) $\displaystyle=-\sum_{k=1}^{K}y^{k}\frac{\nabla\hat{y}_{\bm{\phi}}^{k}}{\hat{y} _{\bm{\phi}}^{k}},$ (23) which due to optimality of the classifier we have $y^{k}=\hat{y}_{\bm{\phi}}^{k}$ $\displaystyle=-\sum_{k=1}^{K}\nabla\hat{y}_{\bm{\phi}}^{k},$ (24) due to the properties of the softmax, the labels must sum to 1 $\displaystyle=-\nabla\sum_{k=1}^{K}\hat{y}_{\bm{\phi}}^{k}=\bm{0}.$ (25) $$

### A.2 SoftCVI normalization term gradient

When $p^{-}(\bm{\theta})=q_{\bm{\phi}}(\bm{\theta})$, we claim in eq. 14 that the normalization term gradient can be written as $\frac{1}{K}\sum_{k=1}^{K}\nabla_{\bm{\phi}}\log q_{\bm{\phi}}(\bm{\theta}_{k})$. We can show this as follows

$$ $\displaystyle\nabla_{\bm{\phi}}\log\sum_{k=1}^{K}\frac{q_{\bm{\phi}}(\bm{ \theta}_{k})}{p^{-}(\bm{\theta}_{k})}$ $\displaystyle=\frac{\nabla_{\bm{\phi}}\sum_{k=1}^{K}q_{\bm{\phi}}(\bm{\theta}_ {k})/p^{-}(\bm{\theta}_{k})}{\sum_{k=1}^{K}q_{\bm{\phi}}(\bm{\theta}_{k})/p^{- }(\bm{\theta}_{k})},$ $\displaystyle=\frac{1}{K}\sum_{k=1}^{K}\frac{\nabla_{\bm{\phi}}q_{\bm{\phi}}( \bm{\theta}_{k})}{p^{-}(\bm{\theta}_{k})},$ $\displaystyle=\frac{1}{K}\sum_{k=1}^{K}\frac{\nabla_{\bm{\phi}}q_{\bm{\phi}}( \bm{\theta}_{k})}{q_{\bm{\phi}}(\bm{\theta}_{k})},$ $\displaystyle=\frac{1}{K}\sum_{k=1}^{K}\nabla_{\bm{\phi}}\log q_{\bm{\phi}}( \bm{\theta}_{k}),$ (26) $$

where on the first and last line we make use of $\nabla_{\bm{x}}\log f(\bm{x})=\frac{\nabla_{\bm{x}}f(\bm{x})}{f(\bm{x})}$.

### A.3 Gradient signal-to-noise ratio

By evaluating an objective gradient over a set of random seeds (different sets of $\{\bm{\theta}_{k}\}_{k=1}^{K}$), it is possible to empirically inspect the signal-to-noise ratio (SNR)ratio of the gradient for different objectives. Following Rainforth et al. (), we compute the SNR as

$$ $\operatorname{SNR}(\nabla_{\bm{\phi}}\mathcal{L}(\bm{\phi}))=\left|\mathbb{E}[ \nabla_{\bm{\phi}}\mathcal{L}(\bm{\phi})]\right|\ \big{/}\sigma[\nabla_{\bm{ \phi}}\mathcal{L}(\bm{\phi})]$ $$

where $\sigma[\cdot]$ denotes the element-wise standard deviation for each parameter gradient. A low SNR indicates that the gradient estimation is dominated by noise, making stochastic optimization challenging. However, it is important to recognize that whilst a high SNR is preferable, it does not alone guarantee the objective itself is practically useful (or even sensible), for example it may be biased or heavily favor mode-seeking solutions.

Figure: Figure 3: The signal, noise and signal-to-noise ratio of the objective gradients on a toy normal task of varying dimensionality. When $d=50$, the gradient properties are computed parameter-wise and averaged. The vertical dashed line shows the true parameter values from the closed form posterior solution.
Refer to caption: x3.png

We consider a toy normal task from Glöckler et al. (), with the distinction that we vary the dimensionality of the task $d\in\{1,50\}$ and the parameterization of variational distribution. The task is defined through the model

$$ $\bm{\theta}\sim N(\bm{0},4\cdot\bm{I}_{d}),\quad\bm{x}\sim N(\bm{\theta},\bm{I }_{d}),$ $$

where the aim is to infer the mean vector $\bm{\theta}$, given an observation ${\bm{x}_{\text{obs}}}=\bm{1}_{d}$, where $\bm{1}_{d}$ is the d-dimensional vector of ones. The variational distribution is parameterized as a normal distribution, with mean vector $\bm{\mu}\in\mathbb{R}^{d}$ and log standard deviations $\log\bm{\sigma}\in\mathbb{R}^{d}$. When assessing the gradient properties of $\bm{\mu}$, we hold $\log\bm{\sigma}$ fixed at the closed form posterior solution, $\log\bm{\sigma}=\log(\sqrt{4/5})\cdot\bm{1}_{d}$. Similarly, when assessing the gradient properties of $\log\bm{\sigma}$, we hold $\bm{\mu}$ fixed at the true closed form solution $\bm{\mu}=4/5\cdot\bm{1}_{d}$. The signal, noise and SNR for each objective are shown in fig. 3.

SoftCVI with $\alpha=1$, and the SNIS-fKL objective yield very similar signals, which results from the two objectives producing the same gradients in expectation (section 3.2). However, the SNIS-fKL objective has positive noise when the variational distribution approaches the true posterior, meaning the SNR degrades to zero. In contrast, for the SoftCVI objectives, the gradient noise approaches zero (see section A.1), and a reasonable SNR is present as the variational distribution approaches the true posterior.

When $d=50$ the signal and consequently the SNR, deteriorates larger values of $\log\bm{\sigma}$, and lower choices of $\alpha$. In this case, both expanding the dimension of the problem, and decreasing the $\alpha$ value, increases the variance of $z_{k}=\log[p(\bm{\theta}_{k},{\bm{x}_{\text{obs}}})/p^{-}(\bm{\theta}_{k})]$ and $\hat{z}_{k}=\log[q_{\bm{\phi}}(\bm{\theta}_{k})/p^{-}(\bm{\theta}_{k})]$. This can lead to degeneracy in the labels and predictions, meaning very few samples meaningfully contribute to the loss function, reducing the SNR. In an extreme case, e.g. setting $\alpha=0$ and further expanding the dimensionality of the problem, only a single label and prediction will be significantly non-zero. In this regime, negligible signal would exist for the $\log\bm{\sigma}$ parameter, as the variational distribution, despite the incorrect $\log\bm{\sigma}$, would still result in the correct degenerate predicted labels.

### A.4 Alternative negative distribution choices

The main results focus on parameterizing the negative distribution as a function of the proposal distribution $p^{-}(\bm{\theta})=\pi(\bm{\theta})^{\alpha}$. However, another possible choice is to use the unnormalized posterior to parameterize the negative distribution, for example $p^{-}(\bm{\theta})=p(\bm{\theta},{\bm{x}_{\text{obs}}})^{\alpha}$. This choice, when $\alpha=1$, implies equality between the assumed negative and positive distributions, meaning through eq. 2 the ground truth labels become constant $y_{k}=\frac{1}{k}$ for $k=1,...,K$. Inputting these labels into eq. 7 yields the objective

$$ $\displaystyle\mathcal{L}(\bm{\phi};\{\bm{\theta}_{k}\}_{k=1}^{K},\bm{y})$ $\displaystyle=-\frac{1}{K}\sum_{k=1}^{K}\log q_{\bm{\phi}}(\bm{\theta}_{k})+ \log\sum_{k=1}^{K}\frac{q_{\bm{\phi}}(\bm{\theta}_{k})}{p^{-}(\bm{\theta}_{k}) }+\text{const}.$ (27) $$

Assuming the proposal distribution is chosen to match the variational distribution, since $\mathbb{E}_{\pi(\bm{\theta})}[\nabla_{\bm{\phi}}\log q_{\bm{\phi}}(\bm{\theta}
)]=\mathbb{E}_{q_{\bm{\phi}}(\bm{\theta})}[\nabla_{\bm{\phi}}\log q_{\bm{\phi}
}(\bm{\theta})]=0$, the first term in eq. 7 has an expected gradient of zero. In contrast to using the proposal distribution as the negative distribution, which results in the normalization term which punishes the ratios acting as a control variate, here, the first term instead acts as a control variate, with the normalization term providing the signal by punishing the ratio punishing the ratio $q_{\bm{\phi}}(\bm{\theta})/p(\bm{\theta},{\bm{x}_{\text{obs}}})\propto q_{\bm{
\phi}}(\bm{\theta})/p(\bm{\theta}|{\bm{x}_{\text{obs}}})$. As such this choice heavily punishes $q_{\bm{\phi}}(\bm{\theta})\gg p(\bm{\theta}|{\bm{x}_{\text{obs}}})$, favoring overconfident posteriors. We report in fig. 4 the metrics, using $p^{-}(\bm{\theta})=p(\bm{\theta},{\bm{x}_{\text{obs}}})^{\alpha}$, with $\alpha=0.75$ and $\alpha=1$.

Figure: Figure 4: Additional metrics analogous to fig. 1, using $p^{-}(\bm{\theta})=p(\bm{\theta},{\bm{x}_{\text{obs}}})^{\alpha}$ as the negative distribution choice in the SoftCVI objectives.
Refer to caption: x4.png

### A.5 Additional figures

Figure: Figure 5: The run times including compilation time for each task using $k=8$ (CPU, 8GB RAM).
Refer to caption: x5.png

Figure: Figure 6: Additional metrics analogous to fig. 1, using $K=2$, instead of $K=8$ samples when approximating the objectives. The SNIS-fKL objective performance degrades substantially when $K$ is small.
Refer to caption: x6.png

Figure: Figure 7: The multimodal 2-dimensional posterior marginals for the first six runs of the SLCP task for each method, with the reference posterior samples shown in green.
Refer to caption: x7.png

Figure: Figure 8: The distribution of the posterior marginals for a single run of the eight schools, linear regression and GARCH (1,1) tasks. For the linear regression task, we restrict the plot to the first 10 parameters of the model.
Refer to caption: x8.png

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. In line 57, the authors mention proposing **a family of variational objectives**. Could you clarify this point? It seems that there is only one objective—the cross-entropy between true soft labels and the predicted categorical distribution. Is there indeed a family of objectives, or is it a single objective?

2. [questions from Weaknesses] How does the choice of $K$ influence the training curves and final results? Is there an optimal value for $K$? What criteria should be used to select $K$?
soundness: 3
strengths: The paper presents an elegant method for fitting the variational distribution by training the model to solve a classification problem. 

I identify two key benefits of the proposed approach:

1. There is no need for samples from either the negative or the target distribution.
2. The densities of the target and negative distributions need only be known up to a normalizing constant.

Overall, the paper is well-written and easy to follow.
summary: In this paper, the authors propose an approach that reformulates the inference problem as a classification task using contrastive learning and the generation of soft labels. 

These soft labels, derived from an unnormalized target distribution, are employed in a cross-entropy loss to fit the variational distribution. The authors demonstrate the efficacy of their methodology through a series of experiments on Bayesian inference.
weaknesses: I note the following weaknesses in the paper:

### Major:

1. Deep Learning examples:

The experimental evaluation feels modest and misses an opportunity to show the benefits of the proposed methodology in deep learning scenarios. The authors focus solely on "classical" examples of Bayesian inference. I think it would be very illustrative to consider deep learning applications.

For example, in the related work section, the authors mention papers on Variational Auto-Encoders (VAE). It would be interesting to see how the proposed methodology performs in this context.

I see several ways your approach could be applied:

- End-to-End VAE training: Instead of optimizing the standard ELBO (or alternatives like IWAE [1]), incorporate the SoftCVI training procedure and evaluate the negative log-likelihood.
- Fitting variational distribution for pretrained VAE: For an already pretrained VAE (decoder) with a complex, multimodal posterior, apply your training procedure to fit a variational distribution to this posterior. Then demonstrate that the variational distribution fitted with SoftCVI covers more modes and provides more diverse samples (see e.g. [2]).

I appreciate the example with the SLCP task. Similarly, an example with a "circular" posterior in [3] could be considered for demonstration.

2. Training curves with different $K$:

Since $K$ is a hyperparameter we can choose as we wish, it's important to understand how it influences the results. In Figures 1 and 6, two different choices of $K$ are considered, but there is no discussion on which is better. It would be beneficial to see not only the final evaluations but also the training curves for different $K$ values. Given that $K$ can range widely (from 2 to infinity), practical guidance on selecting $K$ would be valuable.

Additionally, training curves are interesting in the context of comparison with SNIS-fKL, as the variance of estimation will be greater there, and it is interesting to see it.

### Minor:

1. Motivation in Section 2.1:

Initially, I found the beginning of Section 2.1 a bit confusing. It mentions a scenario with $K−1$ negative samples and only 1 positive, which doesn't seem to match the scenario in the paper. As I understand it, we have samples only from the proposal distribution and assign probabilities to them as if they come from the true distribution to generate soft labels. We don't assume that only one sample comes from the true distribution. If my understanding is incorrect, please clarify. Otherwise, consider rephrasing this section for clarity.


2. Typo:

Line 329: "in contrast **to to**".

----

I am positive about the methodological contributions. With the incorporation of extensions to the experimental section, I am willing to significantly increase my evaluation score.


----

References:

[1] Burda, Y., Grosse, R., & Salakhutdinov, R. (2015). Importance Weighted Autoencoders. arXiv preprint arXiv:1509.00519.

[2] Thin, A., Kotelevskii, N., Denain, J. S., Grinsztajn, L., Durmus, A., Panov, M., & Moulines, E. (2020). MetFlow: A New Efficient Method for Bridging the Gap Between Markov Chain Monte Carlo and Variational Inference. arXiv preprint arXiv:2002.12253.

[3] Thin, A., Kotelevskii, N., Doucet, A., Durmus, A., Moulines, E., & Panov, M. (2021). Monte Carlo Variational Auto-Encoders. In International Conference on Machine Learning (pp. 10247-10257). PMLR.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: It would be helpful if the paper is a bit clearer what kind of reference posteriors are used. Are they from MCMC algorithms that can be assumed to closely approximate the true posterior? What about the SBI posteriors?

Re line 127, can you elaborate what is the exact parameterisation of the classifier in terms of the ratio of two distributions. Are the parameters (infinite-dimensional) probability densities?

Can you clarify why you compare against simulation-based reference models (SBI)? How is the proposed approach applicable without having access to the likelihood function?
soundness: 2
strengths: The paper addresses the classic and important question of Bayesian inference, particularly for the challenging setting where the approximate posterior should be close to the target with respect to the forward KL to yield well calibrated inferences. The approach avoids previous objectives that yield mass-covering behaviour, but are more challenging to optimise, particular in high dimensions. The suggested objective is now as far as I am aware. It is also very interesting that the suggested approach can be used to construct a previously constructed objective (SNIS-fKL), but with a gradient estimator that has zero variance in the idealised setting when the approximate posterior coincides with the true posterior. Experiments do indeed show that the suggested method improves the coverage and leads to approximate posterior that are closer to the target in terms of the forward KL.
summary: The paper considers a variational inference approach that builds upon classifying ‘noisy’ samples from a learned proposal distribution as negative samples, similar to a noise contrastive estimation approach, and where ‘ground truth soft labels’ depend on a self-normalised importance sampling estimate for the log ratio of the target density relative to the proposal density. In contrast to a classical ELBO minimizing the reverse KL to the target, the suggested approach appears to better approximate the forward KL to the target. Numerical experiments on benchmark Bayesian inference problems illustrate that suggested approach yields te better coverage probabilities and yields to approximations with less bias for the posterior mean estimate.
weaknesses: The empirical verification seems a bit limited. In particular, the only baseline (beyond a standard ELBO) is the SNIS-fKL objective. However, as the paper shows, they have the same expected gradients. It is thus not clear how well the suggested approach compares to other (approximate) Bayesian methods that have different averaged gradients, e.g. [1-3]. 

Also as the suggested approach has a higher computation complexity requiring sampling and density-evaluations of K particles, evaluations with a simple IWAE approach instead of the ELBO baseline appears appropriate. 

The authors consider (in Section 2.1) that they have access to samples from the true posterior. However, is having access to samples from the true posterior not the aim of the method? In (4), is $y$ equal to $\hat{y}$ as in Algorithm 1 when one does not have true posterior samples? 

It is not clear why eq. (4) is the right objective to optimise when one does not have access to the true labels $y$. Does optimizing it correspond to some proper scoring rule or divergence [4] that are minimized for this objective? 

Can you clarify why the optimal classifier recovers the true posterior and that the true posterior is equal to the approximate variational distribution as per eq. (6)? The ‘optimal classifier parameterized as in eq. (3)’ depends (apparently?) on the number of samples $K$ from the proposal and on the proposal distribution. Do all the arguments on page 3 from line 141 not depend on these choices? I understand that the classifier is optimal with respect to a cross-entropy loss if it recovers the Bayesian posterior [4]. But is this optimality even realisable in the suggested form (3)?

The experiments target posteriors that are not that high-dimensional (<=200). Does the method scale to high-dimensional problems? Does it overcome scalability issues commonly faced when optimizing the forward KL or for importance sampling?


[1] Dieng, Adji Bousso, et al. "Variational Inference via $\chi $ Upper Bound Minimization." Advances in Neural Information Processing Systems 30 (2017).\
[2] Li, Chengrui, et al. "Forward $\chi^ 2$ Divergence Based Variational Importance Sampling." arXiv preprint arXiv:2311.02516 (2023).\
[3] Naesseth, Christian, Fredrik Lindsten, and David Blei. "Markovian score climbing: Variational inference with KL (p|| q)." Advances in Neural Information Processing Systems 33 (2020): 15499-15510.\
[4] Mohamed, Shakir, and Balaji Lakshminarayanan. "Learning in implicit generative models." arXiv preprint arXiv:1610.03483(2016).

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 2
questions: * 50k "optimisation steps": this sounds a lot, any explanation for this? What is the definition of step?
* K=8: Any ablation of the effect of K on accuracy vs wall-clock time? i.e. trade off in slower steps vs number of steps
     * Fig 6 (K=2) is hard to compare to Fig 1 (K=8) due to their separation and the SNIS result dominates the scale, squashing up the rest.
soundness: 4
strengths: * The proposed method makes sense and appears straight-forward to implement. 
* The method appears general to problems with a parametric model where $p(x,\theta)$ can be computed.
* Comparison to SNIS is revealing and clear.
summary: The paper proposes a new way to approximate the posterior distribution over parameters for a parametric model $p(x;\theta)$. 

The approach:
* defines 2 multinomial distributions:
   *  $y(\\{\theta_k\\})$ over the index $k$ of $K$ parameter samples $\theta_k$ as a function of the model $p(x,\theta)$ and noise distribution; 
   * $\hat{y}(\\{\theta_k\\})$ as a function of the parameterised approximate posterior $q_\phi(\theta)$ and the noise distribution;
* $\hat{y}$ is trained to match $y$ under cross entropy loss;
* at optimality $q_\phi(\theta)$ matches the true posterior $p(\theta|x)$ (subject to the choice of the parametric family $q_\phi$).

Results compare the proposed method to typical amortised variational inference (maximising the ELBO) and a more similar method, showing improved outcomes on those datasets.
weaknesses: Overall I like the paper, find it to be of high quality and soundness and the proposed method is a clever trick that converts a difficult problem into a simpler classification-type task. Areas in which the paper could improve (from good to excellent) are: 

**Clarity**: as someone familiar with this field but not immersed in it, the paper could be clearer to read. As specific examples:
* [033, 078] From the outset, the paper refers to *latent variables*, which often refer to variables that are one-to-one with data samples, whereas the term *parameters* typically relates to global variables of the model, i.e. common across all samples. (Of course parameters are often treated as unobserved, hence latent, variables). I appreciate that this may be subjective, but I find it to be fairly universal and a useful distinction, so this phrase put the wrong starting point in my mind. I was wondering throughout when "parameters of the model" get updated and only when I reached the experiments did it become clear that it is about parameters and it made more sense. I would suggest referring to "parameters" (treated as latent variables).
* The preamble to Eq 2 is not directly relevant and could be more clear: labels y are constructed analytically and the approximate posterior is trained by learning to predict them. The backstory of predicting k* doesn't seem so relevant (as it is in NCE, say) so this bit threw me off until I realised that.
* [202-213] This part stands out as less precise relative to preceding arguments. The method trains one ratio of ratios to match another ratio of ratios using cross entropy loss. Reference to "classification being too easy" and "punishing" seems out of place and unclear. The issue seems to be the relative range of values in the softmax function (+/- numerical instability?), which is mentioned but blended with loose wording. This part could be improved, particularly as computing these ratios/choosing the negative sampling distribution is the crux and potential pitfall of ratio estimation.
* [161] Ensuring $\pi$ covers the true posterior - I can't see how this is "guaranteed", unless pi is diffuse, which sounds problematic for ratios.
* Experiments: to compute a posterior, you need a prior, but the priors don't seem clearly specified.
* Fig 1: how are some log probability values positive?
* GARCH experiment: The descriptions of $\alpha_1$ and $\beta_1$ are rather vague.

**Analysis/Experiments**: For such a widely applicable method, the experiments to showcase its efficacy/analyse performance etc seem thin. 
* Understanding how well it works in higher dimensions is of interest, where potentially finding a negative distribution close to the true posterior could be less straightforward.
* Understanding how the methods compare as number of data points increases is of interest to know when the method is most applicable.
* Understanding how the methods compare when the posterior is mis-specified is of interest to understand robustness (mode seeking may be preferable to distribution covering?).

Minor
* [134] it would aid clarity to explicitly say optimisation is "with respect to $\phi$"
* [134] it would be much clearer to state labels y "(from Eq 2)".
* [370] "a different observation" of what?
* [421-424] the paper does well to be relatively self-contained, this part could be made more clear to the less familiar reader.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 4
questions: 1. Given the connection between SNIS-fKL and SoftCVI and that MC estimation of the SNIS-fKL objective is a biased estimation of the forward KL, does SoftCVI give a biased estimation of the posterior when $p^-(\theta)=\pi(\theta)=q_w(\theta)$? Could you clarify whether SoftCVI with $p^-(\theta)=q(\theta)^\alpha$ gives a biased estimation of the posterior?
2. What is claimed when stating that SoftCVI enables a stable objective (line 064)? Is numerical stability, stable convergence or something else? What evidence is there for the type of stability claimed?
3. Why does the equation after eq.4 hold (line 146)? It seems like there is a population limit ($\lim_{K \rightarrow \infty}$) missing. Is this the case? Otherwise, please explain why this is not necessary.
4. The article establishes that $\alpha=0$ leads to leakage and that $\alpha=0.75$ is an good choice for the experiments. How do you suggest deciding $\alpha$? Could you provide theoretical justification or guidelines for choosing in practice $\alpha$? And, is there a threshold under which $\alpha$ will lead to leakage, or is this only the case for $\alpha=0$?
5. Could you explain the second equality in the derivation leading to eq. 26? For example, by a step-by-step derivation.
6. The claim that the parametric variational distribution can recover the posterior exactly under only a support constraint (lines 143-144) seems odd. Is there a condition that the posterior is contained within the parametric variational family missing? If so, please add this and any other assumption to the paragraph. If not, explain the flaw in the following case: Let the posterior $p(\theta|\mathcal{D})$ be multi-modal (for example, a mixture posterior) and the variational distribution parameterized by $\phi$ be unimodal, e.g., Gaussian, both supported on the real line. By construction, we would have that $\forall \phi.p(\theta|\mathcal{D}) \not = q_\phi(\theta)$ contradicting the claim in lines 143-144 that there $\exists \phi. p(\theta|\mathcal{D}) = q_\phi(\theta)$ namely the $\phi$ that parameterizes the optimal classifier from eq. 3.
soundness: 3
strengths: ## Reason for score
- The presentation is lucid.
- The derivations look good,  aside from a couple of technicalities (see questions).
- Using soft labels in contrastive learning for Bayesian inference is exciting and novel.
- The connection to SNIS-fKL is clear, and viewing the gradient of eq. 7 as a control variate is insightful and supported by evidence.
- Experimenting with normal and normalizing flow variational distribution demonstrates SoftCVI's flexibility, and the results support that SoftCVI is more mass-covering than ELBO and SNIS-fKL.
summary: The paper proposes SoftCVI, a variational inference method based on contrastive learning that uses soft labels, i.e., binary classification probabilities. SoftCVI reframes fitting a variational distribution as finding a classifier to identify actual posterior samples among a set of samples. The paper claims that:
  1. SoftCVI works without having access to samples from the posterior.
  2. SoftCVI forms objectives that are mass-covering.
  3. SoftCVI yields better posterior approximation than baselines with simple and flexible variational distributions.
weaknesses: - It is unclear whether the variational distribution is a biased estimator of the posterior. Please see the questions.

Meta-review: This paper proposes a noise contrastive estimation (NCE) type of objective for variational inference. The idea is to use the unnormalised posterior and a noise distribution to assign a "soft label" (i.e., the classification probability vector) on samples from a proposal distribution, and then fit another classifier to this "dataset" where the classifier itself is defined using density ratios between an approximate posterior and the same noise distribution used in the "soft label" generation process. 
The idea is clear and the presentation is sound for people who knows variational inference and NCE. Initially reviewers had questions regarding some details, and they also requested further experiments. These concerns were mostly addressed during author rebuttal.

After a brief read, I was a bit surprised that, although the idea looks very obvious in hindsight, to the best of my knowledge no previous work has formalised this idea and made it work. Multi-class NCE has been around for sometime regarding mutual info estimation but it doesn't work very well in high dimensions, so I would be curious to read more about what's the key change that make things work (I am guessing the improved SNR), and whether that key technique can also benefit mutual info estimation. 

As someone who has worked extensively in both VI and mutual info estimation, I think this paper worths a highlight for the probabilistic modelling community at ICLR. I think the paper did a reasonably good job beyond just presenting the idea -- their algorithmic choices are also well justified with both mathematical and empirical analyses. The paper also positioned its contribution fairly well by a detailed comparison with many of existing works, except that -- 

There is another line of work in VI literature that also used the classification loss idea, e.g., see https://arxiv.org/abs/1701.04722 and other related papers published around the same time. It would be useful if the authors further discuss this line of work in camera ready, making their positioning effort complete.

Decision: Accept (Spotlight)