SimVAE: Narrowing the gap between Discriminative & Generative Self-Supervised Representation Learning

• "Omitting important related work": Thank you for raising these works.
  • [1-4]: for clarity (and benefit of other reviewers), we note that these are highly related to one another by the same authors:
    • [4] is a workshop submission which further investigates a graphical model proposed in [3]. It appeared after the ICLR submission date (due at a NeurIPS 23 workshop), so we omit this per ICLR guidelines.
    • [1, 2] are a workshop paper and abstract based on [3], so we restrict comments to [3].
    • [3] is on Arxiv (not yet cited) and not peer-reviewed, and we were unaware of it. (Please note ICLR Guidelines: "Authors ... may be excused for not knowing about papers not published in peer-reviewed conference proceedings or journals, which includes papers ... on arXiv.") [3] is related to Sim-VAE, we provide a detailed comparison below and cite it in our work (section 2).
  • [5] introduces a VAE with a diffusion prior and a discriminative SSL loss (see Eq 4, 5 & Fig 2). This does not compare closely to our generative model for SSL which unifies existing discriminative SSL under a unique ELBO. [5] does combine SSL with VAEs, we therefore cite it as related work in section 2.
  • [6] trains Energy Based Models with a contrastive component and is less comparable to our generative latent variable approach.
• "Similarity & differences between SimVAE and [3]": SimVAE and [3] both set out to improve self-supervised learning and develop a more principled approach by understanding existing existing SSL methods from a probabilistic perspective. We highlight main differences:
  • Model:
    • [3] proposes a different graphical model for each class of SSL method considered (Fig 1): contrastive (CL), cluster-based (CB) and negative-free (NF). We note that these are not generative models of the data $x$ (e.g. arrows point away from $x$) leading to objective functions distinct from SimVAE's.
    • In SimVAE, we consider the distribution over representations ($z$) induced by several SSL methods and formalise the intuition in a single generative hierarchical latent variable model (LVM): $p_{\theta, \psi, \pi}(x)\!=\!\int_{y,z} p_\theta(x|z)p_\psi(z|y)p_\pi(y)$ (Fig 2). This model is general, and it is assumed that $y$ conditions lower-variance distributions $p(z|y)$, e.g. clusters, in the latent space. We claim that all of the SSL methods considered (and related methods, by extension) implicitly assume this one latent variable model (Fig 2), i.e. representations are learned by assuming that those of semantically related data share the same $y$ and are therefore clustered together, while clusters are kept apart/preventing collapse. Specifically:
      • instance discrimination (ID) treats the data index $i$ as the latent $y$, hence representations of each sample and its augmentations form clusters. CL methods (e.g. SimCLR) are shown to approximate ID, so exploit the same latent assumption (not their relationship to mutual information, as often assumed). NF models, e.g. BYOL, impose similar instance-level latent clustering, and differ algorithmically in how representations are prevented from collapsing.
      • CB methods (e.g. DeepCluster) directly perform clustering in the latent space, using a large (arbitrary) number of clusters indexed by $y$, relying on the inductive bias of the encoder to map representations of semantically related data close together, even at random initialisation (see DeepCluster).
  • Objective:
    • [3] proposes a new SSL objective, GEDI (Eq 10), that maximises the likelihood under an energy-based model (EBM), $p_\psi(x) \propto e^{u^\top enc(x)}$ and minimises several KL divergences derived from CB & NF methods. This amalgamates different aspects of SSL approaches and a model of $p(x)$. GEDI appears to assume (in $p_\psi(x)$ and $L_{DI}$ terms of Eq 10) that the number of classes $c$, for some ground truth "class" is known, which contrastive SSL methods do not assume.
    • In SimVAE, the proposed generative LVM unifies the considered SSL methods under one probabilistic model and the proposed objective is the ELBO for that LVM derived from first principles.
  • Implementation:
    • [3] uses a 2-stage process to first train the EBM using SGLD sampling and then train a discriminative model.
    • SimVAE follows a single optimisation process (e.g. SGD).
    • [3] also introduces a novel augmentation strategy ("DAM") that improves results, as does the assumption that $c$ is known (Figs 3 and 7). These provide GEDI with additional samples and information relative to other SSL methods (inc SimVAE), obfuscating like-for-like comparison.

• "Originality: the authors did not discuss prior VAE SSL work.": Thank you for raising these works. After carefull consideration, we cited [3] and [4] in the updated manuscript, in section 2 while we do not believe [1] is closely related to our work for the following reasons: [1] does not tackle the typical self-supervised learning problem we consider, of learning representations from semantically related data. Instead, they target richer and higher-quality generation using VAEs, by applying a "self-supervised" lossy augmentation to the data (e.g. downscaling or edge detection) that is assumed part of the generative process, allowing that process to be considered in stages (resembling a diffusion-like approach). [3] combines several loss components from SSL and VAEs to arrive at a comparable training objective to SimVAE's but does not propose a hierarchical latent variable model for SSL and does not derive its associate lower bound. [4] is related to our work but differs from SimVAE as it does not propose a generative latent variable model for SSL in the same way, rather their model $p(x|z)p(z)$ is defined in terms of a posterior and the data distribution itself (their Eq 9).

• Clarity & Quality: we have reworded key sections of the paper to improve clarity and we address specific questions below.

1. "but do not fit its posterior due to their discriminative nature": this has been rephrased as it is an unecessary level of detail for the introduction.
2. "$p(x|z)$ is implicity defined by ..." we have reworded this for clarity.
3. "geometric properties of representations implicitly correspond to latent probabilistic assumptions": we have reworded this for clarity and more clearly defined the relationship between an encoder $f$ and the latent distribution it induces $p_f(z)$.
4. "$z$ ... only identifiable up to certain symmetries" this has been noted previously for generative models (c.f., Locatello et al. 2019), in particular VAEs, and we have added citations.
5. "insight into the information captured by representations": we mean this less quantitatively than suggested and have amended the wording to make this clear.
6. "variance of each $p(z|y)$ is low relative that of $p(z)$": at this point in the paper, we aim to be intuitive and motivate the steps we take. Note that we do not prove properties of the latent space learned by SSL methods, but draw intuition from the latent structure they impose, design a latent variable model based on that and aim to justify our choices empirically. We have rephrased "low relative to that of $p(z)$" as "more concentrated" and a concrete example of what we mean is if $p(z|y)$ and $p(z)$ are both Gaussian and the variance of former is lower than the latter.
7. "the model (Equation 4) justifies representation learning methods that heuristically perform clustering": we have rephrased this section to hopefully make the explanation more clear. Maximising Eq 4 learns the model in Eq 3, and is a prinicipled means of training a VAE with mixture model prior. Methods that take a similar but heuristic approach can therefore be interpreted as also (approximately) learning the model in Eq 3.
8. "samples of each class differ only in style (and classes are mutually exclusive) this collapse leads to style-invariant representations": we have reworded for clarity. We reference "Variational Classification" [Dhuliawala et al., 2023,] regarding the latent perspective and "collapse" of representations of each class under softmax. We agree that Dosovitskiy et al. do not reference content/style, we meant to refer to their reference to "transformation invariance" and have made this more clear.
9. "representations of a class $y$ converge to class parameter": we have reviewed this wording and removed this part as it is unnecessary for the point we hope to make.
10. " $z^T z’$ approximates $z^T w_y$, ... what “store $w_y$” means." we have reviewed this wording and removed this part as it is unnecessary for the point we hope to make.
11. "representations are comparable to those learned by softmax": we have re-worded to make this more clear. We aim to highlight common high-level clustering of representations of semantically related samples across various SSL methods.
12. typos: thank you, these have been fixed.

We hope we have answered all your queries and improved the paper. Let us know if we can address any additional questions or concerns and please consider updating your score if appropriate.

• Impact of $p(x|z)$ variance Thank you, we agree with the role this parameter has in governing the amount of information representations learn. We had considered this parameter in preliminary experiments but not fine tuned it for final experiments. We note that, for unsupervised representation learning, one should be cautious about tuning hyperparameters to a particular downstream task. However, following the reviewer's suggestion, we have re-run experiments for several values of $\sigma^2=Var[X|Z]$ ($\sigma^2 = 1.0$ in the paper) and include results in Figure 7 for downstream CIFAR10 classification. This shows a slight performance improvement ($+2\%$) from using $\sigma^2=0.6$. We are testing lower values to see if there is a better value "in general" (i.e. across various datasets) and if so will update the results. (We do not expect this change the overall picture, e.g. to "brige the gap" to discriminative performance.)
• Hierarchical VAE missing references: Thank you for your feedback. We agree SimVAE fits within the wider hierarchical VAE literature and this was an omission. The related work (section 2) has been updated to include prior relevant hierarchical latent variable models ( Valpola, 2015; Ranganath et al., 2016; Rolfe, 2017; He et al., 2018; Sønderby et al., 2016; Edwards & Storkey, 2016) and how they relate to SimVAE. The closest work by Edwards & Storkey (2016), proposes a very similar hierarchical graphical model for modelling datasets. The main differerences to our work are (a) that we present a hierarchical latent variable to unify existing SSL methods (purpose); and (b) how the posterior is factorised, which is crucual for representations $z$ to be inferred for each sample $x$ independently of related $x'$ (note that Edward & Storkey first learn a global statistic $c$ from a dataset via $q(c|D)$ and only then infer a representation of the data $q(z|x,c)$).
• Model specifications: We have clarified the distribution assumptions used in the model in §5. Section A.3 has been reworded to tie the ELBO formulation to Algorithm 1.

We hope we have answered all your queries and improved the paper. Let us know if we can address any additional questions or concerns and please consider updating your score if appropriate.

• "Authors propose variational approaches for performing instance discrimination, deep clustering"/"Equation 7 is the primary contribution of the paper" Thank you for your feedback. As described in Section 1, our contributions are two-fold:

  • [Section 3 & 4.1] we perform an analysis of existing discriminative SSL methods (i.e., contrastive learning, instance discrimination, latent clustering). By taking a latent variable perspective, we show how these methods can be tied back to a common framework.
  • [Section 4.2] we leverage this analysis to formally define the graphical latent variable model (i.e., Figure 2) which unifies the aforementioned methods and for which we perform experimental validation to verify its soundness and how it can help improve downstream prediction over existing generative and SSL methods.

  We do not propose novel variational approaches for instance discrimination or deep clustering, but rather a latent variable model to unify these methods. §3 & §4.1 formalise a relationship between discriminative and generative methods and define the latent variable model to motivate SimVAE.

• "Hard to read": Thank you for this feedback, which we have taken seriously and reworded the manuscript (in particular §1-4) to improve clarity. Additional care was given to the detailed derivation of the proposed objective as well as its connection to the computational steps performed in practice (c.f., section A.3).

• "Claims that do not belong to the paper": As suggested, part of the representation learning section (§3) has been moved to Background & Related Work (§2). Sections 3 & 4.1 have been reworded to make clearer their contribution and connection to the rest of the paper. §3 focuses on the relationship between generative and discriminative approaches to representation learning, makes this more formal than previously and is hopefully more clear. This shows that both methods can induce latent structure, hence we analyse various discriminative SSL methods (§4.1) to show that they induce comparable latent structure, then take a generative approach to achieve the same (§4.2). Without this link, we wouldn't provide a rationale for those existing SSL methods, we would simply be proposing a new approach.

• "Algorithm in the main paper": Agreed, we have reformatted the paper to accommodate this.

• "Purpose of equation 9" This was included for clarity since many SSL methods compare pairs of data ($J=2$). However, we agree this is not necessary and have removed it to the appendix.

• "Equation used during training" Thank you for your feedback. We reworded section 4.2, 5 and A.3 for greater clarity regarding the SimVAE objective, the involved distribution (including $q(y|z)$) and the connection between the ELBO and Algorithm 1, respectively.

  • Equation 7 refers to the ELBO used to train SimVAE in the general case.
  • In the paragraph below Eq. 7 we specify assumptions that we make (which for different data sets may be varied).
  • In particular, we assume following:
    • $q(z|x)$ and $p(x|z)$ are Gaussian, as in in typical VAE.
    • $p(z|y) = \mathcal{N}(z; \mu_{y}, \sigma^2)$ are Gaussian, with (small) fixed common variance $\sigma^2$.
    • that $y$ is continuous and uniformly distributed, allowing $y$ to be integrated out (as described in §4.2).

  These assumptions are reflected in Algorithm 1, which we implement in the experiments.

We hope we have answered all your queries and improved the paper. Let us know if we can address any additional questions or concerns and please consider updating your score if appropriate.