A Probabilistic Model behind Self- Supervised Learning
Authors: Alice Bizeul
Bernhard Schölkopf
Carl Allen
Venue: Accepted by TMLR ()

Abstract: In self-supervised learning (SSL), representations are learned via an auxiliary task without annotated labels. A common task is to classify augmentations or different modalities of the data, which share semantic _content_ (e.g. an object in an image) but differ in _style_ (e.g. the object's location). Many approaches to self-supervised learning have been proposed, e.g. SimCLR, CLIP and VicREG, which have recently gained much attention for their representations achieving downstream performance comparable to supervised learning. However, a theoretical understanding of the mechanism behind self-supervised methods eludes. Addressing this, we present a generative latent variable model for self-supervised learning and show that several families of discriminative SSL, including contrastive methods, induce a comparable distribution over representations, providing a unifying theoretical framework for these methods. The proposed model also justifies connections drawn to mutual information and the use of a ``projection head''. Learning representations by fitting the model generatively (termed SimVAE) improves performance over discriminative and other VAE-based methods on simple image benchmarks and significantly narrows the gap between generative and discriminative representation learning in more complex settings. Importantly, as our analysis predicts, SimVAE outperforms self-supervised learning where style information is required, taking an important step toward understanding self-supervised methods and achieving task-agnostic representations.

Published in Transactions on Machine Learning Research (09/2024)
A Probabilistic Model behind Self-Supervised Learning
Alice Bizeul'
alice. cer.rbbizeel@@i..ceh....
Deportment of Computer Science & ETH Al Center, ETH Zurich
Bernhard Schölkopf
bs@lucbingen.
Mar Planck Institute for Intelligent Systems, Täbingen
Carl Allen
earl. Sai etha.ch
Deportment of Computer Science ETH Al Center, ETH Zurich
Reviewed on OpenReview: https: // openreview. net/ forum? id= QEw274476R
Abstract
In self-supervised learning (SSL), representations are learned via an auxiliary task without annotated labels. A common task is classify augmentations or different modalities of the data, which share semantic content (e.g. an object in an image) but differ in style (e.g. the object's location). Many approaches to self-supervised learning have been proposed, c.g. SimCLR, CLIP, and DINO, which have recently gained much allention for their represen- tations achieving downstream performance comparable to supervised learning. However, theoretical understanding of self-supervised methods eludes. Addressing this, we present a generative latent variable model for self-supervised learning and show that several families of discriminative SST, including contrastive methods, induce a comparable distribution over representations, providing a unifying theoretical framework for these methods. The proposed model also justifies connections drawn to mutual information and the use of "projection head". Learning representations by fitting the model generatively (termed Sim VAE) improves performance over discriminative and other VAE-based methods on simple image benchnarks and significantly narrows the gap between generative and discriminative representation learn- ing in more complex settings. Importantly, as our analysis predicts, SimVAE outperforms self-superviscr learning where style information is required, taking an important step toward understanding sule-supe methods and achieving task-agnostic representations.
1 Introduction
In self-supervised learning (SSL), a model is trained to perform an auxiliary task without class labels and, in the process, learns representations of the data that are useful in downstream tasks. Of the many approaches to SSL (e.g. see Ericsson al., 2022), contrastive methods, such as InfoNCE (Oord et al., 2018), SimCLR (Chen et al., 2020a), DINO (Caron al., 2021) and CLIP (Radford al., 2021), have gained attention for their representations achieving downstream performance approaching that of supervised lcarning. These methods exploit semantically related obscrvations, such as different parts (Mikolov el al., 2013; Oord et al., 2018), augmentations (Chen et al., 2020a; Misra & Maaten, 2020), or iodaaliiie//iiew (Baevski et al., 2020; Radford et al., 2021) of the data, considered to share latent semantic content (e.g. the object in a scene) and differ in style (e.g. the object's position). SSL methods observed to "pull together" representations of semantically related samples relative to those chosen at random and there has been growing interest in formalising this to explain why SSL methods learn useful representations, but mathematical mechanism to justify their performance remains unclear.
Z
N
Figure 1: SSL Model for J semantically re- lated samples (terms explained in $3)
*Correspondence alice. www.:/@inf.tttaac.c The code reproduce SAEVAE he found at https: htgp:://g/tttttttg http://///o.iicbiiee///sseeeee
carl.al allennai.ethh.c

---

Published in Transactions on Machine Learning Research (09/2024)
p(zly)
Augmontation: shape
Instance: color
E RN
zf E IR"
Figure 2: The representation space Z of self-supervised representation learning. (left) common view: representations of pairs of semantically related data (e.g. augmentations of an image) are "pulled together" (grey arrows), while other pairs "pushed apart" (red arrows); (right) probabilistic view: representations of semantically related data are samples from common mixture component, 2 p(zly = =t). Maximising ELBOssL pulls zf towards a mode for each i (grey dots), clusters are held apart by the need to reconstruct.
We propose a principled rationale for multiple self-supervised approaches, spanning instance discrimination (Dosovitskly et al., 2014), deep clustering (Caron et al., 2018) and contrastive learning, including the popular InfoNCE loss (Chen et al., 2020a; Radford et al., 2021) (referred to as predictive SSL). We draw connection between these thcoretically opaque methods and fitting latent variable model by variational inference. More specifically, we: i) treat representations as latent variables and vice versa; ii) propose the SSL Model (Fig. 1) as generative latent variable model for SSL and derive its evidence lower bound (ELBOssL): and iii) show that the discriminative loss functions of prodictive SSL algorithms approximate the gencrative ELBOSsL. up to its generative reconstruction term.
Under the SSL Model, data is assumed to be generated by sampling: i) high-level latent variable y~-(y), which determines the semantic content of the data; ii) latent variable 2 p<z(() which governs the style of the data and gives rise to differences between semantically related samples (i.e. with same y); and iii) the data p(x22) Under simple Gaussian assumptions, latent variables semantically related data form clusters, 2² ~ p(zl)) j [1, (Fig. 2, right), mirroring in a principled way how SSL objectives "pull together" representations semantically related data and "push apart" others (Wang & Isola, 2020) (Fig. 2, left).
We derive the evidence lower bound, ELBOssL. as training objective to fit the SSL Model and show that it closely relates to the loss functions of discriminative predictive SSL methods. Thus, predictive SSL methods induce comparable latent structure to the mixture prior p(z) of the SSL Model, but dilfer in that they encourage latent clusters p(zly) to "collapse" and so lose style information that distinguishes semantically rolated data samples (sce Figures 3 and 4). Given that a downstream task may require style information, c.g. for object localisation, this could limit the generality of SSL-trained represenlatious.
Having proposed the SSL Model and its ELBO as a theoretical rationale behind self-supervised learning, anticipate that directly maximising ELBOsst. (termed SimVAE) should give comparable or improved representations. However, it is well known that, at present, generative modelling brings additional challenges relative to discriminative modelling (He et al., 2022), making results of discriminative methods a benchmark to aspire to and VAE-based generative methods a relevant point of comparison. As such, we show that SimVAE representations i) consistently outperform other VaE-based representations on downstream tasks; ii) compare to or even outperform predictive SST. methods at downsttrea classification on simple benchmark datasots (MNIST, FashionMNIST and CclebA); iii) significantly bridge the gap between generative and discriminative methods for content classification on more complex datasets; and iv) consistent with our analysis, outperform discriminative methods on tasks that require style information.
Overall, our results provide empirical support for the SSL Model as a mathematical basis for self-supervised learning and suggest that SSL methods may overfit to content classification tasks. They also indicate that generative SSL may be promising approach to task-agnostic representation learning if distributions can be well modelled, particularly given the potential added benefits of: uncortainty estimation from the posterior, a means to generate novel samples and qualitatively assess captured information from reconstructions, and model arbitrary sets of semantically related data, not only pairs.
2

---

Published in Transactions on Machine Learning Research (09/2024)
To summarise our main contributions:
we propose the SSL Model (Fig. 1) and its ELBO as a theoretical basis to justify and unify predictive SSL, several families of self-supervised learning algorithms,, including popular contrastive methods ($3);
we show that the SSL Model predicts that discriminative methods lose style information and rationalises: the view that SSL methods "pull together"/ ""push apart" representations, perceived link to mutual information, and the use of projection head (5-1); and
we show that SimVAE representations, learned generatively by maximising ELBOssL. outperform predictive SSL methods at content prediction on simpler benchmarks and at tasks requiring style information across multiple benchmarks (+14.8% for CelebA); and significantly outperform previous generative ( (VAE-based) methods on content classification tasks (+15% on CIFART0), including one tailored to SST. (87).
2 Background and Related Work
Representation Learning aims to learn a mapping X' (an encoder) from data x E to representations == f(x) (typically | <|x]) that perform well on downstream tasks. Representation learning is not "well defined" in that downstream tasks are arbitrary and good performance one may not mean good performance another (Zhang et. al., 2022). For instance, image representations are commonly evaluated by predicting semantic class labels, but the downstream task could instead be to detect lighting, position or orientation, which representations useful for predicting class may not capture. This suggests that generul-purpos or task-agnostic representations should capture as much information about the data as possible, as supported by recent. works that evaluate on a range of downstream tasks (e-g. Balažević et al., 2023).
Self-Supervised Learning includes many approaches that can be categorised in several ways (e.g. Balestriero et al., 2023; Garrido et al., 2022). The SSL methods that we focus on (predictive SSL) are defined below:
Instance Discrimination (Dosovitskiy et al., 2014; Wu et al., 2018) treats each data point a, and any of its augmentations samples of distinct class labelled by the index i. A softmax classifier (encoder + softmax layer) is trained to predict the "class" label (i.e. index) and encoder outputs are taken as representations.
Latent Clustering performs clustering on representations. Song et al. (2013); Xie et al. (2016); Yang et al. (2017) apply K-means, or similar, to the hidden layer of a standard auto-encoder. DeepCluster (Caron et al., 2020) iteratively clusters ResNet encoder outputs by K-means and uses cluster assignments as "pseudo-labels" to train a classifier. DINO (Caron et al., 2021), a wransormer-based model, can be interpreted similarly as clustering representations (Balestriero et al., 2023).
Contrastive Learning encourages representations of semantically related data (positive samples) to be "close" " relative to those sampled at random (negative samples). Early SSL approaches include energy-based models (Chopra et al., 2005; Hadsell et al., 2006); and word2vee (Mikolov et al., 2013) that predicts co-occurring words and encodes their pointwise mutual information (PMI) in its embeddings (Levy & Goldberg, 2014; Allen & Hospedales, 2019). InfoNCE (Oord et al., 2018; Sohn, 2016) extends word2vec to other data domains. For positive pair of semantically related samples (x,x x+ and randomly selected negative samples X- (2g 1 the InfoNCE objective is defined by:
e simi(,,
LINC = E En x log Σ Ex'E(at)ux- (siu(2,22)
(1)
where sim(..)) a similarity function, o.g. dot product. The InfoNCE objective (Eq. 1) optimisod if + e, for some constant (Oord et al., 2018). Many works build on InfoNCE, e.g. SimCLR (Chen et al., 2020a) uses synthetic augmentations and CLIP (Radford et al., 2021) uses different modalities as positive samples; DIM (Hjelm et al., 2019) takes other encoder paramcters as representations; and MoCo (He et al., 2020), BYOL (Grill et al., 2020) and VicREG (Bardes et al., 2022) find alternative strategies to negative sampling to prevent representations from collapsing.
sim(2, =
Due to their wide variety, we do not address all SSL algorithms, in particular those with regression-style auxiliary tasks, e.g. reconstructing data from perturbed versions (He et al., 2022; Xie et al., 2022); or predicting perturbations, e.g. rotation angle (Gidaris et al., 2018).
3

---

Published in Transactions on Machine Learning Research (09/2024)
Variational Auto-Encode (VAE). For generative latent variable model x, parameters 0 of model po(x)) po p)) can be learned by maximising the evidence lower bound (ELBO)
[. go (0(2)2) log po(x)=))o(=) (0(2/z)
Eg log p(x)]
2 Ez [o))
Ex
(ELBO)
where (0(2(z) learns to approximate the model posterior pe (=(2)))= po(ala))o(a)
J. po(a)=)po(s) Latent variables z can be
used representations ($1.1). A VAE (Kingma & Welling, 2014) maximises the ELBO, with pa, do modelled Gaussians parameterised by neural networks. A 3-VAE (Higgins et al., 2017) weights ELBO terms to increase disentanglement of latent factors. A CR-VAE (Sinha & Dieng, 2021) considers semantically related samples through an added regularisation term. Further VAE variants are summarised in Appendix A.I.
Variational Classification (VC). Dhuliawala et al. (2023) define a latent variable model for classifying labels y, p(yl=r/]_ (<))))p(yl=)) that generalises softmax neural network classifiers. The encoder is interpreted as parameterising (2(xr;; and the softmax layer encodes p(y|z) by Bayes' rule (VC-A). For continuous data domains X, e.g. images, q(2 of softmax classifier (parametersised by the encoder) shown overfit to a delta-distribution for all samples of each class, and representations of class "collapse" together (recently termed "neural collapse", Papyan et al., 2020), losing semantic and probabilistie information that distinguishes class samples and so harming properties such as calibration and robustness (VC-B).²
Prior theoretical analysis of SSL. There has been considerable interest in underst understanding the mathemat- ical mechanism behind self-supervised learning (Arora et al., 2019; Tsai et al., 2020; Wang & Isola, 2020; Zimmermann al., 2021; Lee al., 2021; Von Kügglgen et al., 2021; HaoChen et al., 2021; Wang al., 2021; Saunshi et al., 2022; Tian, 2022; Sansone & Manhaeve, 2022; Nakamura et al., 2023; Shwartz- Ziv et al.,
2023; Ben-Shaul et al., 2023), as summarised by HaoChen et al. (2021) and Saunshi et al. (2022). A thread of works (Arora et al., 2019; Tosh et al., 2021; Lee et al., 2021; HaoChen et al., 2021; Saunshi et al., 2022) aims to prove that auxiliary task performance translates to downstream classification accuracy, but Saunshi et al. (2022) prove that to be impossible for typical datasets without also considering model architecture. Several works propose an information theoretic basis for SSL (Hjelm et al., 2019; Bachman et al., 2019; Tsai et al., 2020; Shwartz-Zit Ziv el. al., 2023), e-g. maximising mutual information between representations, but Tschannen et al. (2020); McAllester & Stratos (2020); Tosh al. (2021) raise doubts about this. We show that this apparent connection to mutual information is justified more fundamentally by our model for SSL. The previous works most similar to our own are those that take probabilistic view of self-supervised learning, in which the encoder approximates the posterior of generative model and so can be interpreted to reverse the generative process (e.g. Zimmermann et al., 2021; Von Kügelgen et al., 2021; Daunhawer et al., 2023).
7
Lee
MNIST
VAE
B-VAE
CR-VAE
SIMVAE
SimCLR
VicReg
Fashion
VAE
B-VAE
CR-VAE
SimVAE
SimCLR
VicReg
Figure 3: Assessing the information in representations. Original images (left cols) and reconstructions from representations learned by generative unsupervised learning (VAE, 8-VAE, CR-VAE), generative SSL (our SimVAE) and discriminative SSL (SimCLR, VicREG) on MNIST (1), Fashion MNIST (r). Discriminative methods lose style information (e.g. orientation).
2In practice, constraints such regularisation and early stopping arbitrarily restrict this optimum being fully attained.
4

---

Published in Transactions on Machine Learning Research (09/2024)
3 A Probabilistic Model for Self-supervised Learning
Here, we propose a latent variable model for self-supervised learning, referred to as the SSL Model, and its evidence lower bound, ELBOssL, as the rationale behind predictive SSL methods.
We consider the data used for predictive SSL, X X X,, as a collection of subsets x; x== of semantically related data, where a! EX' may, for example, be different augmentations, modalities or snippets (indexed j) of data observations zg 3 All of Ex; are considered to have some semantic information in common, referred to content, while varying in what we refer to style, e.g. mirror images of an object reflect the same object (content) in different orientations (stylc). A hicrarchical generative process for this data, the SSL Model, is shown in Fig. 1. Under this model, a subset x of J semantically related data samples is generated by sampling: (i) y>-p(y, which determines the common semantic content, (ii) z ~p(zly) conditionally independently for E [1, J), which determine the style of each x (given the same y), and (iii) each x² ~p(x|22); hence
p(xly) = /[I 1()))=)) MP(>'W)
(2)
where z= {22) )/= and y, which defines which x are semantically related, is known, i.e. observed. If distributions in Eq. 2 are modelled parametrically, their parameters be learned by maximising the evidence lower bound to the conditional log likelihood (analogous to the standard ELBO),
E [log((((v) x.y
Σ
(= (xd) (loggo((oo))) - log (o(((())))) + +f.os [ (xx(xx))logppoll))
(ELBOss1.)
X.9
where the approximate posterior q(z () II is assumed to factorise. 4 A derivation of ELBOssL is given in Appendix A.2. Similarly to the standard ELBO, ELBOsst. comprises: a roconstruction term, the approximate posterior entropy and the (now conditional) log prior over all 2 Ez.
By assuming that p(zly) are unimodal and concentrated relative to p(z) (i.e. Varizly] < Var(z), Vy), latent variables of semantically related data z/ E Ez form clusters and are closer on average than random pairs just SSL representations are described. Thus, fitting the SSL Model to the data by maximising ELBOssL- induces a mixture prior distribution over latent variables, Pssu(z) = E p(zly)p(v) (Fig. 2, right), comparable to the distribution over representations induced by predictive SSL methods (Fig. 2, left). This connection goes to the heart of our main claim: that the SSL Model and ELBO underpin predictive SSL algorithms ($2). To formalise this, we first establish how discriminative and generative methods relate so that the loss functions of discriminative predictive SSL methods can be compared to the generative ELBOssL objective.
4 A Probabilistic Model behind Predictive SSL
So far, we have assumed that latent variables make useful representations without justification. We now justify this by considering the relationship between discriminative and generative representation lcarning (64.1) in order to understand how the generative ELBOssL objective can be emulated discriminatively ($4.2). We then review predictive SSL methods and show how they each emulate ELBOssL (84.3).
4.1 Discriminative vs Generative Representation Learning
Standard classification tasks can be approached discriminatively or generalively, defined by whether or not distribution over the data space X' is learned in the process. We make an analogous distinction for representation lcarning where, whatever the approach, the aim is to train an encoder f X + so that representations z= fz) follow a distribution useful for downstream tasks, e.g. clustered by classes of interest.
Under a generative model z xx latent variables are assumed determine semantic properties of the data x, which are predicted by the posterior p(z]x). Thus, learning an approximate posterior q(=|<r) (parameterised by
3 xxy may vary and domains X² 3r) can differ with modality.
4 Expected to be reasonable for 2d that carry high information about æ!, such that observing related g. its representation 2 provides negligible extra information, f.e. p(z] 22,, *²) p=)))
5

---

Published in Transactions on Machine Learning Research (09/2024)
f) by variational inference offers a generative route to learning semantically meaningful data representations. Meanwhile, the predictive SSL hods of interest train deterministic encoder under a loss function without a distribution over X' and are considered discriminative, We aim to reconcile these two approaches.
To discuss both approaches in common terms, we note that a deterministic encoder f can be viewed as a (very low variance) posterior Py (2(x) R fc=)( which, with p(x), defines the joint distribution ps(x; ====)())r))((( (x))))(()) and the marginal over representations Pf [= pr(x, z). Thus, in principle, an encoder can be trained so that representations follow a "useful" distribution p" (z) in one of two ways:
generatively, by optimising the ELBO for generative model with prior p"(=) or discriminafiuely, by optimising loss function that is minimised if (and only if) (()))>>(()))
This defines a form of equivalence between discriminative and generative approaches allowing our claim to be restated as: prodictive SSL methods induce a distribution (z(z) over representations comparable to the prior (ssr(z) (z) under the SSL Model.
4.2 Emulating a Generative Objective Discriminatively
Having seen in $4.1 that, in principle, an encoder can be trained generatively or discriminatively that representations follow target distribution p=(),, we now consider how this can be achieved in practice for the prior of the SSL Model, Pa=P-
While ELBOSst. gives a principled generative objective to learn representations that fit, Pasu.2), a discriminative loss that is minimised when p, pz(z)= psss (z) is unclear. Therefore, using ELBOşsL, as a template, we consider the individual effect each term has on the optimal posterior (<2[2) (parameterised by encoder f) to understand how a discriminative (or "p(xlz)-frree" objective might induce a similar distribution over representations.
Entropy: as noted $4.1, discriminative methods can be considered to have a posterior with very low fixed variance, q(=[x)) s 82-f(=e; for which H[q] is constant and is omitted from a discriminative objective.
log P (zly)] is optimal w.r.t. q iff all related samples x² Ex map to modes of p(zly). For uni- modal p(zly), this means representations 2 Ez collapse to a point, losing style information that distinguishes x² Ex.5 Since the prior term is not generative it is included directly in a discriminative objective.
Prior:
Reconstruction: Po [=] terms are maximised w.r.t. q iff each x maps to a distinct representation 23, which po (x maps back to x²; countering the prior to prevent collapse. In a discriminative objective, this generative term is excluded, and is considered emalated by term that requires 2 to be distinct.
We see that ELBOssL includes a term acting to collapse clusters p(zly) and another preventing such collapse, which counterbalance one another when the full ELBOsst. objective is maximised. Since the reconstruction term, p(x|²), 2), must be excluded from a discriminative objective, a discriminative objective ÊSSL is considered to emulate ELBOssL. if it combines the prior term with (p(x)2))frro)) substitute for the reconstruction term, Rt-), that prevents representation collapse (both between and within clusters), i.e.
fssl. = Fxx, If ((22)))(ogg((l))) + R(z (x) x))]
(3)
We can now restate our claim as: predictive SSL objectives have the form of Eq. and so emulate ELBOssL. to induce a distribution over representations comparable to the prior pest of the SSL Model (Fig. 1).
The analysis of ELBO terms shows that objectives of the form of Eq. 3 fit the view that SSL "pulls together" representations of related data and "pushes apart" others (o.g. Wang & Isola, 2020). That is, the prior "pulls" representations to modes and the R(-) term "pushes" to avoid collapse. Thus ELBOssI. which underpins Eq. 3, provides the underlying rationale for this "pull/push" perspective of SSL, depicted in Fig. 2.
Lastly, we note that the true reconstruction ELBOss1. not only avoids collapse but is necessary component in order to approximate the posterior and learn meaningful representations. We will see, as may be expected, that substituting this term by R(-) can impact representations and their downstream performance.
5This assumes classes x; are distiint, as is the case for empirical self-supervised learning datasets of interest..
6

---

Published in Transactions on Machine Learning Research (09/2024)
4.3 Emulating the SSL Model with Predictive supervidee Learning
Having defined ELBOssI, as generalive objective for self-supervised representation learning and ESSI. as discriminate approximation to it, we now consider classes of predictive SSL, specifically instance discrimination, latent clustering and contrastive mothods, and their relationship to ESSL: and thus the SSL Model.
Instance Discrimination (ID) trains softmax classifier on data samples labelled by their index {(.) 3s }). From VC-A (62), this softmax cross-entropy objective can be interpreted under the SSL Model as maximising a variational lower bound given by
Ezy (o.googp( 2 Ex.. [[ ((2)2) logg)/)))) = Ee,y [[
logple)) tc.
(4)
Eq. (RHS) matches Eq. 3 with J=1 and R())=(((()) the entropy of p(z), where 2 f(xx are outputs of the classifier encoder. Intuitively, maximising entropy (in lieu of the reconstruction term) might be expected to avoid collapse, but although representations (2) of distinct classes are spread apart, those within the same class, z Ezi, collapse together (VC-B, $2).
Deep Clustering (DC) iteratively assigns temporary labels = c e [1, to data z by clustering representations output by ResNet encoder; and trains the encoder together with softmax head to predict those labels. While subsets x of data with the same label are now defined according to a ResNet's 's "inductive bias", the same loss is used as Eq. 4, having the form of Eq. 3 with J=1. Each cluster of representations z Ez again collapses together.
Contrastive Learning based on the InfoNCE objective (Eq. 1), contrasts the representations of positive pairs x= [x,x²} with those sampled at random. The InfoNCE objective is known to be optimised when sim(z, z') = PMI(x,,'') x²) + c ($2). From this, it can be shown that for X'= {xt, 22 .T., X x= X'U(,,,
CINCE( x') Ex 9o (Z(
sim(s,at)
Ex
/ 9 (o(2
plat(a)/p(at)
Eeegr sim(z,=")
Eeezz p(=11)//())
x Ex
a (z) log
pa,e p(=)p(=1)
log(l: 1)
/ as (oogo)))))))(lz)
Ex
Σ log (ogg((z))) - (5)
(See Appendix A.3 for a full derivation.) Here, p(z,=t) is the probability of sampling 2, zt assuming that 2, x are semantically related. Under the SSL Model this is denoted p(2, (t,2=1))) (ly = p(zly), hence the change in notation for ease of comparison to Fq. 3. Thus, the InfoNCE loss lower bounds an objective of the form of Eq. with [())]][((=)] as in Eq. 4, but with J=2 samples.
As brief aside, we note that ID and DC consider J J=1 sample of at time, computing p(=/un vhs) from parameters th stored for each class (weights of the softmax layer), which could be memory intensive. However, for J 2 samples, one can estimate p(zzlyi) without stored parameters as p(zzlys) - [fg p(th) II, p(2/ yes v/)== = s(2,) if s/ be integrated out (see Appendix A.4 for an example). Under this "V/-free" approach, joint distributions over representations depend on whether they are semantically related:
s(2,) I(=.)((E.))
if i; Vr (i.e. if H, Vr,
p(zin o. 2i,)
InfoNCE implicitly employs this "trick" within sim(-.-) to avoid parameterisation of each cluster.
To summarise, we have considered the loss functions of each category of predictive SSL and shown that:
1. despite their differences, each has the form of fSSL and so emulates ELBOssL to induce a distribution over representations comparable to the prior of the SSL Model; and
cach method substitutes the reconstruction term of ELBOsst. by entropy H[p(z)], which causes represen- tations of semantically related data 2f Ez, to form distinct clusters (indexed by i), as in Psal, but those clusters collapse and lose style information as representations (indexed by j) become indistinguishabl..
7

---

Published in Transactions on Machine Learning Research (09/2024)
Style information is important in many real-world tasks, e.g. detecting an object's location orientation; and is the main focus of representation learning elsewhere (Higgins et al., 2017; Karras et al., 2019). Thus, contrary to the aim of gemeral representation learning, our analysis predicts that discriminative SSL over-fit to content-related tasks (discussed further in Appendix A.6).
Even restricting to predictive SSL,, an exhaustive analysis of methods is infeasible, however, many other methods adopt or approximate aspects analysed above. For example, SimCLR (Chen et al., 2020a) and CLIP (Radford et al., 2021) use the InfoNCE objective. MoCo (He et al., 2020), BYOL (Grill et al., 2020) and VicREG (Bardes et al., 2022) replace negative sampling in different ways but ultimately "pull together" representations of semantically related data, modelling the prior; and "push apart" others, avoiding collapse via a mechanism (e.g. momentum encoder, stop gradients or (co-)variance terms) in lieu of reconstruction. DINO (Caron et al., 2021) assigns representations of pairs of semantically related data to the same cluster, as in DC (Balestriero et al., 2023) but with J
Lastly, we discuss how the SSL Model also provides principled explanation for the connection between SSL and mutual information drawn in prior works, and a plausible rationale for using a projection head in SSL.
4.3.1 Relationship to Mutual Information
InfoNCE is known to optimise lower bound on matual information (MI), T(x,=²) = B Blog p(a.a')
pl=)p(z')) (Oord
et al., 2018) and Eq. shows it lower bounds 2'). It has been argued that maximising MI is in fact the underlying rationale that explains contrastive learning (Hjelm et al., 2019; Ozsoy et al., 2022). Previous works have challenged this showing that better MT estimators do not give better representations (Tschannen et al., 2020), that MI approximation is inherently noisy (McAllester & Stratos, 2020) and that the InfoNCE estimator is arbitrarily upper-bounded (Poole et al., 2019). Note also that for disjoint x, (e.g. non-overlapping augmentation sets), the range of pointwise mutual information values is unbounded, [-o0, k), k> 1, yet using the bounded cosine similarity function, sim(2z=')= |>]
=T=' E[-1, (e.g. Chen et al., 2020a) to model PMI values outperforms use of the unbounded dot-product (see $A.5). 6 Our analysis supports the idea that MI is not the fundamental mechanism behind SSL and explains the apparent connection: MI arises by substituting entropy H[p(z)] for the reconstruction term in ELBOssL as in predictive SSL methods. Rather than aide representation learning, this shown to collapse representations togother and lose information.
4.3.2 Rationale for Projection Head
Downstream performance is often found to improve by adding several layers to the encoder, termed a "projection head", and using encoder outputs as representations and projection head outputs in the loss function. Our analysis has shown that the ELBOsst. objective clusters representations of semantically related data while ensuring that all representations are distinct, whereas predictive SSL objectives collapse those clusters. However, ncar-final layers are found to exhibit similar clustering, but with higher intra-cluster variance (Gupta et al., 2022), as also observed in supervised classification (Wang et al., 2022). We conjecture that representations from near-final layers outperform those used in the loss function because their higher intra-cluster variance, or lower collapse, gives distribution closer to Past. of the SSL Model.
5 Generative Superveee Learning (SimVAE)
The proposed SSL Model (Fig. 1) has been shown to justify: (i) the training objectives of predictive SSL methods; (ii) the more general notion that SSL "pulls logether"/ ""pushes apart" representations; (iii) the connection to mutual information; and (iv) the use of a projection head.
As a proposed basis for self-supervised learning, we look to provide empirical validation by maximising ELBOsst. to (generatively) learn representations that perform comparably to predictive SSL. Maximising ELBOssL can be viewed as training VAE with mixture prior (z) E, p(zly)p(y) where representations e z of semantically related data are conditioned on the same y, which we refer to as SimVAE.
6In $A.5, we show that cosine similarity gives representations comparable using softmax cross entropy (ID) but without stored class parameters di for each 24- Wang & Isola (2020) note this conclusion but do consider the known PMI minimiser.
8

---

Published in Transactions on Machine Learning Research (09/2024)
Cifar10
VAE
B-VAE
CR-VAE
SimVAE
SmCLR
VicReg
Celeb-A
VAE
B-VAE
CR-VAE
SimVAE
SimCUR
VieReg
Figure 4: Assessing information in representations. Original images (left cols) and reconstructions from representations Iearned by generative unsupervised learning (VAE, 3-VAE, CR- VAE), generative SSL (our SimVAE) and discriminative SST. (SimCLR, VicREG) Cifar10 (I), CelebA (r). Discriminative methods lose style information (e.g. orientation/colour))
In practice, it is well known that training a generative model is more challenging, at present, than training a discriminative model (Radford et al., 2021; Tschannen et al., 2024). not least because it requires modelling a complex distribution in the high dimensional data space. SimVAE performance is therefore anticipated to most comparable to other VAE-basod approaches. Thus we add empirical support for the SSL Model by validating the following hypotheses: [H1] SimVAE achieves sule-super learning if distributions can be well modelled; [H2] SimVAE retains more style information than discriminative objectives; [H3] SimVAE learns better performing representations than other VAE-based objectives.
Testing hypotheses [H1-3] requires instantiating distributions in ELBOssL- As for a standard VAE, we assume p(x[z) and (<z|z) are Gaussians parameterized by neural networks. The prior p(y) assumed uniform over {1...., N], for N training samples and p(z + th, o²) are Gaussian with small fixed variance o2. Assuming a uniform distribution p(t) over their means (over a suitable space), & integrates out (see Appendix A.4) give:
Implementing SimVAE.
p(zyy) x expl- 1. (
(6)
20
a Gaussian centred on the mean representation Σ 25. Maximising this within ELBOssL can be considered to "pinch together" representations of semantically related data.
While contrastive methods typically compare pairs of related representations (J=2), ELBOssL allows any number J to be compared. In practice a balance is struck between better approximating p(zly) (high J) and diversity in mini-batch (low J). Algorithm in Appendix A.7.1 details the steps to optimise ELBOssL-
6 Experimental Setup
Datasets and Evaluation Metrics. We evaluate SimVAE representations on four datasets including two with natural images: MNIST (LeCun, 1998), MNSHT (Xiao et al., 2017), CelebA (Liu et al., 2015) and CIFARIO (Krizhevsky et al., 2009). We augment images following the SimCLR protocol (Chen al., 2020a) which includes cropping and flipping, and colour jitter for natural images. Frozen pre-trained representation are evaluated by k-Nearest Neighbors (k-NN; Cover & Hart, (K67)(k-NN) and a non-linear MLP probe classification tasks (Chen al., 2020a; Caron et al., 2020). Downstream performance is mcasured in terms of classification accuracy (Acc). Generative quality is evaluated by FID score (Heusel el al., 2017) and reconstruction error. For further experimental details and additional results for clustering under a gaussian mixture model and linear probe see Appendices A.7 and A.8.
Baselines methods. We compare SimVAE to other VAE-based models including the vanilla VAE (Kingma & Welling, 2014), 3-VAE (Higgins et al., 2017) and CR-VAE (Sinha Dieng, 2021), as well as to state- of-the-art sele-super discriminative methods including SimCLR (Chen et al., 2020a), VicREG (Bardes et al., 2022), MoCo (He et al., 2020) and its extension MoCo v2 (Chen et al., 2020b). As lower bound,

---

Published in Transactions on Machine Learning Research (09/2024)
Ac
Acc-k-NN
MNIST
Fashion
CelebA
CIFARIO
MNIST
Fashion
CelebA
CIFAR10
Random
38.1 3.8
49.8
83.5
16.3
46.1 2.5
66.5
0.4
80.0
13.1 =
SimCLR
97.2 = 0.0
74.9 = 0.2
93.7 + 0.4
67.4 = 0.1
97.2 = 0.1
76.0 0.1
91.6 = 0.3
64.0 = 0.0
MoCo
78.6 1.2
65.1
91.2 =
56.4
94.6 0.3
76.9
87.9
54.0 1
MoCov2
94.6 + 0.4
71.2 = 0.1
91.7 = 0.1
56.4 = 1.6
94.6 1 0.3
76.9 = 0.2
88.8 + 0.4
54.0 1 2.0
VicREG
96.7 6.0
73.2
94.7 +
69.7 +
97.0 0.0
76.0
92.7 0.4
68.3
VAE
97.8 0.1
80.2 + 0.3
89.0 # U.S
30.3 # u.4
98.0 + 0.1
83.7 0.2
86.9 + u.7
25.6 + 0.5
3-VAE
98.0 0.0
82.2
93.4 +
36.6 0.1
98.3 0.0
86.1 0.0
92.0
28.5 +
CR-VAE
97.7 0.0
82.6 0.0
93.1 = 0.4
36.8 0.0
98.0 + 0.0
86.4 0.0
91.6 = 0.6
28.1 = 0.1
SimVAR
98.4 0.0
82.1 = 0.0
95.6 0.4
51.8 #
98.5 = 0.0
86.5 = 0.0
93.2 0.1
47.1 = 0.0
Table 1: Content retrieval. Top-1% classification accuracy(1) for MNIST, PashionMNIST, CIFARIO, and CelebA (gender classification) using MLP probe (MP) and k-Nearest Neighbors (k-NN) classification methods; We report mean and standard errors over three runs; Bold indicates best scores in each method class: generative (teal), discriminative methods (red).
also provide results for randomly initialized encoder. For fair comparison, the augmentation strategy, representation dimensionality, batch size, and encoder-decoder architectures are invariant across methods. To enable a qualitative comparison of representations, decoder networks were trained for each discriminative baseline on top of frozen representations using the reconstruction error. See Appendix A.7 for further details on training baselines and decoder training.
Implementation Details. We use MLP and Resnet18 (He et al., 2016) network architectures for simple and natural image datasets respectively. For all generative approaches, we adopt Gaussian postcriors, (<2(x), priors, p(z), and likelihoods, p(æ[2z, with diagonal covariance matrices (Kingına & Welling, 2014). For SimVAE, we adopt Gaussian p(=ly) as described in $5 and, we fix the number of augmentations to 10 (see Fig. 8 for an ablation). Ablations were performed for all sensitive hyperparameters for each method and parameter values were selected based on the best average MLP Acc across datasets. Further details regarding hyperparameters and computational resources can be found Appendix A.7.
7
Results
We assess content and style information retrieval by predicting attributes that are preserved (e.g. object class, face gender) or disrupted (e.g. colour, position and orientation) under augmentation, respectively.
Content Retrieval. Table reports downstream classification performance across datasets using class labels. SimVAE oulperforms predictive SSL for simple datasets (MNIST, PashionMNIST, CelebA), which empirically supports H1, demonstrating that SimVAR achieves self-supervised learning where data distributions can be well-modelled. SimVAE also materially reduces the performance gap (A) between generative and discriminative SSL, by approximately half (A 32.8% + 17.6%) for more complex settings (CIFAR10). The performance improvement. by SimVAE over other generative met supports H3
Style Retrieval. We reconstruct MNIST, FashionMNIST, CIFAR-10, and CelebA images from represen- fations to gain qualitative insights into the style information captured. Figures 3 and 4 show a significant loss of orientation and colour in reconstructions from predictive SSL representations across datasets. CelebA has 40 mixed labeled attributes, some of which clearly pertain to style (e.g. hair color). Removing those with high-class imbalance, we evaluate classification of 20 attributes and report results in Figure 5. SimVAE outperforms both generative and predictive SSL baselines at classifying hair color (left, middle). Across all attributes (right), SimVAE outperforms or performs comparably to both generative and predictive SSL basclines. These results support H2 qualitatively and quantitatively, confirming that discriminative mothods lose more style information relative to generative methods (Figure 5, middle). Performance for each CelebA attribute is reported in Figure 6 in Appendix A.8.
10

---

Published in Transactions on Machine Learning Research (09/2024)
Random
56.2 + 0.8
Random
25
SimCLR
52.0 = 0.2
66
VAE
AAVAE
MoCo
57.7
20
MoCo v2
59.4 # 0.5
Style MP Acc
62
SIIVAE
MP
15
VePeg
VieREG
52.8 0.4
MP Accsunuet
111
VAE
64.4 0.3
28-
8-VAE
66.4 +: 0.4
CR-VAE
66.2
34-
SimVAE
67.5 0.3
84
WA McCov2 MoCo CRIVAE (AVAE VieReg SinCUR Method (M)
Content MP Acc
Figure 5: Style retrieval on CelebA. (Acc-MP f) (left) hair colour prediction (mean and standard error over runs, best results in bold); (middle) content vs. style prediction (gender vs hair colour), best performance in top-right; (right) Performance gain of SimVAE vs baselines (M) across all 20 CelebA altributes. From left. to right, lowest to highest mean score;
Image Generation. While generative quality is not relevant to our main hypotheses, out of interest we show randomly generated SimVAE images and quality metrics Appendix A.8. We observe small but significant improvements in FID score and roconstruction error relative previous VAE methods.
8 Conclusion
The impressive performance of self-supervised learning over recent years has spurred considerable interest in understanding the theoretical mechanism underlying SSL. In this work, propose latent variable model (SSL Model, Fig. 1) as the theoretical basis to explain several families of self-supervised learning methods, termed "predictive SSL." including contrastive learning. We show that the ELBO the SST. Model relates to the loss functions predictive SSL methods and justifies the general notion that SSL "pulls torether" representations of semantically related data and "pushes apart" others.
Predictive SSL methods are found to maximise entropy H[p(=)), in lieu of the ELBO's reconstruction term. Although this creates scemingly intuitive link to the mutual information between reprosentations that was thought to explain contrastive methods, it in fact causes represenlations of semantically related data to collapse together, losing style information about the data, and reducing their generality. Our analysis also justifics the use of a "projection head", the existence of which suggests cantion in (over-)analysing representations that are ultimately not used.
We provide empirical validation for the proposed SSL Model by showing that fitting by variational inference, termed SimVAE, learns representations generatively that perform comparably, significantly reduce the performance gap, to discriminative methods at content prediction. Meanwhile, SimVAE outperforms where style information is required, taking step towards task-agnostic representations. SimVAE also outperforms previous generative VAE-based approaches, including CR-V CR-VAE tailored to SSL.
Learning representations complex data distributions generatively, rather than discriminatively, remains a challenging and actively researched area (Tschannen al., 2024). Balestriero LeCun (2024) recently highlighted the scale of the challenge for higher-variance datasets such as ImageNet (Deng et al., 2009). We hope that the theoretical connection to popular SSL methods and the notable performance improvements of SimVAE over other VAE- based methods encourages further investigation into generative representation learning. This seemss particularly justified given the potential for uncertainty ostimation from the posterior, the ability to generate novel data samples, and the potential for fully lask-agnostic representations learned under a principled rather than heuristic approach. The transparency the SSL Model may also allow representations learned that disentangle rather than discard style information (Higgins et al., 2017) through careful choice of model parameters. The SSL Model thus serves as basis for understanding self-supervised learning and offers guidance for the design and interpretation of future representation learning methods.
11

---

Published in Transactions on Machine Learning Research (09/2024)
Acknowledgments
Alice Bizeul gratefully supported by an ETH Al Center Doctoral Fellowship and by Professors Julia Vogt and Bernhard Schölkopf. Carl gratefully supported by an ETH AT Centre Postdoctoral Fellowship, a grant from the Haslerstiftung (no. 23072) and by Professors Mrinmaya Sachan and Ryan Cotterell. We thank Pawel Czyz and Thomas Sutter for their valuable feodback.
References
Carl Allen and Timothy Hospedales. Analogies explained: Towards understanding word embeddings. In
ICML, 2019. 3
Sanjeev Arora, Hrishikesh Khandeparkar, Mikhail Khodak, Orestis Plevrakis, and Nikunj Saunshi. A theoretical analysis of contrastive unsupervised representation learning. In TCML, 2019. 4
Philip Bachman, R Devon Hjelm, and William Buchwalter. Learning representations by maximixing mutual information across views. In NewrIPS, 2019. 4
Alexei Baevski, Yuhao Zhou, Abdelrahman Mohamed, and Michael Auli. wav2vec 2.0: A framework for self-supervised learning of speech representations. In NeurIPS, 2020. 1
Ivana Balažević, David Steiner, Nikhil Parthasarathy, Relja Arandjelović, and Olivier J Hénaff. Towards in-context scene understanding. NeurlPS, 2023.
Randall Balestricro and Yann LeCun. How learning by reconstruction produces uninformative features for perception. In Forty-first International Conference on Machine Learning, 2024. 11
Randall Balestriero, Mark Ibrahim, Vlad Sobal, Ari Morcos, Shashank Shekhar, Tom Goldstein, Florian Bordes, Adrien Bardes, Gregoire Mialon, Yuandong Tian, et al. A cookbook of self-supervised learning. arXiv preprint arXiv.230..12210, 2023. 3, 8
Adrien Bardes, Jean Ponce, and Yann LeCun. Viereg: Varianceiivvrraancccovrriii regularization for self-supervised learning. In ICLR, 2022. 3, 9, 19
Ido Ben-Shaul, Ravid Shwartz-Ziv, Tomer Galanti, Shai Dekel, and Yann LeCun. Reverse engineering self-supervised learning. In NeurIPS, 2023.
Mathilde Caron, Piotr Bojanowski, Armand Joulin, and Matthijs Douze. Deep clustering for unsupervised learning of visual features. In ECCV, 2018. 2
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervisod
learning of visual features by contrasting cluster assignments. In NeurIPS, 2020. 3, 9
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand
Joulin. Emerging properties in self-supervised vision transformers. In ICCV, 2021. 1, 3, 8
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In ICML, 2020a. 1, 2, 3, 8, 9, 21, 22, 23
Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He. Improved baselines with momentum contrastive learning. arXiv preprint arXiv:2003.0422 04497, 2020b.
Sumit Chopra, Rain Hadsell, and Yann LeCun. Learning a similarity metric discriminatively, with application to face verification. In CVPR, 2005. 3
Thomas Cover and Peter Hart. Nearest neighbor pattern classification. In IEEE Transactions on Information Theory, 1967. 9
Imant Daunhawer, Alice Bizeul, Emanuele Palumbo, Alexander Marx, and Julia E Vogt. Identifiability results for multimodal contrastive learning. In ICLR, 2023.
12

---

Published in Transactions on Machine Learning Research (09/2024)
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEFE conference computer vision and pattern recognition, pp. 248 118-255. leee, 2009. 11
Shchzaad Dhuliawala, Mrinmaya Sachan, and Carl Allen. Variational Classification. In TMLR, 2023. 4,
Alexey Dosovitskiy, Jost Tobias Springenberg, Martin Riedmiller, and Thomas Brox. Discriminative unsuper- vised feature learning with convolutional neural networks. In NeurIPS, 2014. 2, 3
Harrison Edwards and Amos Storkey. Towards a neural statistician. In ICLR, 2016. 16
Linus Ericsson, Henry Gouk, Chen Change Loy, and Timothy M Hospedales. Self-supervised representation learning: Introduction, advances, and challenges. IEEE Signal Processing Magazine, 39(3):42-62, 62, 2022.
Quentin Garrido, Yubei Chen, Adrien Bardes, Laurent Najman, and Yann Lecun. On the duality between contrastive and non-contrastive self-supervised learning. arXiv preprint arXiv:2206.02574, 2022.
Spyros Gidaris, Praveer Singh, and Nikos Komodakis. Unsupervised representation learning by predicting image rotations. arXiv preprint
2018. 3
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. In NeurlPS, 2020. 3, 8, 19
Kartik Gupta, Thalaiyasingam Ajanthan, Anton van den Hengel, and Stephen Gould. Understanding and improving the role of projection head in self-supervised learning. arXiv preprint arXiv:2212.11491, 2022.
Raia Hadsell, Sumit Chopra, and Yann LeCun. Dimensionality reduction by learning invariant mapping. In CVPR, 2006. 3
Jeff 7. HaoChen, Colin Wei, Adrien Gaidon, and Tengyu Ma. Provable guarantees for self-supervised deep learning with spectral contrastive loss. In NeurIPS, 2021. 4
Jiawei Hc, Yu Gong, Joseph Marino, Greg Mori, and Andreas Lehrmann. Variational autocncoders with jointly optimized latent dependency structure. In ICLR, 2018. 16
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. CVPR, 2016. 10
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In CVPR, 2020. 3, 8, 9
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders
are scalable vision learners. In CVPR, 2022. 2, 3
Martin Hensel, Hubert Ramsauor, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. In NeurIPS, 2017.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. bela-vae: Learning basic visual concepts with a constrained variational framework. In ICLR, 2017. 4, 8, 9, 11
R Devon Hjelm, Alex Fedorov, Samuel Mavoie-Marchildon, Karan Grewal, Phil Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. In ICLR, 2019. 3,
Tero Karras, Samuli Laine, and Timo Aila. A style-based generalor architecture for generative adversarial networks. In CVPR, 2019.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. In ICLR, 2014. 4, 10
13

---

Published in Transactions on Machine Learning Research (09/2024)
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009. 9, 20
Yann LeCun. The mnist database of handwritten digits. http://yann. lecun. com/cxdb/mma //nssd////is//.....//// 1998. 9, 20
Jason D Lee, Qi Lei, Nikunj Saunshi, and Jiacheng Zhuo. Predicting what you already know helps: Provable sell-supervised learning. In NeurIPS, 2021. 4
Omer Levy and Yoav Goldberg. Neural word embedding as implicit matrix faclorization. In NeurIPS, 2014.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In ICCV, 2015. 9, 20
David McAllester and Karl Stratos. Formal limitations on the mcasurement of mutual information. In ATStats, 2020. 4, 8
Tomas Mikolov, Nya Sutskever, Kai Chen, Greg S Corrado, and Jeff Dean. Distributed representations of words and phrases and their compositionality. In NeurIPS, 2013. 1,
Ishan Misra and Laurens van der Maaten. Self-supervised learning of pretext-invarien representations. In CVPR, 2020. 1
Hiroki Nakamura, Masashi Okada, and Tadahiro Taniguchi. Representation uncertainty in self-supervised learning as variational inference. In ICCV, 2023. 4, 16
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding, arXiv preprint arXiv:1807.0374 03748, 2018. 1, 3, 8, 16, 17, 18
Serdar Ozsoy, Shadi Hamdan, Sercan Arik, Deniz Yuret, and Alper Erdogan. Self-supervised learning with information maximization criterion. In NeurIPS, 2022. 8
Vardan Papyan, XY Han, and David I, Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Procoedings of the National Academy of Sciences, 117 (0)44)))44522225 2020. 4
Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On variational bounds of mutual information. In ICML, 2019. 8, 17
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. Tn ICML, 2021. 1, 2, 3, 8,
Rajesh Ranganath, Dustin Tran, and David Blei. Hierarchical variational models. In ICML, 2016.
Jason Tyler Rolfe. Discrete variational autoencoders. In ICLR, 2017. 16
Emanuele Sansone and Robin Manhaeve. Gedi: Generative and discriminative training for self-supervised learning. arXiv preprint arXiv:2212.13425, 2022.
Nikunj Saunshi, Jordan Ash, Surbhi Goel, Dipendra Misra, Cyril Zhang, Sanjeev Arora, Sham Kakade, and Akshay Krishnamurthy. Understanding contrastive learning requires incorporating inductive biases. ICML, 2022. 4
Ravid Shwartz-Ziv, Randall Balestriero, Kenji Kawaguchi, Tim GJ Rudner, and Yann LeCun. An information- theoretic perspective on wariane-iinvriince-coaaraar regularization. In NeurIPS, 2023. 4
Abhishek Sinha, Jiaming Song, Chenlin Meng, and Stefano Ermon. D2c: Diflusion-decodiig models for
few-shot conditional generation. In NeurIPS, 2021. 16
Samarth Sinha and Adji Bousso Dieng. Consistency regularization for variational auto-encoders. In NeurIPS. 2021. 4, 9
Kihyuk Sohn. Improved deep metric learning with multi-class n-pair loss objective. In NeurIPS, 2016.
14

---

Published in Transactions on Machine Learning Research (09/2024)
Casper Kaac Sonderby, Tapani Raiko, Lars Maalee, Soren Kaac Sønderby, and Ole Winther. Ladder variational autoencoders. In NTPS, 2016. 16
Chunfeng Song, Feng Liu, Yongzhen Huang, Liang Wang, and Tieniu Tan. Auto-encoder based data clustering. In Progress in Pattern Recognition, Image Analysis, Computer Vision, and Applications, 2013. 3
Yuandong Tian. Understanding deep contrastive learning via coordinate-wise optimization. In NeurIPS, 2022.
Christopher Tosh, Akshay Krishnamurthy, and Damiel Hsu. Contrastive learming, multi-view redundancy, and linear models. In Algorithmic Learning Theory, pp. 1179-1206. PMLR, 2021. 4
Yao-Hung Hubert Tsai, Yue Wu, Ruslan Salakhutdinov, and Louis-Philippn Morency, Self-supervised learning from multi-view perspective. In ICLR, 2020. 4
Michael Tschannen, Josip Djolonga, Paul K Rubenstein, Sylvain Gelly, and Mario Lucic. On mutual information maximization for representation learning. In ICLR, 2020. 4, 8, 19
Michael Techannen, Manoj Kumar, Andreas Steiner, Xiaohua Zhai, Neil Houlsby, and Lucas Beyer. Image captioners are scalable vision learners too. NeurIPS, 2024. 9, 11
Harri Valpola. From neural pea to deep unsupervised learning. In Advances in Independent Component Analysis and learning machines, pp. 143-171. Elsevier, 2015. 16
Julius Von Kügelgen, Yash Sharma, Luigi Gresele, Wieland Brendel, Bernhard Schölkopf, Michel Besserve, and Francesco Locatello. Self-supervised learning with data augmentations provably isolates content from style. In NeurIPS, 2021.
Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through alignment and uniformity on the hypersphere. In ICML, 2020. 2, 4, 6, 8
Yifei Wang, Qi Zhang, Yisen Wang, Jiansheng Yang, and Zhouchen Lin. Chaos is a ladder: A new theoretical understanding of contrastive learning via augmentation overlap. In ICLR, 2021. 4
Yizhou Wang, Shixiang Tang, Feng Zhu, Lei Bai, Rui Zhao, Donglian Qi, and Wanli Ouyang. Revisiting the transferability of supervised pretraining: an mlp perspective. In CVPR, 2022. 8
Chengzhi Wu, Julius Pfrommer, Mingyuan Zhou, and Jürgen Beyerer. Generatiie-contrrstiiv learning for self-supervised latent representations of 3d shapes from multi-modal euclidean input. arXiv preprint arXiv:2301.04612. 2023. 16
Zhirong Wu, Yuanjun Xiong, Stella X Yu, and Dahua Lin. Unsupervised feature learning via pon-paraettric instance discrimination. In CVPR, 2018. 3
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.07747, 2017. 9, 20
Junyuan Xie, Ross Girshick, and Ali Farhadi. Unsupervised deep embedding for clustering analysis. In ICML, 2016.
Zhenda Xie, Zheng Zhang, Yue Cao, Yutong Lin, Jianmin Bao, Zhuliang Yao, Qi Dai, and Han Hu. Simmim: A simple framework for masked image modeling. CVPR, 2022. 3
Bo Yang, Xiao Fu, Nicholas D Sidiropoulos, and Mingyi Hong. Towards f-means-friendly spaces: Simultancous deep learning and clustering. In ICML, 2017. 3
Mingtian Zhang, Tim Z Xiao, Brooks Paige, and David Barber. Improving vac-based representation lcarning. arXiv preprint ar 2rXiv:2205.14533 2022. 3
Roland Zimmermann, Yash Sharma, Steffen Schneider, Matthias Bethge, and Wieland Brendel. Contrastive learning inverts the data generating process. In ICML, 2021. 4
15

---

Published in Transactions on Machine Learning Research (09/2024)
A
Appendix
A.1 Background: Relevant VAE architectures
The proposed hierarchical latent variable model for self-supervised learning (Fig, 1) is fitted to the data distribution by maximising the ELBOssL and can be viewed as a VAE with a hierarchical prior.
VAEs have been extended to model hierarchical latent structure (c.g. Valpola, 2015; Ranganath et al., 2016; Rolfe, 2017; He et al., 2018; Sønderby et al., 2016; Edwards Storkey, 2016), which our work relates to. Notably, Edwards & Storkey (2016) propose the same graphical model as Fig. 1, but methods differ in how posteriors are factorised, which is key aspect for learning informative representations that depend only the sample they represent. Wu et al. (2023) and Sinha et al. (2021) combine aspects of VA and contrastive learning but do not propose a latent, variable model for SSL. Nakamura et al. (2023) look to explain SSL methods via the ELBO, but with second posterior approximation not generative model.
A.2 Derivation of ELBOssL and SimVAE Objective
Let (xx [.)".....') be set of J semantically related samples with common label y (c.g. the index of set of augmentations of an image æ;). Let w {0, vi, J be parameters of the generalive model for SSL (Fig. 1) and o be the parameter of the approximate posterior (z(z))) We derive the Evidence Lower Bound (ELBOsst) that underpins the training objectives of (projective) discriminative SSL methods and is used to train SimVAE ($5).
min DKLL (xlv)) = max E
[og P.
x.y
J/ g0 (x<()()) log log Pw (x(W)
max E
with x.y
[/c (z/x) log po (x(x)))p (x)p)))
max E
(x(x)))))) (x(a)))
with x.y
[/ (x)x)) log (a(xle)ps + (orl(0022x)) [2|x,)))
max E
(xax)
w. x;9
// 2o (x(z))x) xx log po(x) ()(xx)
max E
log log (2lw)
(ef Eq. ELBO)
w. x.9
EV 20 (2' log Po(e'1z) (0
+ / (z(z))) (o(zz))loggpe po (zl))(
max E
2 x.y
(= ELBOssL)
Terms of ELBOsst. are analogous to those of the standard ELBO: reconstruction error, entropy of the approximate posterior (qq and the (conditional) prior. Algorithm 1 provides an overview of the computational steps required to maximise ELBOSsL under Gaussian assumptions described in $5, referred as SimVAE. As our experimental setting considers augmentations semantically related samples, Algorithm incorporates a preliminary step augment data samples.
A.3 Detailed derivation of InfoNCE Objective
For data sample I4, E Xigs let 24, E Exi, be a semantically related positive sample and {a be random negative samples. Denote by x= (iig,,' the positive pair, by x- ( ..., a/, all negative samples, and
by X x-X all samples. The InfoNCE objective is derived as follows (by analogy to Oord et. al. (2018)).
Ex [logrly = 0[}]]
(predict y = index of positive sample X-)
Ex log / p(y 012)g(Z(x)]
(introduce latent variables Z: 2 + X)
2 Ex( / ((Z(x) )) logr(u = 0/2)]
(by Jensen's 's incquality)
16

---

Published in Transactions on Machine Learning Research (09/2024)
Ex ((([x) log
Fig = 0)
(Bayes rule, note p(y p(y=r)= 1.- Vr)
p(2 (ig.))))
Ex [// (((1X) log
p(zig |2np) موgا p(zi.))
(from sample similarisy/indeeeedenee)
r=o p(zz, |z4) IIaser pl p(z
Ex [// ((11x) log
platy (:10 )/m(t,,)
(divide through by oo
> p(xi, 1810 p/p((.)))
Ex [/ʃ (Z|x) log
pltig (1g)) (10gg11))// Σ plrl mmo (ool
The final expression is parameterised using a similarity function sim(2,z') give the objective.
LINCE = =Ex / ((( log sim(2/n =20))
- sim( =250
Oord el al. (2018) show, and Poole el al. (2019) confirm, that this loss is a lower bound on the mutual information, which improves with the number of negative samples k.
LINCE
=Ex[ [// (() log
p / pp(z),)) E. p(z/, (102 p(=
Exx
(multiply Eq. 7 through by p(zio)
&
Ex 1/8 (Z log plzio(zoooooo (mo) - log (pl Zzo (2(,)
=Ex
Σ
p(zi (x4,))
Ex [// a(21x) log (1+ Σ
p( (o(
(divide through by ((2)ol=())
p Flo
r=1
Ex [/, ((z)x) log (1 (k- - 1)Bug L ((2g/az))
plano Fio }]
plang
Fx / (((2x) log (1 + (I: - 1) p(21o Finy"
p(zin)
Ex / ((z/x) log(k 1) p(zto 1%
plain)
Ex / ((z)x) log p(angrat)) ] + +log (k: 1)
1
p(ano)l(zig)) { Ex / (o((=)x)(oogp(()b)))
Σ log (lgpp())))
+
In the last step, we revert to the terminology used in the main paper for ease of reference.
A.4 Derivation of parameter-free p(z,(y)) s(z;)
Instance Discrimination methods consider J=1 sample 2; at a time, labelled by its index yi =i, and computes p(a,ly= = i;),) from stored instanco-specific parameters 0,, This requires parameters proportional to the dataset size, which could be prohibitive, whereas parameter number is often independent of the dataset size, or grows slowly. We show that contrastive methods (approximately) optimise the same objective, but without parameters, and here explain how that is possible. Recall that the "label" i is semantically meaningless and simply identifics samples of a common distribution =)(xx=) For J semantically related samples x; {x]))=> a! ~ p(x(y,), their latent variables are conditionally independent, hence p(z,)u)) = f p(t)))p(z.l) th)) = J p())) p(x/ [ys: vi)) s(z,), function of the latent variables that non-paraetrrcaally approximates the joint distribution of latent variables of semantically related data. (Note that unsemantically related data are independent and the joint distribution over their latent variables is a product of marginals).
17

---

Published in Transactions on Machine Learning Research (09/2024)
We assume Gaussian prior p())) (N and class-conditionals p(z/|1))= = (2/1100,, (for fixed variance x²).
p(zslus)= = / p(z: (14ilss)p((1) ((bp))(())) = -/1, p((1)s) [(((/)))))
j
expl 17pl-77700 v/ II expl 20/ - )
41
/ exp( 20²
( + ( - 01)
41
/ expl
2/ - 2( Σ vi + ( + S)
=
20
/
exp(
(Vi
Σ
x } + expl 20
Σ
12
Σ 2))3
(+)
20
0
4/4
8
x expl
1
12
Σ +1
20
jfk
=xxpf
((1 02/72-11 1)
Σ
20
d2/
1
x exppl 20-
(if ||:||2 = 1)
J
The result can be rearranged into a Gaussian form (a well known result when all distributions are Gaussian), but the last line also shows that, under the common practice of setting embeddings to unit length (1==|2=1), s(-) can be calculated directly from dot products, or cosine similarities (up to a proportionality constant, which does not affect optimisation).
If we instead assume a uniform prior, we can take the limit of the line marked (+) as 7 + oo:
expf
1
52
1
E 2()3)
202
02172+1
eexp
(13)
1
2)99
20
j
=cxpf
222
202
= expl 20
22
21 + +1,211
= eepp- - 1 (1
(2) - 27 Σ 2 + Σ );²)
2
3
=exp1
1
(2)2 - 22(222111
20
expl- - [x1 -
(8)
A.5 Relationship between InfoNCE Representations and PMI
For data sampled x ~ p(x) and augmentations x' ~ (x'lx) sampled under synthetic augmentation strategy, Oord et al. (2018) show that the InfoNCE objective for a sample is optimised if their respective representations 2, z' satisfy
(xp(sim(2z2")) 2²)) p(x,z)) p(z)o(z)))
(9)
18

---

Published in Transactions on Machine Learning Research (09/2024)
where sim(;s) is the similarity function (e.g. dot product), and is a proportionality constant, specific to 2. Since may differ arbitrarily with x it can be considered an arbitrary function of x; but for simplicity we consider a particular a and fixed c. Further, e > 0 is strictly positive since it is a ratio between positive (exponential) and non-negative (probability ratio) terms. Accordingly, representations satisfy
sim(z, 2²) = PMI(z,x') +c',
(10)
where e' = log IR and (M, x') is the pointwise mutual information between samples x and x'. Pointwise mutual information (PMI) is an information theoretic term that reflects the probability of events occurring jointly versus independently, For an arbitrary sample and augmentation this is given by:
Pr (x'1z) p(x, x') = log p(x²)
())(2,2')= = log p(a)p(a')
(11)
We note that (x if x can not be augmented to produce x'; and that, in continuons domain, such as images, two augmentations are identical with probability zero. Thus augmentations of different samples are expected to not overlap and the marginal is given by p(x')= I, (x' (x)((x))))))))))))))))))))))([ = p+ (x' [a'l="")0)0)"") px)))))) where x² is the sample augmented to give x'. Thus
(r(211z) Pr (x' plx²) (+(1'22* (x(=))))))))))))))))))))))))))
1/D(a*)
if === (i.e. x' is an augmentation of x) otherwise
(12)
0
and PMI(x, x²) -lg 2 k > 0 or PMI -'= = -00 respectively. Here = - log arg max, p(z) is a finite value based on the most. likely sample. For typical datasets, this can be approximated empirically by N where N is the size of the original dataset (since that is how often the algorithm observes each sample), hence k = log N, often of the order 5 10 (depending the dataset).
If the main objective were to accurately approximate PMT (subject to a constant d') in Eq. 10, e.g. to approximate mutual information, or if representation learning depended on it, then, at the very least, the domain of sim(-..) must span its range of values, seen above as from -o0 for negative samples to a small positive value (e.g. 5-10) for positive samples. Despite this, the popular bounded cosine similarity function (cossim(zzz") T=laa|eTT E 1]) is found to putperform the unbounded dot product, even though the cosine similarity function necessarily cannot span the range required to reflect true PMT values, while the dot product can. This strongly suggests that representation learning does not require reprosentations to specifically learn PMI, or for the overall loss function to approximale mutual information.
_T
Instead, with the cosine similarity constraint, the InfoNCE objective is as optimised as possible if represen- tations of a data sample and its augmentations are fully aligned (cossim(2, z/) = 1) and representations of dissimilar data are maximally misaligned cossim((2, =') = 11, since these minimise the error from the true PMI values for positive and negative samples (described above). Constraints, such as the dimensionality of the representation space vs the number of samples, may prevent these revised theoretical optima being fully achieved, but the loss function is optimised by clustering representations of sample and its angmentations and spreading apart those clusters. Note that this is the same geometric structure as induced under softmax cross-entropy loss (Dhuliawala et al., 2023).
We note that our theoretical justification for representations not capturing PMI is supported by the empirical observation that closer approximations of mutual information do not appear to improve representations (Tschannen et al., 2020). Also, more recent contrastive self-supervised methods increase the cosine similarity between semantically related data but spread apart represeniation the without negative sampling of InfoNCE, yet outperform the InfoNCE objective despite having no obvious relationship to PMI (Grill et al., 2020; Bardes et al., 2022).
A.6 Information Loss due to Representation Collapse: discussion
While may seem appealing to lose information by way of representation collapse, e.g. obtain representations invoriant to nuisance factors, this is problematic notion from the perspective of ceneral-pupposs representatin ation
19

---

Published in Transactions on Machine Learning Research (09/2024)
learning, where the downstream task is unknown there may be many, since what is noise for task may be of use in another. For example, "blur" is often considered noise, but camera on an autonomous vehicle may be better to detect blur (e.g. from soiling) than be invariant to it and eventually fail when blurring becomes too severe. We note that humans can observe a scene including many irrelevant pieces of information, e.g. we can tell when an image is blurred or that we are looking through a window, and "disentangle" that from the rest of the image. This suggests that factors can be preserved and disentangled.
To stress the point that representation collapse is not desirable in and of itself, we note that collapsing together representations of semantically related data x, would be problematic if subsets x, overlap. For example, in the discrete case of wordêver, words are considered semantically related if they co-occur within a fixed window. Representation collapse, here, would mean that co-occurring words belonging to the same x; would have the same representation, which is clearly undesirable.
A.7 Experimental Details
A.7.1 SimVAE Algorithm
Algorithm SimVAE
Require:
data {xi)M1: batch size N; data dim D; latent dim L; augmentation set T; number of augmenta- tions J; encoder føi decoder go; variance of zly, a²; for randomly sampled mini-batch {x,} do for augmentation Ø ~ T do
x = (x:);
augmeet samples
n² E = (x();
# forward pass: z ~ px(2(x)
z ~ N (1, E));
x = go (z));
# x E(x 0)
end for
Li 1 J
2 ||x
- x/112
# minimize loss
CH = 2 2j==
1 Σ j=1
log( (1
prior = 1= [|( - ト 2//0
min Em= 1= Cive + L + ) w.r.t. Φ, 0 by SGD;
end for return Φ
A.7.2 Datasets
MNIST. The MNIST dataset (LeCun, 1998) gathers 60'000 training and 10'000 testing images representing digits from 0 to 9 in various caligraphic styles. Images were kept to their original 28x28 pixel resolution and were binarized. The 10-class digit classification task was used for evaluation.
The FNshionMNIST dataset (Xiao et al., 2017) is a collection of 60'000 training and 10'000 test images depicting Zalando clothing items (i.e., t-shirts, trousers, pullovers, dresses, coats, sandals, shirts, sneakers, bags and ankle boots). Images were kept to their original 28x28 pixel resolution. The 10-class clothing type classification task used for evaluation.
FashionMNIST.
The CIFAR10 dataset (Krizhevsky et al., 2009) offers compact dataset of 60,000 (50,000 training and 10,000 testing images) small, colorful images distributed across ten categories including objects like airplanes, cats, and ships, with various lighting conditions. Images were kept to their original 32x32 pixel resolution.
CIFAR10.
CelebA. The CelebA dataset (Liu et al., 2015) comprises vast collection of celebrity facial images. It encompasses diverse set of 183'000 high-resolution images (i.e., 163 163'000 training and 20'000 test images),
20

---

Published in Transactions on Machine Learning Research (09/2024)
each depicting distinet individual. The dataset showeases wide range of facial attributes and poses and provides binary labels for 40 facial attributes including hair & skin colour, presence absence of attributes such as eyeglasses and facial hair. Each image was cropped and resized to 64x64 pixel resolution. Attributes referring to hair colour were aggregated into a 5-class attribute (i.e., bald, brown hair, blond hair, gray hair, black hair). Images with missing or ambiguous hair colour information were discarded at evaluation.
All datasets were sourced from Pytorch's dataset collection.
A.7.3 Data augmentation strategy
Taking inspiration from SimCLR's (Chen el al., 2020a) augmentation stralegy which highlights the importance of random image cropping and colour jitter on downstream performance, our augmentation strategy includes random image cropping, random image flipping and random colour jitter. The colour augmentations are only applied to the of the datasets we explored, hyperparameters such as the cropping strength were adapted to each dataset. to ensure that semantically meaningful features remained after augmentation. The augmentation strategy hyperparameters used for each dataset are detailed in table 2.
gray-scale datasets (i.e., CIFAR10 & CelebA datasets). Due to the varying complexity
Dataset
Crop
Vertical Flip
Colour Jitter
scale
ratio
prob.
b-s-c
hue
prob.
MNIST
0.4
10.3
0.5
Fashion
0.4
[0.75,1.3]
0.5
CIFAR10
0.6
[0.75,1.3
0.5
0.8
0.2
0.8
CelebA
0.6
[0.75,1.3]
0.5
0.8
0.2
0.8
Table 2: Data augmentation strategy for each dataset: (left to right) cropping scale, cropping ratio, probability of vertical flip, sigghttess--aaturationo-ooo colour jitter
jitter, hue jitter, probability of
A.7.4 Training Implementation Details
This section contains all details regarding the architectural and optimization design choices used to train SimVAE and all basclincs. Method-specifie hyperparanetors are also reported below.
Network Architectures. The encoder network architectures used for SimCLR, MoCo, VicREG, and VAE-based approaches including SimVAE for simple (i.e., MNIST, FashionMNIST) and complex datasets (i.e,, CIFAR10, CelebA) are detailed in Table 3a, Table 4a respectively. Generative models which include all based methods also require decoder networks for which the architectures are detailed in Table 3b and Table 4b. The latent dimensionality for MNIST and FashionMNIST is fixed at 10 and increased to 64 for the CelobA and CIFAR10 datasets. The encoder and decoder architocture networks are kept constant across methods including the latent dimensionality to ensure a fair comparison.
Layer Name
Output Size
Block Parameters
Layer Name
Output Size
Block Parameters
fel
500
784x500 fc, relu
fel
2000
10x2000 fe, relu
fc2
500
500x500 fc, relu
fc2
500
2000x500 fc, relu
fc3
2000
500x2000 fc, relu
fc3
500
500x500 fc, relu
lo4
10
2000x10
fe4
784
500x784 fc
(a) Encoder
(b) Decoder
Table 3: Multi-layer perceptron network architectures used for MNIST & FashionMNIST training
Optimisation & Hyper-parameter tuning. All methods were trained using an Adam optimizer until
training loss convergence. The batch size fixed to 128. Hyper-parameter tuning was performed based
21

---

Published in Transactions on Machine Learning Research (09/2024)
Layer Name
Output
Block Parameters
Layer Name
Output
Block Parameters
convl
32x32
4x4, 16, stride 1 batchnorm, relu
fc
256x4x4
64x4096 fc
3x3 maxpool, stride
conv2_xx x
32x32
3x3, 32, stride 1
convl
8x8
3x3, 128, stride 2
3x3, 82, stride 1
3x3, 128, stride
conv3
16x16
3x3, 64, stride 2
conv2_
16x16
3x3, 64, stride 2
3x3, 64, stride 1
3x3, 64, stride 1
conv x
3x3, 128, stride
conv3 x
32x32
3x3, 32, stride
3x3, 128, stride
3x3, 32, stride
conv5
4x4
3x3, 256, stride 2
conv4 X
64x64
3x3, 16, stride 2
3x3, 256; stride
3x3, 16, stride
lc
64
4096x64 fc
conv5
64x64
5x5, 3, stride
(a) Encoder
(b) Decoder
Table 4: Resnet18 network architectures used for CIFARIO & CelebA datasets
on the downstream MLP classification accuracy across datasets. The final values of hyperparameters were selected to reach the best average rownstreeam performance across datasels. While we observed stable performances across datasets for the VAE family of models, VicREG and MoCo, SimCLR more sensitive, leading to difficulties when having to define unique set of parameters across datascts. For VAEs, the learning rate was set to 85-S, and the likelihood probability, p(æ[[=), variance parameter was set to 0.02 for 3-VAE, CR-VAF and SimVAF. CR-VAF's À parameter was set to 0.1. SimVAE's prior probability, p(=ly), variance was set to 0.15 and the number of augmentations to 10. VicREG's parameter was set to 25 and the learning rate to lc-4. SimCLR's temperature paramcter, т, was set to 0.7 and learning rates were adapted for each dataset due significant performance variations across datasets ranging from -e-5 to le-3 For MoCo, the temperature parameter was fixed to 0.7, for MoCov2 to 0.2. To ensure a more fair comparison between methods, we kept the augmentation strategy identical across methods and did not add blurring for MoCov2 training.
A.7.5 Evaluation Implementation Details
Following common practices (Chen et al., 2020a), downstream performance is assessed using a linear probe, a multi-layer rerceptron probe, a k-nearest neighbors (kNN) algorithm, and a Gaussian mixture model (GMM). The linear probe consists of fully connected layer whilst the mlp probe consists of two fully connected layers with relu activation for the intermediate layer. Both probes were trained using Adam optimizer with learning rale of 34 4 for 200 epochs with batch size fixed to 128. Scikit-learn's Gaussian Mixture model with a full covariance matrix and 200 initialization was fitted to the representations using the ground truth cluster number. The k-NN algorithm from Python's Scikit-learn library was used with k spanning from 1 to 15 neighbors. The best performance was chosen as the final performance measurement. No augmentation strategy was used at evaluation.
Computational Resources. Models for MNIST, FashionMNIST and CIFAR10 were trained on a RTX2080ti GPU with 12G RAM. Models for CelebA were trained on an RTX3090 GPU with 24G RAM. We observe that while the family of generative models requires more time per iteration, the loss converges faster while discrininative meth converge at a slower rate when considering the optimal set of hyperparameeers. As a consequence, generative baselines and SimVAE were trained for 400 epochs while discriminative methods were trained for 600 to 800 epochs.
A.7.6 Generation Protocol
Here detail the image generation protocol and the quality evaluation of generated samples.
22

---

Published in Transactions on Machine Learning Research (09/2024)
Acc-LP
Acc-GMM
MNIST
Fashion
CelebA
CIFAR10
MNIST
Fashion
CelebA
CIFAR10
Random
39.7 2.4
51.2 = 0.6
64.4 = 0.9
15.7 0.9
42.2 1.2
48.6 0.2
59.2 + 0.3
13.1 = 0.6
SimCLR
96.8 1 0.1
73.0 0.3
94.2 +: 0.2
65.4 : 0.1
83.7 + 0.8
53.6 0.3
71.6 0.6
28.2 :: 0.2
MoCo
88,6 1.7
65.0 i: 1.3
53.3 1.3
70.5 4.0
56.6 1.1
52.4 0.3
VieREG
96.7 + 0.0
71.7 + 0.1
94.3 + 0.3
68.2 # 0.0
79.8 + 0.6
60.2 = 1.1
53.9 + 0.2
35.0 = 2.8
VAE
97.2
79.0
81.5 =
24.7
96.3 0.4
57.9 0.8
58.8 0.2
23.4 +
3-VAE
97.8 0.0
79.6 + 0.0
81.9 + 0.2
26.9 = U.D
96.2 + 0:2
68.0 0.3
59.5 0.6
31.2 # D.1
CR-VAE
97.5 0.0
79.7 0.0
81.6
26.8
96.9 =
63.4 0.4
58.9 0.4
30.3 =
SimVAE
98.0 = 0.0
80.0 0.0
87.1 + 0.3 40.1 + 0.0
96.6 + 0.0
71.1 + 0.0
58.4 = 0.6
39.3 = 0.0
Table Top-1% self-supervised Acc (t) for MNIST, FashionMNIST, CIFARIO, and CelebA (gender classification) using linear probe (LP) and Gaussian Mixture Model (GMM) classification methods; We report mean and standard errors over three runs; Bold indicate best scores in each method class: generative (teal), discriminative methods (red).
Ad-hoc decoder training VAE-based approaches, including SimVAE, are fundamentally generative mothods aimed at approximating the logarithm the marginal likelihood distribution, denoted as log)(x). In contrast, most traditional self-supervised methods adopt discriminative framework without a primary focus on accurately modeling p(z). However, for the purpose of comparing reprosentations, and asscssing the spectrum of features present in 2, intend to train a decoder model for SimCLR & VicREG models. This decoder model is designed to reconstruct images from the fixed representations initially trained with these approaches. To achieve this goal, we train decoder networks using the parameter configurations specified in Tables 3b and 4b, utilizing the mean squared reconstruction error as the loss function. The encoder parameters remain constant, while we update the decoder parameters using an Adam optimizer with learning rate of 1 until a minimal validation loss is achieved (i.c. 10-80 epochs).
Conditional Image Generation. To allow for a fair comparison, all images across all methods are generated by sampling from multivariate Gaussian distribution fitted to the training samples' representations. More precisely, each Gaussian distribution is fitted to z conditioned on label y. Scikit-Learn Python library Gaussian Mixture model function (with full covariance matrix) is used.
A.8 Additional Results & Ablations
Content retrieval with linear & gaussian mixture model prediction heads. ?? reports the top-1% self-supervised classification accuracy using linear prediction head and gaussian mixture model. From ??, draw similar conclusion with Table 1: SimVAE significantly bridges the gap between discriminative and generative self-supervised learning methods when considering supervised linear predictor and fully unsupervised methods for downstream prediction. Table 6 report the normalized mutual information (NMI) and adjusted rank index (ARI) for the fitting of the GMM prediction head.
Content & Style retrieval for all CelebA attributes. Figure reports average classífication accuracy using a MLP probe (over 3 runs) for the prediction of 20 CelebA facial attributes for SimVAE, generative and discriminative baselines.
Augmentation protocol strength ablation. Figure 7 reports the downstream classification accuracy across methods for various augmentations strategics. More preciscly, we progressively incrcase the cropping scale and color jitter amplitude. Unsurprisingly (Chen et al., 2020a), discriminative methods exhibit high sensitivity to the augmentation strategy with stronger disruption leading to improved content prediction. The opposite trend is observed with vanilla generative methods where reduced variability amongst the data leads increased downstream performance. Interestingly, SimVAE is robust augmentation protocol and performs comparably across settings.
23

---

Published in Transactions on Machine Learning Research (09/2024)
100
VAE
3-VAE
CR-VAE
SimVAE
SimCLR
VicReg
CA MLP
70
hair color
arch. eyebrow
bags e99
lips
big nose
big eyebroa
maleup
high check
moath open
narras eyees
100
VAE
3-VAE
CR-VAE
SimVAE
SimCLR
VicReg
CA MLP
70
gonder
beand
oval lace
pointy nose
amile
sbright hair
wapy lair
eariings
lipttic
young
CelebA Attributes
Figure 6: CelebA 20 facial attributes prediction using a MP. Average scores and standard errors are reported
across 3 random seeds.
24

---

Published in Transactions on Machine Learning Research (09/2024)
Dataset
VAE
B-VAE
CR-VAE
SimVAE
MoCo
VicREG
SimCLR
ART
89.0 + 1.0
93.3 + 0.3
94.0 + 0.0
93.1 + 0.0
58.3 3.8
72.0 +: 0.7
77.4 = 0.2
MNIST
NMI
94.9 +
96.7 +: 0.2
96.9
96.6 +
71.4
86.8 + 0.4
89.6
ARI
44.3 +: 0.9
53.3 4: 0.4
47.6 0.4
56.8 4: 0.0
30.9 0.5
41.2 d: 0.5
33.2 0.3
Fashion
NM1
69.1 + 0.6
75.6 + 0.1
72.6
77.1 + 0.0
50.4 0.6
66.9 + 0.3
62.1 0.2
ARI
5.7 0.2
6.2 0.7
6.6 + 0.9
2.6 + 0.7
18.7 0.8
0.0 0.1
CclebA
NMT
3.9 0.2
4.7 0.9
5.0 + 0.7
2.9 +: 0.7
24.3 0.3
0.0 0.0
ARI
0.6
2.9 0.1
2.0 +
12.2 à
27.2
25.7 0.2
52.2
CIFAR10
NMI
31.7 + 0.0
33.5 + 0.1
32.4 0.0
42.8 + 0.1
16.5 0.4
55.3 = 0.1
21.7 0.1
Table 6: Normalized mutual information (NMI) and Adjusted Rank Index (ARI) for all methods and datascts; Average scores and standard errors are computed across three runs
VAE
VAE
SimVAE
SimVAE
SimCLR
SimCLR
60
VicHeg
VicReg
50-
CA
30
30
20
20
10
Weak
Intermediate
High
Weak
Intermediate
High
Image Cropping Strength
Color Jitter Strength
Figure 7: Ablation experiment across the augmentation strength for (left) image cropping and (right) color jitter strength considered during training of the SimVAE and baseline models using the CIFAR dataset.
# of augmentation ablation. Figure reports the downstream classification accuracy for increasing numbers of augmentations considered simultaneously during the training of SimVAE for MNIST and CIFAR10 datasets. On average, larger number of angmentations result in performance increase. Further exploration is needed to understand how larger sets of augmentations can be effectively leveraged potentially by allowing for batch size increase. From Figure 8, we fix our number of augmentations to 10 across datasets.
Likelihood p(x|2) variance ablation. We explore the impact of the likelihood, p(z|z), variance, o², across each pixel dimension on the downstream performance using the MNIST and CIFAR10 datasets. Figure 9 highlights how the predictive performance inversely correlated with the o² on the variance range considered for the CIFAR10 datasot. A similar ablation was performed on all VAE-based models and led to similar conclusion. We therefore fixed o² 0.02 for 3-VAE, CR- TR-VAE and SimVAE across datasets.
Generated Images. Figure 10 report examples of randomly generated images for each digit class and clothing item using the SimVAE trained on MNIST FashionMNISTT, CIFAR10 and CelebA respectively.
25

---

Published in Transactions on Machine Learning Research (09/2024)
MLP KNN
GMM
61-
Arc
Mh.T
(11
05
Number augmentations
Number augmentations
Figure 8: Ablation experiment across the number of augmentations considered during training of the SimVAE model using the MNIST (left) and CIFAR10 (right) datasets. Two, four, six, eight, and 10 augmentations were considered. The average and standard deviation of the downstream classification accuracy using Linear, MLP probes, and KNN & GMM estimators reported across three seeds. Batch size of 128 for all reported methods and number of augmentations. Mcans and standard errors are reported for three runs.
MLF
GMM
MLP
KNN
53
07.5
00.0
00.3
MNIT Are
OFAIOO Arc
90.0
87.3
130
0.066
0.01
5.55
0.2
0.6
0.01
0.32
3.55
0.1
0.4
0.6
2.6
Likel huod plal=) variance
Likelibood plae [) variance
Figure 9: Ablation experiment across the likelihood p(z[2) variance considered during training of the SimVAE model using the MNIST (left) and CIFAR10 (right) datasets. The average and standard deviation of the downstream classification accuracy using Linear, MLP probes and KNN & GMM estimators are reported across three seeds. Means and standard errors are reported for three runs.
26

---

Published in Transactions on Machine Learning Research (09/2024)
1 550
MNIST
Fashion
Cifar10
Celeb-A
Figure 10: Samples generated from SimVAE model using MNIST, FashionMNIST, Cifar10 and CelebA training datasets
Generative Quality. Table 7 reports the FID scores and reconstruction error for all generative baselines and SimVAE.
MNIST
Fashion
CelebA
CIFAR1O
MSE (4)
VAE
0.029 = 0.0
0.012 = 0.0
0.016 = 0.0
0.008 = 0.0
8-VAE
0.029 0.0
0.008 + 0.0
0.005 + 0.0
0.004 0.0
CR-VAF
0.030
0.008
0.005
0.004
SimVAE
0.026 0.0
0.009 = 0.0
0.004 0.0
0.003 =
FID ())
VAE
150.1
99.4
162.9
365.4
8-VAF
155.3
99.9 = 0.7
163.8
376.7 = 1.7
CR-VAE
153.0 = 0.9
98.7 = 0.0
159.3 + 5.4
374.4 + 0.4
SimVAE
152.7 0.3 96.1 1.0
157.8 2.3
349.9 + 2.1
Table 7: Generation quality evaluated by: mean squared reconstruction error (MSE), Fréchet inception distance (FID). Mean and standard errors are reported across three runs.
27


Reviews:
[{"review_id": "5Ubm2bTQpr", "reviewer": "Reviewer_WADv", "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": "# Strenghts\n\n- The paper is very well-narrated, and gives a excellent overview of modern SSL.\n\n- The generative model is very sensible, and it is always compelling to have unifying perspectives on a topic.\n\n- To the best of my understanding, the mathematical derivations seem sound.\n\n- The experiments are convincing, albeit limited to smaller datasets than the ones used in most important SSL papers.\n\n# Weaknesses\n\n- One of the main claims is that most SSL losses are related to the ELBO of the proposed model. I feel like more details could be given to support this. It would help a lot to say explicitly, for each method, what is $z$, what is $y$, what is $J$, and what are the additional/missing terms. It is sometimes said clearly, and sometimes very unclearly (for deep clustering for instance). \n\n- Modern SSL has been mostly influential for very large datasets, like Imagenet or Places (Zhou et al., Learning Deep Features for Scene Recognition using Places Database, NeurIPS 2014), but the authors consider smaller image datasets (MNIST, Fashion MNIST, CelebA, and Cifar10). This is not terrible, because training on such huge datasets is challenging, and requires expensive compute. However, the reasons for this choice should be more clearly explained, and this limitation should be more clearly acknowledged.\n\n\n\n# Minor details \n\n- This is a matter of personal preference, but I think that the notation that the authors use for integrals, namely $\\int_x f(x)$, is not very popular, compared to more standard ones like $\\int  f(x)dx$ or $\\int  f(x) \\mathrm{d} x$.\n\n- typo in the abstract: Abstract: \"self-supevised\".", "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": "The authors propose a new generative model for self-supervised learning (SSL). The main idea is to assume the existence of two latent variables: one, $y$, that models the \"content\" of data points, and one, $z$ that governs the \"style\". For instance, different augmentations of the same image have the same content, but different styles.\n\nThe model can be trained using a VAE-style evidence lower bound (ELBO), that involves an encoder $q(z |x)$ that can be used as a learned representation.\n\nThe authors argues that several SSL objectives resemble their ELBO, generally with fewer terms. They perform a series of experiments on representation learning, and show that their reprentations compare savourably against SSL alternatives and other VAE models.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Clarifying the two points I mentioned in the \"weaknesses\" section.", "broader_impact_concerns": "I do not have particular concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ox7EvwTz2d", "reviewer": "Reviewer_43pU", "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:\n\n1. The paper provides a comprehensive review of related work and uses the ELBO to connect predictive self-supervised learning with generative learning.\n\n2. The paper is well-structured and easy to follow, with clear and straightforward derivations.\n\n3. The empirical results show SimVAE is competitive with or outperforms several existing benchmarks, particularly on simpler datasets and tasks requiring style information.\n\nWeaknesses:\n\n1. The empirical evaluation is limited to relatively small and simple datasets.\n\n2. While the rationale behind the projection head is interesting, it lacks substantial evidence to support the claim that it brings representations closer to the SSL prior.", "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 proposes a latent variable model as a theoretical framework to understand self-supervised learning (SSL) methods, particularly contrastive learning objectives. The model is learned by maximizing the Evidence Lower Bound (ELBO), which the authors use to connect their framework with existing SSL methods. They introduce SimVAE with a mixture of Gaussian prior to bridge the gap between generative and predictive SSL.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. In SimVAE with Gaussian mixture prior, please clarify how the number of clusters is defined. Is a uniform prior assumed for p(y)?\n\n2. While bridging the gap between generative and predictive SSL is valuable, the paper could benefit from more discussion on fundamental insights to improve SSL beyond simply adding a reconstruction term. Adding reconstruction term will increase the model parameters and also perhaps instable training.", "broader_impact_concerns": "I don't see any broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "YeYnRg5mKN", "reviewer": "Reviewer_3vMd", "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. This paper is well-structured and well-written. \n2. Theoretical analysis makes this study more solid.\n3. SimVAE shows better performance than the considered VAE-based methods.\n\n***Weakness***:\n1. While SimVAE shows theoretical promise, it is currently not as competitive as non-VAE-based baselines such as MoCo, and maybe other stronger methods, such as MoCoV2 [a], and MAE [b]. Given the current experimental results, it may be premature to fully assert the potential of VAE-based methods.\n2. The current experiments are still relatively small scale. It would be interesting to know the performance on more complex tasks (a larger number of classes and/or the distribution is harder to model), such as ImageNet. Does the model perform consistently and stably better than other methods?\n3. It is unclear why MoCo is not analysed in Figure 5. The analysis on \"Content\" vs. \"Style\" is interesting. However, more experiments (more datasets) is needed to validate the generalisability of conclusions.\n3. Referring to \"uncertainty estimation\" is confusing, as there is already a research area dedicated to aligning model confidence with model performance.\n\n4. Typos:\n\n   1. In Abstract: supevised $\\rightarrow$ supervised\n\n   2. abd $\\rightarrow$ and\n\n[a] Improved baselines with momentum contrastive learning, arXiv 2020.\n\n[b] Masked Autoencoders Are Scalable Vision Learners, CVPR 2022.", "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 introduces a generative latent variable model for self-supervised learning, demonstrating that various types of discriminative self-supervised learning (SSL) methods generate a comparable distribution over representations. This offers a unifying theoretical framework for these approaches. Additionally, the proposed model elucidates the relationship with mutual information and the application of a \"projection head.\" Experimental results indicate that SimVAE outperforms other VAE-based models.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See weakness.", "broader_impact_concerns": "No societal concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper provides an interesting theoretical insight for SSL. Their theoretical results look solid and well-motivated.

As many reviewers pointed out, one of the weaknesses could be the small-scale experiments. However, as the rebuttal clarified, the goal of this work is to propose a theoretical basis for families of SSL methods to give insight into how they work. It would be great if this paper could include more large-scale experiments, as with many previous SSL works, but I don't think it is necessary to convince the readers.

I think the current form of the paper sufficiently satisfies the TMLR evaluation criteria and recommend "Accept as is".
Decision: Accept (poster)