Learning Robust Representations for Malicious Content Detection via Contrastive Sampling and Uncertainty Estimation

TL;DR

The paper tackles malicious content detection under Positive-Unlabeled learning by introducing the Uncertainty Contrastive Framework (UCF), which fuses uncertainty-aware contrastive loss, adaptive temperature scaling, and a self-attention LSTM encoder to produce robust, calibrated embeddings. It presents a two-stage training pipeline (ConPU optimization followed by pseudo-negative triplet refinement) and a detailed PU-specific loss that emphasizes uncertain samples. Through synthetic PU-like data and extensive evaluation, UCF demonstrates rapid convergence, clear embedding separation, and strong downstream performance across classifiers, highlighting improved robustness and generalization in noisy, imbalanced settings. The work shows potential for scalable deployment in cybersecurity and biomedical text mining, with future directions toward multimodal PU and cross-dataset adaptation.

Abstract

We propose the Uncertainty Contrastive Framework (UCF), a Positive-Unlabeled (PU) representation learning framework that integrates uncertainty-aware contrastive loss, adaptive temperature scaling, and a self-attention-guided LSTM encoder to improve classification under noisy and imbalanced conditions. UCF dynamically adjusts contrastive weighting based on sample confidence, stabilizes training using positive anchors, and adapts temperature parameters to batch-level variability. Applied to malicious content classification, UCF-generated embeddings enable multiple traditional classifiers to achieve more than 93.38% accuracy, precision above 0.93, and near-perfect recall, with minimal false negatives and competitive ROC-AUC scores. Visual analyses confirm clear separation between positive and unlabeled instances, highlighting the framework's ability to produce calibrated, discriminative embeddings. These results position UCF as a robust and scalable solution for PU learning in high-stakes domains such as cybersecurity and biomedical text mining.

Figures (6)

• Figure 1: The proposed Uncertainty Contrastive Framework (UCF) for malicious content detection. The framework integrates preprocessing and sampling, adaptive temperature scaling, uncertainty-weighted contrastive representation learning, and a self-attention LSTM encoder to generate discriminative embeddings used for final classification.
• Figure 2: Training loss curves for Stage 1 (ConPU loss and raw $\tau$) and Stage 2 (pseudo-negative triplet loss).
• Figure 3: t-SNE of unlabeled_train embeddings after Stage 2.
• Figure 4: t-SNE of validation embeddings after Stage 2.
• Figure 5: Confusion matrices for tuned classifiers: (a) Logistic Regression and (b) Gradient Boosting.