Bayesian Appropriation: General Likelihood for Loss Functions

In this blog post, we explore how some losses could be rewritten as a Bayesian objective using ideas from variational inference—hence, the tongue-in-cheek “Bayesian Appropriation.” This can make it easier to see connections between loss functions and Bayesian methods (e.g. by spotting similar patterns in the wild). We will first provide a brief overview of the necessary background concepts, followed by step-by-step derivations. Overall, this post focuses on a rather simple observation and is mainly intended to be a straightforward reference for myself, but I hope it will be useful to others as well.

Background

In the following background section, we look at the Gibbs/Boltzmann distribution, the information-theoretic notation we are going to use, and the variational inference. These concepts are essential to understanding the main topic of this blog post, so we will provide a brief explanation of each concept to make it more accessible to readers who may not be familiar with them. As always, feel free to skim and skip as is best :-)

Gibbs/Boltzmann Distribution

The Gibbs or Boltzmann distribution is a probability distribution that describes the probability of a system being in a particular state, given its energy. It is based on the principle of maximum entropy, which states that the most likely state of a system is the one with the highest entropy (or disorder). In simple terms, the Gibbs distribution helps us understand how systems tend to distribute their energy among various possible states. The discrete form of the Gibbs distribution is given by the formula: $$P(E)=\frac{e^{-E/kT}}{\sum _{i=1}^N{e}^{-E_i/kT}}$$ where $E$ is the energy of the system, $k$ is Boltzmann’s constant, and $T$ is the temperature. The denominator is a normalizing factor that ensures that the probabilities sum to 1.

Information-Theoretic Notation

Information theory deals with the quantification, storage, and communication of information. In this blog post, we use information-theoretic notation to express various quantities related to probability distributions and their relationships. Some key quantities we will use include the information content of an event $x$, the entropy of a random variable $X$, and the Kullback-Leibler Divergence (🥬) between two probability distributions:

The information content of an event $x$ is $-\log \mathrm{p}(x)$: $$\mathrm{H}[x]\triangleq \mathrm{H}(\mathrm{p}(x))\triangleq -\log \mathrm{p}(x).$$ We commonly use this information content as minimization objective in deep learning (as negative log likelihood, which is just the cross-entropy).

The entropy is the expectation over the information content: $$\mathrm{H}[X]={\mathbb{E}}_{\mathrm{p}(x)}[\mathrm{H}[x]]={\mathbb{E}}_{\mathrm{p}(x)}[\mathrm{H}(\mathrm{p}(x))]={\mathbb{E}}_{\mathrm{p}(x)}[-\log \mathrm{p}(x)],$$ where we use the following convention (introduced in “A Practical & Unified Notation for Information-Theoretic Quantities in ML” (Kirsch et al, 2021)): when we have an expression with (upper-case) random variables and (lower-case) observations, we take an expectation over all the random variables conditional on all the given observations. For example, the joint entropy $\mathrm{H}[X,y|Z]$ is: $$\mathrm{H}[X,y|Z]={\mathbb{E}}_{\mathrm{p}(x,z|y)}[\mathrm{H}[x,y|z]].$$

We will also use the Kullback-Leibler Divergence (🥬) and the cross-entropy: $$\begin{array}{rl}{\mathrm{D}}_{\mathrm{KL}}(\mathrm{p}(X)\Vert \mathrm{q}(X))& =\mathrm{H}(\mathrm{p}(X)\Vert \mathrm{q}(X))-\mathrm{H}(\mathrm{p}(X))\\ & ={\mathbb{E}}_{\mathrm{p}(x)}[\mathrm{H}(\mathrm{q}(x))]-{\mathbb{E}}_{\mathrm{p}(x)}[\mathrm{H}(\mathrm{p}(x))]\\ & ={\mathbb{E}}_{\mathrm{p}(x)}[-\log \mathrm{q}(x)]-{\mathbb{E}}_{\mathrm{p}(x)}[-\log \mathrm{p}(x)].\end{array}$$

Variational Inference

To connect the Gibbs distribution to a Bayesian principle, we start by looking at variational inference.

Variational inference is a technique used in Bayesian statistics to approximate complex probability distributions, such as the posterior distribution of model parameters given observed data. The main idea behind variational inference is to find an approximate distribution that is close to the true distribution by minimizing a divergence measure (usually the Kullback-Leibler divergence) between the two distributions. This is particularly useful when the true posterior distribution is difficult or computationally expensive to compute directly (or even infeasible to compute).

We can side-step computing the actual posterior by using Bayes’ law to express it via a prior distribution and a likelihood: let’s assume, we have model parameters $\mathit{\omega }$ and an observation $x$. We have a prior distribution $\mathrm{p}(\mathit{\omega })$ and a likelihood $\mathrm{p}(x|\mathit{\omega })$.

We are interested in performing Bayesian inference and learning about $\mathrm{p}(\mathit{\omega }|x)$. We can do so by using a different parameterized distribution ${\mathrm{q}}_{\psi }(\mathit{\omega })$ and optimizing $\psi$ to make ${\mathrm{q}}_{\psi }(\mathit{\omega })$ be as close to $\mathrm{p}(\mathit{\omega }|x)$ as possible.

We will use a functional approach for now, which is to say, we ignore $\psi$ and use $\mathrm{q}(\mathit{\omega })$.

Then, we are interested in: $$\arg \underset{\mathrm{q}}{min}{\mathrm{D}}_{\mathrm{KL}}(\mathrm{q}(\mathbf{\Omega })\Vert \mathrm{p}(\mathbf{\Omega }|x))$$ For divergences (and the 🥬 in particular), it is 0 iff $\mathrm{q}(\mathit{\omega })=\mathrm{p}(\mathit{\omega }|x).$ That is, the two distributions match.

For variational inference, we apply Bayes’ law to obtain: $$\begin{array}{rl}& {\mathrm{D}}_{\mathrm{KL}}(\mathrm{q}(\mathbf{\Omega })\Vert \mathrm{p}(\mathbf{\Omega }|x))={\mathrm{D}}_{\mathrm{KL}}(\mathrm{q}(\mathbf{\Omega })\Vert \frac{\mathrm{p}(x|\mathbf{\Omega })\mathrm{p}(\mathbf{\Omega })}{\mathrm{p}(x)})\\ & =\mathrm{H}(\mathrm{q}(\mathbf{\Omega })\Vert \mathrm{p}(x|\mathbf{\Omega }))+\mathrm{H}(\mathrm{q}(\mathbf{\Omega })\Vert \mathrm{p}(\mathbf{\Omega }))\\ & -\mathrm{H}(\mathrm{p}(x))-\mathrm{H}(\mathrm{q}(\mathbf{\Omega }))\\ & =\mathrm{H}(\mathrm{q}(\mathbf{\Omega })\Vert \mathrm{p}(x|\mathbf{\Omega }))+{\mathrm{D}}_{\mathrm{KL}}(\mathrm{q}(\mathbf{\Omega })\Vert \mathrm{p}(\mathbf{\Omega }))-\mathrm{H}(\mathrm{p}(x)).\end{array}$$ This leads directly to the ELBO (though it is an upper-bound in information-theory): $$\begin{array}{rl}& {\mathrm{D}}_{\mathrm{KL}}(\mathrm{q}(\mathbf{\Omega })\Vert \mathrm{p}(\mathbf{\Omega }|x))\ge 0\\ ⟺& \mathrm{H}(\mathrm{q}(\mathbf{\Omega })\Vert \mathrm{p}(x|\mathbf{\Omega }))+{\mathrm{D}}_{\mathrm{KL}}(\mathrm{q}(\mathbf{\Omega })\Vert \mathrm{p}(\mathbf{\Omega }))\ge \mathrm{H}(\mathrm{p}(x))\\ ⟺& {\mathbb{E}}_{\mathrm{q}(\mathit{\omega })}[-\log \mathrm{p}(x|\mathit{\omega })]+{\mathrm{D}}_{\mathrm{KL}}(\mathrm{q}(\mathbf{\Omega })\Vert \mathrm{p}(\mathbf{\Omega }))\ge -\log \mathrm{p}(x),\end{array}$$ where:

• ${\mathrm{D}}_{\mathrm{KL}}(\mathrm{q}(\mathbf{\Omega })\Vert \mathrm{p}(\mathbf{\Omega }))$ is the prior term (regularizer), and
• ${\mathbb{E}}_{\mathrm{q}(\mathit{\omega })}[-\log \mathrm{p}(x|\mathit{\omega })]$ is the (marginal) loss.
• $-\log \mathrm{p}(x)$ is the evidence.

Expressing Losses as a Bayesian Objective

The main idea of this blog post is to show that we rewrite some loss functions and the corresponding optimization objectives into a more Bayesian form by introducing matching likelihoods or joint probabilities. In other words, we can represent the loss using probabilities, which might allows us to connect the loss function to Bayesian principles and variational inference—we do not go into the details of this though but present a few research works that can be a starting point for a more serious exploration of this topic.

General Joint Loss

To demonstrate the idea, we first consider an entropy-regularized general objective of the form: $$\arg \underset{\mathrm{q}}{min}{\mathbb{E}}_{\mathrm{q}(\mathit{\omega })}[L(x,\mathit{\omega })]-\mathrm{H}(\mathrm{q}(\mathbf{\Omega })),$$ where we want to minimize a “general joint loss” $L(x,\mathit{\omega })$ for observation $x$ given model parameters $\mathit{\omega }$ while maximizing the entropy of the approximate distribution. $L(x,\mathit{\omega })$ does not have to be normalized and does not have to be non-negative.

By defining a joint probability distribution that is proportional to the exponential of the negative loss function, we can express this objective as a probabilistic objective and find the optimal solution using variational inference as well.

Concretely, we solve for an equivalent (probabilistic) joint probability likelihood: $L(x,\mathit{\omega })\equiv -\log \mathrm{p}(x,\mathit{\omega })$. We thus let: $$\mathrm{p}(x,\mathit{\omega })\propto \exp (-L(x,\mathit{\omega })).$$With normalization constant $Z_L=\int \exp (-L(x,\mathit{\omega }))dxd\mathit{\omega },$ we have:$$\mathrm{p}(x,\mathit{\omega })=\frac{\exp (-L(x,\mathit{\omega }))}{Z_L},$$ and vice-versa: $$L(x,\mathit{\omega })=-\log \mathrm{p}(x,\mathit{\omega })-\log Z_L.$$

Hence, after dropping the $\log Z_L$ constant term, we observe the equivalent objective: $$\begin{array}{rl}& {\mathbb{E}}_{\mathrm{q}(\mathit{\omega })}[-\log \mathrm{p}(x,\mathit{\omega })]-\mathrm{H}(\mathrm{q}(\mathbf{\Omega }))\\ & \stackrel{(1)}{=}\underset{\mathrm{H}(\mathrm{q}(\mathbf{\Omega })\Vert \mathrm{p}(\mathbf{\Omega }|x))}{\underbrace{{\mathbb{E}}_{\mathrm{q}(\mathit{\omega })}[-\log \mathrm{p}(\mathit{\omega }|x)]}}-\mathrm{H}(\mathrm{q}(\mathbf{\Omega }))-\log \mathrm{p}(x)\\ & \stackrel{(2)}{=}{\mathrm{D}}_{\mathrm{KL}}(\mathrm{q}(\mathbf{\Omega })\Vert \mathrm{p}(\mathbf{\Omega }|x))-\log \mathrm{p}(x)\\ & \ge -\log \mathrm{p}(x),\end{array}$$ where (1) follows from the chain rule ($\mathrm{p}(a,b)=\mathrm{p}(a|b)\mathrm{p}(b)$), and (2) is just the definition of the KL (🥬) divergence. This is following the same steps as variational inference.

Optimal Solution

What does this tell us about the optimal distribution?

From above, we can read off the optimal solution as $\mathrm{p}(\mathit{\omega }|x)$, which is the posterior distribution for appropriately defined $\mathrm{p}(\mathit{\omega })$ and $\mathrm{p}(x|\mathit{\omega }):$ $$\begin{array}{rl}\mathrm{p}(\mathit{\omega })& =\int \mathrm{p}(x,\mathit{\omega })dx\\ \mathrm{p}(x|\mathit{\omega })& =\frac{\mathrm{p}(x,\mathit{\omega })}{\mathrm{p}(\mathit{\omega })}.\end{array}$$ These terms might often be intractable analytically, but when we can solve them, we have a closed form for an equivalent variational inference problem.

Specifically, we have:$$\mathrm{p}(\mathit{\omega }|\mathit{x})=\frac{\mathrm{p}(\mathit{x},\mathit{\omega })}{\mathrm{p}(\mathit{x})}=\frac{\exp (-L(x,\mathit{\omega }))}{\int \exp (-L(x,\mathit{\omega }))d\mathit{\omega }}.$$ Writing this without the normalization constant and keeping in mind that we have to normalize over $\mathit{\omega }$, we have: $$\arg \underset{\mathrm{q}}{min}{\mathbb{E}}_{\mathrm{q}(\mathit{\omega })}[L(x,\mathit{\omega })]-\mathrm{H}(\mathrm{q}(\mathbf{\Omega }))\propto \exp (-L(x,\mathit{\omega })).$$

Aside: Invariance to Variable Transformation

Once, we have mapped a loss to a probabilistic form, the parameterization of $\mathit{\omega }$ does not matter (assuming smooth reversible transformations).

This is because we can re-parameterize the Kullback-Leibler divergence without change in value: $${\mathrm{D}}_{\mathrm{KL}}(\mathrm{q}(X)\Vert \mathrm{q}(X))={\mathrm{D}}_{\mathrm{KL}}(\mathrm{q}(Y)\Vert \mathrm{p}(Y)),$$when $Y=f(X)$ and $f$ is smooth and bijective.

Above optimization problem might attain a different value for $-\log \mathrm{p}(x)$, but the optimal solution for $\mathrm{q}$ is the same.

Maximum Entropy Principle

This connection between loss functions and Bayesian objectives can also be interpreted through the lens of the maximum entropy principle—we already see this from the entropy regularization above.

We can find the optimal solution that minimizes the expected loss while maximizing the entropy of the approximate posterior distribution. This provides an alternative perspective on the relationship between loss functions and Bayesian principles.

Let’s look at a simpler objective: $$\arg \underset{\mathrm{q}}{min}{\mathbb{E}}_{\mathrm{q}(\mathit{\omega })}[L(\mathit{\omega })]-\mathrm{H}(\mathrm{q}(\mathit{\omega })).$$ First, using the same approach, we see that ${\mathrm{q}}_{\ast }\propto \exp (-L(\mathit{\omega }))$.

Importantly, this also tells us two things for ${\mathrm{q}}_{\ast }$:

1. for fixed entropy $\mathrm{H}(\mathrm{q}(\mathit{\omega }))$, this is also the minimum of ${\mathbb{E}}_{\mathrm{q}(\mathit{\omega })}[L(\mathit{\omega })];$ and
2. for fixed ${\mathbb{E}}_{\mathrm{q}(\mathit{\omega })}[L(\mathit{\omega })]$, this is a maximum entropy distribution (maximum $\mathrm{H}(\mathrm{q}(\mathit{\omega }))$).

This simply follows from ${\mathrm{q}}_{\ast }$ minimizing above combined objective: if there was a better value of either while the other remains fixed, we would have found a better optimum.

Another way to argue is that whenever we have an optimization object that can be split into a “loss” and a “regularizer”, we can obviously swap the two: $$\arg \underset{\mathrm{q}}{min}{\mathbb{E}}_{\mathrm{q}(\mathit{\omega })}[L(\mathit{\omega })]-\mathrm{H}(\mathrm{q}(\mathit{\omega }))=\arg \underset{\mathrm{q}}{max}\mathrm{H}(\mathrm{q}(\mathit{\omega }))-{\mathbb{E}}_{\mathrm{q}(\mathit{\omega })}[L(\mathit{\omega })].$$ All of this is not rocket science, but worth pointing out nonetheless.

General Likelihood

We can apply this approach to simpler loss functions that only depend on inputs and a separate regularization term.

Specifically, we consider objectives of the form: $$\arg \underset{\mathrm{q}}{min}{\mathbb{E}}_{\mathrm{q}(\mathit{\omega })}[L(x;\mathit{\omega })]+{\mathrm{D}}_{\mathrm{KL}}(\mathrm{q}(\mathbf{\Omega })\Vert \mathrm{p}(\mathbf{\Omega })),$$ where we use the notation $L(x;\mathit{\omega })$ for the loss for observation $x$ given model parameters $\mathit{\omega }$. $L(x;\mathit{\omega })$ does not have to be normalized and does not have to be non-negative.

We can rewrite the objective as: $${\mathbb{E}}_{\mathrm{q}(\mathit{\omega })}[L(x;\mathit{\omega })]+{\mathrm{D}}_{\mathrm{KL}}(\mathrm{q}(\mathbf{\Omega })\Vert \mathrm{p}(\mathbf{\Omega }))={\mathbb{E}}_{\mathrm{q}(\mathit{\omega })}[L(x;\mathit{\omega })+\mathrm{H}(\mathrm{p}(\mathit{\omega }))]-\mathrm{H}(\mathrm{q}(\mathbf{\Omega })),$$ by using the definition of the KL (🥬) divergence and cross-entropy.

Using the joint loss $L^{\prime }(x,\mathit{\omega })\triangleq L(x;\mathit{\omega })+\mathrm{H}(\mathrm{p}(\mathit{\omega }))$, we immediately see from section “General Joint Loss” that we have the following equivalent objective with optimal solution: $$\arg \underset{\mathrm{q}}{min}{\mathbb{E}}_{\mathrm{q}(\mathit{\omega })}[L(x;\mathit{\omega })]+{\mathrm{D}}_{\mathrm{KL}}(\mathrm{q}(\mathbf{\Omega })\Vert \mathrm{p}(\mathbf{\Omega }))\propto \exp (-L^{\prime }(x,\mathit{\omega }))=\exp (-L(x;\mathit{\omega }))\mathrm{p}(\mathit{\omega }),$$ where $\exp (-L(x;\mathit{\omega }))$ can be viewed as a `generalized likelihood’.

However, this `generalized likelihood’ is not necessarily a probabilistic likelihood, and the prior $\mathrm{p}(\mathit{\omega })$ is not necessarily the prior of an equivalent Bayesian objective.

To see this, let’s follow the steps from the previous section to obtain the probabilistic form of the objective using a model ${\mathrm{p}}^{\prime }(x,\mathit{\omega })$: $$\begin{array}{rl}{\mathrm{p}}^{\prime }(x,\mathit{\omega })& =\frac{\exp (-L(x;\mathit{\omega }))\mathrm{p}(\mathit{\omega })}{\int \exp (-L(x;\mathit{\omega }))\mathrm{p}(\mathit{\omega })dxd\mathit{\omega }}\\ {\mathrm{p}}^{\prime }(\mathit{\omega })& =\int {\mathrm{p}}^{\prime }(x,\mathit{\omega })dx\\ {\mathrm{p}}^{\prime }(x|\mathit{\omega })& =\frac{{\mathrm{p}}^{\prime }(x,\mathit{\omega })}{{\mathrm{p}}^{\prime }(\mathit{\omega })}.\end{array}$$

The equivalent Bayesian objective is then: $$\begin{array}{rl}& \arg \underset{\mathrm{q}}{min}{\mathbb{E}}_{\mathrm{q}(\mathit{\omega })}[L(x;\mathit{\omega })]+{\mathrm{D}}_{\mathrm{KL}}(\mathrm{q}(\mathbf{\Omega })\Vert \mathrm{p}(\mathbf{\Omega }))\\ & =\arg \underset{\mathrm{q}}{min}{\mathbb{E}}_{\mathrm{q}(\mathit{\omega })}[{\mathrm{p}}^{\prime }(x|\mathit{\omega })]+{\mathrm{D}}_{\mathrm{KL}}(\mathrm{q}(\mathbf{\Omega })\Vert {\mathrm{p}}^{\prime }(\mathbf{\Omega })),\end{array}$$ with the posterior ${\mathrm{p}}^{\prime }(\mathit{\omega }|x)\propto \exp (-L(x;\mathit{\omega }))\mathrm{p}(\mathit{\omega })$ as optimal solution.

Importantly, the prior ${\mathrm{p}}^{\prime }(\mathit{\omega })$ might not be equal $\mathrm{p}(\mathit{\omega })$ from the original objective. Substituting the definitions, we have: $${\mathrm{p}}^{\prime }(\mathit{\omega })\propto \int \exp (-L(x;\mathit{\omega }))dx\mathrm{p}(\mathit{\omega }),$$ which is only $=\mathrm{p}(\mathit{\omega })$ when the integral $\int \exp (-L(x;\mathit{\omega }))dx$ is constant (independent of $\mathit{\omega }$).

Additional notes:

1. We need the $\exp (-\dots )$ because of the $-\log \dots$ that is required for combining the prior distribution with the likelihood in variational inference.
2. This is exactly the Boltzmann/Gibbs distribution with a temperature that cancels out Boltzmann’s constant. To be fair, we could also define $kT$ as the constant with $k$ being a conversion factor between a temparature (in Kelvin) and a unitless quantity.
3. We use probability distributions as a tool to find the posterior distribution using an equivalent variational inference objective.

MAP Solution

For a simpler objective, we can choose a Dirac distribution for the approximate posterior distribution and obtain a maximum a posteriori (MAP) solution. This is particularly useful for deterministic neural networks, as it allows us to convert a loss function with a regularization term into an equivalent Bayesian objective.

We can choose a Dirac distribution for ${\mathrm{q}}_{\psi }(\mathit{\omega })={\delta }_{\psi }(\mathit{\omega })$. Then we obtain the simpler objective: $$\arg \underset{\psi }{min}L(x;\psi )+H(\mathrm{p}(\psi ))=\arg \underset{\mathrm{q}}{min}L(x;\psi )-\log \mathrm{p}(\psi ).$$ Therefore, we can convert a loss for a deterministic neural network of the form: $$\arg \underset{\psi }{min}L(x;\psi )+R(\psi )$$into an equivalent “Bayesian” objective by setting: $$\begin{array}{rl}\mathrm{p}(x|\mathit{\omega })& \propto \exp (-L(x;\mathit{\omega }))\\ \mathrm{p}(\mathit{\omega })& \propto \int \exp (-L(x;\mathit{\omega }))dx\exp (-R(\mathit{\omega })).\end{array}$$ For example, if $R(\psi )={\psi }^2$ (L2 regularization) and $\int \exp (-L(x;\mathit{\omega }))dx$ is constant, we can recognize $\mathrm{p}(\mathit{\omega })\propto \exp (-\frac{(\mathit{\omega }-0{)}^2}{1})$ as unit Gaussian distribution.

Conclusion

In conclusion, we can construct Bayesian objectives by expressing losses as likelihoods and joint probability distributions. This allows us to leverage insights from variational inference to write down an optimal solution.

Keep in mind, however, that we do not have to use variational inference to find the optimal solution and that writing down the solution does not mean that we can actually compute it. Actually computing it would require computing the partition constants, too.

On the other hand, we don’t always need the partition constant or the normalized density. For example, we can use the unnormalized density to sample from the distribution.

It might also be helpful to know the form of these solutions to recognize them in other contexts and to use them as a starting point.

Please leave some feedback and send along more papers for the next section :)

More Serious Research Papers

Knoblauch, Jeremias, Jack Jewson, and Theodoros Damoulas. “Generalized variational inference: Three arguments for deriving new posteriors.” arXiv preprint arXiv:1904.02063 (2019). https://arxiv.org/abs/1904.02063

Wild, Veit David, Robert Hu, and Dino Sejdinovic. “Generalized variational inference in function spaces: Gaussian measures meet bayesian deep learning.” Advances in Neural Information Processing Systems 35 (2022): 3716-3730. https://arxiv.org/abs/2205.06342

Kendall, Alex, Yarin Gal, and Roberto Cipolla. “Multi-task learning using uncertainty to weigh losses for scene geometry and semantics.” Proceedings of the IEEE conference on computer vision and pattern recognition. 2018. https://arxiv.org/abs/1705.07115v3 https://openaccess.thecvf.com/content_cvpr_2018/html/Kendall_Multi-Task_Learning_Using_CVPR_2018_paper.html

Cobb, Adam D., Stephen J. Roberts, and Yarin Gal. “Loss-calibrated approximate inference in Bayesian neural networks.” arXiv preprint arXiv:1805.03901 (2018). https://arxiv.org/abs/1805.03901

Bissiri, Pier Giovanni, Chris C. Holmes, and Stephen G. Walker. “A general framework for updating belief distributions.” Journal of the royal statistical society. series b, statistical methodology 78.5 (2016): 1103. https://arxiv.org/abs/1306.6430 https://www.ncbi.nlm.nih.gov/pmc/articles/PMC5082587/

Ganev, Stoil, and Laurence Aitchison. “Semi-supervised learning objectives as log-likelihoods in a generative model of data curation.” arXiv preprint arXiv:2008.05913 (2020). https://arxiv.org/abs/2008.05913

Aitchison, Laurence. “InfoNCE is a variational autoencoder.” arXiv preprint arXiv:2107.02495 (2021).
https://arxiv.org/abs/2107.02495

Acknowledgements. Thanks to Linara Adilova for valuable discussion & feedback, to Laurence Aitchison for spotting a mistake in the initial derivations and interesting discussions, and to Yarin Gal for pointing to additional related literature.