Paper: SimVAE: Narrowing the gap between Discriminative & Generative Self-Supervised Representation Learning (ICLR-2024)

Reviewer comment (Novelty):
The paper omits important related work [1-6]. At a minimum, a discussion about the similarities and differences should be included (**Quality**).

Author response:
* **"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](https://nips.cc/virtual/2023/80847#:~:text=We%20present%20a%20novel%20objective,to%20permutations%20of%20cluster%20assignments.)), 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](https://iclr.cc/Conferences/2024/ReviewerGuide): "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.

Reviewer comment (Novelty):
Some of the main claims of the paper are not well-supported, especially the ones about the unification between SSL and generative approaches. Please refer to the general analysis in [3] and [4]. The novelty and theoretical contribution is somewhat limited, perhaps lying in specializing the existing framework (GEDI in [3] and [4]) to the VAE setting (**Novelty**).

Author response:
* **"Claims not well supported & novelty of theoretical contribution vs [3]"**:
    - Please see the above comparison between SimVAE and [3], showing that the SimVAE model and training objective differs materially to [3] and, we believe, offers novel theoretical insight.

Reviewer comment (Novelty):
While the experimental analysis provides evidence on the benefits of the proposed unification, the conclusions drawn from the experiments are rather limited confirming what has been already observed partly in [8] and in [3,4] (**Significance/Novelty**). Perhaps, the authors should focus on deepening the analysis on the existing gap observed on natural images (CIFAR-10), in order to improve in terms of significance and novelty.

Author response:
* **"Conclusions from experiments limited vs [3,8] ... Missing comparison with [3]."**
    - [8] relates specifically to the separation of *style* and *content* in SSL and runs experiments on  synthetic data, whereas we propose a hierarchical generative latent variable model for SSL and run experiments on real data, so do not see a direct comparison between these experimental results. For example, Fig 1 in [8] does not share the same hierarchical structure as our model (Fig 2).
    - [3] please see the comparison to SimVAE above, which explains why experimental results are not fairly comparable. Also, SimVAE learns parameters of a latent variable model under variational inference (using the ELBO) and compares to similar VAE-based models to show the benefit of our assumptions, whereas [3] is a significantly different EBM-based method.
* **"Reasons for the performance gap between VAE-like and SSL models"**:
    - From our analysis, we conjecture that this may be due, as least in part, to the last pooling layer in standard ResNets. This layer sums the tensor representation over spatial dimensions. 
        - For discriminative SSL this seems beneficial for downstream class prediction as representations of an object in different spatial locations are made similar, in line with the SSL objective. 
        - For generative SSL, the non-invertability of this step loses information required for reconstruction, but removing this step causes representations of the same object to differ, harming the SSL objective of clustering them together.
    - Thus the pooling step in standard ResNets is useful for discriminative SSL but hampers generative SSL and we believe this suggests that a different choice of encoder architecture may be required to narrow the performance gap in future work.
* **"*content* and *style* in natural images"**:
    - we use these terms loosely to provide intuition and in keeping with works that consider the separation of content and style in images [e.g. 8 and $\beta$-VAE], where *content* typically refers to "object class" and *style* refers to attributes such as lighting or pose. (This distinction is subjective, depending on which latent attributes are important to a task.)

Reviewer comment (Experiments):
The experimental analysis is missing a comparison with other existing generative and discriminative models [3] and [4] (**Quality**).

Author response:
* **"Conclusions from experiments limited vs [3,8] ... Missing comparison with [3]."**
    - [8] relates specifically to the separation of *style* and *content* in SSL and runs experiments on  synthetic data, whereas we propose a hierarchical generative latent variable model for SSL and run experiments on real data, so do not see a direct comparison between these experimental results. For example, Fig 1 in [8] does not share the same hierarchical structure as our model (Fig 2).
    - [3] please see the comparison to SimVAE above, which explains why experimental results are not fairly comparable. Also, SimVAE learns parameters of a latent variable model under variational inference (using the ELBO) and compares to similar VAE-based models to show the benefit of our assumptions, whereas [3] is a significantly different EBM-based method.

Reviewer comment (Writing):
In Section 4.2, all conditional densities should be explicitly defined.

Author response:
* **"definition of conditionals in §4.2"**:
    - §4.2 describes the general model, which can be used for different types of data for which different distributional choices might be made. For our experiments on image data, conditionals $p(x|z)$ and $p(z|y)$ are modelled by Gaussians. This was made more explicit in section 5.

Reviewer comment (Novelty):
Originality: the authors did not discuss prior VAE SSL work, such as [1, 4].

Author response:
**"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).

Reviewer comment (Writing):
Clarity: coupled with the last point, reading certain parts of the draft can be challenging as some terms are not clearly defined or certain steps are missing. The details are also in the Questions section.

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

Reviewer comment (Writing):
Page 1, “but do not fit its posterior due to their discriminative nature”: It is unclear how to define “fit” and why due to the discriminative nature.

Author response:
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.

Reviewer comment (Writing):
Page 5, “and for any \( z \in \mathcal{Z}, \) a distribution \( p(x|z) \) is implicitly defined by the probabilities of samples mapping to it \(\{ x \in \mathcal{X} | f(x) = z \}\)”: It is unclear what it means to be “implicitly defined by the probabilities of samples mapping to it.” This is a vague statement without mathematical backing.

Author response:
2. **"$p(x|z)$ is implicity defined by ..."** we have reworded this for clarity.

Reviewer comment (Theory):
Page 5, “Hence geometric properties of representations implicitly correspond to latent probabilistic assumptions.”: The authors could have shown theoretically rigorous proof to show this. And it is unclear what specific geometric properties the latent probabilistic assumption induced.

Author response:
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)$.

Reviewer comment (Theory):
Page 5, “and insight into the information captured by representations from their regenerations”: The reviewer is not sure this claim is valid without further explanation; why do generative models have better insights into the information captured by the representation? It is better to define the “information” here, as in SSL literature, there are numerous works studying the information the representation captures (some of which the authors rightfully cited) [5-9].

Author response:
5. **"insight into the information captured by representations"**: we mean this less quantitatively than suggested and have amended the wording to make this clear.

Reviewer comment (Theory):
Page 5, “Note that if the variance of each \( p(z|y) \) is low relative that of \( p(z) \), this fits with the notion that contrastive methods ‘pull representations of related samples together and push those of random samples apart.’”: It would be much more valid if the authors showed proof of this. Also, it is not clear how to define rigorously “fits with the notion,” e.g., via asymptotic analysis. And how to quantify “low relative that of \( p(z) \)” is unclear.

Author response:
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.

Reviewer comment (Theory):
Page 5, “Thus, the model (Equation 4) justifies representation learning methods that heuristically perform clustering”: It is unclear how it is justified. Factorizing \( p(x) \) into the form of Equation 4 is a good start, but it did not justify why heuristic clustering methods are working well or necessarily capturing \( p(x) \) well, e.g., through a tight error bound.

Author response:
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.

Reviewer comment (Theory):
Page 6, “samples of each class differ only in style (and classes are mutually exclusive) this collapse leads to style-invariant representations.”: Despite correct intuition, this statement is, in general, very strong; Dosovitskiy et al. did not explicitly claim anything about the style vs. semantic information in the representations, and the authors did not cite any other work supporting this claim nor specify any assumptions.

Author response:
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.

Reviewer comment (Experiments):
Page 6, “In effect, representations are comparable to those learned by softmax, subject to unit length constraint.”: The authors may clarify how to define “comparable.” It may be helpful to at least cite related work directly, or show empirical evidence to show under what tasks the representations are comparable.

Author response:
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.

Reviewer comment (Writing):
Typos: Eq.(5) the support could be simply \( y \) for the last integral, and in the paragraph below Eq.(5) the lower bound should be \(\log p_{\theta}(z) \geq \int_{y} \mathbf{q_{\phi}(y|z)} \log \frac{p_{\theta}(z|y)p(y)}{q_{\phi}(y|z)}\) (the main result in Eq.(5) is correct).

Author response:
13. typos: thank you, these have been fixed.

Reviewer comment (Presentation):
The paper is very hard to read. The primary contribution of the paper is equation (7) defined over J semantically related samples. The rest of the paper is filled with a lot of claims that do not belong to the paper. For instance, section 4.1 has a latent variable approach to instance discrimination. It is neither interesting nor surprising that a latent-variable version of instance discrimination or any other model can be created. Unless it serves some purpose or offers extra insights, it should be removed. Everything except 4.2 needs to be removed from section 4.

Author response:
* **"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.

Reviewer comment (Presentation):
Having an entire section for representation learning is again wasteful. The representation learning section needs to be moved to related work.

Author response:
* **"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.

Reviewer comment (Writing):
What is the purpose of adding equation 9) since J=6 is used during training and J=2 is never used?

Author response:
* **"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.

Reviewer comment (Reproducibility):
Which equation is used during training? Which equation corresponds to Algorithm 1? If it is equation 8), what is q(y|z1, ..., zJ). In fact it is necessary to show how each of the distribution is represented.

Author response:
* **"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.