urlBERT:A Contrastive and Adversarial Pre-trained Model for URL Classification

In 2018, Google introduced the BERT[3], which following pre-training on a substantial corpus of document-level text through Masked Language Modeling (MLM) and Next Sentence Prediction (NSP) tasks, demonstrated outstanding performance when fine-tuned for various downstream tasks such as text classification, named entity recognition, and question-answering. Subsequently, teams across numerous research domains were inspired by this "pre-training + fine-tuning" approach, recognizing the advantages of the Transformer architecture in representing textual information. They have pre-trained BERT models using domain-specific corpora and deployed them in various downstream tasks through fine-tuning, achieving significant advancements.BioBERT[12], trained on a large biomedical corpus, outperformed traditional NLP models in understanding biomedical literature. CodeBERT[5], based on RoBERTa[15], was pre-trained on a substantial dataset containing natural language-programming language pairs, exhibiting remarkable performance in natural language code search and code documentation generation tasks. PhishBERT[21] pre-trained BERT on a substantial amount of unlabeled URL data and achieved excellent accuracy in fine-tuning for phishing link detection tasks. Today, facing a multitude of downstream tasks such as URL webpage topic classification, phishing link recognition, and advertisement URL identification, pre-training a publicly available model based on the Transformer architecture remains valuable.

URL classification tasks, including malicious URL detection, webpage topic classification, and advertisement URL recognition, have always been crucial components of cybersecurity governance and improvements in network services optimization, among which the identification of malicious URLs has garnered particular attention from researchers in the field of cybersecurity. In the past, traditional identification methods relied on blacklist-based approaches, where URLs were categorized based on predefined lists. However, the timeliness and accuracy issues associated with this method prompted researchers to adopt methods involving the creation of rich datasets [18, 1, 19] and training various neural network models. As URLs represent a unique form of textual data, classification tasks are typically handled by commonly used neural networks in the NLP domain. Particularly focusing on the task of malicious URLs detection, researchers have designed and trained specialized neural network architectures tailored to URL data features. Models such as Grambeddings [1], URLTran [16], URLNet [11], and URL2Vec [23] have demonstrated excellent performance in malicious URL identification tasks. Moreover, these networks, when adjusted and fine-tuned for other tasks, have also shown promising accuracy. When deploying these models for downstream tasks, it is necessary to fine-tune the model’s structure and undergo multiple epochs of training on large datasets to achieve convergence. In URL classification, following the approach outlined in BERT [3], fine-tuning a pre-trained model that has learned abundant semantic features of URLs, along with task-specific fully connected layers, for a few epochs, can lead to higher convenience and efficiency.

Represented by MoCo[8] and SimCLR[2], contrastive learning aims to enhance performance by simultaneously maximizing the consistency between different transformations (positive samples) of the same instance and minimizing the consistency between different instances (negative samples), demonstrating significant performance improvements on ImageNet. And self-supervised contrastive learning utilizes training data itself as supervisory information, generating positive and negative samples for the model’s contrastive learning. In the field of NLP, prior work has explored the application of self-supervised contrastive learning in BERT training. ConSERT[22], during the fine-tuning stage, employs sample augmentation on the generated data representations at the Embedding layer to obtain corresponding positive and negative samples. After training on the contrastive loss represented by the shared BERT model, performance improvements are achieved on the original model. On the other hand, CERT[4] enhances pre-training data through multilingual back-translation in the pre-training phase, incorporating contrastive learning, resulting in performance improvements in some downstream tasks.

The concept of Virtual Adversarial Training (VAT)[17] is derived from adversarial training, enabling semi-supervised learning. For unlabeled data, the model’s predictions on the data from the previous stage can be used as virtual labels. Adversarial perturbations are then computed based on these virtual labels to generate adversarial samples. During training, the goal is to minimize the discrepancy between the model outputs for the original and adversarial samples, thereby enhancing the model’s robustness. ALUM[13] applies VAT to both the pre-training and fine-tuning stages of the BERT model. Whether introducing VAT in both stages or applying it in one stage, significant improvements in both accuracy and robustness of the model on downstream tasks are observed.

The pre-training of urlBERT consists of three tasks: Masked Language Modeling (MLM), sentence-pair self-supervised contrastive learning, and token-level virtual adversarial training. The training process structure for this stage is depicted in Figure. 1.

URL data typically consists of various components, including network protocols, domain names or IP addresses, resource paths, query parameters, etc., organized in a specified format. Each URL is uniquely identified under this structural specification, and even slight differences between two URLs can result in significant distinctions in their functionalities. Taking phishing links as an example, harmful and legitimate links may only exhibit subtle differences in their query statements. In response to this scenario, we introduce a self-supervised contrastive learning(SSCL) task to enable the model to learn the subtle differences and structural features between URLs at the sentence level during the pre-training phase. We denote the contrastive learning loss as ${Loss}_{con}$, and detailed training specifics and loss functions for this task will be elaborated in Section 3.2.

We retained the Masked Language Modeling (MLM) pre-training task from BERT, consistent with the settings outlined in the original paper[3]. This task involves predicting the actual values of tokens masked with "[MASK]", facilitating the model in learning the compositional structure and semantic information of URLs. In addition to the MLM task, we introduced the Virtual Adversarial Training(VAT) task. VAT involves generating adversarial samples corresponding to the data samples and performing adversarial training. We compute the adversarial training loss, ${Loss}_{VAT}$, at the MLM task’s output of the shared BERT model. This enhances the robustness of the model’s predictions for URL fluctuations at the token level. The detailed implementation of this task will be presented in Section 3.3.

The losses for each task are weightedly summed, with the combination of the VAT task training loss following the configuration outlined in ALUM[13]. The final loss is as follows:

Inspired by the work of ConSERT[22] and SimCSE[6], the self-supervised contrastive learning task in urlBERT incorporates data augmentation at the token embedding level to generate positive and negative samples. Differing slightly from ConSERT[22], urlBERT utilizes the dropout module within BERT’s Embedding layer to produce two augmented samples for a given input. Subsequently, the fast gradient sign method[7] is applied separately to both samples to generate adversarial samples, thereby enhancing the effectiveness of data augmentation. The sample generation scheme is illustrated in Figure. 2.

A batch of $N$ token embeddings generates $2N$ augmented data samples. For a single data sample within a batch, there is only one positive sample, while the rest consist of $2(N-1)$ negative samples. The objective of the contrastive learning task is to minimize the distance between positive samples and maximize the distance between negative samples. Therefore, similar to ConSERT[22], urlBERT employs the Normalized Temperature-scaled Cross-Entropy Loss (NT-Xent):

$r_{i}$ is the representation obtained from the final layer of the transformer block output in the shared BERT encoder after encoding an augmented data sample. $r_{j}$ signifies the representation for another augmented data sample corresponding to the same input. The function $sim\left(\bullet\right)$denotes the cosine similarity, and $\tau$ represents the temperature coefficient. After calculating this loss for each data point, taking the average yields the loss for this task, denoted as ${Loss}_{con}$. The pseudocode for the SSCL task is illustrated in Algorithm 1.

The VAT task in urlBERT is primarily inspired by ALUM[13], introducing regularized perturbations in the token embedding space to generate adversarial samples corresponding to unlabeled URL data. Similar to ALUM[13], urlBERT employs the gradient forward propagation approach, introducing perturbations with a mean of 0 and a variance of 1 during training initialization. Subsequently, the data samples obtain corresponding model output logits and ${logits}_{adv}$ through encoding by the shared BERT encoder. After computing the virtual adversarial loss (${Loss}_{VAT}$) using the KL divergence function, a regularization perturbation is added in the reverse gradient direction to maximize the adversarial loss. This process can be repeated multiple times to enhance perturbation effects, but urlBERT introduces perturbation only once. The final Loss is obtained by adding the weighted sum of ${Loss}_{VAT}$ and the MLM task loss. Algorithm 2 outlines the training process for the VAT task.

During the pre-training phase, urlBERT utilized the Web Tracking Datasets released by the Broadband Communications Systems and Architectures Research Group at UPC, which comprises approximately 3 billion web URLs. Our team extracted unlabeled URLs from this extensive dataset and employed them for urlBERT’s pre-training. The pre-training of urlBERT was based on the BERTForMaskedLM model, conducted on four A100-40G GPUs. AdamW optimizer was employed, with a learning rate of 2e-6 for the SSCL task and 1e-5 for the other two tasks. The batch size was set to 64, and the temperature coefficient $\tau$ for the objective loss function of the SSCL task was set to 0.1.

We conducted fine-tuning experiments on three downstream tasks: phishing link detection, webpage topic classification, and advertisement link identification. Traditional fine-tuning approaches typically involve obtaining the [CLS] data representation from the last layer transformer block of BERT and inputting it into a fully connected layer constructed for the downstream task. During training, all parameters, including those of BERT, are fine-tuned. Another approach involves using BERT as a feature extractor, freezing its parameters, and inputting the data representation into other commonly used deep neural networks for training. And generally the former method tends to yield superior performance. Inspired by the work of Ran Wang[20], we chose to adopt the FT-TM fine-tuning approach to enhance urlBERT’s performance on downstream tasks.

Initially, we froze the parameters of the urlBERT model and used a CNN as the neural network for the downstream tasks. Fine-tuning was performed for each task, followed by unfreezing the BERT parameters for full-parameter fine-tuning. Throughout the experiments, the AdamW optimizer was employed, with a learning rate of 2e-3 in the first stage and 2e-5 in the second stage and both stages involve 5 epochs of training.

Table 1 presents a comparison of the Accuracy metrics on the test set for models trained on various downstream tasks using urlBERT, under both the FT-TM and full-parameter fine-tuning methods. It is evident that urlBERT, when fine-tuned with the FT-TM approach, exhibits a substantial improvement in performance across all three downstream tasks. Particularly noteworthy is the task of recognizing advertisement links, where the Accuracy on the test set has increased from 0.9990 to 0.9995, showcasing further impressive performance in predicting Accuracy for this task.

To demonstrate the effectiveness of urlBERT in URL security-related research tasks, we conducted comparative experiments on three downstream tasks: phishing link detection, webpage topic classification, and advertisement link recognition. We compared urlBERT with commonly used neural network models in the NLP domain, namely BiLSTM[9] and TextCNN[10]. All models were experimented with on the same training and testing sets, utilizing the AdamW optimizer. The learning rate for BiLSTM was set to 0.005, TextCNN’s learning rate was set to 2e-3, and urlBERT’s learning rate was set to 2e-5. The batch size for all models was set to 5, and testing comparisons were performed after 5 training epochs.

Table 2 presents a comparative analysis of urlBERT and other commonly used neural networks for NLP tasks across different downstream tasks. The urlBERT group represents results obtained under full-parameter fine-tuning, while urlBERT(FT-TM) represents results obtained under FT-TM method. It is evident that, compared to other commonly used neural network models, urlBERT consistently demonstrates higher predictive accuracy across various downstream tasks. Particularly noteworthy is the task of webpage topic classification, where urlBERT achieves an accuracy of 0.7448, a 35% improvement over BiLSTM’s test accuracy of 0.5531 and a slight enhancement compared to TextCNN’s metric of 0.7373. In the phishing link recognition task, urlBERT achieves a test accuracy of 0.9739, showing significant improvement over BiLSTM (0.9545) and TextCNN (0.9667).

Due to the significance of malicious link recognition in the current research landscape related to URLs, we have documented various metrics for different models at different training points in this task. The optimization processes for these metrics are illustrated in Figure 4. It is evident that during the 5 training rounds on this downstream task, urlBERT consistently outperforms BiLSTM in terms of Accuracy, Precision, and F1-score. Additionally, compared to TextCNN, urlBERT exhibits smoother and faster convergence in terms of Accuracy and F1-score, surpassing TextCNN after 5 rounds of training. While both models experience a decline in Precision after the second round of training, urlBERT demonstrates greater stability in this aspect.

We evaluated the predictive capabilities of urlBERT models trained on different data scales in the phishing link detection task. Models were trained on datasets consisting of 30,000, 320,000, and 520,000 instances, respectively, and were tested on the corresponding test sets. We compared the performance with TextCNN and BiLSTM to assess urlBERT’s dependency on dataset scale. All three models were trained for 5 epochs on datasets of varying sizes, maintaining consistent training settings with those used in downstream task comparison experiments.

Table 4 presents the metrics of models obtained through training by urlBERT and two other neural networks at various data scales. It is evident that urlBERT exhibits the least dependency on data scale. Across three different data volumes, urlBERT consistently outperforms TextCNN and BiLSTM in terms of Accuracy, Precision, and F1-Score. Only at a data volume of 320,000 does urlBERT’s Recall (0.9645) fall slightly below that of TextCNN (0.9656). TextCNN, on the other hand, demonstrates significant data scale dependence, with noticeable declines in all metrics at a data volume of 30,000. Figure 4 illustrates the ROC curves for models obtained through urlBERT training at different data volumes. It is observed that, despite a decrease in training data volume, urlBERT’s inference performance does not exhibit significant deterioration. At the fixed FPR of 0.01, the model’s test TPR still reaches at approximately 0.9, indicating that urlBERT does not heavily rely on large-scale training data.

To investigate the feature extraction capabilities of different transformer blocks in urlBERT, we extracted the [CLS] representations from outputs at various layers of the transformer during fine-tuning. These representations were then input into a fully connected layer for training in the phishing link recognition task. Additionally, different pooling methods were employed in subsequent experiments to assess the impact of this factor on urlBERT’s performance in downstream tasks. Throughout the training process, each experiment group utilized the AdamW optimizer with a learning rate set to 2e-5 and a weight decay set to 1e-4. The models were trained for 5 rounds, and the metrics were compared on the test set thereafter.

Table 3 presents the test results obtained from training urlBERT on different layers’ feature outputs in the downstream task of phishing link recognition and also illustrates the impact of various feature pooling methods on the model’s performance. The tested transformer layers are the last four layers of the BertEncoder, and the experimental results confirm the effectiveness of these layers’ data representations. Employing different pooling methods on the output features of the last four layers yields varying effects on model performance. In comparison to the best single-layer test accuracy of 0.9717, WeightedPooling and AttentionPooling result in performance improvements, reaching 0.9718 and 0.9720, respectively. However, MaxPooling disrupts the model’s data representations, causing a decline in metrics to 0.9612.

Multi-task learning is another important direction for fine-tuning urlBERT for various downstream tasks. MT-DNN[14] uses BERT as a shared data encoder and assigns separate neural network layers for different downstream tasks during training. This training approach allows the model to learn knowledge from a large amount of training data from different downstream tasks, particularly beneficial when training data for a single task is limited. It enhances the model’s generalization capabilities and mitigates the negative effects of small training data volumes. We attempted multi-task learning for urlBERT, designing classifier networks for phishing link recognition, webpage topic classification, and advertisement link recognition with the shared urlBERT. During training, to create a suitable scenario for multi-task learning, we set different training data volumes for the three tasks: 40,000 samples for advertisement link recognition and 180,000 samples for the other two tasks. We employed the AdamW optimizer with a learning rate of 2e-5 and a weight decay of 1e-4.

We employed traditional full-parameter fine-tuning methods with the same data volume settings for multi-task learning to compare with multi-task learning, aiming to demonstrate the effectiveness of multi-task fine-tuning for urlBERT in scenarios where data scale is limited. The experimental results are presented in Table 5. It is evident that conducting multi-task learning for fine-tuning urlBERT in situations of limited training data effectively enhances the model’s predictive performance. Particularly noteworthy is the webpage topic classification task, where, compared to the test accuracy obtained through fine-tuning on 180,000 data points (0.7264), urlBERT achieves a test accuracy of 0.7304 through multi-task fine-tuning. In the phishing link recognition task, multi-task learning also increases accuracy from 0.9601 to 0.9612. While the model’s performance in the advertisement link recognition task is slightly lower with multi-task fine-tuning compared to full-parameter fine-tuning, the accuracy still maintains an impressive performance of 0.9981.