CWE Classifier: Improvements Report

CWE Classifier: Improvements Report#

Technical report on changes made to the CWE classification pipeline (dataset generation, CWE knowledge base, and cwe_guesser_patches.py trainer) following the analysis in VulnTrain#17.

Context#

An external review of the CWE suggestions surfaced through VLAgentIc noted that most suggested CWEs were Discouraged for mapping by MITRE (CWE-703, CWE-707, CWE-284, …) while obvious, Allowed candidates (such as CWE-620 for an unverified password change) were missed.

The root cause was in the training labels: the trainer mapped every CWE of the dataset to its deepest ancestor via deep_child_to_ancestor.json, and those deep ancestors are mostly Pillars and high-level Classes — exactly the entries MITRE marks as Discouraged. The model could only ever predict discouraged CWEs.

Changes#

Label set restricted to allowed CWEs#

The CWE knowledge base is now refreshed from the Vulnerability-Lookup API by tools/cwe/update_cwe_knowledge_base.py, which stores each CWE’s Mapping_Notes usage (Allowed, Allowed-with-Review, Discouraged, Prohibited) in tools/cwe/cwe_usage.json. As of July 2026: 749 Allowed, 93 Allowed-with-Review, 44 Discouraged, 83 Prohibited.

tools/cwe/build_child_to_ancestor.py then regenerates vulntrain/data/deep_child_to_ancestor.json: each CWE is mapped to its highest ancestor whose usage is Allowed or Allowed-with-Review. The label set went from 26 discouraged deep ancestors to 303 allowed CWEs. Since a classifier can only predict labels it was trained on, the model structurally cannot suggest a Discouraged or Prohibited CWE anymore.

90 CWEs have no allowed entry anywhere on their ancestor path (including common ones like CWE-20, CWE-200, CWE-287, CWE-400 — this is MITRE’s actual guidance). Examples carrying only such CWEs are excluded from training rather than being mapped to a discouraged label.

Dataset generation fixes#

Two bugs in vulntrain/datasets/cwe-guesser-dataset.py reduced the quantity and quality of the training data:

  • Stale CWE labels (cvelistv5): the cwe variable was only assigned when a record had problemTypes descriptions, so records without them silently inherited the CWE of the previously processed vulnerability. Each record now gets a fresh list, all problemTypes descriptions are collected (not just the first of the first), and the structured cweId is preferred over the free-text description.

  • GHSA examples silently dropped: GHSA records stored their CWEs in a separate cwes column while the trainer only reads cwe, so no GHSA example ever contributed to training. Both sources now emit a cwe column as a list of strings.

Because the cwe column type changed, the CIRCL/vulnerability-cwe-patch dataset must be rebuilt once with the new --from-scratch flag:

python vulntrain/datasets/cwe-guesser-dataset.py --sources cvelistv5,github,pysec,csaf_redhat --repo-id=CIRCL/vulnerability-cwe-patch --from-scratch

(The flag was added together with a fix to the incremental push: the previous concatenation call used a non-existent Dataset.concatenate method, whose AttributeError was silently swallowed — appends actually replaced the dataset.)

Top-k accuracy metrics#

The model is consumed as a suggester of candidate CWEs, so whether the right CWE appears among the top suggestions matters more than exact top-1 accuracy — especially with 303 fine-grained classes. compute_metrics now reports accuracy_top3 and accuracy_top5 alongside accuracy and macro F1, and the best checkpoint is selected by f1_macro instead of eval loss (metric_for_best_model).

Long-context input and ModernBERT#

The input is a vulnerability description plus commit messages and diffs, but only the first patch was used and everything was truncated at 512 tokens (roberta-base’s limit), so most of the patch content was never seen by the model. Changes:

  • all patches of a record are concatenated into the input text;

  • the tokenizer truncation length defaults to the model’s maximum input length (capped at 8192), overridable with --max-length;

  • padding is now dynamic per batch (DataCollatorWithPadding) instead of padding everything to the maximum length.

This makes ModernBERT-base (8192-token context, code-aware pretraining) the recommended base model:

vulntrain-train-cwe-classification --base-model answerdotai/ModernBERT-base --dataset-id CIRCL/vulnerability-cwe-patch --repo-id CIRCL/vulnerability-cwe-classification-modernbert-base --batch-size 8

Lower the batch size compared to roberta-base: sequences are up to 16× longer.

Class weighting strategies#

With 303 heavily imbalanced classes, fully balanced class weights (the previous hard-coded behaviour) over-boost classes with a handful of examples. The --class-weights flag from the CNVD severity trainer was ported: none (uniform loss), sqrt (sqrt-dampened weights), balanced (previous behaviour, still the default) and focal (focal loss, gamma=2.0, balanced alpha). Worth comparing sqrt and focal against the balanced baseline.

Baseline to beat#

First model trained with the allowed-only label set (roberta-base, 40 epochs, balanced weights, 512-token context, before the dataset regeneration):

Metric

Value

Accuracy

0.5683

F1 macro

0.2101

Labels

303

Train examples

6,460

The accuracy/F1-macro gap shows the long tail of rare CWEs is barely learned. Expected levers, in order: dataset regeneration (recovers all GHSA examples and fixes mislabeled cvelistv5 rows), long-context ModernBERT (sees the full patches), then class-weighting experiments.

Possible future work#

  • Multi-label training: records often carry several CWEs; the trainer currently keeps only the first mappable one. Sigmoid multi-label targets would use all of them and match reality (a vulnerability can legitimately be CWE-620 and CWE-640).

  • Minimum label support: collapse labels with fewer than ~5 training examples out of the label set instead of teaching the model classes it cannot learn.

  • Hierarchy-aware evaluation: report a soft accuracy where predicting an ancestor or descendant of the reference CWE counts as partially correct.