From Evaluation to Optimisation: Hierarchy-Aware Training Signals for CWE Prediction in Python

2026-07-23

AI Summary

This study addresses the performance degradation of CWE-level vulnerability prediction in Python code under distribution shift. It presents the first empirical validation that the hierarchical evaluation penalty defined in the ALPHA benchmark can be effectively leveraged as a training signal. The authors systematically compare three training mechanisms: supervised fine-tuning, dual-head classification loss, and GRPO reinforcement learning incorporating normalized ALPHA penalties. Experimental results demonstrate that GRPO significantly outperforms supervised approaches under distribution shift. On the Qwen2.5-Coder-7B model, the best GRPO policy reduces cumulative ALPHA penalty by 27.9% with greedy decoding and 25.5% with sampling decoding on the SVEN dataset, achieving statistically equivalent performance to a zero-shot teacher model 4.5 times larger in scale.

Abstract

The original ALPHA benchmark introduced a taxonomy-aware penalty for evaluating CWE-level vulnerability prediction in Python and proposed that the penalty could theoretically also serve as a training signal. This paper provides that validation. We compare three delivery mechanisms: supervised fine-tuning, a dual-head classification loss, and reinforcement learning with a dense reward derived from the normalised penalty. We find that supervised approaches consistently regress below the zero-shot baseline under distribution shift, while GRPO succeeds. Our best policy reduces the cumulative ALPHA penalty of Qwen2.5-Coder-7B on Security Hardening and Adversarial Testing (SVEN) dataset by 27.9% under greedy decoding, and by 25.5% under sampled decoding(p = 0.005, Welch's t-test), reaching statistical parity with its 4.5x larger zero-shot teacher. We conclude that the value of a hierarchical penalty as a training signal depends largely on the directness of its delivery.

Problem

• CWE prediction
• hierarchy-aware penalty
• training signal
• distribution shift
• vulnerability detection

Innovation

• hierarchy-aware training
• CWE prediction
• reinforcement learning
• ALPHA penalty
• distribution shift