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

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 tests that proposal empirically, comparing 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 and generalises across independent training seeds. Using a default, non-test-informed configuration, GRPO reduces the mean cumulative ALPHA penalty of Qwen2.5-Coder-7B-Instruct on Security Hardening and Adversarial Testing (SVEN) by 13.5%. A further test-selected configuration reduces the mean penalty by 22.9% (27.9% in the best-observed run), reaching statistical parity with its 4.5× larger zero-shot teacher. Our experiments show that reinforcement learning outperforms the alternative delivery mechanisms we test. We conclude that delivery mechanism, rather than the specific hierarchical structure of the reward, is the primary variable—though whether taxonomy-aware structure contributes further remains an open question.

1. Introduction

Large language models (LLM) are now widely used to generate code [1,2,3], and a growing body of evidence indicates that the code they produce frequently contains exploitable vulnerabilities [4,5]. The prior ALPHA work [6] argued that binary vulnerability detection—labelling a function as vulnerable or not—provides insufficient specificity for the iterative repair pipelines increasingly used to improve generated code [1,3], since effective repair requires specific feedback [7]. ALPHA [6] therefore evaluates predictions at the level of the common weakness enumeration (CWE) and introduces a taxonomy-aware penalty under which a prediction that falls near the true CWE in the MITRE hierarchy is penalised less than one that falls in an unrelated region of the taxonomy. Because this penalty is continuous and direction-aware, it is a natural candidate for a training objective rather than only an evaluation metric. Furthermore, the penalty’s direction-aware asymmetry explicitly encodes downstream utility by penalising the loss of diagnostic information (over-generalisation) more heavily than unwarranted specificity (over-specification).

While the original ALPHA paper theorised this penalty could function as a loss function, the empirical viability of this approach was unexplored. Crucially, an effective evaluation metric does not automatically translate into an effective loss function. Even if a signal accurately measures a target, the optimisation mechanism can distort or discard that information before it reaches the model parameters. Therefore, realising the theoretical promise of this penalty depends substantially on how it is architecturally delivered during training. We explore three delivery mechanisms, which differ in whether the penalty enters the training loss and in whether the output vocabulary is restricted to classes observed during training,

• Supervised fine-tuning (SFT): The penalty does not enter the loss; the model imitates the reasoning and CWE labels produced by a stronger teacher model, and the penalty is used only to score the result.
• A dual-head supervised loss: This is the architecture proposed in the original ALPHA paper. A classification head is trained with a differentiable expected-penalty term, ${\mathcal{L}}_{\mathrm{ALPHA}}$, added to cross-entropy (CE). The penalty enters the loss, but the head is restricted to the CWE vocabulary observed in training.
• Reinforcement learning (RL): The ${\mathcal{L}}_{\mathrm{ALPHA}}$ penalty defines the reward, which is one minus the normalised penalty; and the output vocabulary is unrestricted.

To evaluate these delivery mechanisms, we train on the SecurityEval [8] and CVEfixes [9] datasets, and evaluate on the Security Hardening and Adversarial Testing (SVEN) [10] dataset (Section 4). On a matched in-distribution split (SecurityEval 70/30% train–test split), the classification head is highly effective when trained with CE and mean pooling. However, under the distribution shift to SVEN, both supervised mechanisms degrade significantly, regressing below the base model’s (Qwen2.5-Coder-7B-Instruct) zero-shot performance. In contrast, group relative policy optimisation (GRPO) is the only mechanism that demonstrates robustness under shift. Across every base-backbone configuration, GRPO successfully reduces the cumulative ALPHA penalty relative to the baseline. Notably, the default GRPO configuration, without test-informed tuning, achieves a mean 13.5% reduction in the penalty across three independent training seeds, while a further, test-selected configuration achieves a mean 22.9% reduction (27.9% in the best-observed run under greedy decoding, 25.5% under matched sampled decoding; p = 0.005, Welch’s t-test) and reaches statistical parity with a model 4.5× its size. We make two contributions:

• Evidence that RL can deliver the ALPHA penalty as an effective training signal under distribution shift, and that the delivery mechanism, not merely the penalty’s availability, determines whether its benefit is realised. Supplying the same penalty through supervised mechanisms does not improve, and often harms, CWE prediction under shift. Whether specifically the penalty’s hierarchical structure, as opposed to reinforcement learning with an unrestricted output vocabulary, is responsible for this advantage remains open (Section 7).
• RL configurations that improve CWE prediction under distribution shift—a base-backbone GRPO policy with a dense, penalty-derived reward that reduces the SVEN ALPHA penalty of Qwen2.5-Coder-7B-Instruct by a mean of 13.5% across 3 different seeds using only default, pre-specified hyperparameters; and by a mean 22.9% (up to 27.9% best-observed) under a further test-selected configuration that reaches parity with its 4.5× larger zero-shot teacher, together with the data, hyperparameter, and seed-stability analysis that characterise it.

2. Background

LLMs now assist a substantial share of everyday software development, yet a recurring body of evidence shows that the code they produce is frequently insecure [4,5,7]. This has prompted two broad responses: steering models towards secure code and detecting vulnerabilities after generation so they can be repaired. SVEN [10] and SafeCoder [11] exemplify the former using SFT, which raises the proportion of secure completions from 59.1% to 92.3% without degrading functional correctness. For detection, comparative studies report that LLMs and static application security testing (SAST) tools occupy complementary regimes: SAST tools achieve low detection rates but relatively few false positives, whereas LLMs detect far more vulnerabilities at the cost of higher false-positive rates [6,12]. Prompting-augmented frameworks such as DLAP [13] attempt to combine these strengths in a feedback loop setting.

Most SAST-versus-LLM comparisons reduce the task to binary classification [12,14], which is too coarse to drive targeted repair [7]. Our prior work, ALPHA [6], addressed this limitation by introducing the first function-level Python benchmark that scores predictions at CWE-level granularity for both LLMs and SAST tools. Rather than treating every misprediction as equally wrong, ALPHA exploits the hierarchical structure of the MITRE CWE taxonomy [15]. This taxonomy-aware perspective aligns with a broader line of hierarchy-aware learning, in which penalising errors by their semantic distance within a class hierarchy yields “better mistakes” than a flat loss [16].

RL has become the standard mechanism for aligning LLMs with objectives that are not directly differentiable. Building on preference-based RL [17] and the RL-from-human-feedback pipeline popularised by InstructGPT [18], most implementations optimise a reward using proximal policy optimisation [19]. GRPO [20] simplifies this recipe by discarding the learned value network and instead estimating the baseline from the relative scores of a group of sampled outputs, thereby reducing the computational overhead.

Applying RL to code generation is now well established—CodeRL [21], PPOCoder [22], RLTF [23], and StepCoder [24] all treat execution, unit-test, or compiler feedback as a reward to improve functional correctness. Reward-driven approaches aimed specifically at security, however, remain scarce. The closest examples train models to repair vulnerabilities using PPO with syntactic and semantic similarity rewards [25,26], arguing that ordinary cross-entropy underweights the few security-critical lines relative to the bulk of functional code. Concurrent work has begun to pair online RL with a vulnerability reward model [27]. To our knowledge, none of these approaches use a hierarchy-aware CWE penalty as an RL reward. This is the gap the present work addresses: we take the original ALPHA penalty and optimise it directly as a GRPO reward, testing whether this yields a principled, taxonomy-aware training signal for CWE-specific vulnerability detection.

3. Methodology

3.1. The Training-Corpus Problem and Label Withholding

Function-level, human-annotated Python-CWE datasets are small—SecurityEval [8] contains 121 samples across 69 CWE classes, and the Python subset of SVEN [10] contains 342 samples across four base classes. Larger CWE-labelled corpora exist [9] but trade reliability for scale [28]. Their labels are typically derived from static-analysis tools, whose coverage and precision are known to be poor [6,28,29], making them unsuitable as a source of trustworthy supervision. This scarcity of reliable data creates two distinct problems for training. First, training on a classification objective alone, with no LM signal at all, risks eroding the pretrained backbone’s generative competence and collapsing it into a narrow classifier. Second, independently of data quantity, raw Python code itself provides a weak LM signal since short Python functions contain very little that a code-pretrained backbone has not already seen.

The natural remedy is to supplement the LM and classification head with reasoning and explanations generated by a teacher model that has access to the true label. This, however, introduces label leakage into the classification head—if the explanation names the CWE, the language head can read the answer directly from its input. Our protocol resolves this through label withholding. The teacher (Qwen2.5-Coder-32B-Instruct) receives the code and the true CWE label and generates a mechanistic explanation under an explicit constraint to omit any CWE identifier. This yields two corpora: an SFT corpus, in which the teacher’s output retains the CWE identifier; and a classifier corpus, in which the identifier is withheld so that the classification head predicts under genuine uncertainty. Both corpora share the same schema; the gold label is used only in loss computation.

3.2. Dual-Head Architecture and Supervised Objective

We adopt Qwen2.5-Coder-7B-Instruct as a shared backbone—a decoder-only, open-weight model with a native causal LM head and a published zero-shot ALPHA baseline, enabling direct comparison. We adapt the backbone with LoRA [30] adapters (Appendix C), keeping the base transformer weights frozen throughout. Two heads read from this backbone: a causal LM head that generates a mechanistic vulnerability explanation, and a classification head that predicts a CWE from pooled hidden states. The heads are trained in two sequential stages, each optimising a single objective. We employ two-stage training: In Stage 1, we train the LM head with next-token CE only on the teacher-explanation tokens, yielding the SFT backbone; its role is to regularise the shared representation. In Stage 2, we load the SFT adapter, apply a second LoRA adapter, and attach the classification head to the pooled last-layer hidden states. Only Stage 2’s LoRA adapter and the head are updated under the classification loss. The head is a multi-layer perceptron (MLP) mapping the pooled representation to the $n=69$ CWE classes observed in the training corpus, with three hidden layers of width 2056, 1024, and 512 before the output projection. The LoRA-adapted backbone supplies the features; this deeper head captures non-linear structure in the pooled embeddings. We evaluate both mean and last-token pooling [31] to form the pooled representation. The Stage 2 objective combines a standard CE loss with the hierarchy-aware ${\mathcal{L}}_{\mathrm{ALPHA}}$ term,

$${\mathcal{L}}_{\mathrm{total}}={\mathcal{L}}_{\mathrm{CE}}+\lambda ·{\mathcal{L}}_{\mathrm{ALPHA}},{\mathcal{L}}_{\mathrm{ALPHA}}\left(\theta \right)=\sum _{v\in V_{\mathrm{obs}}}{p}_{\theta }(v\mid x)\widehat{P}(v,c_{true}).$$

(1)

$\widehat{P}(v,c_{true})=P(v,c_{true})/P_{max}$ is the ALPHA graph penalty [6] normalised to $[0,1]$. Both terms are differentiable and read the same class distribution. ${\mathcal{L}}_{\mathrm{ALPHA}}$ is an expected penalty weighting each candidate misclassification by its probability and its taxonomic severity. For $n=69$, ${\mathcal{L}}_{\mathrm{CE}}\in [0,lnn]\approx [0,4.23]$ while normalised ${\mathcal{L}}_{\mathrm{ALPHA}}\in [0,1]$. Therefore, CE would dominate by 3–$4\times$ at $\lambda =1$, so we set $\lambda =3$ to place the terms on a commensurate scale.

A property left unanalysed in the original ALPHA proposal is that ${\mathcal{L}}_{\mathrm{ALPHA}}$ is insufficient to train the classifier on its own in the supervised regime. The two logit gradients (derived in Appendix A) expose why

$${\displaystyle \frac{\partial {\mathcal{L}}_{ALPHA}}{\partial z_k}}=p_{\theta }\left(k\right)[\widehat{P}\left(k\right)-{\mathcal{L}}_{\mathrm{ALPHA}}],{\displaystyle \frac{\partial {\mathcal{L}}_{\mathrm{CE}}}{\partial z_k}}=p_{\theta }\left(k\right)-\mathbf{1}[k=c_{true}].$$

(2)

CE is corrective through the indicator $\mathbf{1}[k=c_{true}]$. It applies an upward push of magnitude $1-p_{\theta }\left(c_{true}\right)$ on the true class, strongest exactly when the model is most wrong and unconditional on the training phase. ${\mathcal{L}}_{\mathrm{ALPHA}}$, on the other hand, is confirmatory. Since $\widehat{P}\left(c_{true}\right)=0$, its push on the true class is $p_{\theta }\left(c_{true}\right)·{\mathcal{L}}_{ALPHA}$, which vanishes if $p_{\theta }\left(c_{true}\right)\to 0$. It goes silent precisely when correction is most needed and favours the true class only while that class already leads (Appendix B). It can extend a lead but not create one.

This pathology is specific to the supervised expected-penalty setting. Under GRPO, the ${\mathcal{L}}_{\mathrm{ALPHA}}$ penalty enters through sampled rollouts and normalised advantages, rather than a belief-weighted expected-penalty gradient, so it retains a corrective signal even when the model is confidently wrong; this is why the ALPHA penalty can drive learning on its own under RL but must be paired with CE under SFT. The classification head assigns probability only to the classes represented in the training corpus; consequently, ${\mathcal{L}}_{\mathrm{ALPHA}}$ can only redistribute probability mass within this fixed simplex. When the deployment distribution is dominated by a class seen only a handful of times in training, hierarchical shaping cannot compensate for a dimension the model has barely learned.

This analysis makes three predictions for SFT:

• First, when the shared representation is weak and the model’s distribution is already committed, ${\mathcal{L}}_{\mathrm{ALPHA}}$ can only rescale existing mass. So, toggling it on or off should leave the score unchanged.
• Second, when the representation is strong, its effect should be small and refining rather than decisive.
• Third, whichever class leads at initialisation should be reinforced, so a biased starting point should collapse toward that class rather than away from it.

3.3. Reinforcement Learning with an ALPHA Reward

SFT encourages adaptation to the training distribution at the expense of the base model’s broader capabilities [32], and sufficient adaptation can lead to overspecialisation [33] and catastrophic forgetting [34]. RL removes these constraints. The model generates a free-form analysis, the predicted CWE is extracted from the text, and the reward is computed from the penalty. We view this as a contextual bandit setting [35], in which the entire response is treated as a single action with a terminal reward, thereby avoiding token-level credit assignment. The state s is the prompt, the action a is the complete response, and the policy ${\pi }_{\theta }\left(a\right|s)$ is the language model factorised autoregressively. The objective is

$$J\left(\theta \right)={\mathbb{E}}_{x\sim D,a\sim {\pi }_{\theta }(·|x)}\left[R(a,c_{true})\right]$$

(3)

For a response a whose extracted CWE is $\widehat{c}$,

$$R(a,c_{true})=1-{\displaystyle \frac{P(\widehat{c},c_{true})}{P_{max}}}\in [0,1]$$

(4)

We optimise $J\left(\theta \right)$ with GRPO [20], chosen for three reasons: it is suited to a single-GPU budget; it requires no separate value network, estimating the advantage baseline from a group of completions for the same prompt; and a Kullback–Leibler (KL) divergence [36] (${\mathbb{D}}_{\mathrm{KL}}$) anchor to a frozen reference policy to keep training stable. For each prompt, we sample G completions, score each, and form group-relative advantages A while skipping any group whose reward standard deviation falls below ${10}^{-6}$,

$$A_{i,t}={\displaystyle \frac{R_i-\mathrm{mean}(R_1,\dots ,R_G)}{\mathrm{std}(R_1,\dots ,R_G)}}.$$

(5)

Although the optimisation problem is formulated as maximising $J\left(\theta \right)$ in Equation (3), the implementation minimises its negation (Equation (6)), consistent with gradient-descent optimisers such as AdamW [37],

$${\mathcal{L}}_i=-(A_i·\overline{log{\pi }_{\theta }}\left(y_i\right|x)-\beta ·{\mathbb{D}}_{\mathrm{KL}}({\pi }_{\theta }\left|\right|{\pi }_{\mathrm{ref}}\left)\right).$$

(6)

The contrast with the supervised penalty is structural. The dual-head gradient scales with $p_{\theta }\left(k\right)$, attenuating updates to low-probability classes, whereas the GRPO policy gradient scales with ${\nabla }_{\theta }log{\pi }_{\theta }\left(y_i\right|x)$, introducing the inverse dependence. This distinction concerns gradient scaling rather than providing a guarantee that reinforcement learning is unconditionally superior at escaping local optima. Additionally, the KL anchor independently constrains policy movement (Section 6.1), and improvement still requires the sampled group to contain a better action. Beyond constraining policy movement and stabilising training, the KL penalty prevents degenerate solutions, such as the model defaulting to the most frequent class regardless of input (e.g., CWE-89, which alone accounts for 54% of SVEN samples).

3.4. Use of AI Tools in Manuscript Preparation

Generative AI tools (Claude Sonnet 4.5 and Claude Sonnet 5, Anthropic, San Francisco, CA, USA) were used during manuscript preparation for two purposes: First, to assist in drafting and refining explanatory text, including identifying inconsistencies between reported claims and their supporting data. Second, to assist with debugging during implementation. The authors reviewed and verified all AI-assisted output and take full responsibility for the accuracy and integrity of the analysis and content.

4. Experimental Setup

4.1. Datasets

SecurityEval [8] (121 Python functions, 69 CWEs) is the primary training source; a 70/30 train–test split was carved out before any corpus generation for SFT training, so the teacher never saw the held-out samples, which are reserved for in-distribution evaluation. CVEfixes [9] (1548 samples, 71 classes) supplies mined patched functions that augment the RL pool. SVEN [10] (342 samples; CWE-022, CWE-078, CWE-079, CWE-089) is the out-of-distribution (OOD) benchmark: its four classes are common enough to be well represented in the base model’s pretraining corpora but are thinly covered by the training corpora used here. Table 1 summarises sizes, class counts, and overlap.

Table 1. “Shared” counts distinct CWEs also present in SecurityEval. CVEfixes^⋆ denotes CVEfixes restricted to the 23 CWEs shared with SecurityEval; ${\mathcal{D}}_{\mathrm{comb}}$ additionally applies deduplication and length filtering to CVEfixes^⋆ before union with SecurityEval (7).

Corpus	Role	Samples	CWEs	Shared
SecurityEval	SFT source, RL pool	121	69	—
CVEfixes	raw	1548	71	23
CVEfixes⋆	class-filtered	710	23	23
${\mathcal{D}}_{\mathrm{comb}}$, capped	RL corpus	435	69	23
${\mathcal{D}}_{\mathrm{comb}}$, uncapped	RL corpus	584	69	23
SVEN	OOD benchmark	342	4	4

CWE labels are normalised across all three corpora: identifiers are mapped to the integer node form used by the ALPHA taxonomy graph [6]. Labels retain their original abstraction level (Base, Class, or Variant) rather than being collapsed to a common level, since the ALPHA penalty is direction-aware and this distinction separates over-generalisation from over-specification. These three corpora differ in label provenance: SecurityEval and SVEN each carry one hand-annotated CWE per sample; CVEfixes labels are inherited from the National Vulnerability Database (NVD) record of the associated common vulnerabilities and exposures (CVE) via the CVEfixes relational schema and are consequently noisier. A further source of label noise affects the CVEfixes-derived data specifically: NVD sometimes assigns more than one CWE to a single CVE. This affects 53 of the 463 CVEfixes samples (11.45%) that remain after the preprocessing described below. For these samples, we kept the first CWE label returned by the database.

Preprocessing: CVEfixes^⋆ (710 samples, 23 CWEs—Table 1) passes through exact deduplication followed by a 20–400 word length bound. The lower bound removes stubs, bare imports and trivial accessors that carry no vulnerability semantics and contribute negligible gradient signal while the upper bound guarantees every retained snippet fits entirely within the 512-token encoder window alongside its instruction wrapper, so no sample is scored on truncated code. Both are standard for CVEfixes-derived subsets, where audits report substantial duplication and length noise in the raw mining output [9,38,39]. Preprocessing leaves 463 usable Python samples over the 23 shared CWEs, which combine with SecurityEval to form ${\mathcal{D}}_{\mathrm{comb}}$, consisting of 584 samples over 69 CWEs.

$${\mathcal{D}}_{\mathrm{comb}}={\mathcal{D}}_{\mathrm{SecurityEval}}\cup \mathrm{preprocess}\left({\mathcal{D}}_{\mathrm{CVE}}^{\star }\right).$$

(7)

This augmentation deepens classes the classifier already targets rather than adding new ones. The filter is load-bearing: CVEfixes is dominated by CWE-918 (125), CWE-352 (120) and CWE-444 (85), none of which appear in SecurityEval or SVEN, and training GRPO on unfiltered CVEfixes yields the worst SVEN score of any RL configuration (see Section 6.2).

Within the 23 shared classes, the preprocessed CVEfixes’ contribution is uneven—the largest class supplies 74 samples while half supply 6 or fewer—so a few classes would otherwise dominate the RL gradient signal. To test whether this matters, we also build a variant with a per-class ceiling of 31 samples (mean per-class count ≈ 30.87 of the CVEfixes^⋆ pool), which binds on 7 of the 23 classes and reduces the three largest classes’ share of the shared-class mass from 42% to 29%. This gives two variants of ${\mathcal{D}}_{\mathrm{comb}}$: capped and uncapped. The uncapped variant is used throughout the main experiments, with the capped variant reported in the ablation in Section 6.2.

4.2. Policy and Training

The policy and backbone is Qwen2.5-Coder-7B-Instruct [40]; the teacher that generates the explanation corpora is Qwen2.5-Coder-32B-Instruct [40]. All fine-tuning is parameter-efficient (PEFT) [41] via LoRA [30]. Unless stated otherwise, all reported scores represent cumulative ALPHA values. In selected instances, we report mean per-sample scores and label them accordingly. We report two complementary settings:

• Out-of-distribution: Training on SecurityEval or ${\mathcal{D}}_{\mathrm{comb}}$ and evaluation on SVEN—the deployment-realistic regime.
• In-distribution: A $70/30$ SecurityEval split, with both teacher corpora generated only from the training portion so that the held-out samples are unseen.

The in-distribution setting is diagnostic, not a target in its own right—it isolates whether the dual-head loss behaves as the gradient analysis in Section 3.2 predicts. Deployment is where a trained model meets inputs unlike its training distribution, so distribution shift is this work’s actual concern. We evaluate SFT and the dual-head loss in both regimes because their in-distribution behaviour explains why they fail under shift (Section 5). GRPO is evaluated only under shift: we adopted it because SFT showed no distribution-shift gains despite adequate in-distribution performance, so we did not run an in-distribution GRPO evaluation.

GRPO: GRPO sweeps group size $G\in \{8,16\}$, KL regularisation coefficient $\beta \in \{0.02,0.04\}$, and 2–6 epochs. Completions are sampled at temperature 0.8 (Ollama default) during training. Unless otherwise stated, GRPO checkpoints are evaluated under greedy decoding, matching typical deployment practice; Section 5.2 additionally reports a sampled-decoding evaluation of the best configuration for direct statistical comparison with the sampled-decoding baseline. Continued training runs reduce the learning rate from $5\times {10}^{-6}\mathrm{to}2\times {10}^{-6}$ and retain the original base model as the frozen KL reference so that the KL bound accumulates over epochs rather than relative to the most recent checkpoint. All experiments used a single NVIDIA A100 (80 GB) GPU.

5. Results

5.1. Supervised Results: In-Distribution and Under Shift

On the matched $70/30$ SecurityEval split, mean pooling substantially outperforms last-token pooling, and the dual-head marginally outperforms the SFT generative baseline. Table 2 reports the held-out SecurityEval evaluation.

Table 2. In-distribution evaluation on the held-out SecurityEval split (36, 30%). Lower is better.

Condition	Cumulative ALPHA	Mean ALPHA (per Sample)
Qwen2.5-Coder-7B-Instruct Baseline	230.7	6.41
${\mathcal{L}}_{\mathrm{CE}}$ only (last-token pooling)	233.6	6.49
${\mathcal{L}}_{\mathrm{CE}}$ only (mean pooling)	145.2	4.03
${\mathcal{L}}_{\mathrm{CE}}+{\mathcal{L}}_{\mathrm{ALPHA}}$ (last-token pooling)	233.6	6.49
${\mathcal{L}}_{\mathrm{CE}}+{\mathcal{L}}_{\mathrm{ALPHA}}$ (mean pooling)	141.6	3.93

Bold indicates the best (lowest) cumulative ALPHA score.

Three observations follow: first, pooling is the dominant factor—mean pooling reduces the penalty by roughly 40% relative to last-token pooling. Second, CE fine-tuning alone adds value over parsing the model’s own free-form output when the distribution matches. The mean-pooled classification head trained with only ${\mathcal{L}}_{\mathrm{CE}}$ reaches 145.2, improving on the zero-shot baseline of 230.7. Third, adding ${\mathcal{L}}_{\mathrm{ALPHA}}$ to the CE objective improves the mean-pooled score only marginally ($145.2\to 141.6$), a difference of 0.1 per sample over 36 samples, which we do not claim as significant and on which we rest no conclusion. This marginal improvement is the second prediction realised in Section 3.2—once CE has fixed the ranking, the ${\mathcal{L}}_{\mathrm{ALPHA}}$ term that only redistributes residual mass has limited scope to alter it. Under last-token pooling, switching the penalty term ${\mathcal{L}}_{\mathrm{ALPHA}}$ on or off changes nothing (233.6 in both cases), consistent with the gradient analysis. This is the first prediction of Section 3.2 realised—because the ${\mathcal{L}}_{\mathrm{ALPHA}}$ gradient scales with $p_{\theta }\left(k\right)$, it can only amplify the distribution that the representation already produces; under a weak (last-token) representation, that distribution is fixed, and the penalty term has no distribution to reshape, and the score is identical.

Re-running the supervised mechanisms with evaluation on SVEN reverses the picture (Table 3). Both SFT and the dual-head classifier exceed the 7B zero-shot baseline’s cumulative penalty of 753.2.

Table 3. Out-of-distribution evaluation (SecurityEval → SVEN). Cumulative ALPHA, lower is better.

Condition	Cumulative ALPHA
Qwen2.5-Coder-7B-Instruct Baseline	753.2
${\mathcal{L}}_{\mathrm{CE}}$ only (last-token pooling)	2123.5
${\mathcal{L}}_{\mathrm{CE}}$ only (mean pooling)	3063.1
${\mathcal{L}}_{\mathrm{CE}}+{\mathcal{L}}_{\mathrm{ALPHA}}$ (last-token pooling)	2312.7
${\mathcal{L}}_{\mathrm{CE}}+{\mathcal{L}}_{\mathrm{ALPHA}}$ (mean pooling)	2799.2

Bold indicates the best (lowest) cumulative ALPHA score.

Two points are notable: first, the ordering of the pooling strategies inverts relative to Table 2—mean pooling, the strongest in-distribution, is the weakest under shift, because a representation tuned to the training distribution does not transfer when the test distribution is concentrated elsewhere. Second, under mean pooling, the effect of ${\mathcal{L}}_{\mathrm{ALPHA}}$ is also larger under shift than in-distribution—an 8.6% reduction (3063.1 → 2799.2)—though this does not reverse the classifier’s overall degradation, and the effect is not consistent across pooling strategies (${\mathcal{L}}_{\mathrm{ALPHA}}$ increases the penalty by 8.9% under last-token pooling).

The classifier conditions are the weakest overall—it can predict only one of its 69 training classes, and SVEN’s dominant class appeared twice in training. Table A3 reports per-class recall for each pooling-loss combination on SVEN’s four evaluation classes, together with each condition’s most frequently predicted class overall. Last-token pooling never recovers the correct class on any of the four SVEN CWEs regardless of loss, collapsing almost entirely onto CWE-020; mean pooling performs unevenly across classes and, like last-token pooling, defaults most often to a class absent from SVEN’s evaluation taxonomy. This pattern is sharpest on SVEN’s dominant class CWE-89—CE with mean pooling predicts it correctly in only 2 cases, most often predicting CWE-522 instead; CE with last-token pooling never predicts CWE-89 correctly on this subset. This is the third prediction in Section 3.2 realised—with only two training instances of the dominant class, $p_{\theta }\left(\mathrm{CWE}-89\right)$ begins negligible, so the penalty gradient towards it is negligible and shrinks further as the model commits elsewhere. The loss cannot lift a class the representation has not surfaced.

5.2. GRPO Under Distribution Shift

Across all GRPO-initialised base backbone configurations in Table 4, every configuration improves on baseline; under a null hypothesis of no directional effect, this outcome has probability ${\left(0.5\right)}^6=0.0156$ (excluding checkpoints continued from an earlier epoch of the same run). The exceptions, which do not improve, are policies initialised from the SFT checkpoint and policies trained on unfiltered out-of-domain CVEfixes data; both are explained in Section 6.

Table 4. All GRPO configurations on SVEN, reporting cumulative ALPHA. Lower is better. “Cont.” marks continued training from the preceding checkpoint. * signifies the capped training corpus.

Backbone	G	$\mathit{\beta }$	Epochs	Training Data	SVEN ↓	$\mathbf{\Delta }$ (%)
7B Baseline	-	-	-	-	753.2	0%
SFT	8	0.04	2	SecurityEval	1450.5	$-92.58\%$
SFT	8	0.04	2	Combined	1377.8	$-82.93\%$
Base	8	0.04	2	CVEfixes	1763.0	$-134.07\%$
Base	8	0.04	2	SecurityEval	690.6	$8.31\%$
Base	8	0.04	4 (cont.)	SecurityEval	671.6	$10.83\%$
Base	8	0.04	6 (cont.)	SecurityEval	658.7	$12.55\%$
Base	8	0.04	2	Combined *	659.7	$12.41\%$
Base	8	0.04	2	Combined	614.1	$18.47\%$
Base	8	0.04	4 (cont.)	Combined	609.7	$19.05\%$
Base	8	0.04	6 (cont.)	Combined	624.1	$17.14\%$
Base	8	0.02	2	Combined	584.3	$22.42\%$
Base	8	0.02	4 (cont.)	Combined	571.4	$24.14\%$
Base	16	0.04	2	Combined	570.1	$24.31\%$
Base	16	0.04	4 (cont.)	Combined	708.2	$5.97\%$
Base	16	0.02	2	Combined	546.2	$27.48\%$
Base	16	0.02	3 (cont.)	Combined	543.1	27.89%
Base	16	0.02	4 (cont.)	Combined	642.1	$14.75\%$

↓ indicates that lower values correspond to better (less vulnerable) performance. Bold indicates the best (lowest) cumulative ALPHA score.

To establish a result that does not depend on any test-informed choice, we first report the pre-specified, default configuration (base backbone; $G=8,\beta =0.04$, 2 epochs; SecurityEval training data), repeated across three independent training seeds (42, 43, 44) to quantify variability from initialisation, data order, and policy optimisation, distinct from the decoding-time variance reported below. This configuration reduces SVEN’s ALPHA penalty to $651.8\pm 33.8$ across three seeds, a 13.5% reduction relative to baseline (per-seed reductions 8.3%, 16.4%, 15.7%), equivalent to 2.57 baseline standard deviations (SD) below the baseline mean (Table 5). All three seeds individually improve on baseline.

Table 5. Training-seed variability for the default (base backbone; $G=8,\beta =0.04$, 2 epochs; SecurityEval training data) and best-observed GRPO (base backbone; $G=16,\beta =0.02$, 3 epochs; ${\mathcal{D}}_{\mathrm{comb}}$ training data) configurations on SVEN, evaluated under greedy decoding. Cumulative ALPHA, lower is better.

Configuration	Seed 42	Seed 43	Seed 44	Mean	SD
Default	690.6	629.4	635.3	651.8	33.8
Best	543.1	615.1	583.0	580.4	36.1

Hyperparameter and data-recipe selection, evaluated against SVEN, identify an additional configuration (base backbone; combined data ${\mathcal{D}}_{\mathrm{comb}};G=16,\beta =0.02$; epoch 3) that further improves on this result. Repeated across the same three seeds, it reduces the SVEN penalty to $580.4\pm 36.1$, a 22.9% reduction (per-seed reductions 27.9%, 18.3%, 22.6%), equivalent to 4.37 baseline SD below the baseline mean (Table 5). Because the configuration was selected using SVEN performance, the seed-42 result of 543.1 and its associated 27.9% reduction is treated as the best observed value within this configuration’s seed distribution rather than as an unbiased estimate of expected performance. The seed-averaged 22.9% is reported as the more representative estimate.

For the best configuration, since the published baseline was evaluated under sampled decoding, we additionally evaluated matched-sampled decoding (temperature 0.8) from the seed-42 checkpoint across six independent decoding runs (526.0, 541.5, 548.2, 577.8, 586.4, 587.8; mean $561.3\pm 26.1$) and recorded a 25.5% reduction. A Welch’s two-sample t-test against the baseline’s three published runs (mean $753.2\pm 39.5$) gives t = −7.62 (df = 2.9), p = 0.005 (two-tailed); even the weakest of the six sampled runs clears the baseline mean by 4.2 baseline standard deviations. Both the seed-averaged best-configuration score (580.4) and sampled-decoding figures (561.3) lie within, or close to, the 4.5× larger 32B model’s reported performance ($574.7\pm 42.4$).

Because cumulative ALPHA depends on test-set size and composition, Table A3 and Table A4 report standard prediction measures alongside it. The default GRPO configuration achieves 77.5% exact accuracy and 52.5% macro-averaged recall across SVEN’s four classes, against 0.0–8.8% accuracy and 0.0–18.4% macro-averaged recall for the four classifier conditions (Table A4). Ancestor and descendant errors—predictions that are hierarchically close to the true CWE—are rare across every condition (at most 0.3% of predictions); almost all non-exact predictions for the classifier and GRPO are lateral—landing in a sibling branch of the CWE taxonomy rather than a hierarchically nearby one. GRPO’s advantage therefore comes overwhelmingly from converting near-misses into exact matches. GRPO’s max penalty share (3.5%) is driven primarily by outputs indicating no vulnerability, for which no CWE identifier could be extracted from the model’s free-form response (3.2% of predictions)—a category that does not apply to the classifier heads, which always output one of their 69 trained classes via softmax.

A different comparison, between the two GRPO configurations, isolates a separate mechanism. Comparing default and best-observed GRPO at seed 42 shows both achieve identical exact-match counts (265/342), so the reduction from 690.6 to 543.1 does not come from getting more predictions right. It comes from 11 samples moving out of the max penalty category (12 to 1) and into lateral (65 to 76). The best configuration’s advantage over the default therefore comes primarily from fewer max penalty errors among its incorrect predictions, rather than from resolving substantially more cases correctly.

6. Ablations

The GRPO study comprises many configurations; three factors determine the outcome: the reference policy, the training data, and the optimisation hyperparameters, together with the number of epochs.

6.1. Reference Policy: Base Versus SFT Backbone

The largest single effect is the choice of reference policy. Initialising GRPO from the SFT backbone records 1450.5 ALPHA penalty on SVEN; initialising from the base Qwen2.5-Coder-7B-Instruct records $651.8\pm 33.8$ on the same SecurityEval training data—better by more than a factor of two. A likely cause is the KL anchor. If the SFT checkpoint is biased towards any particular class as a result of its teacher-supervised training, the KL term keeps the policy anchored near that biased reference, so RL cannot fully correct the bias even with a correct reward signal—consistent with the pattern in Table 4, where both SFT-initialised configurations (SecurityEval: 1450.5; Combined: 1377.8) score far above every base-initialised configuration.

6.2. Training Data Composition

SecurityEval provides only 121 training samples, so we tested whether enriching the RL training set with CVEfixes improves the policy and whether the composition of that enrichment matters. Training with unfiltered CVEfixes does not help (1763.0, worse than SecurityEval alone at 690.6)—its dominant classes are absent from SVEN, so RL teaches the wrong distribution. Restricting the enrichment to CVEfixes samples whose CWE also appears in SecurityEval (Section 4) improves the score. The ${\mathcal{D}}_{\mathrm{comb}}$ uncapped variant performs substantially better than the capped one at matched settings: 614.1 ($18.47\%$) against 659.7 ($12.41\%$), a gap that persists under continued training (609.7, $19.05\%$ at 4 epochs; 624.1, $17.14\%$ at 6). We therefore use the uncapped corpus throughout the main experiments.

6.3. Group Size and KL Coefficient

Increasing the group size G from 8 to 16 sharpens the advantage estimate, yielding more usable gradient updates per epoch. Lowering $\beta$ from 0.04 to 0.02 increases the reward’s influence relative to the KL anchor. The two effects are approximately additive (Table 6).

Table 6. Hyperparameter ablation (combined data, 2 epochs, base backbone, seed 42). The final column reports the change in cumulative ALPHA relative to the default GRPO configuration; negative values indicate a lower penalty and, hence, better performance.

G	$\mathit{\beta }$	SVEN ↓	vs. $\mathit{G}=8,\mathit{\beta }=0.04$
8	0.04	614.1	0
8	0.02	584.3	$-29.8$
16	0.04	570.1	$-44.0$
16	0.02	546.2	$-67.9$

↓ indicates that lower values correspond to better (less vulnerable) performance.

We did not reserve a validation set for the RL runs: we identified the peak epoch, like the other hyperparameters, by its score on SVEN, the test benchmark. A validation set disjoint from the test data would allow principled early stopping without selecting on the test set; we identify this in Section 9.

The group-size and KL-coefficient effects reported above are drawn from single training runs at each setting. We therefore read the ordering in Table 6 as suggestive of direction rather than as precisely quantified independent contributions, pending a fully reseeded ablation (Section 9).

7. Discussion

The key findings of this study are the contrast between the results in Table 2 and Table 3. The dual-head penalty loss itself is not flawed; rather, the supervised learning mechanisms that support it fail to deliver robust performance. Specifically, SFT learns to imitate a teacher whose labels and vocabulary may embed systematic biases; the dual-head loss delivers the penalty into a classifier whose vocabulary is fixed at training time. Both deliveries are mediated, and each mediator is a point of failure under shift.

GRPO is the only mechanism that improves under shift, and its advantage is consistent with the KL property—anchoring to a base reference allows the policy to reinforce knowledge already present from pretraining.

We have not isolated the contribution of the ALPHA penalty’s hierarchical structure within reinforcement learning; part of the improvement may be attributable to reinforcement learning with a free output vocabulary rather than to the graded penalty specifically. This remains an open question that this study does not resolve (Section 9). We therefore frame the contribution more narrowly: the penalty can serve as an effective training signal when delivered directly, and the delivery mechanism is the primary variable distinguishing success from failure across our experiments.

A related observation from Section 5.2 is suggestive on this point, though it is not a substitute for that baseline. The improvement from the default to the best-observed GRPO configuration (Section 5.2, Appendix D) occurs entirely through samples moving from the max penalty category into lateral, with the exact-match count unchanged. A binary or exact-match reward cannot distinguish between these two outcomes, since neither is an exact match and both would receive identical, zero, reward; only a reward that assigns different values to different degrees of wrongness could provide a gradient favouring this shift. This pattern is consistent with the graded structure of the reward doing genuine work, but it is evaluated only across two configurations of the same graded-reward setup rather than against an actual binary-reward run, and we do not treat it as resolving the question posed above.

8. Threats to Validity

We selected the best-observed GRPO value using the test results; therefore, it is an optimistic estimate. We address this by leading our results with a default configuration selected using no test-informed choices (Section 5). We further address this by repeating training for both configurations across three seeds, confirming the default configuration’s improvement is consistent across all three without any SVEN-informed tuning (Table 5). The best configuration achieves a seed-averaged improvement of 22.9%, while the single best-observed run, obtained with seed 42, achieves a 27.9% reduction; we report both statistics throughout, with the seed-averaged improvement treated as the more representative estimate. The training-seed variability is measured here for two of the six independent configurations in Table 4; the remainder, including the full hyperparameter sweep for Table 6, remain open (Section 6.3).

The in-distribution comparison of ${\mathcal{L}}_{\mathrm{ALPHA}}$ and CE rests on 36 samples and a 0.1-per-sample difference; accordingly, we treat the two losses as comparable in-distribution and draw no conclusion about the difference. Furthermore, the classifier cannot predict unobserved classes and no nearest-ancestor fallback is implemented other than applying max penalty.

9. Limitations and Future Work

The most acute limitation is data sparsity: 121 training samples spread across 69 distinct CWEs (a mean of fewer than two samples per class, with a long tail of singletons) and as few as two samples for SVEN’s most common class. This is the binding constraint on all three mechanisms and the principal motivation for future work—larger labelled corpora augmentation to balance CWE coverage or an inference-time nearest-ancestor fallback for classes the classifier has not observed. Such a fallback is relevant only to a deployed system, which is outside the scope of the present study, and its own viability—reliably predicting the correct ancestor class for an unobserved label—is not established here.

Furthermore, the evaluation is function-level Python on two datasets, evaluated against a single base backbone model family. The central finding, that direct delivery via RL survives distribution shift while mediated supervised delivery does not, is a claim about delivery mechanism rather than about Python, CWE taxonomy, or this model family specifically; we would expect the same asymmetry wherever a supervised objective must pass through a fixed-vocabulary classifier or an imitated teacher, but this is untested. Three axes of generalisation remain open and should not be assumed: other programming languages, where CWE prevalence and code idioms differ substantially from Python; repository or commit-level analysis, where the unit of prediction and available context differ from the isolated functions used here; and CWE distributions substantially different from training class, particularly distributions with a long tail of singleton classes, which our own analysis (Section 5.1, Appendix D) identifies as the regime most damaging to the supervised mechanisms specifically.

Three further extensions follow directly from the analysis in Section 6 and Section 7: a binary-reward GRPO baseline to isolate the contribution of the graded penalty; a validation set disjoint from the test benchmark for principled early stopping; and extending our seed-repetition analysis to the remaining configurations, particularly the hyperparameter sweep underlying Table 6 and possibly to more seeds for tighter confidence intervals.

Beyond these results, the longer-term direction is to deploy the trained policy within an iterative repair pipeline and evaluate remediation success directly. Such deployment would require several safeguards. Predictions and free-form explanations should be verified against the CWE taxonomy and, where appropriate, complementary SAST outputs, with invalid or off-graph predictions rejected or flagged rather than surfaced as actionable. A human reviewer should approve any downstream code change, with the model providing ranked, explained suggestions rather than applying fixes autonomously. Predictions should also remain traceable to the model checkpoint, reward configuration, and supporting taxonomic evidence, given the sensitivity of GRPO outcomes to training choices. Finally, deployment would require maintaining the ALPHA graph as the MITRE CWE taxonomy evolves and treating generated explanations as decision support rather than ground truth. This is consistent with approaches in other safety-critical domains that integrate LLM-generated domain knowledge into structured knowledge-management systems while retaining explicit human validation checkpoints [42].

We also did not directly compare the computational cost of the three delivery mechanisms against their predictive gains. All experiments were run on a single NVIDIA A100 (80 GB) GPU (Section 4), but the mechanisms differ substantially in the compute required to reach their reported result: SFT and the dual-head classifier are each a single fixed-cost training run, whereas the best-performing GRPO policy was selected from a sweep of fifteen configurations over group size, KL coefficient, and epoch count (Table 6), several extended through continued training. We did not record wall-clock time or GPU-hours per configuration, and a direct cost-accuracy comparison across delivery mechanisms, analogous to complexity–accuracy trade-offs examined in comparative machine-learning studies in other engineering domains [43], is left for future work.

10. Conclusions

This paper validated the proposal that the ALPHA penalty could serve as a training signal, and the validation succeeded unexpectedly. The proposed supervised mechanism fails: the expected-penalty loss leaves predictions essentially unchanged in-distribution, and the full supervised pipeline regresses below zero-shot under distribution shift for reasons distinct from the penalty itself—a fixed vocabulary that converts class sparsity into worst-case penalties is directly evidenced in our analysis (Section 5.1); mode collapse and label noise introduced by the teacher-supervision pipeline are plausible contributing mechanisms, though we have not isolated their individual contribution. Reinforcement learning with the penalty as a direct, dense reward eliminates all three limitations by construction and delivers a mean $13.5\%$ improvement under the default configuration across three seeds. A further test-selected configuration delivers a mean $22.9\%$ improvement, bringing the 7B model to parity with one $4.5\times$ its size. The ablations localise the gains: reference-policy choice dominates decisively, data curation comes second, group size and KL strength contribute near-additively, and early stopping is mandatory. This study demonstrates the value of direct delivery under a graded, hierarchy-derived reward but does not, by itself, establish that the hierarchical structure of the reward, rather than free-form generation under RL, is responsible for GRPO’s advantage; disentangling these remains open (Section 9). The design principle for hierarchy-aware training that follows is therefore not that the penalty matters little but that a useful penalty is squandered unless it is delivered directly to the model’s parameters; for CWE prediction under realistic distribution shift, reinforcement learning provides that directness.

Appendix A.1. Notation

• $z_k\in \mathbb{R}$: the raw logit for class k, produced by the linear classification head.
• $S={\sum }_{j=1}^n{e}^{z_j}$: the softmax normalisation constant.
• $p_{\theta }\left(k\right)={\displaystyle \frac{e^{z_k}}{S}}$: the model’s predicted probability for class k after softmax.
• $c_{\mathrm{true}}$: the ground-truth CWE class.
• $\widehat{P}(v,c_{\mathrm{true}})\in [0,1]$: the normalised ALPHA penalty for predicting class v when the true class is $c_{\mathrm{true}}$; a constant with respect to model parameters.
• $\mathbf{1}[·]$: the indicator function, equal to 1 if the condition holds and 0 otherwise.

Appendix A.2. Softmax Derivative

Both derivations rely on the derivative of the softmax output $p_{\theta }\left(v\right)$ with respect to logit $z_k$. We derive this first using the quotient rule.

• Case 1: $v=k$.

$$\begin{array}{cc}\hfill {\displaystyle \frac{\partial p_{\theta }\left(k\right)}{\partial z_k}}=& {\displaystyle \frac{\partial }{\partial z_k}}{\displaystyle \frac{e^{z_k}}{S}}\hfill \\ \hfill =& {\displaystyle \frac{e^{z_k}·S-e^{z_k}·e^{z_k}}{S^2}}\hfill \\ \hfill =& {\displaystyle \frac{e^{z_k}}{S}}-{\displaystyle \frac{e^{z_k}}{S}}·{\displaystyle \frac{e^{z_k}}{S}}\hfill \\ \hfill =& p_{\theta }\left(k\right)-p_{\theta }{\left(k\right)}^2\hfill \end{array}$$

(quotient rule)

$$\begin{array}{cc}\hfill \hspace{1em}\hspace{1em} =& p_{\theta }\left(k\right)\left(1-p_{\theta }\left(k\right)\right).\hfill \end{array}$$

(A1)

The numerator step uses ${\displaystyle \frac{\partial }{\partial z_k}}{e}^{z_k}=e^{z_k}$ and ${\displaystyle \frac{\partial S}{\partial z_k}}=e^{z_k}$ (since only the k-th term in S depends on $z_k$).

• Case 2: $v\ne k$.

$$\begin{array}{cc}\hfill \hspace{1em} \hspace{0.25em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em} {\displaystyle \frac{\partial p_{\theta }\left(v\right)}{\partial z_k}}=& {\displaystyle \frac{\partial }{\partial z_k}}{\displaystyle \frac{e^{z_v}}{S}}\hfill \\ \hfill =& {\displaystyle \frac{0·S-e^{z_v}·e^{z_k}}{S^2}}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}(\mathrm{quotient} \mathrm{rule};{\displaystyle \frac{\partial }{\partial z_k}}{e}^{z_v} = 0)\hfill \\ \hfill =& -{\displaystyle \frac{e^{z_v}}{S}}·{\displaystyle \frac{e^{z_k}}{S}}\hfill \end{array}$$

$$\begin{array}{cc}\hfill \hspace{0.25em}\hspace{0.25em} =& -p_{\theta }\left(v\right)·p_{\theta }\left(k\right).\hfill \end{array}$$

(A2)

Appendix A.3. Gradient of ${\mathcal{L}}_{\mathrm{CE}}$

$${\mathcal{L}}_{\mathrm{CE}}=-log{p}_{\theta }\left(c_{\mathrm{true}}\right)=-log{\displaystyle \frac{e^{z_{c_{\mathrm{true}}}}}{S}}=-z_{c_{\mathrm{true}}}+logS$$

Derivative of $logS$ with respect to $z_k$.

$$\begin{array}{c}{\displaystyle \frac{\partial logS}{\partial z_k}}={\displaystyle \frac{\partial logS}{\partial S}}·{\displaystyle \frac{\partial S}{\partial z_k}}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{1em}\hspace{0.25em}\end{array}$$

(chain rule)

$$\begin{array}{cc}\hfill =& {\displaystyle \frac{1}{S}}·{\displaystyle \frac{\partial S}{\partial z_k}}\hfill \\ \hfill =& {\displaystyle \frac{e^{z_k}}{S}}\hfill \\ \hfill =& p_{\theta }\left(k\right)\mathrm{or}softmax\left(k\right).\hfill \end{array}$$

(only the k-th term survives)

• Case 1: $k\ne c_{\mathrm{true}}$.

The term $-z_{c_{\mathrm{true}}}$ does not depend on $z_k$, so,

$$\begin{array}{cc}\hfill {\displaystyle \frac{\partial {\mathcal{L}}_{\mathrm{CE}}}{\partial z_k}}& ={\displaystyle \frac{\partial }{\partial z_k}}(-z_{c_{\mathrm{true}}})+{\displaystyle \frac{\partial logS}{\partial z_k}}\hfill \\ \hfill & =0+p_{\theta }\left(k\right)\hfill \\ \hfill & =p_{\theta }\left(k\right).\hfill \end{array}$$

Since $p_{\theta }\left(k\right)>0$, the gradient is positive, and the logit is pushed down.

• Case 2: $k=c_{\mathrm{true}}$.

Now $-z_{c_{\mathrm{true}}}$ does depend on $z_k$, contributing $-1$:

$$\begin{array}{cc}\hfill {\displaystyle \frac{\partial {\mathcal{L}}_{\mathrm{CE}}}{\partial z_k}}& ={\displaystyle \frac{\partial }{\partial z_k}}(-z_{c_{\mathrm{true}}})+{\displaystyle \frac{\partial logS}{\partial z_k}}\hfill \\ \hfill & =-1+p_{\theta }\left(k\right)\hfill \\ \hfill & =p_{\theta }\left(k\right)-1.\hfill \end{array}$$

Since $p_{\theta }\left(k\right)\in (0,1)$, the gradient is negative, and the logit is pushed up.

• Unified form:

Both cases combine into

$${\displaystyle \frac{\partial {\mathcal{L}}_{\mathrm{CE}}}{\partial z_k}}=p_{\theta }\left(k\right)-\mathbf{1}[k=c_{\mathrm{true}}].$$

(A3)

This gradient is always well-defined and is strongest early in training when $p_{\theta }\left(c_{\mathrm{true}}\right)$ is small, providing the largest upward push to the correct logit precisely when the model is most uncertain.

Appendix A.4. Gradient of ${\mathcal{L}}_{\mathrm{ALPHA}}$

$${\mathcal{L}}_{\mathrm{ALPHA}}\left(\theta \right)=\sum _{v\in V_{\mathrm{obs}}}{p}_{\theta }(v\mid x)\widehat{P}(v,c_{\mathrm{true}}).$$

The penalties $\widehat{P}(v,c_{\mathrm{true}})$ are constants with respect to $\theta$; only $p_{\theta }\left(v\right)$ carries dependence on $z_k$.

$$\begin{array}{cc}\hfill {\displaystyle \frac{\partial {\mathcal{L}}_{\mathrm{ALPHA}}}{\partial z_k}}& =\sum _v\widehat{P}(v,c_{\mathrm{true}})·{\displaystyle \frac{\partial p_{\theta }\left(v\right)}{\partial z_k}}.\hfill \end{array}$$

Split the sum into the $v=k$ term and all $v\ne k$ terms, then substitute the softmax derivatives from (A1) and (A2):

$$\begin{array}{cc}\hfill {\displaystyle \frac{\partial {\mathcal{L}}_{\mathrm{ALPHA}}}{\partial z_k}}& =\widehat{P}(k,c_{\mathrm{true}})·p_{\theta }\left(k\right)\left(1-p_{\theta }\left(k\right)\right)+\sum _{v\ne k}\widehat{P}(v,c_{\mathrm{true}})·\left(-p_{\theta }\left(v\right)p_{\theta }\left(k\right)\right).\hfill \end{array}$$

Factor out $p_{\theta }\left(k\right)$,

$$\begin{array}{cc}\hfill & =p_{\theta }\left(k\right)\left[\widehat{P}(k,c_{\mathrm{true}})\left(1-p_{\theta }\left(k\right)\right)-\sum _{v\ne k}\widehat{P}(v,c_{\mathrm{true}})p_{\theta }\left(v\right)\right].\hfill \end{array}$$

Expand $\widehat{P}(k,c_{\mathrm{true}})(1-p_{\theta }\left(k\right))$,

$$\begin{array}{cc}\hfill & =p_{\theta }\left(k\right)\left[\widehat{P}(k,c_{\mathrm{true}})-\{\widehat{P}(k,c_{\mathrm{true}})p_{\theta }\left(k\right)+\sum _{v\ne k}\widehat{P}(v,c_{\mathrm{true}})p_{\theta }\left(v\right)\}\right].\hfill \end{array}$$

Observe that

$$\widehat{P}(k,c_{\mathrm{true}})p_{\theta }\left(k\right)+\sum _{v\ne k}\widehat{P}(v,c_{\mathrm{true}})p_{\theta }\left(v\right)=\sum _v\widehat{P}(v,c_{\mathrm{true}})p_{\theta }\left(v\right)={\mathcal{L}}_{\mathrm{ALPHA}}$$

$$\begin{array}{cc}\hfill \therefore {\displaystyle \frac{\partial {\mathcal{L}}_{\mathrm{ALPHA}}}{\partial z_k}}& =p_{\theta }\left(k\right)\left[\widehat{P}(k,c_{\mathrm{true}})-{\mathcal{L}}_{\mathrm{ALPHA}}\right].\hfill \end{array}$$

(A4)