SlowGuess/NAACL_2021_OCR huggingface dataset:
Dejiao Zhang $^{1}$  Feng Nan $^{1}$  Xiaokai Wei $^{1}$  Shang-Wen Li $^{1}$  Henghui Zhu $^{1}$   
Kathleen McKeown $^{1,2}$  Ramesh Nallapati $^{1}$  Andrew O. Arnold $^{1}$  Bing Xiang $^{1}$
$^{1}$ AWS AI  $^{2}$ Columbia University
dejiaoz, nanfen, xiaokaiw, shangwel, henghui
mckeownk, rnallapa, anarnld, bxiang@amazon.com
Unsupervised clustering aims at discovering the semantic categories of data according to some distance measured in the representation space. However, different categories often overlap with each other in the representation space at the beginning of the learning process, which poses a significant challenge for distance-based clustering in achieving good separation between different categories. To this end, we propose Supporting Clustering with Contrastive Learning (SCCL) – a novel framework to leverage contrastive learning to promote better separation. We assess the performance of SCCL on short text clustering and show that SCCL significantly advances the state-of-the-art results on most benchmark datasets with  $3\% - 11\%$  improvement on Accuracy and  $4\% - 15\%$  improvement on Normalized Mutual Information. Furthermore, our quantitative analysis demonstrates the effectiveness of SCCL in leveraging the strengths of both bottom-up instance discrimination and top-down clustering to achieve better intracluster and inter-cluster distances when evaluated with the ground truth cluster labels<sup>1</sup>.
Clustering, one of the most fundamental challenges in unsupervised learning, has been widely studied for decades. Long established clustering methods such as K-means (MacQueen et al., 1967; Lloyd, 1982) and Gaussian Mixture Models (Celeux and Govaert, 1995) rely on distance measured in the data space, which tends to be ineffective for high-dimensional data. On the other hand, deep neural networks are gaining momentum as an effective way to map data to a low dimensional and hopefully better separable representation space.
Many recent research efforts focus on integrating clustering with deep representation learning
Instance-CL
SCCL
Figure 1: TSNE visualization of the embedding space learned on SearchSnippets using Sentence Transformer (Reimers and Gurevych, 2019a) as backbone. Each color indicates a ground truth semantic category.
by optimizing a clustering objective defined in the representation space (Xie et al., 2016; Jiang et al., 2016; Zhang et al., 2017a; Shaham et al., 2018). Despite promising improvements, the clustering performance is still inadequate, especially in the presence of complex data with a large number of clusters. As illustrated in Figure 1, one possible reason is that, even with a deep neural network, data still has significant overlap across categories before clustering starts. Consequently, the clusters learned by optimizing various distance or similarity based clustering objectives suffer from poor purity.
On the other hand, Instance-wise Contrastive Learning (Instance-CL) (Wu et al., 2018; Bachman et al., 2019; He et al., 2020; Chen et al., 2020a,b) has recently achieved remarkable success in self-supervised learning. Instance-CL usually optimizes on an auxiliary set obtained by data augmentation. As the name suggests, a contrastive loss is then adopted to pull together samples augmented from the same instance in the original dataset while pushing apart those from different ones. Essen
<sup>1</sup>We plan to open source our implementation. Please visit https://arxiv.org/abs/2103.12953 for the release updates.
Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 5419-5430 June 6-11, 2021. ©2021 Association for Computational Linguistics
tially, Instance-CL disperses different instances apart while implicitly bringing similar instances together to some extent (see Figure 1). This beneficial property can be leveraged to support clustering by scattering apart the overlapped categories. Then clustering, thereby better separates different clusters while tightening each cluster by explicitly bringing samples in that cluster together.
To this end, we propose Supporting Clustering with Contrastive Learning (SCCL) by jointly optimizing a top-down clustering loss with a bottom-up instance-wise contrastive loss. We assess the performance of SCCL on short text clustering, which has become increasingly important due to the popularity of social media such as Twitter and Instagram. It benefits many real-world applications, including topic discovery (Kim et al., 2013), recommendation (Bouras and Tsogkas, 2017), and visualization (Sebrechts et al., 1999). However, the weak signal caused by noise and sparsity poses a significant challenge for clustering short texts. Although some improvement has been achieved by leveraging shallow neural networks to enrich the representations (Xu et al., 2017; Hadifar et al., 2019), there is still large room for improvement.
We address this challenge with our SCCL model. Our main contributions are the following:
- We propose a novel end-to-end framework for unsupervised clustering, which advances the state-of-the-art results on various short text clustering datasets by a large margin. Furthermore, our model is much simpler than the existing deep neural network based short text clustering approaches that often require multi-stage independent training.
- We provide in-depth analysis and demonstrate how SCCL effectively combines the top-down clustering with the bottom-up instance-wise contrastive learning to achieve better intercluster distance and intra-cluster distance.
- We explore various text augmentation techniques for SCCL, showing that, unlike the image domain (Chen et al., 2020a), using composition of augmentations is not always beneficial in the text domain.
Self-supervised learning Self-supervised learning has recently become prominent in providing effective representations for many downstream tasks.
Early work focuses on solving different artificially designed pretext tasks, such as predicting masked tokens (Devlin et al., 2019), generating future tokens (Radford et al., 2018), or denoising corrupted tokens (Lewis et al., 2019) for textual data, and predicting colorization (Zhang et al., 2016), rotation (Gidaris et al., 2018), or relative patch position (Doersch et al., 2015) for image data. Nevertheless, the resulting representations are tailored to the specific pretext tasks with limited generalization.
Many recent successes are largely driven by instance-wise contrastive learning. Inspired by the pioneering work of Becker and Hinton (1992); Bromley et al. (1994), Instance-CL treats each data instance and its augmentations as an independent class and tries to pull together the representations within each class while pushing apart different classes (Dosovitskiy et al., 2014; Oord et al., 2018; Bachman et al., 2019; He et al., 2020; Chen et al., 2020a,b). Consequently, different instances are well-separated in the learned embedding space with local invariance being preserved for each instance.
Although Instance-CL may implicitly group similar instances together (Wu et al., 2018), it pushes representations apart as long as they are from different original instances, regardless of their semantic similarities. Thereby, the implicit grouping effect of Instance-CL is less stable and more data-dependent, giving rise to worse representations in some cases (Khosla et al., 2020; Li et al., 2020; Purushwalkam and Gupta, 2020).
Short Text Clustering Compared with the general text clustering problem, short text clustering comes with its own challenge due to the weak signal contained in each instance. In this scenario, BoW and TF-IDF often yield very sparse representation vectors that lack expressive ability. To remedy this issue, some early work leverages neural networks to enrich the representations (Xu et al., 2017; Hadifar et al., 2019), where word embeddings (Mikolov et al., 2013b; Arora et al., 2017) are adopted to further enhance the performance.
However, the above approaches divide the learning process into multiple stages, each requiring independent optimization. On the other hand, despite the tremendous successes achieved by contextualized word embeddings (Peters et al., 2018; Devlin et al., 2019; Radford et al., 2018; Reimers and Gurevych, 2019b), they have been left largely unexplored for short text clustering. In this work, we leverage the pretrained transformer as the back
Figure 2: Training framework SCCL. During training, we jointly optimize a clustering loss over the original data instances and an instance-wise contrastive loss over the associated augmented pairs.
bone, which is optimized in an end-to-end fashion. As demonstrated in Section 4, we advance the state-of-the-art results on most benchmark datasets with  $3\% - 11\%$  improvement on Accuracy and  $4\% - 15\%$  improvement on NMI.
We aim at developing a joint model that leverages the beneficial properties of Instance-CL to improve unsupervised clustering. As illustrated in Figure 2, our model consists of three components. A neural network  $\psi (\cdot)$  first maps the input data to the representation space, which is then followed by two different heads  $g(\cdot)$  and  $f(\cdot)$  where the contrastive loss and the clustering loss are applied, respectively. Please refer to Section 4 for details.
Our data consists of both the original and the augmented data. Specifically, for a randomly sampled minibatch  $\mathcal{B} = \{x_i\}_{i=1}^M$ , we randomly generate a pair of augmentations for each data instance in  $\mathcal{B}$ , yielding an augmented batch  $\mathcal{B}^a$  with size  $2M$ , denoted as  $\mathcal{B}^a = \{\tilde{x}_i\}_{i=1}^{2M}$ .
3.1 Instance-wise Contrastive Learning
For each minibatch  $\mathcal{B}$ , the Instance-CL loss is defined on the augmented pairs in  $\mathcal{B}^a$ . Let  $i^1 \in \{1, \ldots, 2M\}$  denote the index of an arbitrary instance in augmented set  $\mathcal{B}^a$ , and let  $i^2 \in \{1, \ldots, 2M\}$  be the index of the other instance in  $\mathcal{B}^a$  augmented from the same instance in the original set  $\mathcal{B}$ . We refer to  $\tilde{x}_{i^1}, \tilde{x}_{i^2} \in \mathcal{B}^a$  as a positive pair, while treating the other  $2M - 2$  examples in  $\mathcal{B}^a$  as negative instances regarding this positive pair. Let  $\tilde{z}_{i^1}$  and  $\tilde{z}_{i^2}$  be the corresponding outputs of the head  $g$ , i.e.,  $\tilde{z}_j = g(\psi(\tilde{x}_j)), j = i^1, i^2$ . Then for  $\tilde{x}_{i^1}$ , we try to separate  $\tilde{x}_{i^2}$  apart from all negative instances in  $\mathcal{B}^a$  by minimizing the following
$$
\ell_ {i ^ {1}} ^ {I} = - \log \frac {\exp \left(\operatorname {s i m} \left(\tilde {z} _ {i ^ {1}} , \tilde {z} _ {i ^ {2}}\right) / \tau\right)}{\sum_ {j = 1} ^ {2 M} \mathbb {1} _ {j \neq i ^ {1}} \cdot \exp \left(\operatorname {s i m} \left(\tilde {z} _ {i ^ {1}} , \tilde {z} _ {j}\right) / \tau\right)}. \tag {1}
$$
Here  $\mathbb{1}_{j\neq i^1}$  is an indicator function and  $\tau$  denotes the temperature parameter which we set as 0.5. Following Chen et al. (2020a), we choose  $\mathrm{sim}(\cdot)$  as the dot product between a pair of normalized outputs, i.e.,  $\mathrm{sim}(\tilde{z}_i,\tilde{z}_j) = \tilde{z}_i^T\tilde{z}_j / \| \tilde{z}_i\| _2\| \tilde{z}_j\| _2$
The Instance-CL loss is then averaged over all instances in  $\mathcal{B}^a$ ,
$$
\mathcal {L} _ {\text {I n s t a n c e - C L}} = \sum_ {i = 1} ^ {2 M} \ell_ {i} ^ {I} / 2 M. \tag {2}
$$
To explore the above contrastive loss in the text domain, we explore three different augmentation strategies in Section 4.3.1, where we find contextual augmenter (Kobayashi, 2018; Ma, 2019) consistently performs better than the other two.
3.2 Clustering
We simultaneously encode the semantic categorical structure into the representations via unsupervised clustering. Unlike Instance-CL, clustering focuses on the high-level semantic concepts and tries to bring together instances from the same semantic category together. Suppose our data consists of  $K$  semantic categories, and each category is characterized by its centroid in the representation space, denoted as  $\mu_{k}, k \in \{1, \dots, K\}$ . Let  $e_{j} = \psi(x_{j})$  denote the representation of instance  $x_{j}$  in the original set  $\mathcal{B}$ . Following Maaten and Hinton (2008), we use the Student's t-distribution to compute the probability of assigning  $x_{j}$  to the  $k^{th}$  cluster,
$$
q _ {j k} = \frac {\left(1 + \| e _ {j} - \mu_ {k} \| _ {2} ^ {2} / \alpha\right) ^ {- \frac {\alpha + 1}{2}}}{\sum_ {k ^ {\prime} = 1} ^ {K} \left(1 + \| e _ {j} - \mu_ {k ^ {\prime}} \| _ {2} ^ {2} / \alpha\right) ^ {- \frac {\alpha + 1}{2}}} \tag {3}
$$
Here  $\alpha$  denotes the degree of freedom of the Student's t-distribution. Without explicit mention, we follow Maaten and Hinton (2008) by setting  $\alpha = 1$  in this paper.
We use a linear layer, i.e., the clustering head in Figure 2, to approximate the centroids of each
cluster, and we iteratively refine it by leveraging an auxiliary distribution proposed by Xie et al. (2016). Specifically, let  $p_{jk}$  denote the auxiliary probability defined as
$$
p _ {j k} = \frac {q _ {j k} ^ {2} / f _ {k}}{\sum_ {k ^ {\prime}} q _ {j k} ^ {2} / f _ {k ^ {\prime}}} . \tag {4}
$$
Here  $f_{k} = \sum_{j=1}^{M} q_{jk}, k = 1, \ldots, K$  can be interpreted as the soft cluster frequencies approximated within a minibatch. This target distribution first sharpens the soft-assignment probability  $q_{jk}$  by raising it to the second power, and then normalizes it by the associated cluster frequency. By doing so, we encourage learning from high confidence cluster assignments and simultaneously combating the bias caused by imbalanced clusters.
We push the cluster assignment probability towards the target distribution by optimizing the KL divergence between them,
$$
\ell_ {j} ^ {C} = \mathrm {K L} [ p _ {j} \| q _ {j} ] = \sum_ {k = 1} ^ {K} p _ {j k} \log \frac {p _ {j k}}{q _ {j k}}. \tag {5}
$$
The clustering objective is then followed as
$$
\mathcal {L} _ {\text {C l u s t e r}} = \sum_ {j = 1} ^ {M} \ell_ {j} ^ {C} / M \tag {6}
$$
This clustering loss is first proposed in Xie et al. (2016) and later adopted by Hadifar et al. (2019) for short text clustering. However, they both require expensive layer-wise pretraining of the neural network, and update the target distribution (Eq (4)) through carefully chosen intervals that often vary across datasets. In contrast, we simplify the learning process to end-to-end training with the target distribution being updated per iteration.
Overall objective In summary, our overall objective is,
$$
\begin{array}{l} \mathcal {L} = \mathcal {L} _ {\text {I n s t a n c e - C L}} + \eta \mathcal {L} _ {\text {C l u s t e r}} \\ = \sum_ {j = 1} ^ {M} \ell_ {j} ^ {C} / M + \eta \sum_ {i = 1} ^ {2 M} \ell_ {i} ^ {I} / 2 M. \tag {7} \\ \end{array}
$$
$\ell_{j}^{C}$  and  $\ell_{i}^{I}$  are defined in Eq (5) and Eq (2), respectively.  $\eta$  balances between the contrastive loss and the clustering loss of SCCL, which we set as 10 in Section 4 for simplicity. Also noted that, the clustering loss is optimized over the original data only. Alternatively, we can also leverage the augmented data to enforce local consistency of the cluster assignments for each instance. We discuss this further in Appendix A.3.
Implementation We implement our model in PyTorch (Paszke et al., 2017) with the Sentence Transformer library (Reimers and Gurevych, 2019a). We choose distilbert-base-nli-stsb-mean-token as the backbone, followed by a linear clustering head  $(f)$  of size  $768 \times K$  with  $K$  indicating the number of clusters. For the contrastive loss, we optimize an MLP  $(g)$  with one hidden layer of size 768, and output vectors of size 128. Figure 2 provides an illustration of our model. The detailed experimental setup is provided in Appendix A.1. We, as in the previous work Xu et al. (2017); Hadifar et al. (2019); Rakib et al. (2020), adopt Accuracy (ACC) and Normalized Mutual Information (NMI) to evaluate different approaches.
Datasets We assess the performance of the proposed SCCL model on eight benchmark datasets for short text clustering. Table 2 provides an overview of the main statistics, and the details of each dataset are as follows.
- **SearchSnippets** is extracted from web search snippets, which contains 12,340 snippets associated with 8 groups Phan et al. (2008).
- StackOverflow is a subset of the challenge data published by Kaggle², where 20,000 question titles associated with 20 different categories are selected by Xu et al. (2017).
- Biomedical is a subset of the PubMed data distributed by BioASQ<sup>3</sup>, where 20,000 paper titles from 20 groups are randomly selected by Xu et al. (2017).
- AgNews is a subset of news titles (Zhang and LeCun, 2015), which contains 4 topics selected by Rakib et al. (2020).
- Tweet consists of 2,472 tweets with 89 categories (Yin and Wang, 2016).
- GoogleNews contains titles and snippets of 11,109 news articles related to 152 events (Yin and Wang, 2016). Following (Rakib et al., 2020), we name the full dataset as GoogleNews-TS, and GoogleNews-T and GoogleNews-S are obtained by extracting the titles and the snippets, respectively.
$^{2}$ https://www.kaggle.com/c/predict-closed-questions-onstackoverflow/download/train.zip
<sup>3</sup>http://participants-area.bioasq.org
AgNews	SearchSnippets	StackOverflow	Biomedical
ACC	NMI	ACC	NMI	ACC	NMI	ACC	NMI
BoW	27.6	2.6	24.3	9.3	18.5	14.0	14.3	9.2
TF-IDF	34.5	11.9	31.5	19.2	58.4	58.7	28.3	23.2
STCC	-	-	77.0	63.2	51.1	49.0	43.6	38.1
Self-Train	-	-	77.1	56.7	59.8	54.8	54.8	47.1
HAC-SD	81.8	54.6	82.7	63.8	64.8	59.5	40.1	33.5
SCCL	88.2	68.2	85.2	71.1	75.5	74.5	46.2	41.5
GoogleNews-TS	GoogleNews-T	GoogleNews-S	Tweet
ACC	NMI	ACC	NMI	ACC	NMI	ACC	NMI
BoW	57.5	81.9	49.8	73.2	49.0	73.5	49.7	73.6
TF-IDF	68.0	88.9	58.9	79.3	61.9	83.0	57.0	80.7
STCC	-	-	-	-	-	-	-	-
Self-Train	-	-	-	-	-	-	-	-
HAC-SD	85.8	88.0	81.8	84.2	80.6	83.5	89.6	85.2
SCCL	89.8	94.9	75.8	88.3	83.1	90.4	78.2	89.2
For each dataset, we use Contextual Augmenter (Kobayashi, 2018; Ma, 2019) to obtain the augmentation set, as it consistently outperforms the other options explored in Section 4.3.1.
Table 1: Clustering results on eight short text datasets. Our results are averaged over five random runs.
Dataset	|V|	Documents	Clusters
ND	Len	NC	L/S
AgNews	21K	8000	23	4	1
StackOverflow	15K	20000	8	20	1
Biomedical	19K	20000	13	20	1
SearchSnippets	31K	12340	18	8	7
GooglenewsTS	20K	11109	28	152	143
GooglenewsS	18K	11109	22	152	143
GooglenewsT	8K	11109	6	152	143
Tweet	5K	2472	8	89	249
Table 2: Dataset statistics.  $|V|$ : the vocabulary size;  $N^D$ : number of short text documents; Len: average number of words in each document;  $N^C$  number of clusters; L/S: the ratio of the size of the largest cluster to that of the smallest cluster.
4.1 Comparison with State-of-the-art
We first demonstrate that our model can achieve state-of-the-art or highly competitive performance on short text clustering. For comparison, we consider the following baselines.
- STCC (Xu et al., 2017) consists of three independent stages. For each dataset, it first pretrains a word embedding on a large in-domain corpus using the Word2Vec method (Mikolov et al., 2013a). A convolutional neural network is then optimized to further enrich the repre
sentations that are fed into K-means for the final stage clustering.
- Self-Train (Hadifar et al., 2019) enhances the pretrained word embeddings in Xu et al. (2017) using SIF (Arora et al., 2017). Following Xie et al. (2016), it adopts an autoencoder obtained by layer-wise pretraining (Van Der Maaten, 2009), which is then further tuned with a clustering objective same as that in Section 3.2. Both Xie et al. (2016) and Hadifar et al. (2019) update the target distribution through carefully chosen intervals that vary across datasets, while we update it per iteration yet still achieve significant improvement.
- HAC-SD (Rakib et al., 2020) $^4$  applies hierarchical agglomerative clustering on top of a sparse pairwise similarity matrix obtained by zeroing-out similarity scores lower than a chosen threshold value.
- BoW & TF-IDF are evaluated by applying K-means on top of the associated features with dimension being 1500.
To demonstrate that our model is robust against the noisy input that often poses a significant chal
They further boost the performance via an iterative classification trained with high-confidence pseudo labels extracted after each round of clustering. Since the iterative classification strategy is orthogonal to the clustering algorithms, we only evaluate against with their proposed clustering algorithm for fair comparison.
Figure 3: Ablation study of SCCL. In SCCL-Seq, we first train the model using Instance-CL, and then optimize the clustering objective. We exclude Biomedical for better visualization, full plot can be found in Appendix A.4.
lenge for short text clustering, we do not apply any pre-processing procedures on any of the eight datasets. In contrast, all baselines except BoW and TF-IDF considered in this paper either preprocessed the Biomedical dataset (Xu et al., 2017; Hadifar et al., 2019) or all eight datasets by removing the stop words, punctuation, and converting the text to lower case (Rakib et al., 2020).
We report the comparison results in Table 1. Our SCCL model outperforms all baselines by a large margin on most datasets. Although we are lagging behind Hadifar et al. (2019) on Biomedical, SCCL still shows great promise considering the fact that Biomedical is much less related to the general domains on which the transformers are pretrained. In contrast, Hadifar et al. (2019) learn the word embeddings on a large in-domain biomedical corpus, followed by a layer-wise pretrained autoencoder to further enrich the representations.
Rakib et al. (2020) also shows better Accuracy on Tweet and GoogleNews-T, for which we hypothesize two reasons. First, both GoogleNews and Tweet have fewer training examples with much more clusters. Thereby, it's challenging for instance-wise contrast learning to manifest its advantages, which often requires a large training dataset. Second, as implied by the clustering performance evaluated on BoW and TF-IDF, clustering GoogleNews and Tweet is less challenging than clustering the other four datasets. Hence, by applying agglomerative clustering on the carefully selected pairwise similarities of the preprocessed data, Rakib et al. (2020) can achieve good performance, especially when the text instances are very short, i.e., Tweet and GoogleNews-T. We also high
light the scalability of our model to large scale data, whereas agglomerative clustering often suffers from high computation complexity. We discuss this further in Appendix A.5.
To better validate our model, we run ablations in this section. For illustration, we name the clustering component described in Section 3.2 as Clustering. Besides Instance-CL and Clustering, we also evaluate SCCL against its sequential version (SCCL-Seq) where we first train the model with Instance-CL, and then optimize it with Clustering.
As shown in Figure 3, Instance-CL also groups semantically similar instances together. However, this grouping effect is implicit and data-dependent. In contrast, SCCL consistently outperforms both Instance-CL and Clustering by a large margin. Furthermore, SCCL also achieves better performance than its sequential version, SCCL-Seq. The result validates the effectiveness and importance of the proposed joint optimization framework in leveraging the strengths of both Instance-CL and Clustering to compliment each other.
4.2.1 SCCL leads to better separated and less dispersed clusters
To further investigate what enables the better performance of SCCL, we track both the intra-cluster distance and the inter-cluster distance evaluated in the representation space throughout the learning process. For a given cluster, the intra-cluster distance is the average distance between the centroid and all samples grouped into that cluster, and the inter-cluster distance is the distance to its closest neighbor cluster. In Figure 4, we report each type