🤖 machine learning

Mitigating The Effect of Class Imbalance in Data with Hierarchical and Dependable Structure

This paper proposes a Hierarchy-Aware RoBERTa framework that leverages learnable parent-class embeddings to effectively mitigate class imbalance in CWE vulnerability classification, demonstrating that incorporating hierarchical structure outperforms traditional oversampling techniques which often degrade model performance.

Original authors: Bipin Chhetri, Deepika Giri, Avishek Kadel, Rabin Kumar Karki, Akbar Siami Namin

Published 2026-07-15

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Bipin Chhetri, Deepika Giri, Avishek Kadel, Rabin Kumar Karki, Akbar Siami Namin

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a detective trying to sort a massive pile of clues about computer security flaws. These flaws are organized in a giant family tree called the Common Weakness Enumeration (CWE). At the top of the tree are broad categories like "Base" (the big picture), and as you go down, the branches get more specific, ending in tiny, rare leaves like "Compound" or "Pillar."

The problem? The detective's evidence bag is totally unbalanced. There are hundreds of clues for the big, common categories, but only a handful for the rare, specific ones. It's like having a library with 500 books about "Fruit" but only 5 books about "Dragonfruit." If you try to teach a computer to sort these, it gets lazy and just guesses "Fruit" every time, because that's what it sees most often.

The "Fake Clue" Experiment (What Didn't Work)

To fix this, many experts tried a trick called oversampling. They took the few rare clues and tried to invent new, fake ones to make the numbers look equal. They used two popular methods: SMOTE and ADASYN.

Think of this like a chef trying to make a soup taste like it has more rare spices. Instead of finding more real spices, they take two existing spice grains, blend them together, and hope the new mixture tastes authentic.

The paper tested this on different types of "detectives" (computer models):

The Old-School Detectives (Random Forest and SVM): These models are like detectives who look at simple lists of facts. When fed the fake blended spices, they got a tiny boost. Their accuracy went from 0.65 to 0.69 for the Random Forest, and the Support Vector Machine (SVM) stayed steady around 0.71–0.72. It helped a little, but not much.
The High-Tech Detectives (CNN and BiGRU): These are smarter, deep-learning models that understand how words flow together. When the researchers fed them the fake blended spices, the results were a disaster. The CNN's accuracy crashed from 0.71 down to 0.55 with SMOTE and 0.51 with ADASYN. The BiGRU fell from 0.70 to 0.53 and 0.44.

Why? The paper argues that these high-tech models are like chefs who can taste the difference between real spice and a fake blend. When you mix two different computer "words" together to make a fake one, you break the rules of the family tree. You might create a "Variant" that claims to be a child of a "Base," but the fake blend doesn't actually respect that parent-child relationship. It's like trying to make a "Dragonfruit" by blending an "Apple" and a "Banana." The result isn't a Dragonfruit; it's a confusing mess that confuses the detective.

The "Family Tree" Solution (What Actually Worked)

Instead of making fake clues, the authors built a new detective called Hierarchy-Aware RoBERTa.

Imagine this detective has a special map of the family tree in their pocket. They don't just read the clue; they also check the map to see, "Wait, if this clue is about a 'Base' weakness, the answer must be related to that parent."

The model works like this:

It reads the text description of the flaw (using a powerful tool called SecureBERT).
It grabs a "Parent ID" from the family tree (like knowing the clue belongs to the "Base" branch).
It combines the text reading with the map location to make a final guess.

The Results:
This new detective didn't need any fake clues at all. It achieved a weighted F1-score of 0.76 without any data augmentation.

Compare this to the standard BERT model, which scored 0.74.
Most importantly, look at the rare "Class" category. The standard BERT model only got an F1-score of 0.49 for this rare group. The new Hierarchy-Aware model boosted that to 0.60.

The Bottom Line

The paper suggests that when you have a structured family tree of data, trying to "fake" more data by blending existing pieces (oversampling) is a bad idea. It works okay for simple models but breaks the more advanced ones.

Instead, the best approach is to teach the model to respect the family tree structure from the start. By giving the model a "map" of the parent-child relationships, it can figure out the rare, tricky cases much better than by just stuffing the training data with synthetic noise.

However, the authors are careful to note that even their best detective still struggles with the rarest categories, like "Compound" and "Pillar," which had only 8 and 5 samples respectively. For these extremely rare groups, the F1-score remained 0.00 across all models, suggesting that when there is almost no data at all, even a family tree map isn't enough to solve the mystery.

Technical Summary: Mitigating The Effect of Class Imbalance in Data with Hierarchical and Dependable Structure

Problem Statement

Classifying cybersecurity vulnerabilities using the Common Weakness Enumeration (CWE) taxonomy presents two primary challenges: extreme class imbalance and strong hierarchical dependencies among weakness categories. The CWE taxonomy organizes weaknesses into abstraction levels (Base, Class, Pillar, Variant, Compound), where upper-level nodes represent broad concepts and lower-level nodes represent specific details. In practice, a small number of weakness types dominate the corpus, while rare but high-impact categories remain significantly underrepresented.

While Transformer-based encoders like BERT and RoBERTa have established themselves as strong baselines for CWE text classification, they struggle with rare categories due to class imbalance. Furthermore, CWE labels possess explicit parent-child dependencies (e.g., a "Variant" is always a child of a "Base"), meaning misclassifications at higher abstraction levels cascade downward. Traditional approaches to imbalance, such as random oversampling, undersampling, and synthetic generation techniques like Synthetic Minority Oversampling Technique (SMOTE) and Adaptive Synthetic Sampling (ADASYN), are widely adopted but their effectiveness on hierarchical, structured text embeddings remains largely unexplored. Specifically, it is unclear whether linear interpolation in high-dimensional embedding spaces respects the structural constraints of the CWE hierarchy.

Methodology

The study evaluates a range of models under different data augmentation conditions and proposes a novel architecture that integrates hierarchical structure directly into the learning process.

1. Experimental Baselines and Oversampling

The authors empirically evaluated Machine Learning (Random Forest, SVM), Deep Learning (CNN, BiGRU), and Transformer-based (BERT) models on an imbalanced hierarchical CWE dataset.

Oversampling: SMOTE and ADASYN were applied exclusively to the training set. SMOTE equalized all classes to 393 samples using $k$ -neighbors=2, while ADASYN concentrated synthesis on harder boundary regions to achieve near-parity counts.
Models:
- Classical ML: Random Forest (100 trees) and SVM (linear kernel).
- Deep Learning: CNN (128 filters, kernel sizes [4,6,6]) and BiGRU (128 hidden dimension).
- Transformers: BERT-base (uncased) and the proposed Hierarchy-Aware RoBERTa.
- Training: All deep learning models were trained with a batch size of 8 for 15 epochs.

2. Proposed Approach: Hierarchy-Aware RoBERTa

The paper introduces a Hierarchy-Aware RoBERTa framework that explicitly incorporates CWE structural information through learnable parent-class embeddings.

Architecture: The model utilizes a pretrained SecureBERT encoder to generate a 768-dimensional [CLS] representation from input text.
Hierarchical Injection: A learnable parent-category embedding (24-dimensional) is introduced to capture coarse-grained structural relationships. This embedding is trained jointly with the network.
Fusion: The semantic representation ( $h_{cls}$ ) and the hierarchical context ( $h_p$ ) are fused via concatenation: $h_f = [h_{cls}, h_p]$ .
Classification: The fused representation is regularized with dropout (0.3) and passed through a linear classification head to output logits.
Formulation: The task is framed not as an independent prediction $P(Y_{child}|X_{text})$ , but as a conditional probability problem $P(Y_{child}|X_{text}, Y_{parent})$ , utilizing the "Parent ID" as a structural prior.

Key Contributions

Empirical Evaluation: The study provides a comprehensive evaluation of ML, DL, and Transformer models under SMOTE and ADASYN resampling on an imbalanced hierarchical CWE dataset.
Analysis of Synthetic Interpolation: The authors demonstrate that synthetic interpolation in high-dimensional embedding spaces violates inherent parent-child constraints. While offering marginal benefits for classical ML models, these methods consistently degrade the performance of deep learning architectures.
Hierarchy-Aware Model: The introduction of a Hierarchy-Aware RoBERTa model that injects CWE structural priors via learnable parent-class embeddings, enabling end-to-end joint modeling of semantic content and hierarchical structure.
Performance Gains: The proposed model achieves a weighted F1-score of 0.76 without data augmentation, outperforming all baselines. Notably, it improves the F1-score for the "Class" category from 0.49 (BERT baseline) to 0.60.

Results

Oversampling Impact:
- Classical ML: SMOTE and ADASYN provided weak but positive improvements for Random Forest (accuracy increased from 0.65 to 0.69) and SVM (accuracy remained stable around 0.71–0.72).
- Deep Learning: Oversampling consistently degraded performance. Vanilla CNN accuracy dropped from 0.71 to 0.55 (SMOTE) and 0.51 (ADASYN). BiGRU dropped from 0.70 to 0.53 and 0.44, respectively.
- Transformers: BERT achieved 0.74 accuracy without augmentation. The proposed Hierarchy-Aware RoBERTa achieved the highest accuracy of 0.76.
Per-Class Performance:
- The proposed model achieved the highest weighted F1-score (0.76) compared to BERT (0.73), CNN (0.68 baseline), and BiGRU (0.69 baseline).
- Significant gains were observed in minority classes; for instance, the "Class" category F1-score improved from 0.49 (BERT) to 0.60 (Hierarchy-Aware RoBERTa).
- Extremely sparse categories (Compound and Pillar) remained at 0.00 F1 across all models, highlighting the limits of supervised learning with severe data sparsity.

Significance and Claims

The paper argues that hierarchy-aware representation learning is a more principled alternative to oversampling for structured vulnerability classification. The authors claim that synthetic interpolation techniques (SMOTE/ADASYN) fail in this domain because:

They produce samples that lack linguistic coherence in high-dimensional contextual spaces.
They violate the structural dependencies of the CWE taxonomy (e.g., a "Variant" cannot exist without a "Base"), creating structurally invalid embeddings that confuse the model.
They treat samples as independent entities in a flat feature space, ignoring the tree-structured graph nature of CWEs where errors cascade from parent to child.

By explicitly modeling the parent-child relationship through learnable embeddings, the proposed framework enforces structural consistency and improves discrimination among minority and semantically dependent classes without the noise introduced by synthetic data generation.

Limitations

The authors modestly acknowledge several limitations:

Inference Dependency: The model assumes parent-level metadata is available during inference, which aligns with coarse-to-fine pipelines but may not suit all scenarios.
Static Hierarchy: The hierarchy is treated as fixed; the model may not generalize to new or rearranged nodes without retraining, given the continuous updates to the CWE taxonomy.
Scope of Hierarchy: Only direct parent information is captured, disregarding multi-hop ancestors, sibling relationships, and tree depth, which limits performance on deeply nested classes like Compound and Pillar.
Parameter Overhead: Learnable parent embeddings add modest parameter overhead, which could hinder scaling to very large taxonomies.
Domain Specificity: The approach was evaluated only within the CWE domain, leaving its generalizability to other hierarchical tasks empirically unvalidated.