Cyber-Attack Technique Classification Using Two-Stage Trained Large Language Models

TL;DR

This work tackles the challenge of extracting fine-grained attack techniques from unstructured CTI reports by framing it as sentence-level classification and addressing data sparsity with a two-stage training pipeline. It combines similarity-based auxiliary data selection from MITRE with a domain-specific LLM (CTI-BERT) pretrained on cybersecurity text, followed by a second stage that re-focuses training on the primary TRAM data. The approach yields Macro-F1 gains of about 5–9 percentage points and maintains competitive Micro-F1, with notable improvements when augmenting only rare classes using the most similar auxiliary examples (up to $k=10$). The method is domain-agnostic and can be adapted to other low-resource text classification tasks in cybersecurity and beyond.

Abstract

Understanding the attack patterns associated with a cyberattack is crucial for comprehending the attacker's behaviors and implementing the right mitigation measures. However, majority of the information regarding new attacks is typically presented in unstructured text, posing significant challenges for security analysts in collecting necessary information. In this paper, we present a sentence classification system that can identify the attack techniques described in natural language sentences from cyber threat intelligence (CTI) reports. We propose a new method for utilizing auxiliary data with the same labels to improve classification for the low-resource cyberattack classification task. The system first trains the model using the augmented training data and then trains more using only the primary data. We validate our model using the TRAM data1 and the MITRE ATT&CK framework. Experiments show that our method enhances Macro-F1 by 5 to 9 percentage points and keeps Micro-F1 scores competitive when compared to the baseline performance on the TRAM dataset.

Figures (2)

• Figure 1: System Overview. For the given primary data ($D^P$) and auxiliary data ($D^A$), we first select a subset of the auxiliary samples ($D^{A\_sim}$) based on their similarity with $D^P$. The model training is done using $D^P$ and $D^{A\_sim}$ at the first stage and then with only $D^{A\_sim}$ at the second stage.
• Figure 2: Distribution of classes based on the number of member sentences. The TRAM data is highly imbalanced, having 14 classes with only 3 sentences (one each in train, validation and test splits), while the MITRE data is more balanced, with only 3 classes with less than 3 sentences.