SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jian Meng\* Li Yang; Kyungmin Lee; Jinwoo Shin; Deliang Fan; and Jae-sun Seo\*
*Cornell Tech, USA †University of North Carolina at Charlotte, USA
$^{\ddagger}$ KAIST, South Korea John Hopkins University, USA
{kyungmnlee, jinwoos}@kaist.ac.kr† lyang50@uncc.edu†
dfan10@jhu.edu  $\S$  {jm2787, js3528}@cornell.edu*
Contrastive learning (CL) has been widely investigated with various learning mechanisms and achieves strong capability in learning representations of data in a self-supervised manner using unlabeled data. A common fashion of contrastive learning on this line is employing large-sized encoders to achieve comparable performance as the supervised learning counterpart. Despite the success of the labelless training, current contrastive learning algorithms failed to achieve good performance with lightweight (compact) models, e.g., MobileNet, while the requirements of the heavy encoders impede the energy-efficient computation, especially for resource-constrained AI applications. Motivated by this, we propose a new self-supervised CL scheme, named SACL-XD, consisting of two technical components, Slimmed Asymmetrical Contrastive Learning (SACL) and Cross-Distillation (XD), which collectively enable efficient CL with compact models. While relevant prior works employed a strong pre-trained model as the teacher of unsupervised knowledge distillation to a lightweight encoder, our proposed method trains CL models from scratch and outperforms them even without such an expensive requirement. Compared to the SoTA lightweight CL training (distillation) algorithms, SACL-XD achieves  $1.79\%$  ImageNet-1K accuracy improvement on MobileNet-V3 with  $64\times$  training FLOPs reduction. Code is available at https://github.com/mengjian0502/SACL-XD.
To overcome the labeling bottleneck for supervised training of deep neural networks (DNNs), self-supervised learning has been widely investigated to learn representations without intensive labeling. In particular, contrastive learning (CL) has demonstrated its capability of representation learning in various machine learning domains, e.g., image [7, 34], video [31], language [2], speech [11], and medical imaging [30]. The success of CL is built upon different data augmentations from the original (training) samples, and the representation is learned by maximizing the latent common knowledge between contrastive embeddings [8, 7, 35], which are separately encoded from the augmented images by DNN models [25]. Despite the various contrastive learning techniques, learning the latent knowledge and representations requires wide and deep encoders. In particular, current SoTA CL algorithms [3, 18] have to employ a large-sized encoder (e.g., ResNet-50) to achieve comparable performance as the supervised learning counterpart.
On the other hand, training the lightweight models [23, 22, 29] from scratch is largely under-explored in CL. Almost none of the prior works have reported the CL performance with directly-trained lightweight models. In fact, the prior success of the lightweight models in supervised learning cannot be smoothly transferred to CL, where the performance gap between the lightweight models and large-sized models is largely amplified under recent CL methods [7, 19]. For example, with
supervised learning, MobileNet-V3Large (Mob-V3) [22] can achieve  $74.04\%$  on ImageNet-1K, which is comparable to ResNet-50  $(76.15\%)$ . However, the  $\sim 2\%$  accuracy difference is amplified to  $>30\%$ $(75.2\% \rightarrow 36.3\%)$  in CL, as reported in [16] with MoCo-V2. Training MobileNet-V3 with more recent Barlow Twins [35] improves the accuracy to  $\sim 52\%$ , but this is still unsatisfactory.
To overcome the limited trainability of the lightweight DNN models in CL, contrastive lightweight model learning has been investigated as an unsupervised knowledge distillation task. For example, SEED [16] divides the entire training process into a two-step process of (1) teacher pre-training with CL and (2) unsupervised knowledge distillation from the frozen teacher to the student lightweight encoder (e.g., MobileNet-V3 [22]). In addition to the two-step process of "pretraining-and-tuning", the authors employ different input schemes between two steps, which further elevates the complexity of the entire training process. On the other hand, ReKD [36] introduces the latent relation knowledge into online distillation with a
single-step training process. However, the large-sized encoder is still required as the online teacher. To that end, both ReKD [36] and SEED [16] focus on unsupervised distillation with the large-sized teacher, which actually amplifies the training cost for lightweight CL, as shown in Fig. 1. Meanwhile, the performance of training lightweight models with CL from scratch remains unsatisfactory [27] or requires dedicated input scaling design with low generality [28] on normal CL. This is quite meaningful to explore since training a high-performance compact encoder ( $\leq 5\mathrm{M}$  parameters) [22, 23, 29] by necessitating a large-sized ResNet ( $>20\mathrm{M}$ ) [20] teacher largely degrades the "time-to-deploy" of the model due to the magnified training cost. Given the challenges of CL and the limitations of prior works, we raise the following question as our motivation:
Figure 1: ImageNet-1K accuracy of the proposed method (SACL+XD) vs. state-of-the-art methods. "RN" represents the ResNet teacher of the prior works [16, 36]. Following the settings in [36, 16], the number of FLOPs is computed based on 200 epochs of training.
Is there a contrastive learning algorithm that can train the high-performance lightweight model without using a large-sized teacher?
To answer this question, we propose Slimmed Asymmetrical Contrastive Learning (SACL) together with Cross-Distillation (XD), a novel self-supervised contrastive learning algorithm designed for efficient and high-performance CL with lightweight encoders. Specifically, SACL considers lightweight contrastive learning as an asymmetrical sparse training problem based on the shared weights from a dense host encoder. A lightweight encoder can be treated as a “subset” model (sub-model) sliced from a wide host model, which naturally formulates the asymmetrical encoding in CL. Different from the knowledge distillation-based methods [16, 36], the proposed SACL-XD algorithm completely eliminates the large-sized ResNet teacher from the lightweight contrastive learning. On top of that, we introduce cross-distillation (XD) to facilitate SACL. XD minimizes the decorrelation between the latent information distorted by asymmetrical CL, elevating the training stability and performance. As shown in Fig. 1, the proposed method achieves  $64 \times$  and  $8.3 \times$  training cost reduction along with  $1.79\%$  and  $2.09\%$  ImageNet accuracy improvements, compared to SOTA [16] and [36], respectively. The major contributions of the proposed SACL-XD algorithm are:
1. Simplicity: SACL-XD does not require teacher pre-training (and input scaling) before training a lightweight student encoder.
2. Efficiency: SACL-XD does not require the large-sized teacher during training. The training cost of MobileNet-V3 [22] model is reduced by  $64 \times$  and  $8.3 \times$  compared to [36, 16].
3. Performance and generality: SACL-XD achieves up to  $2.09\%$  ImageNet-1K accuracy improvement compared to the previous SoTA method. In particular, the proposed cross-distillation (XD) algorithm solely achieves superior performance with normal contrastive learning settings on ResNet-50 without SACL.
4. Transferability: The SACL-XD-trained lightweight encoder shows high transferability to the downstream tasks. With minimum fine-tuning on top of the ImageNet-1K-trained
Self-supervised contrastive learning. With the absence of deterministic labels, the major objective of contrastive learning is minimizing the distance between the embeddings separately encoded from the augmented input samples. Early research works [7, 19] define "positive" and "negative" sample pairs, and the learning process maximizes the similarity between positives and repels the negative samples. The popular InfoNCE loss [26, 19, 8] or NT-Xent loss [7] has been proposed as the learning objective. The wide and deep DNN encoder is also the key factor of success in contrastive learning. In addition to the cross-reference between positive and negative pairs, contrastive learning is also considered a DNN-based clustering problem, where the samples and their embeddings are grouped into clusters based on similarity metrics [4, 32, 1]. SwAV [5] introduces online cluster learning with the reduced complexity of computation. The entropy-based similarity matching and clustering CL algorithms share the same nature, but they all require an extensive amount of negative samples to generate salient contrastiveness.
Recent methods consider contrastive learning as a knowledge distillation [21] task between the contrastive encoders with separately augmented inputs. BYOL [18] and SimSiam [9] consider the student network as the online model, while the teacher is consistently staying offline with no gradient propagation [9]. The teacher is updated as the moving averaged weights of the student encoder. The "student-teacher" relationship focuses on the distance minimization between latent information. Another perspective is exploring the content hidden inside the embeddings. Barlow Twins [35] pushes the cross-correlation between the encoded embeddings toward the identity matrix. VICReg [3] collectively optimizes the variance, covariance, and distance between and inside the embeddings to achieve better training stability and even enables the CL with different encoders. Among all the previous CL algorithms, the empirical findings show the improved training stability and performance of CL together with the relaxed requirement of batch sizes, but the large-sized encoder (e.g., ResNet-50  $(1\times ,2\times ,\ldots)$  [20]) is a persistent and almost mandatory requirement to achieve high performance, which hinders the development of CL with high energy consumption and extensive training effort.
Contrastive learning with lightweight models. Under the context of supervised learning, various lightweight architecture designs [23, 22, 29] are proposed to minimize the performance difference caused by the model size gap. The intricate architecture and the compact model sizes enable energy-efficient and hardware-friendly computer vision applications. However, the architectural efficiency of the lightweight models becomes invalid in the contrastive learning domain. Employing a lightweight encoder in CL leads to poor performance due to the insufficient trainability caused by the limited model sizes. To resolve this challenge, the mainstream method is introducing a giant model as the teacher of the lightweight student encoder. SEED [16] uses a pre-trained large-sized ResNet [20] as the frozen teacher of the lightweight student model. The similarity-based teacher distribution is generated as the target soft label for student learning. In particular, accurate distillation requires both teacher and student to encode the same input, which is inconsistent with the augmentation strategy of CL. As a result, self-supervised training of the lightweight model becomes a two-step process. However, the dependence on the large-sized teacher and the expensive pre-training complicates the overall learning process. ReKD [36] incorporates the online distillation during learning with the relation knowledge-based loss between embeddings. DisCo [17] complicates the training cost of the distillation even further. The inputs are encoded separately by the teacher and two separate students. Regardless of using an online or frozen teacher, the large-sized ResNet model is always required to train a lightweight encoder. The recent efforts on directly-trained lightweight model require fine-grained efforts on input scaling [28], but the improvements on lightweight model has low generality and has no improvements on the popular ResNet-50-based CL.
To resolve the training dilemma between contrastive learning and lightweight models, this section introduces 1) SACL: Slimmed Asymmetrical Contrastive Learning and 2) XD: Cross Distillation to achieve lightweight contrastive learning from scratch.
MobileNet-V3, ours achieves  $94.80\%$  accuracy on CIFAR-10, outperforms the supervised learning baseline and SEED [16] with  $1.83\%$  and  $14.80\%$  improvements, respectively.
(c)
Figure 2: Overview of (a) contrastive learning with shared encoder, (b) Proposed Slimmed Contrastive Learning (SACL) algorithm and (c) Cross-Distillation-aided SACL.
Similar to prior contrastive learning (CL) algorithms [35, 3, 7], our method performs contrastive learning based on the dual augmentation that is transformed from the clean input image. Given a clean input batch  $X$  sampled from the dataset  $\mathcal{D}$ , the distorted views are generated from a combination of data augmentations  $\mathcal{T}$ , leading to the augmented pairs  $X_{A}$  and  $X_{B}$ , which are fed into the encoders for different contrastive paths. In this work, we follow the recent contrastive learning algorithms [35, 3] and use one shared encoder  $f_{\theta}$  for different contrastive branches, as shown in Fig. 2(a). The encoded outputs are fed into the subsequent projector  $h_{\phi}$ , resulting in the latent embeddings  $z_{A}$  and  $z_{B}$ .
3.1 Slimmed Asymmetric Contrastive Learning (SACL)
Under the standard supervised training, it has been well evidenced that extracting a subset lightweight model out of a wider/larger (e.g.,  $2 \times$ ) encoder often leads to better performance compared to directly train such a lightweight model from scratch [37]. We are inspired by the fact that a lightweight model can be considered as a subset model "sliced" from the original full-sized host model. The asymmetrical relationship between the slimmed model and the original model naturally fits the independent encoding path of contrastive learning.
We propose Slimmed Asymmetric Contrastive Learning (SACL), for obtaining lightweight models under self-supervised CL training. SACL treats the lightweight model contrastive learning as a sparse training optimization problem based on shared weights from a large-sized encoder. During the forward pass of each iteration, the augmented input pair  $(X^{A}, X^{B})$  are separately encoded by the dense host model  $(f_{\theta})$  and the slimmed  $(f_{\theta}^{s})$  encoders with the disabled weight filters (i.e., output channels), as shown in Fig. 2(b).
The host encoder model  $f_{\theta}$  is slimmed by removing a unified amount of weight filters (output channels) of each layer based on weight magnitude score. With the slimmed and dense model  $f_{\theta}^{s}$  and  $f_{\theta}$ , the relationship between  $\theta^{s}$  and  $\bar{\theta}$  is defined as:
$$
\theta_ {s} \subset \theta \quad \text {a n d} \theta_ {s} = \theta \cdot \mathcal {M} \tag {1}
$$
Where  $\mathcal{M}$  is the weight mask that disables the channels of  $f_{\theta}$ , and the subset model  $f_{\theta}^{s}$  is selected when  $\mathcal{M}$  is enabled. As shown in Fig. 2(b),  $f_{\theta}$  and  $f_{\theta}^{s}$  separately encode the input pair in different branches. With the shared weights, the mask  $\mathcal{M}$  is alternatively enabled between branches to formulate the asymmetrical encoding during the forward propagation. The slimming ratio (sparsity) of  $\mathcal{M}$  is defined by the Slimmed Asymmetry (SA) between  $f_{\theta}$  and  $f_{\theta}^{s}$  with the style of “K×-1×”. “K” is the width of the wide host model ( $f_{\theta}$ ) that is employed to slice the  $1\times$  model ( $f_{\theta}^{s}$ ) out of it:
$$
s = 1 - 1 / \mathrm {K} = \operatorname {c a r d} \left(\theta_ {s}\right) / \operatorname {c a r d} (\theta) \tag {2}
$$
where  $\operatorname{card}(\cdot)$  returns the number of nonzero elements of the tensor,  $s$  is the desired slimmed ratio (sparsity) controlled by K.
Starting from the initialization of the training, SACL uniformly generates the masks to slice out the weight filters (output channels) with the least magnitude score from each layer. The masks will be
updated after each epoch to maintain the performance of the slimmed model based on the filter-wise  $L_{1}$  norm score. The resultant slimmed  $(1\times)$  model will be deployed as the final trained encoder. Specifically, SACL drives contrastive learning with the following properties:
1. SACL holds a consistent channel-wise architecture difference (e.g.,  $1.5 \times$  vs.  $1 \times$ ) between contrastive branches throughout the entire training process.
2. SACL removes a unified amount of channels (with the lowest magnitude score) of all the layers. By doing so, the resultant slimmed model will have the exact width as the target lightweight model.
Formally, given the augmented input pair  $(X_A,X_B)$ , the forward pass is characterized as:
$$
z _ {A} = h _ {\phi} \left(f _ {\theta} \left(X _ {A}\right)\right), \quad z _ {B} = h _ {\phi} \left(f _ {\theta} ^ {s} \left(X _ {B}\right)\right). \tag {3}
$$
And the optimization target of SACL is:
$$
\min  _ {\theta , \theta_ {s}} \mathcal {L} _ {\mathrm {S A C L}} \left(z _ {A}, z _ {B}\right), \quad \text {s u c h t h a t} \quad \theta_ {s} \subset \theta \text {a n d} \frac {\operatorname {c a r d} \left(\theta_ {s}\right)}{\operatorname {c a r d} (\theta)} = s, \tag {4}
$$
Inside each mini-batch, the gradient is collectively computed all at once, and the optimizer will update the whole set parameter  $\theta$  of the dense host encoder  $f_{\theta}$ . The slimmed model architecture  $(f_{\theta}^{s})$  will be updated after every epoch based on the magnitude score of each filter. The model is completely online and stop gradient is excluded from learning. Minimizing the contrastive loss  $\mathcal{L}_{\mathrm{SACL}}$  between  $z_{A}$  and  $z_{B}$  is equivalent to overcoming the distortion caused by 1) data augmentation  $\mathcal{T}$  and 2) consistent and structural architecture difference between  $f_{\theta}^{s}$  and  $f_{\theta}$ . And the architecture asymmetry caused by 2) motivates us to explore the enhancement of SACL from the perspective of knowledge distillation, which is presented in the following section.
3.2 Cross Distillation (XD) on top of SACL
We propose Cross-Distillation (XD) on top of SACL, which treats the teacher-student relationship as an interconnected knowledge distillation with the correlation-based optimization on top of the proposed SACL learning scheme. Given the asymmetrical contrastive encoders  $f_{\theta}$  and  $f_{\theta}^{s}$ , we first encode  $X^{A}$  and  $X^{B}$  based on SACL ( $X^{A} \rightarrow f_{\theta}$ ;  $X^{B} \rightarrow f_{\theta}^{s}$ ), leading to the embeddings  $z^{A}$  and  $z^{B}$ . Subsequently, we freeze both  $f_{\theta}$  and  $f_{\theta}^{s}$  while reversing the order of the inputs for encoding ( $X^{B} \rightarrow [f_{\theta}] \rightarrow [\tilde{z}_{B}]$ ;  $X^{A} \rightarrow [f_{\theta}^{s}] \rightarrow [\tilde{z}_{A}]$ ), and characterize the resultant embeddings as  $[\tilde{z}_{A}]$  and  $[\tilde{z}_{B}]$ , where “[·]” represents the frozen encoder for the forward pass only.
As a result, each forward pass will generate two pairs of latent vectors resulting from two groups of SACL forward pass. We first compute the online contrastive loss  $\mathcal{L}_{\mathrm{SACL}}$  based on  $z_{A}$  and  $z_{B}$ . Such online embeddings contain the distortion caused by both data augmentations and asymmetrical encoders. We empirically find out that directly optimizing such high-sparsity difference via single contrastive loss leads to collapsed training. Motivated by that, we compute the cross-distillation loss  $\mathcal{L}_{\mathrm{CD}}$  as the average loss between the pair of  $(z_{A},\tilde{z}_{A})$  and  $(z_B,\tilde{z}_B)$ :
$$
\mathcal {L} _ {\mathrm {C D}} = \frac {\mathcal {L} _ {\mathrm {C D}} ^ {A} \left(z _ {A} , \tilde {z} _ {A}\right) + \mathcal {L} _ {\mathrm {C D}} ^ {B} \left(z _ {B} , \tilde {z} _ {B}\right)}{2}, \tag {5}
$$
where  $\mathcal{L}_{\mathrm{CD}}^A, \mathcal{L}_{\mathrm{CD}}^B$  will be defined formally later (see Eq. 9). As shown in Fig. 2, each term of the cross-distillation loss is computed between the asymmetrical SACL encoders with the same input. In other words, optimizing the cross-distillation loss is equivalent to minimizing the distortion in embeddings caused by the asymmetrical sparsity only. We define the total loss of the training as the weighted sum between  $\mathcal{L}_{\mathrm{SACL}}$  and  $\mathcal{L}_{\mathrm{CD}}$  and weight is defined as  $\alpha$ :
$$
\mathcal {L} = \alpha \mathcal {L} _ {\mathrm {S A C L}} + (1 - \alpha) \mathcal {L} _ {\mathrm {C D}}, \tag {6}
$$
Where  $n$  represents the index of batch, and  $i$  and  $j$  represent the dimensionality indices across the latent output. Correspondingly,  $C_{i,j}$  represents the  $i,j$  element of the correlation matrix.
Regarding contrastive loss design, recent works treat contrastive learning as the optimization problem of principle correlation maximization [35] or the decorrelation minimization [35, 3] of the encoded latent vectors. The cross-correlation between  $z_{A}$  and  $z_{B}$  is computed as:
$$
C _ {i j} = \frac {\sum_ {n} z _ {n , i} ^ {A} z _ {n , j} ^ {B}}{\sqrt {\sum_ {n} \left(z _ {n , i} ^ {A}\right) ^ {2}} \sqrt {\sum_ {n} \left(z _ {n , j} ^ {B}\right) ^ {2}}} \tag {7}
$$
Table 1: ImageNet-1k test accuracy with linear evaluation protocol based on MobileNet-V3 [22] trained by different contrastive learning/distillation methods.
Method	Linear Eval. Acc. (%)	Training Epochs	Pre-training
XD. (Ours)	57.16	100	X
ReKD [36]	56.70	200	X
SEED [16]	55.20	200	✓
In this work, we first adopt the BT-loss [35] as the starting point of the online contrastive loss  $\mathcal{L}_{\mathrm{SACL}}(z_A,z_B)$ , where the principle correlation (diagonal) is maximized, and the cross-decorrelation (off-diagonal) is minimized during training.
$$
\mathcal {L} _ {\mathrm {S A C L}} \left(z _ {A}, z _ {B}\right) = \sum_ {i} \left(1 - C _ {i i} ^ {A B}\right) + \lambda \sum_ {i} \sum_ {i \neq j} \left(C _ {i j} ^ {A B}\right) ^ {2} \tag {8}
$$
Unlike the prior work that uses the SoftMax-based probability [16] as unsupervised soft labels for distillation, we find out the proposed cross-distillation task between  $z$  and  $\tilde{z}$  can still be considered as a correlation-based optimization problem, and our finding shows the proposed cross-distillation is a strong facilitator for contrastive learning even without SACL. First, by duplicating  $L_{\mathrm{SACL}}$  for  $L_{\mathrm{CD}}$ , each term in Eq. 5 becomes:
$$
\mathcal {L} _ {\mathrm {C D}} ^ {A} = \underbrace {\sum_ {i} \left(1 - C _ {i i} ^ {A \tilde {A}}\right)} _ {\text {i n n e r - c o r r e l a t i o n l o s s}} + \lambda \underbrace {\sum_ {i} \sum_ {i \neq j} \left(C _ {i j} ^ {A \tilde {A}}\right) ^ {2}} _ {\text {i n n e r - d e c o r r e l a t i o n}} \tag {9}
$$
By encoding the same input with the identical dense model, the inner correlation  $C_{ii}^{A\tilde{A}} \rightarrow 1$ , and inner de-correlation  $\sum_{i}(1 - C_{ii}^{A\tilde{A}}) \rightarrow 0$ , which makes the cross-distillation loss  $\mathcal{L}_{\mathrm{CD}}^A$  and  $\mathcal{L}_{\mathrm{CD}}^B$  equivalent to the internal decorrelation loss between different dimensions inside  $z_A$  and  $z_B$ , respectively. As a result, minimizing  $\mathcal{L}_{\mathrm{CD}}$  avoids the aliasing feature across different dimensions. We would like to highlight the fact that the combination in Eq. 10 with both cross-decorrelation  $(L_{\mathrm{SACL}})$  and the inner decorrelation  $(L_{\mathrm{CD}})$  to superior performance compared to [3, 35]. Table 1 s lightweight MobileNet-V3-Large [22] trained on the In only 100 epochs, the proposed Cross-distillation (XD) to [16], which uses ResNet-50  $(1\times)$  as the teacher. XD algorithm with other CL methods on ViT [13] an performance of XD proves the importance of the co-op and the inner decorrelation  $(L_{\mathrm{CD}})$ .