Towards a Unified Framework of Contrastive Learning for Disentangled Representations

NIPS 2023

check 22 24 27 papers

Introduction

• Contrastive learning has recently emerged as a promising approach for learning data representations that discover and disentangle the explanatory factors of the data. (advantage)

  • Learning to disentangle the explanatory factors of the observed data is valuable for a variety of machine learning (ML) applications.
  • Contrastive methods (approximately) invert the data generating process and thus recover the generative factors.

• We aim to develop a unified framework of statistical priors on the data generating process to improve our understanding of CL for disentangled representations. (motivation and issue)

  • To deal with nonuniform marginal distributions of the latent factors.
  • To deal with situations when these factors are conditionally dependent to some degree.

• Contributions:

  • We extend and unify theoretical guarantees of disentanglement for a family of contrastive losses under relaxed assumptions about the data generating process.
  • The theoretical findings are empirically validated on several benchmark datasets and we quantitatively compare the disentanglement performance of four contrastive losses.
  • We analyze the impact of partially violated assumptions and investigate practical limitations of the proposed framework.

Related Work

Disentangled Representation Learning

• Historical Roots: Blind source separation, factorial codes.
• Nonlinear Data Relationships: Shift in focus to nonlinear relationships.
• Identifiability Challenge: Fundamental difficulties in identifying latent factors in i.i.d. data.
• Mutual Independence: Recovery of factors under mutual independence and specific regularities.
• Dependent Variables and Auxiliary Information: Use of co-observed dependent variables, time-series data, interventions, and augmentations.
• Exponential Family and VAE: Modeling latent factor distributions with conditional independence given an auxiliary variable.
• Extensions to Distributions: Considering Laplace, von Mises–Fisher distributions, and distance-based conditional distributions.
• Shared Latent Factors: Recovery of latents with known shared factors or interventions, and proving identifiability with general causal dependencies.

Contrastive Learning

• It has been observed that contrastive objectives with tighter mutual information (MI) bounds do not necessarily enhance downstream performance.
• Zimmermann et al. showed that InfoNCE approximately inverts the data generating process, leading to the identification of the true latent factors.

Contrastive Learning for Disentangled Representations

Framework

Data Generating Process

$S\subseteq R^n$ denotes the space of latent factors.

$X\subseteq R^m$ denotes the space of observations.

$$p(\tilde{\mathbf{s}}\mid \mathbf{s})=\frac{Q(\tilde{\mathbf{s}})}{Z(\mathbf{s})}{e}^{-d(\mathbf{s},\tilde{\mathbf{s}})}\text{ with }Z(\mathbf{s})={\int }_{\mathcal{S}}Q(\tilde{\mathbf{s}})e^{-d(\mathbf{s},\tilde{\mathbf{s}})}\mathrm{d}\tilde{\mathbf{s}}$$

$Q$ and

$Z$ are scalar-valued functions.

$d$ is a distance function.

Contrastive Learning Approaches

$$\delta (\mathbf{z},\tilde{\mathbf{z}})=\hat{d}(\mathbf{z},\tilde{\mathbf{z}})+\alpha (\mathbf{z})+\tilde{\alpha }(\tilde{\mathbf{z}})$$

• $\hat{d}$ is a fixed expression that describes the interaction between related examples.

$${\mathcal{L}}_{\delta -\mathrm{NCE}}(f,\delta )=\underset{(\mathbf{x},\tilde{\mathbf{x}})\sim p_{\text{pos }}}{\mathbb{E}}-\log [\mathrm{sig}(-\delta (f(\mathbf{x}),f(\tilde{\mathbf{x}})))]\underset{\mathbf{x},{\mathbf{x}}^{-}\sim p}{-}\mathbb{E}\log [1-\mathrm{sig}(-\delta (f(\mathbf{x}),f\left({\mathbf{x}}^{-}\right)))]$$

$${\mathcal{L}}_{\delta \text{ - }\mathrm{INCE}}(f,\delta ;K)=\underset{\begin{array}{c}(\mathbf{x},\tilde{\mathbf{x}})\sim {\text{ posp }}_{\text{pos }}\\ {\left\{{\mathbf{x}}_i^{-}\right\}}_{i=1}^K\sim {\text{ i.d. }}^{\sim }pp\end{array}}{\mathbb{E}}-\log \frac{e^{-\delta (f(\mathbf{x}),f(\tilde{\mathbf{x}}))}}{e^{-\delta (f(\mathbf{x}),f(\tilde{\mathbf{x}}))}+\sum _{i=1}^K{e}^{-\delta (f(\mathbf{x}),f\left({\mathbf{x}}_i^{-}\right))}}$$

• To fit this method into our theoretical framework, we use a slightly modified version in our analysis, where we substitute
  $f(x{)}^{T}f(x˜)$ with
  $e^{-\delta (f(x),f(x˜)}$

$${\mathcal{L}}_{\delta -\mathrm{SCL}}(f,\delta )=\underset{(\mathbf{x},\tilde{\mathbf{x}})\sim p_{\text{pos }}}{\mathbb{E}}-2e^{-\delta (f(\mathbf{x}),f(\tilde{\mathbf{x}}))}+\underset{\mathbf{x},{\mathbf{x}}^{-}\sim p}{\mathbb{E}}{e}^{-2\delta (f(\mathbf{x}),f\left({\mathbf{x}}^{-}\right))}$$

$${\mathcal{L}}_{\delta -\mathrm{NWJ}}(f,\delta )=\underset{(\mathbf{x},\tilde{\mathbf{x}})\sim p_{\text{pos }}}{\mathbb{E}}\delta (f(\mathbf{x}),f(\tilde{\mathbf{x}}))+\underset{\mathbf{x},{\mathbf{x}}^{-}\sim p}{\mathbb{E}}{e}^{-\delta (f(\mathbf{x}),f\left({\mathbf{x}}^{-}\right))}$$

• We here use it as a loss instead of a lower bound on MI.
• To the best of our knowledge, neither the SCL nor the NWJ objective have been employed to learn disentangled representations or for Independent Component Analysis (ICA).

Identifiability

We derive precise conditions on the data generating process under which the relationship between the learned representation and true latent factors can be described by a simple function.

If they only match up to other simple transformations (e.g., invertible element-wise) of the true or learned latents, the recovered relation can be extended by just those transformations.

Theorem 1 (Weak identifiability). Let

$\mathcal{S}\subseteq {\mathbb{R}}^n$ be open and connected,

$\mathcal{X}\subseteq {\mathbb{R}}^m$, and

$g:\mathcal{S}\to \mathcal{X}$ invertible and differentiable. Let us further assume that the observed data satisfy the generative model given in Eq. (1). If

$d=\hat{d}$ has one of the following properties:

• (i) there exists a function
  $\xi :{\mathbb{R}}^{+}\to {\mathbb{R}}^{+}$, such that
  $\xi \circ d$ is a norm-induced metric
• (ii)
  $d(\mathbf{s},\tilde{\mathbf{s}})=\sum _i{d}_i\left(|s_i-{\tilde{s}}_i|\right)$, where each
  $d_i$ is continuous and strictly increasing

then the optimal estimator of any of the contrastive losses presented above identifies the true latent factors up to affine transformations, i.e.,

$h=f◦g$ is an affine mapping.

• Learned representation can be approximated by a simple affine mapping (linear mapping with translation).
• Allows recovery of true latent factors up to affine transformations.
• Does not involve complex permutations or non-linear operations.

Theorem 2 (Strong identifiability). Assume that all conditions in Theorem 1 are satisfied. Let the function d from

$Eq$. (1) be defined by

$$d(\mathbf{s},\tilde{\mathbf{s}})=\sum _i{\left(|s_i-{\tilde{s}}_i|/{\sigma }_i\right)}^{\beta }$$
with

$\beta \in (0,2)\cup (2,\mathrm{\infty })$ and

${\sigma }_i>0$ for all

$i$, then

$h=f\circ g$ is a generalized permutation matrix, i.e., a composition of a permutation and element-wise scaling and sign flips.

• Requires learned representation to be a generalized permutation matrix (composition of permutation, element-wise scaling, and sign flips).
• Guarantees preservation of relationships between latent factors in a structured manner.
• Involves more complex, non-linear transformations.

Summary:

• Weak Identifiability: Simple transformations (affine mappings).
• Strong Identifiability: Complex, structured transformations (generalized permutations).

Experiments

$$p(\tilde{\mathbf{s}}\mid \mathbf{s})=\frac{Q(\tilde{\mathbf{s}})}{Z(\mathbf{s})}{e}^{-\sum _i{\left(\frac{|s_i-{\tilde{s}}_i|}{\sigma }\right)}^{\beta }}$$

• The results in Table 4 likely aim to assess the identifiability of latent factors under different conditions, such as weak and strong identifiability scenarios.
• Identifiability refers to the ability of the model to accurately distinguish and represent the underlying factors of the data in a disentangled manner.

• Difficulty with Increasing Latent Factors
• Impact of Conditional Distribution Concentration

Conclusion

Our framework accounts for nonuniform marginal distributions of the factors of variation, allows for nonconvex latent spaces, and does not assume that these factors are statistically independent or conditionally independent given an auxiliary variable.

This study provides further evidence that contrastive methods learn to approximately invert the underlying generative process, which may explain their remarkable success in many applications.

Appendix

Nonconvex Latent Spaces:
In the context of disentangled representation learning, the term "nonconvex latent spaces" likely refers to the latent space where the disentangled factors of variation are represented. This means that the manifold or space in which these factors reside is not convex.