WC-SBERT: Zero-Shot Text Classification via SBERT with Self-Training for Wikipedia Categories

SBERT (Sentence-BERT) (Reimers and Gurevych 2019) is a method based on BERT (Bidirectional Encoder Representations from Transformers) that calculates sentence-level vector representations. It utilizes a siamese architecture consisting of two BERT networks with shared weights, where each network processes an input sentence. These two BERT networks learn representations of the two sentences through parameter sharing, enabling comparison of sentence similarity using methods such as cosine similarity. Furthermore, in the fine-tuning process of SBERT for different downstream tasks, the choice of the loss function is crucial. In this study, we employ the multiple ranking loss (MNR Loss), which is particularly suitable when the training dataset contains only positive pairs. The MNR Loss brings positive data closer together in the vector space and pushes negative data further apart, thus forming clusters of similar sentences. Leveraging the architecture and training/inference optimizations of SBERT compared to traditional BERT, this study utilizes the pre-trained model all-mpnet-base-v2 from SBERT as the primary base pre-trained model.

We introduces a novel training approach for BERT-based models, shifting from the conventional training method to an association training based on model labels (categories). The purpose of the basic model fine-tuning stage is to fine-tune the SBERT-based model using the Wikipedia dataset to construct training data and obtain a generalized base model. Based on the provided URLs in the dataset, we identify the corresponding category for each data entry’s text. Each data entry in the dataset is defined by three components: id $i$, text $t_{i}$, and the set $C_{i}$ of categories to which the text belongs. Assume that $C_{i}=\left\{c_{1},c_{2},...,c_{n_{i}}\right\}$, $n_{i}$ represents the size of the category set for the $i^{th}$ data entry. We use each category in the set $C_{i}$ for pairwise combinations. For each data entry, we can obtain $C^{n_{i}}_{2}$ pairs of classes as training data. The set $P_{i}$ comprises pairs of classes for each data entry. Its definition is shown below. We select all $P_{i}$ where $n_{i}\geq 2$ to construct the initial training set and perform fine-tuning on the pre-trained model.

$P_{i}=\left\{(c_{j},c_{k})|c_{j},c_{k}\in C_{i},c_{j}\neq c_{k}\right\}$

Regarding the dataset, we utilize the Wikipedia dataset from Huggingface and assume that the content of each wiki, along with its associated categories, exhibits positive correlations based on their respective wiki page IDs. As a result, we modify the training approach of most BERT-based models for the given text and adopt associated training using model labels (categories). Since the length of labels is significantly shorter compared to the content of the text, this reduces training costs. We use the pre-trained SBERT model as the base model and treat the categories of each wiki page in wiki-cate as labels. We combine these labels in pairs (non-repetitive combinations using Python’s itertools.combinations) and convert them into the SBERT training data format (InputExample) to create the training dataset train-samples. Pseudo code is shown below:

Through SBERT, we train on the aforementioned train-samples to obtain a general purpose foundational model called wiki-cate SBERT model: WC-SBERT.

After the aforementioned data processing, the size of the dataset $D$ is 6,458,670. We utilize the general base model WC-SBERT to encode the first 200 words of the text in each data point from $D$. The resulting embeddings are stored in an h5 file for subsequent fine-tuning.

In the self-training fine-tune stage, we define the target label set as $L$. Initially, we use WC-SBERT to calculate the similarity between each text and all elements $l$ in the label set $L$. Then, for each text $t_{i}$, we select the label $l_{j}$ from $L$ that has the highest similarity score with $t_{i}$ to obtain the similarity score $similarity(t_{i},l_{j})$ between $t_{i}$ and $l_{j}$, and compare the similarity between them. Next, we compare the similarity with the set threshold. If the similarity is greater than the threshold, we pair all elements ($c_{1},c_{2},...,c_{n_{i}}$) of the class set $C_{i}$ of $t_{i}$ with the target label $l_{j}$ to get $(c_{1},l_{j}),(c_{2},l_{j}),...,(c_{n_{i}},l_{j})$, all text categories-target label pairs are used as training data for the fine-tuning stage. Since we have already encoded the text and stored it in an h5 file in the previous step, during the similarity calculation stage, we only need to calculate the embeddings of the target labels, and the text embeddings can be directly retrieved from the h5 file. With the advantage of SBERT’s fast similarity inference, we can quickly obtain $similarity(t_{i},l_{j})$ for each million data points.

As shown in the Figure 1, we use the method described above to obtain the training data and fine-tune the original model $M_{base}$, resulting in model $M_{1}$. We then continue to use the same method, using $M_{1}$ to calculate similarity scores and filter out the data for further fine-tuning of $M_{base}$, obtaining model $M_{2}$. We conduct this process iteratively, until we obtain the final model $M_{final}$.

This study utilized the train set of Wikipedia dataset from Huggingface as the training set for the pre-trained model. The dataset consists of a total of 6,417,006 records, with fields including wiki page ID, URL, title, and text. Since this dataset does not include the categories for each page, a separate web scraping program was designed in this study to retrieve the categories for each page. A total of 1,563,193 categories were collected, and the combination of these categories with the original dataset was named the wiki-cate dataset. Here are some examples of the data:

• •

  AG News (AG’s News Corpus): The AG News dataset (Zhang, Zhao, and LeCun 2015) consists of news titles and descriptions, categorized into four classes: World, Sports, Business, and Sci/Tech. This study utilizes the AG News dataset from Huggingface, with a total of 7,600 data points in the test set. Sci/Tech is further divided into two classes (Science and Technology) for classification purposes.

• •

  Yahoo! Answers: The Yahoo! Answers dataset (Zhang, Zhao, and LeCun 2015) mainly comprises questions and answers from the Yahoo! platform. In this study, Yahoo! Answers dataset from Huggingface is employed, containing 60,000 data points in the test set. There are ten main categories for the topics: Society & Culture, Science & Mathematics, Health, Education & Reference, Computers & Internet, Sports, Business & Finance, Entertainment & Music, Family & Relationships, and Politics & Government. Each label is individually treated as a single label for classification (splitted by the ’&’ character), and the output is mapped back to the original label.

• •

  DBpedia: The DBpedia dataset (Auer et al. 2007) is derived from structured information extracted from Wikipedia. In this study, the DBpedia dataset dbpedia_14 from Huggingface is used as the experimental subject, consisting of 70,000 data points in the test set. It includes 14 distinct categories, namely Company, EducationInstitution, Artist, Athlete, OfficeHolder, MeanOfTransportation, Building, NaturalPlace, Village, Animal, Plant, Album, Film, and WrittenWork. We will split the partial classification into words by separating them with a space as input: EducationInstitution to Education institution, OfficeHolder to Office holder, MeanOfTransportation to Mean of transportation, NaturalPlace to Nature place and WrittenWork to Written work.

In this study, the labels of the test dataset are pre-combined with prompts to perform actual classification. Table 1 shows the prompts for each target dataset.

The parameters used for SBERT experiments are as follows:

• •

  Pre-trained model: all-mpnet-base-v2

• •

  Train batch size: 128

• •

  Sequence length (token length): 128

• •

  Epoch: 1

• •

  Loss function: Multiple Negative Ranking (MNR) Loss

In recent research, Prompt tuning (Schick and Schütze 2020) has shown to significantly improve the accuracy of zero-shot or few-shot classification. We aim to investigate whether Prompt tuning can also enhance the similarity matching of the SBERT-based sentence similarity model, WC-SBERT, which is trained solely on labels. In our experiments, we compare the performance of WC-SBERT using a simple hard prompt (Liu et al. 2022) (”This topic is talk about [Label]”) with that of using only the label for comparison. Figure 2 shows the evaluation results of WC-SBERT on the 7,600 test set of AG News dataset.

WC-SBERT is pre-trained using wiki categories as labels. We are particularly interested in whether direct comparison of the target test set’s queries and labels should be conducted, or if an intermediate step is required. This intermediate step involves mapping the queries to the closest category within the wiki categories and then performing the category-label matching. Intuitively, the presence of seen categories during WC-SBERT’s training process should potentially improve accuracy.

In the experiment, we compare each text (query) in the AG News test set with all wiki categories using cosine similarity. We obtain the category closest to this query as input (for example, query: Fears for T N pension after talks Unions representing workers at Turner Newall say they are ’disappointed’ after talks with stricken parent firm Federal Mogul. corresponds to category: Transport and General Workers’ Union). We then perform a second round of cosine similarity between this category and the labels that should be output for the test. The experimental results can be referred to the Figure 3.

The self-training experiments can be divided into three parts to assess the impact of their configurations on accuracy:

• •

  $i$: The number of iterations

• •

  $m$: The choice of the model for fine-tuning

• •

  $t$: The setting of the similarity score threshold

This study conducted relevant experiments on the AG News test set. Detailed results can be referred to Figure 4.

Continuing from the previous experiment, Self-training relies on iterative training to increase accuracy. The iterative process consists of two main parts: inference on the target training and testing sets, and fine-tuning based on the inference results. We conducted the experiments using WC-SBERT on different datasets and reproduced the current best model in Gera et al. (2022) approach. The experiment involved 2 rounds of fine-tuning and 3 rounds of inference to compare the time cost of inference and fine-tuning between the two methods, thereby evaluating their performance.

To ensure the experiment’s fairness, all tests were conducted using the same hardware environment with the following specifications: CPU: Intel(R) Xeon(R) Gold 6154 (8 cores), GPU: NVIDIA Tesla V100 * 2, RAM: 128 GB, OS: Ubuntu 20.04 LTS. The experimental results are presented in Table 2.

We have also used the currently popular GPT-3.5 model and the GPT-4+fine-tune model provided by the API to test the effects of zero-shot text classification. Compared to the GPT-4 model, the GPT-3.5 model is more affordable, allowing us to test on several datasets. For GPT-3.5, the prompt template we used is as follows:

Using the aforementioned prompt for inference on the GPT-3.5 model, the results are presented in the experiment result section.

Moreover, we also tried using the GPT-4 + fine-tune provided by OpenAI API to address the zero-shot text classification problem. However, due to its high cost, we only used 6,700 Wikipedia entries (split into 5,700 training data and 1,000 validation data) to fine-tune the GPT-4 model. We then randomly selected 300 entries from the Yahoo topic as a test set. If we fine-tuned the GPT-4 model with all 5,700 training data entries at once, the resulting test accuracy was 0.55, far from state of the art. However, OpenAI API also offers the feature to fine-tune an already fine-tuned model, meaning you can fine-tune a model a second time with a smaller dataset. By dividing the 5,700 entries into 5,000 for the first round of fine-tuning and 700 for the second round, we preliminarily achieved a test accuracy of 0.62. Of course, due to the high cost of using the Curie model, we could only preliminarily test on a few hundred entries. This GPT fine-tuning approach is also challenging to apply to large-scale datasets for zero-shot text classification.

From the above experiments, several inferences can be drawn:

• •

  SBERT-based similarity models, when combined with labels as the pre-training and fine-tuning objective, still show sensitivity to the use of prompts. Using prompts leads to better accuracy.

• •

  From the label mapping experiment, we have shattered conventional intuition. Directly comparing the query with the target dataset’s labels using SBERT’s similarity yields better results than mapping the query first.

• •

  Self-training effectively improves accuracy within 1 to 2 iterations.

• •

  Compared to the current best model in Gera et al. (2022), WC-SBERT demonstrates significant improvements in both inference and fine-tuning performance. During the inference phase, WC-SBERT achieves a remarkable reduction in time across different datasets, ranging from 10 to 40 times faster. In the fine-tuning phase, Self-Gera only utilizes 800 samples from the target set, whereas WC-SBERT employs 6.45 million samples from WIKI for fine-tuning. Despite the difference in dataset sizes, WC-SBERT’s overall fine-tuning time remains lower than that of Self-Gera, resulting in a performance boost of 5,000 to 9,700 times.

• •

  GPT-3.5 model does not achieve state-of-the-art performance on several datasets, shows poor results, and has higher costs compared to other models. Currently, it is not suitable for use in zero-shot text classification problems.

Based on the above inferences, we adopt this framework as our foundational model architecture and conduct experiments on different target datasets as shown in the Table 4.