InfoNCE is variational inference in a recognition parameterised model
Authors: Laurence Aitchison
Stoil Krasimirov Ganev
Venue: Accepted by TMLR ()

Abstract: Here, we develop a new class of Bayesian latent variable model, the recognition parameterised model (RPM). RPMs have an implicit likelihood, which is defined in terms of the recognition model. Therefore, it is not possible to do traditional "generation" with RPMs. Instead, RPMs are designed to learn good latent representations of data (in modern parlance, they solve a self-supervised learning task). Indeed, the RPM implicit likelihood is specifically designed so that it drops out of the VI objective, the ELBO. That allows us to learn an RPM without a "reconstruction" step, which is believed to be at the root of poor latent representations learned by VAEs. Indeed, in a very specific setting where we learn the optimal prior, the RPM ELBO becomes equal to the mutual information (MI; up to a constant), establishing a connection to pre-existing self-supervised learning methods such as InfoNCE.

Published in Transactions on Machine Learning Research (03/2024)
InfoNCE is variational inference in a recognition parame- terised model
Laurence Aitchison University of Bristol
lourence.citchhsom@@aaa..o
Stoil Ganev University of Bristol
stoil.gamen@@bristoll.cco
Reviewed on OpenReview: https: // openreview. net/ met/ forum? id- chbRswujan
Abstract
Here, we develop a new class of Bayesian latent variable model, the recognition parameterised model (RPM). RPMs have an implicit likelihood, which is defined in terms of the recognition model. Therefore, it is not possible to do traditional "generation" with RPMs. Instead, RPMs are designed to learn good latent representations of data (in modern parlance, they solve self-supervised lcarning task). Indeed, the RPM implicit likelihood is specifically designed so that it drops out of the VI objective, the ELBO. That allows us to learn an RPM without a "reconstruction" step, which believed to be at the root of poor latent representations lcarned by VAEs. Indeed, a very specific sctting where we learn the optimal prior, the RPM ELBO becomes equal to the mutual information (MI; up to constant), establishing connection to pre-existing self-supervised learning methods such as InfoNCE.
1 Introduction
Self-supervised learning (SSL) involves learning structured, high-level representations of data useful for downstream tasks such as few-shot classification. One common self-supervised approach is to define "pretext" classification task (Dosovitskly et al. 2015 Noroozi & Favaro 2016 Doersch et al. 2015 Gidaris et al. 2018). For instance, we might take a number of images, rotate them, and then ask the model to determine the rotation applicd (Gidaris et al. 2018 The rotation can be identified by looking at the objects in the image (e.g- grass typically on the bottom of the image, while birds are nearer the top), and thus a representation useful for determining the orientation may also extract useful information for other high-level tasks. We are interested in altcrnative class of SSL objectives known InfoNCE (NCE standing for noise contrastive estimation; (Oord et al. 2018 Chen et al. 2020 These methods take two inputs (e.g. two different frames of video), encode them to form two latent representations, and use classification task to maximize a bound on the mutual information between latent representations. As the shared information should concern high-level propertics such as objects, but not low-level details of each patch, this should again extract a useful representation, which is also invariant data augmentation.
Cltimately, the goal of SSL is to extracting useful, structured, high-level representations of data. Inter- estingly, such representations have classically been obtained using a probabilistic latent variable models (PLVMs; Murphy 2022). For instance, consider Latent Dirichlet Analysis (LDA; Blei et al. 2003). LDA clusters words and documents into topics. Given a new document, LDA allows us to assign a combination of topics to that document. Importantly, LDA, as a probabilislic latent variable model, allows us to sample "fake data". However, in the case LDA, it is not at all clear why you would ever want to sample fake data. Specifically, LDA treats documents as "bags of words"; while you could sample bags of words, it unclear why you ever would ever want to. Alternative classic PLVMs include probabilistic PCA/ICA (Pearlmutter Again, in probabilistic PCA/ICA the point
& Parral 1996 MacKay 1996 Tipping & Bishop
1999.

---

Published in Transactions on Machine Learning Research (03/2024)
always to extract interpretable, high-level factors of variation. While you could sample "fake data", it is again not at all clear why you would ever want, to.
Perhaps the most obvious modern incarnation of traditional PLVMs is the variational autoencoder (VAEs Kingma& & Welling 2013 Rezende et al.) 2014 Kingma et al., 2019 The VAE extracts latent variables, and one would hope that these latents would form a good high-level representation, so that VAFs could be used in SSL. However, VAE's typically perform poorly when evaluated on SSL tasks such as few-shot classification or disentanglement (Karaletsos et al., 2015 Higgins et al. 2016 Zhao et al. 2019). The issues are commonly understood to arise because VAE needs to reconstruct the data, the latent representation is forced to encode low-level details (Chen et al.
2020:
Balestriero & LeCun 2024). This raises an important question:
is it possible to develop new PLVMs that give useful high-level representations useful in SSL by eliminating the need to "reconstruct??
Here, we answer the affirmative by developing a new family of recognition paramctorised model (RPM) (note that the "recognition parameterised model" name first appeared in follow-up work Walker et al. 2023b
but the actual idea was introduced here; see Related work for further details). We perform variational inference in this RPM, the resulting objective is the ELBO. Next, connect the RPM to modern SSL methods, in particular InfoNCE (Oord et al. 2018 Chen et al.) 2020). We show that the RPM ELBO is equal (up to a constant) to the MI under an optimized prior (Sec. 4.4], and equal (up to a constant) to the infinite-sample limit of the InfoNCE objective under a different choice of prior (Sec. 4.5). We give an example in a toy system with a latent space without unit-norm constraints in which the usual choice of InfoNCE objective fails completely, but modified system that exploits our prior knowledge about dynamics in the latent space succeeds (Sec. 5]-
This aligus with recent work (Locatello et al., 2019) which argues that good, problem-specifie inductive biases are critical for effective disentanglement. However, it is difficult to introduce such inductive biases in traditional InfoNCE like settings. In contrast, it is straightforward to introduce these inductive biases into the priors of an RPM.
At a high-level, RPMs function by having observations that consist of multiple components, e.g. following a classic SSL setting, we could have two components, where each component is different augmentation of the same underlying image. RPMs then use latent space which models only dependencies between components, and avoids modelling any information about the marginal distribution over each component. Information about potentially complex marginals is instead captured in a complex likelihood. This likelihood is difficult to specify, and lies at the root of problems around "reconstruction" in classical VAEs. Critically, in the RPM framework the likelihood cancels out of the ELBO (see Sec. 4- In this way, an RPM able to specily a full generative model, but able to avoid the problematic reconstruction step typical VAFs.
Interestingly, show that particular PLVM, the RPM, can be effective for SSL (i.e. learning a good high-level representation useful for downstream tasks such as few-shot classification). We would argue that this effectivenoss arises preciscly because we drop the requirement for casily decoding/sampling "fako-data"; This seems to mirror the manner in which state-of-the-arl, generative models such as diffusions arose from the VAE framework by dropping the requirement for potentially interpretable latent space (Sohl- Dickstickt
et al. 2015 Ho et al.] 2020 Kingma et al. 2021 Song et al. 2021
Our results have important implications for the interpretation of methods such as InfoNCE. Specifically, InfoNCE was thought to learn good representations by (approximately) maximizing mutual information (Oord et al. 2018). However, recent work has argued that maximizing the true mutual information could lead to arbitrarily entangled representations, as the mutual information is invariant under arbitrary invertible transformations (Tschannen et al.) 2019). Instead, they argue that InfoNCE, learns good representations because it uses a highly simplified lower bound on the information estimator (Oord et al. 2018)) which forms only loose bound on the true MI. This is highly problematic: TTscannen et al. (2019) argue that better MI estimators give worse representations. Thus, InfoNCE appears to be successful not because it is maximizing MI, but because of ad-hoc choice of simplified mutual information estimator. So what is InfoNCE doing? And how can the success of its simplified mutual information estimator be understood? We show that the InfoNCE objective is equal (up to constant) to the RPM ELBO. Further, with a deterministic encoder, the RPM ELBO becomes equal to the log marginal likelihood, This would argue that the InfoNCE objective is
2

---

Published in Transactions on Machine Learning Research (03/2024)
better motivated in terms of the RPM ELBO or log marginal likelihood (as they are equal up to constant in the infinite sample setting), as opposed to the mutual information (as the InfoNCE objective only forms a bound in the same setting).
2 Related work
The original version of this paper introduced the key recognition paramcterised model idea in 2021, albeit under a different name. However, due to the vagaries of the conference publishing system, follow-up work (Walker et al. 2023b) was published first (at AISTATs). Walker et al. (2023b) introduced the name "Recog- nition parameterised model", and we agree that this is the right name, so to avoid confusing the literature, we have adopted their terminology. Of course, downstream of the key RPM idea (embodied in the definition of the likelihood in Eq. 10) the papers go in quite different directions, as Walker et al. (2023b) is primarily interested in probabilistic modelling, while we are primarily interested in linking to self-supervised learning. This is particularly evident in three aspects. First, we originally set the approximate posterior to be equal to the recognition distributions, while Walker et al. (2023b) allowed more flexiiiiitty in the approximate posteriors, which can slightly improve the quality of their posterior inferences. Second, we made slightly dif- ferent choices in the "base" distribution (see Appendix A) for details. Third, there are considerable notional differences, in the sense that defined RPMs in terms of a "recognition joint", while Walker et al. (2023b) defined RPMs in terms of factor graph. We believe that the recognition joint approach is more intuitive when generalising RPMs to more complex settings, and when using contrastive losses.
A number of papers have taken forward the ideas originally introduced here forward, not just Walker et al. (2023b). First, Walker et al. (2023a) used the RPM in an out-of-distribution setting. Specifically, they are able retain accurate predictions in an distr--iisttibuttion sctting whether both the input distribution and the mapping from inputs to labels can change. Second, Möllers el al. (2023) used this approach to obtain uncertainty estimates in graph SSL setting. Third, Wang et al. (2022) used these ideas to improve sequential recommendation, by improving modelling of user behaviour in the face of sparsity of user-item interactions, uncertainty and a long-tail of items. Fourth, Hasanzadeh et al. (022)) used our approach to introduce principled Bayesian uncertainty estimates in contrastive learning on graph data. In particular, they were able to show improvements in uncertainty estimation, interpretability and predictive performance. Finally, Jelley et al. (2023) used recognition parameterised model to introduce contrastive methods for metalearning in the few-shot setting (specifically, Eq. 2 in arXiv v0). That said, their work differs in that while they have recognition parameterised generative model, they do not optimize using VI. Instead, they optimize using the predictive log-likelihood.
Moreover, there are interesting connections to work on connecting self-supervised learning to learning in a lincar state-space model (Eysenbach et al. 2024h.
Perhaps the closest prior (as opposed to follow-up) work is Zimmermann et al. (202)), p which also identifies an interpretation of InfoNCE as inference in a principled generative model. Our work differs from (Zim- mermann et a 2021) in that we introduce RPM and explicitly identify novel conncction between the InfoNCE objective and the RPM ELBO. In addition, their approach requires four restrictive assumptions. First, they assume deterministic encoder. In contrast, all our theory applies to stochastic and deterministic encoders. While we do explicitly consider deterministic encoders in Appendix B this is only to show that with deterministic encoders, the ELBO bound is tight - all the derivations outside of this very small section (which includes all our key derivations) use fully general encoders. Second, they assume that the encoder is invertible, which not nccessary in our framework. This is particularly problematic as practical encoders commonly used in contrastive SSL are not invertible. Third, they assume that the latent space is unit hyper- sphere, while in our framework there is no constraint on the latent space. Fourth, they assume the ground truth marginal of the latents of the generative process is uniform, whereas our framework accepts any choice of ground-truth marginal. As such, our framework has considerably more flexibility to include rich priors on complex, structured latent spaces.
Other work looked at the specific case of isolating content from style (von Kügelgen et al. 2021 This work used a similar derivation to that in Zimmermann et al. (2021) with slightly different assumptions. While they still required deterministic, invertible encoders, they relax e.g. uniformity in the latent space.
3

---

Published in Transactions on Machine Learning Research (03/2024)
But because they are working in the specific case of style and content variables, they make mimber of additional assumptions on those variables. Importantly, they again do not connect the InfoNCE, objective with the ELBO or log marginal likelihood.
Very different methods use noise-contrastive methods to update a VAE prior (Aneja et al. 2020 Impor- tantly, they still use an explicit decoder.
There is large class of work that seeks use VAΣs extract useful, disentangled representations (e.g. Burgess et al. 2018 Chen al. 2018 Kim & Mnih 2018; Mathieu al. [019 Joy al.) 2020). Again, this work differs from our work in that it uses explicit decoders and thus does not identily an explicit link to self-supervised learning.
Likewise, there work on using GANs to learn interpretable latent spaces (e.g. Chen et al. 2016a Im- portantly, GANs learn a decoder (mapping from the random latent space to the data domain). Morcover, GANs use a classifier estimate a density ratio. However, GANs estimate this density ratio for the data, and x', whereas InfoNCE, like the methods described here, uses classifier to estimate a density ratio the latent space, and z'.
There is work on reinterpreting classifiers as energy-based probabilistic generative models (e.g. Grathwohl et al. 2019), which is related we view SSL methods as being analogous to a classifier. Our work is very different, if for no other reason than because it not possible to sample data from a RPM (even using a method like MCMC), because the decoder written in terms of the unknown true data distribution.
3 Background
3.1 Variational inference (VI)
In VT, we have observed data, x, and latents, and we specify prior, P P()), a likelihood, P (z|22, and an approximate posterior, Q (=[x). When the approximate posterior, (z|æx) is parameterised by a neural network, the resulting model is known as a variational autoencoder (Kingma & Welling 2013 Rezende
jet al., 2014p. We then jointly optimize parameters of the prior, likelihood and approximate posterior using the ELBO as the objective,
((x)z)
P()) 2 L(x) = EE((a)))
log Q (2|z)
(1)
which bounds the log marginal likelihood, log P (x) (as can be shown using Jensen's incquality).
We can rewrite the ELBO as an expected log-likelihood minus a KI-divergence,
((x) = EQ(=)=) [log P (æ|=)] DKL (Q (z]=)]IP(2))
(2)
That KL-divergence is very closely related (with particular choice of prior, (z)) to the MI Alemi et al. 2018 Chen et al. 2016b), so the KT-divergence can be understood intuitively as reducing the MI between data, x and latent, 2. However this connection between the ELBO and MI is not relevant for our work. In particular, the ELBO can be interpreted as minimizing the MI between data and latents, whereas InfoNCE maximizes the MI between different latent variables in structured model.
3.2 InfoNCE
In InfoNCE (Oord et al. 2018 there are two data items, and x'. Oord al. (2018) initially describes time-scries setting where for instance x is the previous datapoint and x' is the current datapoint. But one can also consider other contexts where a and x' are different. augmentations or patches of the same underlying image (Chen et al. 2020). In all cases, there are strong dependencies between æ and x', and the goal is to capture these dependencics with latent representations, 2 and z', formed by passing and x' through neural network encoders. All our derivations consider fully general encoders, R (2|z) and R.g. (2'[2") which could be stochastic deterministic (see Sec. B where we discuss this distinction in-depth). Note that we write
4

---

Published in Transactions on Machine Learning Research (03/2024)
these encoders with explicit & subscripts to emphasise that these are user-specified conditional distributions, often Gaussians over or z' with means and variances specified by applying neural network to or x'.
Importantly, we do not just use R to denote the encoder distributions, R (z|ær) and R& (z'[z'). We gener- alise R. to denote the joint distribution arising from taking the true empirical training data distribution, Raase x") (see Appendix A and encoding using R (z|z) and R (z/|x"). Thus, the joint distribution of all random variables under R can be written,
R (2,,'',
Ro (=/²) R (z'(z') Rhase (x,x²).
We can obtain marginals and conditionals of this joint, such R R 2), R (z')., R R (x'l2) in the usual manner - by marginalising and/or applying Bayes theorem. For instance,
R (2,2)) = f de de' R (z/z) R (2/1a") Rbase 2x,2²).
(4)
Note that R 22) also implicitly depends d. We use the subscript in Rg (2|z) to remind the reader that RS (2|æ² is directly paramcterised by c.g. a neural network with weights o, and we omit the subscript in e.g. R 22) to indicate that its dependence on & is only implicit (through Eq.
The InfoNCE objective was originally motivated as maximizing the mutual information between latent representations,
R (2'12) R (2/)
R (2',22 ER(Z.*)) log (
MI ER(Z,2))
log
(5)
Of course, observed x and x' exhibit dependencies as they are e.g. the previous and current datapoints in a timeseries, or different augmentations of the same underlying image. Thus, 2 and z' must also exhibit dependencies under R (2,='). As the mutual information is difficult to compute directly, InfoNCE uses a bound, (0, Φ), based on a classifier that uses Se with parameters 0 to distinguish positive samples (i.e. the z' paired with the corresponding 2) from negative samples (i.c. 2/ drawn from the marginal distribution and unrelated to z or to the underlying data; see Poole et al. 2019 for further details),
20(2,²") Σn 20(
MI > IN E log fo(2,2')
+ log N.
(6)
Here, the expectation taken over R (z, 2²) I1, R (2/), and we use this objective to optimize (the parameters of Sa) and o (the parameters of the encoder). There two source of slack in this bound, arising from finite reduce but not in general eliminate slack by taking the
N and restrictive choice of f. To start, limit as N goes to infinity, (Oord et al. 2018).
MI > [0, 2 (0,
The bound only becomes light if we additionally optimize an arbitrarily flexible ʃ (Oord et al. 2018). If as usual, we have a restrictive parametric family for f. then the bound does not in general become tight (Oord ct al. 2018). In reality InfoNCE does indced use a highly restrictive class of function for f, which can be expected to give loose bound on the MI dOord al. 2018),
So(2, exp (=²0z'),
where 0 for this particular function is matrix. Note additionally that for this particular functional form for fo to work well, we usually need to restrict and z' to have unit-norm.
This raises a critical question: if our goal is really to maximize the bound on the MI, why not use a more flexible fg? The answer that our goal is not ultimately to maximize the MI. Our goal is ultimately to learn a good ropresentation, and MI is merely a mcans to that end. Further, [Tchannen et al. (2019) argue that optimizing the true MI is likely to lead give poor repesentations, the MI invariant arbitrary invertible transformations that can entangle the representation. They go on to argue that it is
5

---

Published in Transactions on Machine Learning Research (03/2024)
precisely the restrictive family of functions, corresponding to loose bound on the MI, that encourages good representations. TTschannen et al. (2019) thus raise an important question: does it really make sense to motivate an objective that works (the InfoNCE objective) as a loose bound on an objective that does not work (the mutual information)? We offer an alternative motivation by showing that the InfoNCE objective is equal (up to a constant) to the log marginal likelihood under a particular choice of prior and with a deterministic encoder.
Note that this loss has been used in downstream settings, such as SimCLR CChn al.
2020
4 Recognition Parameterised Models
An RPM is specific lype of probabilistic latent variable model (PIVM) where the likelihood is written in terms of recognition model. All RPMs have structured graphical model with multiple observations. We start with perhaps the simplest RPM, where observations are a pair, (x,z²), e-g. two augmentations of the same image, or two different but adjacent frames in a video. In this simple example, we consider a model with latent variable, 2 associated with x and z' associated with x'. The generative probability in our probabilistic generative model can be factorised as,
(xx,x' 22,,) = P (1
(9)
Here, we write the prior as Po (z, z') to emphasise that this could be a user-specified distribution with learned parameters, 0. Importantly, P implies a valid joint distribution (Eq. g over all random variables (z, x', x², 2), so any time we write a distribution with P. we mcan a marginal / conditional of that model joint (Eq. 9
Importantly, we have yet to define the likelihood, and is the likelihood that makes it an RPM. In particular, in RPM, the likelihood is defined in terms of Bayesian inference in the recognition joint, R, formed by taking the true/cmpirical data distribution, and encoding using Rg (2/æ² and R, (z'1z') (Eq. 3 We can use Bayes in the recognition joint to obtain R (x|²) and R (z'lz"). We then decide to use R (xl2) and R (x'l²") as the likelihoods our generative model,
(x]=) R (x|2) Rs (2|2r) Ruase (x)
(10a)
R (2)
R (e'l2=) Rbase (x')
P('l)² R (a'12')
(10b)
R (2')
Here, we have written P (x]22) R (a|2) to denote that we are choosing the generative likelihood, P (x]lz) to be equal to R (z/²) from the recognition joint. We have written Bayes theorem with because that
is not choice. All distributions written with R represent a coherent joint distribution (Eq. 3) and hence
R (z|z) must be given by Bayes theorem applied to that joint distribution.
The normalizing constants, R (2) and R (z'), are,
B(2)) / dx Rs (2/(2) Rbase (z).
(11a)
R R(')) = de' Ris (2'(xe) Riaae (x').
(11b)
See Appendix for further details.
4.1 VI in RPMs
Now, we substitute these generative probabilities into the VT objective (L.e. the ELBO), using an approximate posterior, Qg (2,")
Po (2, 2') R (z) R (2')
Rs (2 1 R (2'lx")
L(x,x') = const + EQ
log
+ log
(12)
Q (z, 2/ 15,2"111211 (r,

---

Published in Transactions on Machine Learning Research (03/2024)
Here,
const = log (Rhase (z) Rhase (z')).
(13)
is constant because Rhase (x) and Rhase (z') are either the true data marginals or the empirical marginals (Appendix A]. In either case, these distributions do not depend on the prior parameters, 0, the recognition parameters, @, or the approximate posterior parameters, tb.
4.2 InfoNCE as a RPM
Now, we choose the approximate posterior to be equal to the corresponding conditional of the recognition joint,
(2,2/10,,)) R (z, (z,2'|z,x'))= R Rz (2'|x')
(14)
Note that. Walker et al. (203b) do not make this choice. Instead, they allow the approximate posterior to be optimized separately; on the probabilistic modelling tasks they consider, their choice is likely to lead to slightly improved performance. Substituting this choice for the approximate posterior, the ELBO simplifies considerably,
Po (2,2') R (
((x,x') const + Eng(sl=) log
(15)
Additionally, consider the expected loss, taking the expectation over the true/cmpirical data distribution, Rhase (I, xx),
Pe (z, (2) R
C = Ressa [Cx x²)] = const + ER(
log
(16)
where R (2,22) is the relevant marginal of the recognition joint (Eq.
4.3 The RPM ELBO can be written as the mutual information minus a KL divergence
To get an intuitive understanding of the ELBO we take Eq. 16 and add and subtract ER(E. (log R (z, =")]
R (2,2")
Po (2,2²)
C = const + ER(E,2") log R (2) R (=')
+ ER(E,)
log
(17)
R (2,2")
The first term is the mutual information between z and x' under the recognition joint, (Eq. 5) and the second term is a KT-divergence,
C const + M1 DKL (2, (2,2"))
(18)
This objective therefore encourages large mutual information between and z' under R (Eq. 4 while encouraging R (z,z') to lie close to the prior, Po (z)
4.4 Under the optimal prior, the RPM ELBO is equal the mutual-information (up to constant)
Looking at Eq. 18 the only term that depends the prior, Pe (2),2), is the negative KL-divergence. As such, maximizing L with respect to the parameters of Po (=',z) is equivalent to minimizing DKL (2, (2)) Of course, the minimal KL-divergence of zero is obtained when,
P" (2,2)) = R (2,2').
(19)
For this optimal prior, the KL-divergence is zero, so the ELBO reduces to just the mutual information between z and z' (and a constant),
LMI const MT.
(20)

---

Published in Transactions on Machine Learning Research (03/2024)
4.5 Under particular prior, the RPM ELBO is equal to the infinite-sample InfoNCE objective (up to constant)
Recent work has argued that the good representation arising from InfoNCE cannot be from maximizing mutual information alonc, bocause the mutual information is invariant under arbitrary invertible transfor- mations [Tschannen et al. 2019 et al., 2021). Instead, the good properties must arise somehow out the fact. that. the InfoNCE objective forms only a. loose bound on the true MI, even the infinite sample limit Eq. (7). In contrast, here we show that the infinite-sample InfoNCE objective is equal (up to a constant) to the ELBO (or log marginal likelihood for deterministic encoders) for specific choice of prior. In particular, choose the prior on implicitly (through R (z)), and we choose the distribution over z' conditioned on z to be given by an energy based model with an unrestricted coupling function, fo(z, 2') (we could of course
use Eq. S for fo),
PlafoNCE (z) = R()
(21a)
PInfONCE (<'1z) Z(z)
1
R (=))
(21b)
The normalizing constant, Z(z), is
Z(z) / dz' R (2)) ER(z) [fo(z,2))]
(22)
Substituting these choices into Eq. 16 and cancelling R (2) R (z'), the average ELBO or log marginal likelihood becomes,
LIniONE const + ER(x.a") [log fo(z, 2²) log Z(z))
(23)
and substituting for Z(z) gives,
LInLONCE const + En(a,27) [log fo(2, ER(=) [log En(E) (o((z,z))]
(24)
Following Wang & Isola (020 and Li et al.
(2021) the right hand side can be identified as the infinite sample
InfoNCE objective that we introduced in Sec. 3.2
MI
Ix = const + LInfONCE.
(25)
Therefore, in this choice of model, the ELBO, LIMONCEE is equal to the infinite-sami InfoNCE objective, Ioo, up to constant. This would argue that the InfoNCE objective has closer link to the ELBO, LINIANCE, than it does to the MI, as the infinite-sample InfoNCE objective, Iz (f, Φ) is equal (up to constant) to log marginal likelihood (with determinstic encoder; Appendix Ch or the ELBO with stochastic encoder. Whereas the infinite-sample InfoNCE objective only gives a bound on the MI.
Of course, this is all in the infinite sample setting: we have an infinite-sample expression for the InfoNCE objective (which forms a bound on the MT) and we have an infinite-sample expression for the ELBOO //oog marginal likelihood. However, the infinite-sample InfoNCE objective and the infinite-sample ELBO are equal up to additive constants, these are really just different interpretations of the same quantity. The question then becomes how to develop finite-sample bound on this single quantity with two slightly different interpretations. As we really only have single quantity, we only need single finite-sample bound. That single finite-aample bound would of course apply equally to both the InfoNCE and ELBO interpretations. Ultimately, we choose to use the finite-sample bound originally described in the InfoNCE framework (Oord et al. 2018), but as discussed above, that bound would of course apply equally to both the InfoNCE, and ELBO interpretations.
5 Experimental results
Our primary results are theoretical: in connecting the ELBO/1op marginal likelihood and mutual information, and in showing that the InfoNCE objective with restricted choice of fo makes more sense bound
8

---

Published in Transactions on Machine Learning Research (03/2024)
A
B
D
E
maan Le distance
60
40
0%
20
Tinarr
supervised
gausisaan
limear
1/N
>= 4/N
P(=1==
Figure 1: Results of the moving balls experiment. A) Example of the motion between consecutive frames. The balls move by full diameter in semi-random direction. B) Locations of the extracted ball centres, after supervised linear decoding. The standard InfoNCE setup fails to extract correct locations. C) The mean distance from the extracted and true centres of the balls for a supervised method, InfoNCE with a Gaussian discriminator after supervised decoding and InfoNCE with a linear discriminator after supervised decoding. D) Probability distribution for the next location of the coral ball in according to an encoder trained with Gaussian discriminator. E) Probability distribution for the next. location of the same ball according an encoder trained with a linear discriminator.
the log marginal likelihood than on the MI. At the same time, our approach encourages a different way of thinking about how to set up contrastive SSL methods, in terms of Bayesian priors. As an example, we considered task in which the goal was extract the locations of three moving balls, based on videos of these balls bouncing around a square (Fig. A Appendix D
Critically, want the latent space, 2 and 2', mirror the underlying true latent space as closely possible. We take the latent spaces to be 6 dimensional, representing the and y positions of the balls. Critically, it does not make sense to impose unit norm constraint on this latent space, such a constraint does not exist in the real latent space. This contrasts with the usual InfoNCE setup, where 2 and z' are usually taken to unit norm.
To resolve this difficulty, we consider the InfoNCE-like setup described in Soc. 4.5 Our prior given by Eq. (21) with R (z) and R (=') defined by Eq- [1 The freedom in this setup is given by the choice of fo- Naively applying the usual InfoNCE choice of fo without unit-norm constraints (using Eq. [ failed (linear in Fig. [BC), because did not correctly encode prior information about the structure of the problem. Critically, our prior is that for the adjacent frames, the locations extracted by the network will be close, while for random frames, the locations extracted by the network will be far apart. The linear estimator in Eq. (8) is not suitable for extracting the proximity of the ball locations, so it fails (linear Fig. BC). In particular, corresponds to non-sensical prior over 2' given 2,
p PInfoNCE (2/12)
R (x')) exp (2"
(26)
Z(2)
(where we have taken Q₂ (=') defined by Fq. 11 to be approximately uniform purely for the purposes of building intuition). This prior will encourage 2' to be very large (specifically, z' should have a large dot- product with zT0. Instead, we would like a prior that encodes our knowledge that z' is likely to be close to z. We can get such a prior by using a Gaussian RBF form for fe.
fo(2,z') exp (- 2 (z z²)2).
(27)
where the only parameter, 0, is scalar learned lengthscale. Critically, this choice of So is natural and obvious if we take a probabilistic generative view of the problem (with a uniform Q₂ (z'), this corresponds Gaussian conditional, PlufoNCE (2'[2)). Because this is the natural and obvious choice of prior, it works well even without unit-norm constraints (which discussed, do not make sense in this setting).
9

---

Published in Transactions on Machine Learning Research (03/2024)
6 Conclusions
In conclusion, have developed new family of probabilistic generative model, the "recognition parame terised model". With determinstic recognition model, the RPM ELBO is equal to the marginal likelihood (Appendix C For the optimal prior, the RPM ELBO is equal (up to a constant) to the mutual information, and with a particular choice of prior, the RPM ELBO is equal (up to a constant) to the infinito-sample InfoNCE objective (up to constants). In contrast, the infinite-sample InfoNCE forms ouly a loose bound on the true MI, which would argue that the InfoNCE objective might be better motivated the RPM ELBO, As such, we unify contrastive semi-supervised learning with generative self-supervised learning (or unsuper- vised learning). Finally, we provide principled framework for using simple parametric models in the latent space to enforce disentangged representations, and our framework allows us to use Bayesian intuition to form richer priors on the latent space.
References
Alexander Alemi, Ben Poole, Ian Fischer, Joshua Dillon, Rif A Saurous, and Kevin Murphy. Fixing a broken elbo. In International Conference on Machine Learning, pp. PMLR, 2018.
Jyoti Aneja, Alexander Schwing, Jan Kautz, and Arash Vahdat. Nep-vae: Variational autoencoders with noise contrastive priors. arXiv preprint arXiv:2010.0000 2020.
Randall Balestriero and Yann LeCun. Learning by reconstruction produces uninformative features for per- ception. arXiv preprint 2xtt:24022118997 2024.
David Blei, Andrew Y Ng, and Michael Jordan. Latent dirichlet allocation. Journal of machine Learning research, 3(Jan):993 1022, 2003.
Christopher P Burgess, Trina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. Understanding disentangling in 8-vae. ar preprint orXiv:1804.03599 2018.
Ricky TQ Chen, Xuechen Li, Roger Grosse, and David Duvenaud. Isolating sources of disentanglement in variational autoencoders. arXiv preprinl XrXi::1802.04 04942, 2018.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. simple framework for contrastive learning visual representations. In International conference on machine learning, pp. 1597-1607. PMLR, 2020.
Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Pieter Abbeel. Infogan: Inter- pretable representation learning by information maximizing generative adversarial nets. arXiv preprint arXiv:1606.0900 2016a.
Xi Chen, Diederik P Kingma, Tim Salimans, Yan Duan, Prafulla Dhariwal, John Schulman, llya Sutskever, and Picter Abbeel. Variational lossy autoencoder. arXiv preprint arXiv:661.009991 2016b.
Carl Doersch, Abhinav Gupla, and Alexei Efros. Unsupervised visual representation learning by context prediction. In Proceedings of the IEEE international conference on computer vision, pp. 1422 1430, 2015.
Alexey Dosovitskiy, Philipp Fischer, Jost Tobias Springenberg, Martin Riedmiller, and Thomas Brox. Dis- criminative unsupervised feature learning with exemplar convolutional neural networks. IEEE transactions on pattern analysis and machine intelligence, 38(9):1734-1747, 2015.
Benjamin Eysenbach, Vivek Myers, Ruslan Salakhutdinov, and Sergey Levine. Inference via interpolation: Contrastive representations provably enable planning and inference. arXiv preprint orXiv:2403.04 04...04082. 2024.
Spyros Gidaris, Praveer Singh, and Nikos Komodakis. Unsupervised representation learning by predicting image rotations. Xiv preprint 1 077288 2018.
10

---

Published in Transactions on Machine Learning Research (03/2024)
Will Grathwohl, Kuan-Chiel Wang, Jörn-Henrik Jacobsen, David Duvenaud, Mohammad Norouzi, and Kevin Swersky. Your classifier is secretly an energy based model and you should treat it like one. arXiv preprint Xr::::1912. 032633, 2019.
Arman Hasanzadeh, Mohammadreza Armandpour, Ehsan Hajiramezanali, Mingyuan Zhou, Nick Duffield, and Krishna Narayanan. Bayesian graph contrastive learning. arXiv preprint arXiv:2112.0783 2021.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. bela-vae: Learning basic visual concepts with constrained variational framework. In International conference on learning representations, 2016.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:684006851,, 2020.
Adam Jelley, Amos Storkey, Antreas Antoniou, and Sam Devlin. Contrastive meta-learning for partially observable few-shot learning. arXiv preprint arXiv:2301.13180 2023.
Tom Joy, Sebastian Schmon, Philip Torr, N Siddharth, and Tom Rainforth. Capturing label characteristies in vacs. In International Conference on Learning Representations, 2020.
Theofanis Karaletsos, Serge Belongie, and Gunnar Rätsch. Bayesian representation learning with oracle constraints. arXiv preprint drXiv:1506.050 05011, 2015.
Hyunjik Kim and Andriy Mnih. Disentangling by factorising, In International Conference on Machine Learning, pp.
PMLR, 2018.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems, 34:21696-21701 2021.
Dicderik P Kingma and Max Welling. Auto-cncoding variational bayes arXiv preprint arXiv:131266114 2013.
Diederik P Kingma, Max Welling, et al. An introduction to variational autoencoders. Foundations and Trends0 in Machine Learning, 12(4):307-38) 332, 2019.
Yazhe Li, Roman Pogodin, Danica J Sutherland, and Arthur Gretton. Self-supervised learning with kernel dependence maximization. arXiv preprint arXiv:2106.0880 083200, 2021.
Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Ractsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled represen- tations. In international conference on machine learning, pp. 4114 4124. PMLR, 2019.
David JC MacKay. Maximum likelihood and covariant algorithms for independent component analysis, 1996.
Emile Mathieu, Tom Rainforth, Nana Siddharth, and Yee Whye Teh. Disentangling disentanglement in variational autocncoders. In International Conference on Machine Learning, pp. 4402-4412. PMLR, 2019.
Alexander Möllers, Alexander Immer, Elvin Isufi, and Vincent Fortuin. Uncertainty in graph contrastive learning with bayesian neural networks. In Fifth Symposium on Advances in Approzimate Bayesian In- [erence, 2023.
Kevin P Murphy. Probabilistic machine learning: introduction. MIT press, 2022.
Didrik Nielsen, Priyank Jaini, Emiel Hoogeboom, Ole Winther, and Max Welling, Survae flows: Surjections to bridge the gap between vaes and flows. Advances in Neural Information Processing Systems, 33:12685 12696, 2020.
Mchdi Noroozi and Paolo Favaro. Unsupervised learning of visual representations by solving jigsaw puzzles. In European conference computer vision, pp. 69 84. Springer, 2016.
11

---

Published in Transactions on Machine Learning Research (03/2024)
Aaron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint Xrxi 1001180..000 180.0037777 2018.
Barak Pearlmutter and Lucas Parra. Maximum likelihood blind source separation: A context-sensitive generalization of ica. Advances in neurl information processing systems, 9, 1996.
Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On variational bounds of mutual information. In International Conference on Machine Learning, pp. 5171 5180. PMLR, 2019.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International conference on machine learning, pp. 1278 1286.
PMLR, 2014.
Jascha Sicl-Dicksteen, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibriun thermodynamics. In International conference on machine learning, pp. PMLR, 2015.
Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score- based diffusion models. Advances in Neural Information Processing Systems, 34:1415-1428, 2021.
Michael E Tipping and Christopher M Bishop. Probabilistic principal component analysis. Journal of the Royal Statistical Sociely Series B: Statistical Methodology, 61(3):611 622, 1999.
Michael Tschannen, Josip Djolonga, Paul K Rubenstein, Sylvain Gclly, and Mario Lucic. On mutual infor- mation maximization for representation learning. arXiv preprint arXin:1907.1660 2019.
Julius von Kügelgen, Yash Sharma, Luigi Grescle, Wieland Brendel, Bernhard Schölkopf, Michel Besserve, and Francesco Locatello. Self-supervised learning with data augmentations provably isolates content from style. arXiv preprint Xix:22106 0ixx2106600469,, 2021.
William Walker, Arthur Gretton, and Maneesh Sahani. Prediction under latent subgroup shifts with high-dimensional obscrvations. arXiu: 2306. 2306.13472, 13472, 2023a.
William Walker, Hugo Soulat, Changmin Yu, and Maneesh Sahani. Unsupervised representational learning with poogguitioon-araanetriee probabilistic models. AISTATS, 2023b.
Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through alignment and uniformity on the hypersphere. In International Conference on Machine Learning, pp. 9929-9939 PMLR, 2020.
Yu Wang, Hengrui Zhang, Zhiwei Liu, Liangwei Yang, and Philip S Yu. Contrastvae: Contrastive variational autoencoder for sequential recommendation. In Proceedings of the 31st ACM International Conference on Information & Knowledge Management, pp. 2056-2066, 2022.
Veit David Wild, Robert Hu, and Dino Sejdinovic. Generalized variational inference in function spaces: Gaussian measures meet bayesian deep learning. Advances in Neural Information Processing Systems, 35: 3716-3738 2022.
Shengjia Zhao, Jiaming Song, and Stelano Ermon. Infovae: Balancing learning and inference variational autoencoders. In Proceedings of the AAAI conference on artificial intelligence, 2019.
Roland Zimmermann, Yash Sharma, Steffen Schneider, Matthias Bethge, and Wieland Brendel. Contrastive learning inverts the data generating process. arXiv preprint Xiv: 08850, 2021.
12

---

Published in Transactions on Machine Learning Research (03/2024)
A Choices for the base distribution
There are three choices of base marginals, Riusee (Ir) and Riaase (x') in an RPM, corresponding to three different variants of the RPM:
1. The empirical marginal RPM.
The true marginal RPM.
The estimated marginal RPM.
Tt will turn out that at least the choice between and makes little difference, as they lead to exactly the same expression for the ELBO (compare Eq. 29 and Eq. 37 albeit with a slightly different interpretation.
A.1 Empirical marginal RPM
The empirical marginal RPM was introduced in Walker et al. 2023b), and uses,
Rhase emp (x) Σ 8(x xz))
(28a)
Riase emp (x') N Σ 8(z' - I/).
(28b)
where indexes training data points. This has the advantage that Ruase emp (x) and Riasee emp (x') are known. Thus, R (2) and R (z') can be evaluated exactly as a mixture distribution,
( + Σ R (2|0s)
(29a)
R = N Σ R (='lx/)
(29b)
In this setting, we know Riase emp (z), R (2) and (of course) R (z(x), so we can exactly evaluate the proba-
bility density for the recognition parameterised likelihood (Eq. 10 As such, we can sample from the model distribution over data, P (xxxx'.
However, the empirical marginal viewpoints does has important issues. Specifically,
/
P ((x,x')=
P
(z'1z')
(30)
Substituting Eq. 10
P (x,x') Rhase emp (x) Rhase emp (x') F (I,
(31)
where,
R (2|z) R (2'122) R (=) (='))
F (xxx')= - dzdz' P (z,z')
(32)
Substituting the empirical marginal form for Rease emp (xr) and Rhase omp (x) (Eq. 28
N 8 8(x-
N
P (xxxx)== =
1 N
mi))
1 N
Σ j=1
8 (z'
-2-2xz) (
F (EE,4))
(33)
1=1
This distribution only places probability density/mass at values of x and x' that were actually observed the data, i.c. {ai and x' E {a/) most cases we do not expect the exact test value of x and o' to be in the training set, and this model assigns those test points zero probability density. Indeed, for
13

---

Published in Transactions on Machine Learning Research (03/2024)
absolutely continuous distributions, we never expect the exact test value of and x' to be in the training set.
Additionally, we can again connect this form to InfoNCE-like objectives. In particular, as the probability density over (x,z') above places mass only at a finite number of training points, we can equivalently write it as a probability mass function over indicies (i,j) describing which datapoint was chosen. The resulting probability mass function is,
(34) Critically, the implied conditional, P(J(), essence embodies a classification problem: classifying which of the N training points, z/j, is associated with a particular z4. And that classification problem almost exactly the InfoNCE classification problem.
P (i,))) - P.j F(r,,
Finally, this implies that even if we do sample from P (z,,,), are definitely going to end up with and x' from the training dataset. As such, "sampling" is just "matching up" potentially consistent x and xf in the training set, and is thus quite different from sampling in a proper generalive model.
A.2 True marginal RPM
While the empirical marginal RPM has many advantages, it does assign zero probability to any lest points that did not turn up in the training data. To resolve this issue, we could instead use the true data distribution for R It seems that this would be problematic as while the true data distribution exists, it almost always unknown (and even unknownable). However, remember that to perform inference and learn the parameters, never need to evaluate probability density of the true data distribution, as the Rhase (x') and Riase (z) terms cancel out of the ELBO (Eq. 12 Instead, we just need to be able to sample R (z, 2'), which can casily be achieved by taking (x', from the real data, and encoding using R (z|xr) and R (z'|z'). Of course, not knowing the true data distribution does prevent us from sampling (x',a)) .)) from the model in a true marginal RPM. But as discussed in the Introduction, we are not interested in sampling data from the model: we were, we would use o.g. a diffusion model. Instead, we are interested in extracting useful structured, high-level representations of data, and for that purpose, all we need to be able to do to encode using Rs (=(æ) and R. (z'(zz) and optimize the parameters of these encoders.
In practice, the empirical and true marginal settings are almost identical. The only minor difference is that in the empirical marginal setting, R (z) and R (z') can be evaluated exactly (Eq. 29), whereas in the true marginal setting, R (2) and R (z') can only be estimated. Critically, though the numerical value of the exact R (z) and R (z') in the empirical marginal setting (Eq, 229] is exactly the same as the true marginal estimate of these quantities. Speicifically, in the true marginal setting,
Jda R = f da Ptrue (1) Ry (2/x)
(35a)
R(=)== = f de'Perue Perue Rz)) R (2'1x').
(35b)
However, we can rewrite these integrals as expectations,
R (2[x)) R (z') Ep (r)) R (2' [x))=')
R (2) = Ep (2)
(36a)
(36b)
We can estimate these expectations using samples from Pirue (x) and Pirue (x'). Of course, the data itself gives samples from Ptrue (x) and Ptrme (x'). Thus, even in the true-narginal setting, we would estimate R (z)
and R (z') using the same sum over datapoints used in the empirical marginal setting (i.e. Eq. 29
R ( ~ (=(xs)
(37a)
R
~ N Σ R (2/14)
(37b)
The only difference is in the interpretation of this sum. Here (in the true marginal setting), the sum estimates the true value of R (=) and R (z'), while in the empirical marginal setting (Eq. 29). the sum is the true value.
14

---

Published in Transactions on Machine Learning Research (03/2024)
A.3 Estimated marginal RPM
We could of course learn Rhase est (a:) and Riase est (z'). That results in model where can both evaluate the probability density, and sample P (z',z). However, this involves training a generative model for complex data, such as images. While that might seem to obviate the whole point of the excrcise, there is the potential for estimated marginal RPMs to be useful. In particular, the estimated marginal RPM only requires training generative model for and x' separately (note that the ELBO in Eq. 12 depends only on the marginals, R (x) and R (z'), and not on 2'). Thus, it may be possible to use the estimated marginal to "stitch together" samples of the joint, (x, from generative model which can sample only from marginals.
B Technical details when using deterministic recognition models
While we do have to, could follow the standard sell-supervised setup, which in effect Dirac-delta recognition models,
R (a (x) 0 [b(()) RS (2'(z²) (2' - g6 (z)))
(38)
(39)
That raises the question of when our objectives make sense. To answer these questions, we need to carefully understand the measure-theory underlying the expressions the main text. In particular, our objectives are all ultimately written in terms of KL-divergences for distributions over the latent variables, (z, z'). We take Z and 2' Z', and consider probability measures μ and v, where the set underlying the measure space is 2 x Z'. Thus, the KT-divergences in the objective can be written in measure-theoretic notation as,
( dv do log dv du
DKL (vllu)= =
(40)
(c.g. see Wild et al. 2022). The critical term is dv/du, which is known as the Radon- Mikon-Nikodym derivative. This derivative is delined when can write in terms of function f 2 Z' + (0)
dv =
(41)
du
v(A) = S fdp
(42)
where A S Z' In that case, f uniquely defined up to p-null set (informally, if = p(A) 0 then = f is not uniquely defined for (z,=') E A). Now, we can consider the two KI-divergence terms in Eq. 11
The first KL-divergence term is,
R (2,2)) log R ( R (21)
ER(2,2)
DRL. (2(2,22))R())RG)
(43)
In this case, we can write,
R =
(44)
the Radon-Nikodyı derivative exists, irrespective of the form of R.g (z|z) and R.g (z'|='). This also implies that in Section 4.4 where we use the optimal prior, the objective (Eq. 20 is always well-defined.
The second KT-divergence term is,
Pe (2,2') = DeL (R (2,,2)))))() log R
ER(Z.E
(45)
If we use an optimal prior, this KL-divergence cancels. However, if we do not use the optimal prior, but instead want. to use parametric prior, this term cause problems. The issue is that if choose
15

---

Published in Transactions on Machine Learning Research (03/2024)
standard form for the prior, such as multivariate Gaussian, then P 222,2') will be absolutely continuous. However if we have deterministic R& (=/ar) and R2 (='[z"), then R (z, x') can be supported on a measure-zero subspace, and hence it will not be absolutely continuous. This for instance might happen if (x,x') x') and (z, x')
are real vector spaces, with x lower dimensional than 2 and/or x' lower dimensional than z'. One way to resolve this issue would be to set,
Rs (z|x) = N (ns(t),
(46)
Rs (2'12²)) N (cb(2'),,I)
(47)
where 0 < is arbilrarily small. In that case, Rs (2/x) Ry (z') and hence R (2,=") are absolutely continuous and non-zero everywhere so the necessary Radon- Nillon-ikodym derivative exists.
Also see Nielsen et al. (2020) which also uses deterministic approximate posteriors in the VAE sctting, and compares them against normalizing flows. Note however that. they work with quantities such Eq(=[=) [log P (x)z) (=|x)] which lack rigorous preaure-theorette formulation, as the distributions in the ratio are over different variables: x in the numerator and in the numcrator. As such, approach presented here which does allow a measure-theoretic formulation is likely to be preferable.
C When we use deterministic recognition models, the ELBO is equal to the (marginal) log likelihood)
The marginal likelihood, P (xx,2') is,
P Pxx,x') dzdz' P (xlz) Pcll)) (2,2').
(48)
Substituting the recognition parameterised likelihoods,
log P (x,x'' = log / dzdz' P 22,22') Ro (z/w) Rease (x) Rop (2'(2e) Rhase (x')
(49)
R (2)
R (2²)
P 2(,,2)
R (x) ((2)
const log (ela)
(50)
where, as usual const = log (Rbase (x) Rease (x')). The recognition models in the expectation are deterministic (Eq. 38 z (1) and x' = 1' (I'), ) so we can easily evaluate the expectation,
P (2,22) R (2) L xx,2').
const log (I, = const + log
(51)
Here, the last equality arises from taking the ELBO in Eq. 15) and substituting the same deterministic encoders. Thus, with a deterministic encoder, the (marginal) log likelihood is equal to the ELBO.
D Experimental details
We generated 900 images in a single continuous video with resolution of 256 x 256 pixels. The three balls had a diameter of 32 pixels. Between consecutive frames the balls moved by a full diameter in a random direction, as illustrated in Fig. DA. The movement trajectory was picked by taking the previous trajectory and adding a uniform noise of -2 +2" If the picked movement resulted in a collision, we sampled new trajectory by doubling the noise range until valid trajectory is found.
We trained the model in classic self-supervised manner. We encoded one "base" frame, one "target" frame (the next frame in a video sequence), along with number of random frames. As usual, the network was trained to distinguish between the target frame (adjacent to the base frame) and random frames. We then trained a linear decoder in a supervised manner to return the (æ, y) locations of the balls.
The encoder itself a simple convolutional neural network, as shown in Fig. It consists of 2 batch normalised convolutional layers with kernel size of 3. The first layer uses ReI.U as the activation function,
16

---

Published in Transactions on Machine Learning Research (03/2024)
Batch Norm
Batch Norm
3x3
3x3
258 258
258 258
256
E
Architecture of the encoder neural network. The first of the two 3x3 convolutional layers outputs 6 feature maps and uses a ReLU activation. The second convolution outputs 3 feature maps and applies sigmoid activation. For each of these 3 maps, we extract their centre of mass, This is done by summing each dimension and normalising it to 1. This then used to perform a weighted average over the axis locations and get the final coordinates.
Figure 2:
while the second layer uses sigmoid. At the output of the convolutional layers, we have 3 feature maps, which we interpret as the locations the 3 different balls. We finally extract these locations by computing the centre of mass of the Teature maps, giving a vector of numbers as output (the x and y locations of the centres of mass of each feature map). The training itself was performed by using stochastic gradient descent with a learning rate of 0.005 over the course of 30 epochs. The batches were made of 30 random pairs of the second
consecutive frames. For any pair, frame of the other pairs in the batch, as the random negative examples, against which we contrast.
the second frame as the positive example and
17


Reviews:
[{"review_id": "C3s3C3rD3s", "reviewer": "Reviewer_c21v", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n1. The paper is clearly written with rigorous definitions and notations, providing a new theoretical concept in understanding self-supervised learning objectives. \n2. In particular, the paper elaborates the equivalence between RPM ELBO and InfoNCE objective, and their relationship to mutual information, which is a common measure that previous works have focused on.\n\n## Weakness\n1. The major weakness of this paper would be the shortage of empirical validation. I suspect that the toy experiment demonstrates that the prior information is important in self-supervised learning, where it was nothing but simply using gaussian RBF for $f_\\theta$. My question is, how does it corroborate your claim that RPM provides a new approach in self-supervised learning? Maybe providing more pragmatic experimental results on vision benchmarks (e.g., self-supervised learning on CIFAR-10, at its minimal approach) might be relevant to better deliver the theoretical finding.\n2. In addition, although the concept of RPM is quite new to the literature, it seems like the equation and derivation is not new from previous literature. Consider $P_\\theta(z,z’)$ as a family of energy-based model with energy $f_\\theta(z,z’)$, then the proposed objective is nothing but Donsker-Varadhna (DV) objective [1], or MINE [2] that is derived for modern neural network architecture. While I agree that writing it into $P_\\theta(z,z’)$ provides a more general formulation, I do not see much advantage in understanding SSL with such a perspective. \n\n[1] Donsker, M. and Varadhan, S. Asymptotic evaluation of certain markov process expectations for large time, iv. Communications on Pure and Applied Mathematics, 36(2):183?212, 1983.\n\n[2] Belghazi, Mohamed Ishmael, et al. \"Mutual information neural estimation.\" International conference on machine learning. PMLR, 2018.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents a novel perspective on self-supervised learning (SSL) by demonstrating that the popular InfoNCE objective is equivalent to Evidence Lower Bound (ELBO) in a new class of probabilistic generative models, dubbed Recognition Parametrized Models (RPMs). Some theoretical investigations are shown as follows:\n1. The RPM ELBO can be written as the subtraction of mutual information with KL divergence between prior and data distribution of pairs, thus the optimal prior reduces to the mutual information up to a constant\n2. Given the infinite sample, ELBO and InfoNCE are equivalent up to constant", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "I believe this paper has a lot of room for improvement by clarifying the advantage of analyzing or implementing self-supervised learning with RPM; the author could do better on theory / experiments by showing some cases where RPM is better than other self-supervised objectives such as InfoNCE.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "lR3qK4AFZJ", "reviewer": "Reviewer_dPq2", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Disclaimer: I have not fully followed all the arguments, but overall the paper appears sensible. It is on topic for TMLR and may be of interest for its audience. My only substantial comment is on the experimental section; apart from that I only have very minor comments.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In the paper \"InfoNCE is variational inference in a recognition parametrised model\", the authors present a different view on the popular InfoNCE loss function, widely used in contrastive learning frameworks such as SimCLR. The authors argue that it is more helpful to see InfoNCE as an approximation to ELBO, rather than an approximation to MI (mutual information).", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "MAJOR COMMENTS\n\n* I must be confused, but I don't quite understand the meaning of \\theta in Equation 8. The way I know InfoNCE, e.g. from SimCLR, the loss function simply uses exp(z^T z' / tau) terms, i.e. there is no matrix multiplication by matrix \\theta, but instead there is a scalar temperature parameter \\tau (which is typically held fixed). Please comment on why you need \\theta matrix here, and whether one can also use parameter-less function f(z, z') that I gave above. I suggest to explicitly refer to SimCLR  here and explain how it fits into your notation/framework. It would also be useful to mention in this section that z and z' are typically normalized to have norm equal 1 (otherwise the objective does not make sense). \n\n* This normalization issue becomes crucial in Section 5, which presents a toy experiment supposedly showing a failure mode of standard InfoNCE setup. However, this looks like a strawman to me, due to missing normalization. SimCLR uses cosine similarity, which can be written as z^T z' but ONLY IF z and z' are normalized to lie on the hypersphere and have norms 1. If they are not normalized, then z^T z' does not give cosine similarity (obviously) and it does not make sense to use that for what the authors call f_theta. What the authors show in Figure 1C as \"linear\" does NOT correspond to SimCLR because it uses unnormalized vectors in z^T z'.\n\n  What about using exp(cos_similarity(z, z')) as f() and show that in Figure 1C instead of \"linear\"? Or at least in addition. Currently this whole section looks like  a strawman.\n\n\nMINOR COMMENTS\n\n* The paper contains many typos and could use a thorough proof reading. Page 2: \"(see ...)\". Page 2: \"This seems to mirrors\". Page 3: \"Third,,\".\n\n* Section 4 title looks incomplete: \"Recognition Parameterised Generative\" -- this sounds like an adjective, not a noun. Is the word \"model\" missing?", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "biZqWU6QZH", "reviewer": "Reviewer_JVxH", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n1. I would say that this paper is excellently written, starting from a intriguing and ambitious goal, and proceeding the arguments with convincing logic flows. It shows that the authors have plenty of knowledge of the literature as well as deep understandings of the essences of different approaches. \n\n2. The paper has a good account of the related work. I fully understand that publishing paper can be quite uncertain in these days, and the authors set up good examples of how we could cope with situations like “because followup works got published earlier, is this work meaningless and hopeless now”. A not that up-to-date work can still give us lots of insights, and we should also account for its contributions in the context of its publishing date. Besides, given the elaborated descriptions of the publication history, to make a good evaluation of these arguments, I search for this paper and browse its submission history. I would say that by shifting the main theme from “InfoNCE is VAE” to “InfoNCE is “RPM”, the authors certainly have done a lot of polishing and it makes this work more understandable and less confusing. \n3. I think that finding new perspective for understanding self-supervised objectives is a meaningful topic, and drawing connection to variational inference is certainly one of the most promising and interesting ways. \n4. The recognition parameterised model (RPM) is an interesting probabilistic generative model, and it kind of inverts the usual convention from “decoder - variational encoder” to “encoder - variational decoder”.  In this way, RPM has an advantage over VAEs by discarding the heavy burden to explicitly reconstruct the outputs, which may help explains the gains of InfoNCE over VAEs for representation learning.\n\n## Weakness\n\n1. **Limited experiments and practical insights.** I would say that although the theory is interesting, the experiments are far from being inspiring or complete to show potential practical insights of this theory. Considering the toy experiment in Sec 5, it is actually a common choice to directly use the cosine similarity between $z$ and $z’$ for parameterizing $f_\\theta(z,z’)$ in InfoNCE. And the new theory here does not add much new information. \n2. **Connection to Poole et al. is not elaborated.** As far as I could see, the RPM model proposed has a close connection to the energy-based model discussed in Poole et al. (2019) which is prior to this work (and its preprint). Specifically, in their energy-based model (Eq. 3), they propose roughly the same formulation as Eq. 21b,  and they also established the connection between InfoNCE and a lower bound Eq.4 (the same lower bound as this paper Eq. 23) and the MI objective. Obviously, the energy-based model in Poole et al. is also a kind of decoder-free probabilistic model that bears great similarities to the definitions to this work (if not the same after addressing Problem 2), and they already show a variational inference of this model gives the InfoNCE loss. So with these similarities, their differences should be elaborated.\n3. **Are MI and ELBO perspectives so different?** I note that the ELBO (as a lower bound of log likelihood) is *nothing but yet another  lower bound of MI*, since its upper bound, log likelihood, is also a variational lower bound of MI (up to constants). To see this, notice that \n    \n    $E_{P(x,y)} \\log P(x,y) = E_{P(x,y)} \\log Q(x,y) + KL(P(x,y), Q(x,y))\\geq E_{P(x,y)} \\log Q(x,y).$\n    \n    I think that this connection should be mentioned to give a holistic view of these different “lower bounds” and derivations. And in view of this limitation, this ELBO perspective is actually not much different from the MI lower bound perspective, e.g., in Poole et al.\n    \n\nGiven these serious concerns, I think this work still lacks enough evidence to show that it is sufficiently different from prior works, or offer concrete new insights to algorithm design.\n\nRef:\n\nPoole et al. On Variational Bounds of Mutual Information. In ICML. 2019.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This papers proposes a new understanding of the InfoNCE from the variational inference perspective, and claims its gain over the conventional mutual information understanding. To achieve this, they propose a recognition parameterized model that in turn, defines the generative probabilistic model as the bayes poster of the recognition model (i.e., the encoder). As a result, unlike VAEs, this model has no explicit decoder (can be a good thing), while they also show the variational lower bound (ELBO) is deeply connected to the mutual information (with a KL divergence accounting for the amortization gap). They also link the InfoNCE objective to the ELBO objective, which offers an alternative view. They show the benefit of this understanding using a rather toy experiment.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The quality of this work could be significantly improved if they could propose new amendments to the manuscript (or show my misunderstandings):\n1. Showing concrete new insights from the ELBO/RPM perspective.\n2. Elaborate the connections and differences to the energy-based model formulation.\n3. Elaborate the real difference between two perspectives.\n\nMinor:\n\nOn page 2, there is an incomplete bracket \"(see ...)\"", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper introduces a novel perspective on the InfoNCE objective by framing it within the context of variational inference, specifically through the lens of Recognition Parametrized Models (RPMs). This approach diverges from the conventional understanding of mutual information in self-supervised learning, offering new conceptual framework that potentially enhances our comprehension of self-supervised learning objectives.

### Major Contributions:
- The paper's primary contribution is the innovative reinterpretation of the InfoNCE objective as an Evidence Lower Bound (ELBO) in the context of RPMs. This reinterpretation broadens the theoretical understanding of self-supervised learning.
- The paper is well-structured, presenting complex ideas in an accessible manner, thereby contributing to the theoretical discourse in the field.

### Points of Consideration:
- Multiple reviewers highlight the need for more empirical evidence to support the theoretical claims made in the paper. While the theoretical framework is robust, additional experiments, particularly on more practical or larger-scale datasets, would bolster the paper's impact and relevance.

### Recommendation:
Despite the concerns regarding empirical validation and the need for a more distinct differentiation from existing literature, the paper presents a significant theoretical contribution that enriches our understanding of self-supervised learning. Therefore, I recommend the acceptance of this paper.
Decision: Accept (poster)