Shuffle & Divide:
Contrastive Learning for Long Text

Self-supervised learning is to learn representation with unlabeled large-scale datasets. Traditionally, most deep learning models have relied on supervised learning. But they have required human supervision. SSL techniques includes MLM(Masked Language Model) [2], NSP(Next Sentence Prediction) [2], and contrastive learning [4]. They uses tasks which can be defined with unlabeled data, i.e, pretext task. A model pre-trained with pretext task is transferred to domain-specific situations. It is called downstream task.

Pre-training methods for SSL includes context prediction [10], image colorization [11], Jigsaw puzzle solving [12], and rotation prediction [13]. As contrastive learning based on InfoMax principle [14], SSL for image have advanced, with the works of CPC [15], DIM [16], AM-DIM [17], MoCo [18], and SimCLR [4]. Pre-trained models for text SSL includes BERT [2] and GPT-3 [19]

Data augmentation aims to generate new data from the original, conserving their labels. For text, it includes text modification and generation. Examples of modification are random noise injection [21], randomly removing, inserting, or replacing words, and TF-IDF based word replacement, modifying less important words which are not in principal component of TF-IDF matrix. Examples of text translation are Back translation [22], translating back and forth, and generative methods [23], using sentence-generating models.

Embedding for documents should be extracted, to cluster them. Traditional ways of embedding are well-known, such as TF-IDF [24], which counts the occurrence of words per document to give weight, LSI (Latent Semantic Indexing) [25], LDA (Latent Dirichlet Allocation) [26], and glove embedding average. Deep learning approaches have added new options. BAT, G-BAT [8] and SS-SB-MT [9] use neural networks to embed documents, to outperform traditional clustering methods. According to [9], ELMo [27] and SBERT [28] extract embedding of sentences or words, then averaging them to get better embedding.

Figure 1 shows the overall architecture of our model. Transformers are initialized with pre-trained weights from the BERT and RoBERTa models. To obtain the embedding of documents required for clustering, we add a mean pooling layer to the Transformers model and use it as an encoder. Mean pooling layer calculates the average of the Transformer output vectors. As a result, the encoder receives a token sequence of a predefined length as input and outputs a 768-dimensional vector.

We use a contrastive self-supervised learning approach to learn latent representations of documents without labels to ultimately cluster them.

During the training process, within each mini-batch, the distance between embeddings of the positive pair is made close to each other, and the distance between embeddings of the negative pair is made farther from each other. Among all document pairs that can be combined in each mini-batch, all pairs that are not positive pairs become negative pairs. Therefore, the same number of positive pairs as the batch size and (batch size * 2 - 2) negative pairs are created for each mini-batch. Then, the token sequences of all documents in the mini-batch are fed to the encoder to obtain output vectors, and the contrastive loss is calculated based on the previously obtained positive and negative pair configurations. The model parameters are updated through the loss calculated in each mini-batch.

In a general contrastive self-supervised method, positive pairs are created through augmentation. In Computer Vision, there are various augmentation options such as rotation, flipping, and color jittering, and by combining them, numerous label-preserving augmented versions of one sample can be obtained. However, In NLP, simple augmentation tricks have little effect in contrastive learning. In order to make meaningful changes to the original sample, it is necessary to use a word dictionary or a trained model such as WordNet Augmenter, Contextual Augmenter, or back translation. However, the effect is limited because the content and sentence structure of the original document and the augmentation result are very similar.

We propose methods suitable for use in contrastive self-supervised learning to replace or improve these existing augmentation methods. The first method is to completely replace text data augmentation. Based on the TF-IDF feature of documents, the pair with the highest cosine similarity is sampled to create a positive pair. Our main approach, the second method, called Shuffle & Divide, simply shuffles the order of the sentences in the documents and splits the document in half to get a positive pair by splitting one document into two.

The first method to solve the limitations of text augmentation is to completely replace augmentation itself with positive sampling using TF-IDF. Even if augmentation is not used, a contrastive learning method can be applied if a positive pair can be created from the dataset.

Instead of text augmentation, we propose a method to perform contrastive learning by making two similar samples as a positive pair. TF-IDF is a method of assigning weights to each word by using the number of documents in which a specific word appears and the number of times in which a specific word appears in a specific document. It is a traditional method often used to compare similarities between documents.

We will create a TF-IDF vector for all documents in the dataset and define it as set $\{D_{n}\}$. For all $D_{n}$, find $D_{m}(m\neq n)$ with the highest cosine similarity and assume $D_{m}$ and $D_{n}$ as a positive pair. Since the positive pair created in this way has a high probability of belonging to the same category but is completely different from each other, it can solve the problem that the result of text augmentation is very weak.

Positive sampling by TF-IDF completely replaces text augmentation, enabling contrastive self-supervised learning even when augmentation is impossible or insufficient. However, there is a fundamental and potential limitation that positive pairs may actually be of different classes. In fact, when top-1 cosine similarity positive sampling is performed based on the TF-IDF vector of documents in the 20NG dataset, the case where the labels of the positive pair match is about 85%.

We propose the Shuffle & Divide algorithm as an alternative to this problem. Figure 2 is an example of applying Shuffle & Divide to document $d$. The sentence order of $d$ is shuffled and then divided in half to make $d^{\prime}$ and $d^{\prime\prime}$. Since they belong to the same category, they can be used as a positive pair. Shuffle & Divide is a very simple and intuitive method, but it showed a very powerful effect in the experimental process described in Experiments.

Figure 3 shows the Shuffle & Divide process. In the representation learning process, only documents consisting of 4 or more sentences are used as training data. Let $D_{n}$ be one of the documents sampled from each mini-batch. When $D_{n}$ is a document composed of $m$ sentences from $s_{1}$ to $s_{m}$, $d_{n}$ is created by randomly shuffled sentence order in $D_{n}$. By dividing $d_{n}$ in half, documents $d^{\prime}_{n}$ and $d^{\prime\prime}_{n}$ composed of $n/2$ sentences are created, and a positive pair ($d^{\prime}_{n}$, $d^{\prime\prime}_{n}$) can be obtained per sample document $D_{n}$.

In other words, we can get as many positive pairs as the batch size in the mini-batch. Also, since the sentence order of the document is randomly shuffled at every epoch and then divided, the augmentation effect of positive pairs is amplified in the contrastive concept.

We use 20 Newsgroups (20NG) [29], Reuters-21578 (Reuters) [30], BBC and BBCSport datasets for our experiments. To compare with SS-SB-MT [9], which is the current state-of-the-art for unsupervised text classification, 20NG and Reuters are preprocessed according to the method of [31], same as SS-SB-MT. Preprocessing the 20NG dataset is summarized as follows. The header and footer of each document were removed, and the URL and email address were also removed. And documents with fewer than 10 words were removed. For the Reuters dataset, multi-labeled data was removed, and after removing documents with empty bodies, duplicate data were also removed. And, only the top 10 categories of documents were extracted.

In case of Reuters, the class imbalance is more severe than 20NG, so the randomness in performance measurement tended to be large. The BBC and BBCSport datasets have not been widely used in previous deep learning studies due to the small number of data, but only the Shuffle and Divide method has been tested experimentally. Table I summarizes the information of each dataset used in the experiment.

In the process of grouping documents expressed as high-dimensional vectors, the definition of distance metric between documents is very important. It is known that cosine distance is more suitable for clustering of high-dimensional vectors than euclidian distance. We used Spherical k-means, an algorithm that is fast and uses cosine distance as a metric, to cluster the 768-dimensional document vectors output from the encoder.

We compared our method with several traditional methods and recent deep learning methods that achieved state-of-the-art. The performance and settings of the models are well summarized in  [32].

We used Accuracy (ACC), an indicator used in most previous studies, to measure and compare clustering performance. Using the Hungarian algorithm [33], it is possible to efficiently calculate the maximum accuracy value among all true label and cluster mappings. We also measured AMI for comparison with previous studies, and measured and referenced the silhouette score at every epoch to pick the best model without a label. formula 1 shows how we calculated ACC.

With the two methods we proposed, the model learns a representation for each dataset without using any labels at all. First, the encoder’s transformers are initialized with “bert-base-uncased” pre-trained weights. Contrastive loss is calculated from positive and negative pairs sampled in each mini-batch and the parameters of the model are updated. As the contrastive loss, the nt-xent loss introduced in SimCLR is used. For representation learning, the learning rate of 3e-5 was used in the AdamW optimizer. The larger the batch size, the more diverse the negative views for positive pairs. Therefore, 320, the maximum usable size, was used as the batch size. The range of the total training epoch was determined for each dataset and method, and the model with the highest silhouette score within the range was used as the final model.