SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Qi Zhang $^{1*}$  Yifei Wang $^{2*}$  Yisen Wang $^{1,3\dagger}$
$^{1}$  National Key Lab of General Artificial Intelligence,
$^{2}$  School of Mathematical Sciences, Peking University
<sup>3</sup> Institute for Artificial Intelligence, Peking University
Existing contrastive learning methods rely on pairwise sample contrast  $z_{x}^{\top}z_{x'}$  to learn data representations, but the learned features often lack clear interpretability from a human perspective. Theoretically, it lacks feature identifiability and different initialization may lead to totally different features. In this paper, we study a new method named tri-factor contrastive learning (triCL) that involves a 3-factor contrast in the form of  $z_{x}^{\top}Sz_{x'}$ , where  $S = \mathrm{diag}(s_1,\dots,s_k)$  is a learnable diagonal matrix that automatically captures the importance of each feature. We show that by this simple extension, triCL can not only obtain identifiable features that eliminate randomness but also obtain more interpretable features that are ordered according to the importance matrix  $S$ . We show that features with high importance have nice interpretability by capturing common classwise features, and obtain superior performance when evaluated for image retrieval using a few features. The proposed triCL objective is general and can be applied to different contrastive learning methods like SimCLR and CLIP. We believe that it is a better alternative to existing 2-factor contrastive learning by improving its identifiability and interpretability with minimal overhead. Code is available at https://github.com/PKU-ML/Tri-factor-Contrastive-Learning.
As a representative self-supervised paradigm, contrastive learning obtains meaningful representations and achieves state-of-the-art performance in various tasks by maximizing the feature similarity  $z_{x}^{\top}z_{x}^{+}$  between samples augmented from the same images while minimizing the similarity  $z_{x}^{\top}z_{x}^{-}$  between independent samples [4, 20, 15, 5, 18]. Besides the empirical success, recent works also discuss the theoretical properties and the generalization performance of contrastive learning [32, 30, 18].
However, there still exist many properties of contrastive learning that are not guaranteed. In this paper, we focus on a significant one: the feature identifiability. Feature identifiability in the representation learning refers to the property there exists a single, global optimal solution to the learning objective. Consequently, the learned representations can be reproducible regardless of the initialization and the optimizing procedure is useful. As a well-studied topic, identifiability is a desirable property for various tasks, including but not limited to, transfer learning [10], fair classification [26] and causal inference [24]. The previous works propose that contrastive learning obtains the linear feature identifiability while lacking exact feature identifiability, i.e., the optimal solutions obtain a freedom of linear transformations [29]. As a result, the different features are coupled, which hurts the interpretability and performance of learned representations.
†Corresponding Author: Yisen Wang (yisen.wang@pku.edu.cn).
(a) 10 most important dimensions
(b) 10 least important dimensions
Figure 1: The visualization of samples on ImageNet-100 that have the largest values in 10 selected dimensions of representations learned by tri-factor contrastive learning (each row represents a dimension).
In this paper, we propose a new contrastive learning model: tri-factor contrastive learning (triCL), which introduces a 3-factor contrastive loss, i.e., we replace  $z_{x}^{\top}z_{x'}$  with  $z_{x}^{\top}Sz_{x'}$  when calculating the similarity between two features, where  $S$  is a learnable diagonal matrix called importance matrix. We theoretically prove that triCL absorbs the freedom of linear transformations and enables exact feature identifiability. Besides, we observe that triCL shows other satisfying properties. For example, the generalization performance of triCL is theoretically guaranteed. What is more, we find that the diagonal values of the importance matrix  $S$  in triCL indicate the degrees of feature importance. In Figure 1, we visualize the samples that have the largest values in the most and least important dimensions ordered by the importance matrix. We find the samples activated in the more important dimensions are more semantically similar, which verifies the order of feature importance in triCL is quite close to the ground truth. Theoretically, we prove that the dimensions related to the larger values in the importance matrix make more contributions to decreasing the triCL loss. With the automatic discovery of feature importance in triCL, the downstream task conducted on the representations can be accelerated by selecting the important features and throwing the meaningless features.
As triCL is a quite simple and general extension, we apply it to different contrastive learning frameworks, such as SimCLR [4] and CLIP [28]. Empirically, we first verify the identifiability of triCL and further evaluate the performance of triCL on real-world datasets including CIFAR-10, CIFAR-100, and ImageNet-100. In particular, with the automatic discovery of important features, triCL demonstrates significantly better performance on the downstream tasks using a few feature dimensions. We summarize our contributions as follows:
- We propose tri-factor contrastive learning (triCL), the first contrastive learning algorithm which enables exact feature identifiability. Additionally, we extend triCL to different contrastive learning methods, such as spectral contrastive learning, SimCLR and CLIP.
- Besides the feature identifiability, we analyze several theoretical properties of triCL. Specifically, we construct the generalization guarantee for triCL and provide theoretical evidence that triCL can automatically discover the feature importance.
- Empirically, we verify that triCL enables the exact feature identifiability and triCL can discover the feature importance on the synthetic and real-world datasets. Moreover, we investigate whether triCL obtains superior performance in downstream tasks with selected representations.
Self-supervised Learning. Recently, to get rid of the expensive cost of the labeled data, self-supervised learning has risen to be a promising paradigm for learning meaningful representations by designing various pretraining tasks, including context-based tasks [14], contrastive learning [20] and masked image modeling [21]. Among them, contrastive learning is a popular algorithm that achieves impressive success and largely closes the gap between self-supervised learning and supervised learning [27, 4, 20, 33]. The idea of contrastive learning is quite simple, i.e., pulling the semantically similar samples (positive samples) together while pushing the dissimilar samples (negative samples) away in the feature space. To better achieve this objective, recent works propose different variants of contrastive learning, such as introducing different training objectives [36, 18, 35], different structures [4, 6, 16, 39], different sampling processes [11, 25, 38, 7] and different empirical tricks [20, 3].
Theoretical Understandings of Contrastive Learning. Despite the empirical success of contrastive learning, the theoretical understanding of it is still limited. [30] establish the first theoretical guarantee on the downstream classification performance of contrastive learning by connecting the InfoNCE loss and Cross Entropy loss. As there exists some unpractical assumptions in the theoretical analysis in [30], recent works improve the theoretical framework and propose new bounds [1, 2, 34]. Moreover, [18] analyzes the downstream performance of contrastive learning from a graph perspective and constructs the connection between contrastive learning and spectral decomposition. Recently, researchers have taken the inductive bias in contrastive learning into the theoretical framework and shown the influence of different network architectures [17, 31]. Except for the downstream classification performance of contrastive learning, some works focus on other properties of representations learned by contrastive learning. [23] discuss the feature diversity by analyzing the dimensional collapse in contrastive learning. [29] prove that the contrastive models are identifiable up to linear transformations under certain assumptions.
Contrastive Pretraining Process. We begin by introducing the basic notations of contrastive learning. The set of all natural data is denoted as  $\mathcal{D}_u = \{\bar{x}_i\}_{i=1}^{N_u}$  with distribution  $\mathcal{P}_u$ , and each natural data  $\bar{x} \in \mathcal{D}_u$  has a ground-truth label  $y(\bar{x})$ . An augmented sample  $x$  is generated by transforming a natural sample  $\bar{x}$  with the augmentations distributed with  $\mathcal{A}(\cdot|\bar{x})$ . The set of all the augmented samples is denoted as  $\mathcal{D} = \{x_i\}_{i=1}^N$ . We assume both sets of natural samples and augmented samples to be finite but exponentially large to avoid non-essential nuances in the theoretical analysis and it can be easily extended to the case where they are infinite [18]. During the pretraining process, we first draw a natural sample  $\bar{x} \sim \mathcal{P}_u$ , and independently generate two augmented samples  $x \sim \mathcal{A}(\cdot|\bar{x})$ ,  $x^+ \sim \mathcal{A}(\cdot|\bar{x})$  to construct a positive pair  $(x, x^+)$ . For the negative samples, we independently draw another natural sample  $\bar{x}^- \sim \mathcal{P}_u$  and generate  $x^- \sim \mathcal{A}(\cdot|\bar{x}^-)$ . With positive and negative pairs, we learn the encoder  $f: \mathbb{R}^d \to \mathbb{R}^k$  with the contrastive loss. For the ease of our analysis, we take the spectral contrastive loss [18] as an example:
$$
\mathcal {L} _ {\mathrm {S C L}} (f) = - 2 \mathbb {E} _ {x, x ^ {+}} f (x) ^ {\top} f \left(x ^ {+}\right) + \mathbb {E} _ {x} \mathbb {E} _ {x ^ {-}} \left(f (x) ^ {\top} f \left(x ^ {-}\right)\right) ^ {2}. \tag {1}
$$
We denote  $z_{x} = f(x)$  as the features encoded by the encoder. By optimizing the spectral loss, the features of positive pairs  $(z_{x}^{\top}z_{x^{+}})$  are pulled together while the negative pairs  $(z_{x}^{\top}z_{x^{-}})$  are pushed apart.
Augmentation Graph. A useful theoretical framework to describe the properties of contrastive learning is to model the learning process from the augmentation graph perspective [18]. The augmentation graph is defined over the set of augmented samples  $\mathcal{D}$ , with its adjacent matrix denoted by  $A$ . In the augmentation graph, each node corresponds to an augmented sample, and the weight of the edge connecting two nodes  $x$  and  $x^{+}$  is equal to the probability that they are selected as a positive pair, i.e.,  $A_{xx^{+}} = \mathbb{E}_{\bar{x} \sim \mathcal{P}_u}[\mathcal{A}(x|\bar{x})\mathcal{A}(x^{+}|\bar{x})]$ . And we denote  $\bar{A}$  as the normalized adjacent matrix of the augmentation graph, i.e.,  $\bar{A} = D^{-1/2}AD^{-1/2}$ , where  $D$  is a diagonal matrix and  $D_{xx} = \sum_{x' \in \mathcal{D}}\bar{A}_{xx'}$ . To analyze the properties of  $\bar{A}$ , we denote  $\bar{A} = U\Sigma V^{\top}$  is the singular value decomposition (SVD) of the normalized adjacent matrix  $\bar{A}$ , where  $U \in \mathbb{R}^{N \times N}$ ,  $V \in \mathbb{R}^{N \times N}$  are unitary matrices, and  $\Sigma = \mathrm{diag}(\sigma_1, \ldots, \sigma_N)$  contains descending singular values  $\sigma_1 \geq \ldots \sigma_N \geq 0$ .
4 Exact Feature Identifiability with Tri-factor Contrastive Learning
In this section, we propose a new representation learning paradigm called tri-factor contrastive learning (triCL) that enables exact feature identifiability in contrastive learning. In Section 4.1, we prove that contrastive learning obtains linear identifiability, i.e., the freedom in the optimal solutions are linear transformations. In Section 4.2, we introduce the learning process of triCL and theoretically verify it enables the exact feature identifiability.
4.1 Feature Identifiability of Contrastive Learning
When using a pretrained encoder for downstream tasks, it is useful if the learned features are reproducible, in the sense that when the neural network learns the representation function on the same data distribution multiple times, the resulting features should be approximately the same. For example, reproducibility can enhance the interpretability and the robustness of learned representations [22]. One rigorous way to ensure reproducibility is to select a model whose representation function is identifiable in function space [29]. To explore the feature identifiability of contrastive learning, we first characterize its general solution.
Lemma 4.1 ([18]). Let  $\bar{A} = U\Sigma V^{\top}$  be the SVD decomposition of the normalized adjacent matrix  $\bar{A}$ . Assume the neural networks are expressive enough for any features. The spectral contrastive loss (Eq. 1) attains its optimum when  $\forall x\in \mathcal{D}$ ,
$$
f ^ {*} (x) = \frac {1}{\sqrt {D _ {x x}}} \left(U _ {x} ^ {k} \operatorname {d i a g} \left(\sigma_ {1}, \dots , \sigma_ {k}\right) R\right) ^ {\top}, \tag {2}
$$
where  $U_x$  takes the  $x$ -th row of  $U$ ,  $U^k$  denotes the submatrices containing the first  $k$  columns of  $U$ , and  $R \in \mathbb{R}^{k \times k}$  is an arbitrary unitary matrix.
From Lemma 4.1, we know that the optimal representations are not unique, due to the freedom of affine transformations. This is also regarded as a relaxed notion of feature identifiability, named linear feature identifiability  $\stackrel{L}{\sim}$  defined below [29].
Definition 4.2 (Linear feature identifiability). Let  $\stackrel{L}{\sim}$  be a pairwise relation in the encoder function space  $\mathcal{F} = \{f : \mathcal{X} \to \mathbb{R}^k\}$  defined as:
$$
f ^ {\prime} \stackrel {L} {\sim} f ^ {*} \Longleftrightarrow f ^ {\prime} (x) = A f ^ {\star} (x), \forall x \in \mathcal {X}, \tag {3}
$$
where  $A$  is an invertible  $k\times k$  matrix.
It is apparent that the optimal encoder  $f$  (Eq. 2) obtained from the spectral contrastive loss (Eq. 1) is linearly identifiable. Nevertheless, there are still some ambiguities w.r.t. linear transformations in the model. Although the freedom of linear transformations can be absorbed on the linear probing task [18], these representations may show varied results in many downstream tasks e.g., in the k-NN evaluation process. So we wonder whether we could achieve the exact feature identifiability. We first further define two kinds of more accurate feature identifiabilities below.
Definition 4.3 (Sign feature identifiability). Let  $\stackrel{S}{\sim}$  be a pairwise relation in the encoder function space  $\mathcal{F} = \{f : \mathcal{X} \to \mathbb{R}^k\}$  defined as:
$$
f ^ {\prime} \stackrel {S} {\sim} f ^ {*} \Longleftrightarrow f _ {j} ^ {\prime} (x) = \pm f _ {j} ^ {\star} (x), \forall x \in \mathcal {X}, j \in [ k ]. \tag {4}
$$
where  $f_{j}(x)$  is the  $j$ -th dimension of  $f(x)$ .
Definition 4.4 (exact feature identifiability). Let  $\sim$  be a pairwise relation in the encoder function space  $\mathcal{F} = \{f:\mathcal{X}\to \mathbb{R}^k\}$  defined as:
$$
f ^ {\prime} \sim f ^ {*} \Longleftrightarrow f ^ {\prime} (x) = f ^ {\star} (x), \forall x \in \mathcal {X}. \tag {5}
$$
4.2 Tri-factor Contrastive Learning with Exact Feature Identifiability
Motivated by the trifactorization technique in matrix decomposition problems [9] and the equivalence between the spectral contrastive loss and the matrix decomposition objective [18], we consider adding a learnable diagonal matrix when calculating the feature similarity in the contrastive loss to absorb
the freedom of linear transformations. To be specific, we introduce a contrastive learning model that enables exact feature identifiability, named tri-factor contrastive learning (triCL), which adopts a tri-term contrastive loss:
$$
\mathcal {L} _ {\operatorname {t r i}} (f, S) = - 2 \mathbb {E} _ {x, x ^ {+}} f (x) ^ {\top} S f \left(x ^ {+}\right) + \mathbb {E} _ {x} \mathbb {E} _ {x ^ {-}} \left(f (x) ^ {\top} S f \left(x ^ {-}\right)\right) ^ {2}, \tag {6}
$$
where we name  $S = \mathrm{diag}(s_1, \ldots, s_k)$  as the importance matrix and it is a diagonal matrix with  $k$  non-negative learnable parameters satisfying  $s_1 \geq \dots \geq s_k \geq 0$ . Additionally, the encoder  $f$  is constrained to be decorrelated, i.e.,
$$
\mathbb {E} _ {x} f _ {i} (x) ^ {\top} f _ {j} (x) = \left\{ \begin{array}{l l} 1, & \text {i f} i = j \\ 0, & \text {i f} i \neq j \end{array} , i, j \in [ k ], \right. \tag {7}
$$
for an encoder  $f: \mathbb{R}^d \to \mathbb{R}^k$ . One way to ensure feature decorrelation is the following penalty loss,
$$
\mathcal {L} _ {\mathrm {d e c}} (f) = \left\| \mathbb {E} _ {x} f (x) f (x) ^ {\top} - I \right\| ^ {2}, \tag {8}
$$
leading to a combined triCL objective,
$$
\mathcal {L} _ {\operatorname {t r i C L}} (f, S) = \mathcal {L} _ {\operatorname {t r i}} (f, S) + \mathcal {L} _ {\mathrm {d e c}} (f). \tag {9}
$$
Similar feature decorrelation objectives have been proposed in non-contrastive visual learning methods with slightly different forms, e.g., Barlow Twins [36].
Since triCL automatically learns feature importance  $S$  during training, it admits a straightforward feature selection approach. Specifically, if we need to select  $m$  out of  $k$  feature dimensions for downstream tasks (e.g., in-time image retrieval), we can sort the feature dimensions according to their importance  $s_i$ 's (after training), and simply use the top  $m$  features as the most important ones. Without loss of generality, we assume  $s_1 \geq \dots \geq s_k$ , and the top  $m$  features are denoted as  $f^{(m)}$ .
Identifiability of TriCL. In the following theorem, we show that by incorporating the diagonal importance matrix  $S$  that regularizes features along each dimension, triCL can resolve the linear ambiguity of contrastive learning and become sign-identifiable.
Theorem 4.5. Assume the normalized adjacent matrix  $\bar{A}$  has distinct largest  $k$  singular values  $(\forall i,j\in [k],\sigma_i\neq \sigma_j$  when  $i\neq j)$  and the neural networks are expressive enough, the tri-factor contrastive learning (triCL, Eq. 9) attains its optimum when  $\forall x\in \mathcal{D},j\in [k]$
$$
f _ {j} ^ {\star} (x) = \pm \frac {1}{\sqrt {D _ {x x}}} \left(U _ {x} ^ {k}\right) _ {j}, S ^ {*} = \operatorname {d i a g} \left(\sigma_ {1}, \dots , \sigma_ {k}\right), \tag {10}
$$
which states that the tri-factor contrastive learning enables the sign feature identifiability.
As shown in Theorem 4.5, the only difference remaining in the solutions (Eq. 10) is the sign. To remove this ambiguity, for dimension  $j$ , we randomly select a natural sample  $\bar{x} \sim \mathcal{P}_u$ , encode it with the optimal solution  $f^{\star}$  of triCL, and observe the sign of its  $j$ -th dimension  $f_{j}^{\star}(\bar{x})$ . If  $f_{j}^{\star}(\bar{x}) = 0$ , we draw another sample and repeat the process until we obtain a non-zero feature  $f_{j}^{\star}(\bar{x})$ . We then store the sample as an original point  $x_{0j}$  and adjust the sign of different learned representations as follows:
$$
\bar {f} _ {j} ^ {\star} (x) = (- 1) ^ {\mathbb {1} \left(f _ {j} ^ {\star} \left(x _ {0 j}\right) > 0\right)} \cdot f _ {j} ^ {\star} (x), \forall x \in \mathcal {D}, j \in [ k ] \tag {11}
$$
By removing the freedom of sign, the solution becomes unique and triCL enables the exact feature identifiability:
Corollary 4.6. Set  $\bar{f}^*$  as the final learned encoder of tri-factor contrastive learning, and then triCL obtains the exact feature identifiability.
5 Theoretical Properties of Tri-factor Contrastive Learning
Besides the feature identifiability, we analyze other theoretical properties of triCL in this section. Specifically, in Section 5.1, we provide the generalization guarantee of triCL and we present another advantage of triCL: triCL can automatically discover the importance of different features. In Section 5.2, we extend triCL to other contrastive learning frameworks.
In practice, we enforce the non-negative conditions by applying the softplus activation functions on the diagonal values of  $S$ . We only enforce the monotonicity at the end of training by simply sorting different rows of  $S$  and different dimensions of  $f(x)$  by the descending order of corresponding diagonal values in  $S$ .
5.1 Downstream Generalization of Tri-factor Contrastive Learning
In the last section, we propose a new contrastive model (triCL) that enables exact feature identifiability. In the next step, we aim to theoretically discuss the downstream performance of the identifiable representations learned by triCL. For the ease of our theoretical analysis, we first focus on a common downstream task: Linear Probing. Specifically, we denote the linear probing error as the classification error of the optimal linear classifier on the pretrained representations, i.e.,  $\mathcal{E}(f) = \min_g\mathbb{E}_{\bar{x}\sim \mathcal{P}_u}\mathbb{1}[g(f(\bar{x}))\neq y(\bar{x})]$ . By analyzing the optimal solutions of triCL when using the full or partial features, we obtain the following theorem characterizing their downstream performance.
Theorem 5.1. We denote  $\alpha$  as the probability that the natural samples and augmented views have different labels, i.e.,  $\alpha = \mathbb{E}_{\bar{x}\sim \mathcal{P}_u}\mathbb{E}_{x\sim \mathcal{A}(\cdot |\bar{x})}\mathbb{1}[y(\bar{x})\neq y(x)]$ . Let  $f_{SCL}^{\star}$  and  $(f_{triCL}^{\star},S^{\star})$  be the optimal solutions of  $SCL$  and triCL, respectively. Then,  $SCL$  and triCL have the same downstream classification error when using all features
$$
\mathcal {E} \left(f _ {S C L} ^ {\star}\right) = \mathcal {E} \left(\left(f _ {t r i C L} ^ {\star}\right)\right) \leq c _ {1} \sum_ {i = k + 1} ^ {N} \sigma_ {i} ^ {2} + c _ {2} \cdot \alpha , \tag {12}
$$
where  $\sigma_{i}$  is the  $i$ -th largest eigenvalue of  $\bar{A}$ , and  $c_{1}, c_{2}$  are constants. When using only  $m \leq k$  features of the optimal features, triCL with the top  $m$  features admits the following error bound
$$
\mathcal {E} \left(f _ {t r i C L} ^ {\star (m)}\right) \leq c _ {1} \sum_ {i = m + 1} ^ {N} \sigma_ {i} ^ {2} + c _ {2} \cdot \alpha := U \left(f _ {t r i C L} ^ {\star (m)}\right). \tag {13}
$$
Instead, without importance information, we can only randomly select  $m$  SCL features (denoted as  $f_{SCL}^{(m)}$ ), which has the following error bound (taking expectation over all random choices)
$$
\mathcal {E} \left(f _ {S C L} ^ {\star (m)}\right) \leq c _ {1} \left(\left(1 - \frac {m}{k}\right) \sum_ {i = 1} ^ {k} \sigma_ {i} ^ {2} + \sum_ {i = k + 1} ^ {N} \sigma_ {i} ^ {2}\right) + c _ {2} \cdot \alpha := U \left(f _ {S C L} ^ {\star (m)}\right). \tag {14}
$$
Comparing the two upper bounds, we can easily conclude that
$$
U \left(f _ {S C L} ^ {\star (m)}\right) - U \left(f _ {t r i C L} ^ {\star (m)}\right) \geq \frac {m (k - m)}{k} \left(\frac {1}{m} \sum_ {i = 1} ^ {m} \sigma_ {i} ^ {2} - \frac {1}{k - m} \sum_ {i = m + 1} ^ {k} \sigma_ {i} ^ {2}\right) \geq 0. \tag {15}
$$
Thus, triCL admits a smaller error when using a subset of features for downstream classification.
Theorem 5.1 shows that triCL is particularly helpful for downstream tasks when we select a subset of features according to the learned feature importance. When using all features, the two methods converge to the same downstream error bound, which also aligns with our observations in practice.
5.2 Extensions to Other Contrastive Learning Frameworks
In the above sections, we propose tri-factor contrastive learning (triCL) which enables the exact feature identifiability and obtains an ordered representation while achieving the guaranteed downstream performance. In the next step, we extend the triCL to a unified contrastive learning paradigm that can be applied in different contrastive learning frameworks.