Learning to Generate Secure Code via Token-Level Rewards

Jiazheng Quan; Xiaodong Li; Bin Wang; Guo An; Like Liu; Degen Huang; Lin Liu; Chengbin Hou

cs.CRFeb 26, 2026

Learning to Generate Secure Code via Token-Level Rewards

Authors: Jiazheng Quan, Xiaodong Li, Bin Wang, Guo An, Like Liu, Degen Huang, Lin Liu, Chengbin Hou

Organizations: Fuyao University of Science and Technology · Xiamen University · Peking University · Huawei

Abstract

Large language models (LLMs) have demonstrated strong capabilities in code generation, yet they remain prone to producing security vulnerabilities. Existing approaches commonly suffer from two key limitations: the scarcity of high-quality security data and coarse-grained reinforcement learning reward signals. To address these challenges, we propose Vul2Safe, a new secure code generation framework that leverages LLM self-reflection to construct high-confidence repair pairs from real-world vulnerabilities, and further generates diverse implicit prompts to build the PrimeVul+ dataset. Meanwhile, we introduce SRCode, a novel training framework that pioneers the use of token-level rewards in reinforcement learning for code security, which enables the model to continuously attend to and reinforce critical fine-grained security patterns during training. Compared with traditional instance-level reward schemes, our approach allows for more precise optimization of local security implementations. Extensive experiments show that PrimeVul+ and SRCode substantially reduce security vulnerabilities in generated code while improving overall code quality across multiple benchmarks.

Figures & tables

Figure 1: Our Methodology. (1) Vul2Safe transforms real-world vulnerable code into high-quality secure repair data with diverse implicit prompts. (2) PrimeVul+ ranks the samples based on four metrics and classifies them into three progressively difficult, curriculum-style tasks. (3) SRCode first applies SFT on detection and repair tasks, and then performs RL on secure code generation with token-level rewards (TLR) for fine-grained safety optimization.

w/o RL Models
Method	Size	CWEval		CodeLMSec	CyberSecEval
Method	Size	FS@1(C/C++)	FS@1(Python)	Sec.Rate	Pass.Rate
Qwen series models
Qwen2.5-Coder-Instruct	7B	28.9	45.0	75.9	61.9
Qwen2.5-Coder-Instruct [SafeCoder*]	7B	19.2	43.5	87.2	70.8
Qwen2.5-Coder-Instruct [PrimeVul+, ours]	7B	$\underline{32.7}^{{\color[rgb]{0,0.6992,0}\uparrow\,3.8\%}}$	$47.8^{{\color[rgb]{0,0.6992,0}\uparrow\,2.8\%}}$	$86.0^{{\color[rgb]{0,0.6992,0}\uparrow\,10.1\%}}$	$70.3^{{\color[rgb]{0,0.6992,0}\uparrow\,8.4\%}}$

Table 1: Main Results. We present the results under the experimental setup described in Section 4.1. FS@1 , i.e., func-sec@1, evaluates both the correctness and security of the code; Sec.Rate and Pass.Rate both represent the proportion of securely implemented code out of the total samples; the higher is better. [PrimeVul+] indicates the model trained using our dataset; [SafeCoder] indicates the model provided by the authors; [SafeCoder*] indicates the model trained by us using the authors’ dataset and setting. We highlight the best and second-best results, and calculate the improvements of our method compared to the baselines without security optimizations.

Qwen2.5-Coder-3B-Instruct
Method	CWEval	CodeLMSec	CyberSecEval
	FS@1(C/C++)	Sec.Rate	Pass.Rate
Base LLM	21.2	73.4	60.9
DiSCo	15.4	76.4	58.4
SafeCoder	13.5	76.7	65.8
SRCode ( Ours )	23.1	79.0	65.4

Table 2: Comparison of security-enhanced code generation methods. The best results are in bold.

Method			CWEval	CodeLMSec	CyberSecEval
SFT	RL	TLR	FS@1	Sec.Rate	Pass.Rate
✔	✔	✔	38.5	88.7	74.8
✗	✔	✔	34.6	86.7	67.3
✔	✗	✗	32.7	85.9	68.8
✔	✔	✗	30.8	83.9	66.8
✗	✔	✗	26.9	84.5	65.8

Table 3: Ablation study on Qwen2.5-Coder-7B-Instruct. SFT refers to supervised fine-tuning on the PrimeVul+, RL refers to reinforcement learning using PPO, and TLR refers to the token-level rewards.

Figure 2: Comparison of Sampling Efficiency Gaps ( $\Delta_{SE}$ ). A smaller $\Delta_{SE}$ indicates that the RL-trained model is closer to the theoretical performance upper bound of the base model under large-k sampling.

Model	HumanEval Pro	MBPP Pro
Qwen2.5-Coder-3B-Instruct	59.1	55.0
w/ SRCode	$60.4^{{\color[rgb]{0,0.6992,0}\uparrow\,1.3\%}}$	$55.6^{{\color[rgb]{0,0.6992,0}\uparrow\,0.6\%}}$
Qwen2.5-Coder-7B-Instruct	65.9	64.8
w/ SRCode	$69.5^{{\color[rgb]{0,0.6992,0}\uparrow\,3.6\%}}$	64.8
DeepSeek-Coder-6.7B-Instruct	55.6	57.1
w/ SRCode	55.6	$57.9^{{\color[rgb]{0,0.6992,0}\uparrow\,0.8\%}}$

Table 4: Evaluation of general code generation capability for the models with SRCode (i.e., w/SRCode).

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Method	CodeLMSec	HumanEval Pro
Method	Sec.Rate	pass@1
SRCode	88.7	69.5
w/o repair task	83.2	68.1
w/o progressive order	85.5	66.6
w/o explicit order	86.2	69.1

Appendix

Table 5: Ablation study for the hierarchical data design. We evaluate the effectiveness of curriculum data design on a security benchmark and a general code generation benchmark. w/o repair task denotes removing the intermediate-difficulty vulnerability repair task; w/o progressive order indicates breaking the progressive difficulty ordering of tasks; w/o explicit order refers to randomly mixing tasks of different difficulty levels without an explicit curriculum structure.

Figure 3: Vulnerability distribution on CodeLMSec across different methods. The number of High, Medium, and Low severity vulnerabilities are reported, representing high-risk, medium-risk, and low-risk security issues, respectively. Each subplot is for a specific base model. Lower vulnerability counts indicate safer model behavior.

Category	Prompt
Vul2Safe framework	Generate Security Patch Code (Listing 3 )
Vul2Safe framework	LLM Self-Reflection (Listing 4 )
PrimeVul+ Dataset	Vulnerable Code Detection (Listing 5 )
	Code Repair (Listing 6 )
	Secure Code Generation (Listing 7 )
RL Training	Teaching Model Reward (Listing 8 )

Appendix

Table 6: Overview of prompt implementations.

Secure Code Generation

Secure code generation studies how to produce code that meets functional requirements while avoiding security vulnerabilities.

See all 48 papers

Enhancing Reliability in LLM-Based Secure Code Generation
May 22, 2026Mohammed F. Kharma, Mohammad Alkhanafseh, Ahmed Sabbah +1LLM Evaluation CoT Reasoning
the Department of Computer Science, Birzeit University, Ramallah, Palestine · the Department of Computer Science, University of Central Florida, Orlando, FL 32816 USA
Learn from Your Mistakes: Tree-like Self-Play for Secure Code LLMs
Jun 2, 2026Wenqi Chen, Ziyan Zhang, Bin Wang +3RL for Code Generation LLM Self-Correction
University of Electronic Science and Technology of China, Chengdu, China · Peking University, Beijing, China · University of Science and Technology of China, Hefei, China +2
DeepGuard: Secure Code Generation via Multi-Layer Semantic Aggregation
Apr 10, 2026Li Huang, Zhongxin Liu, Yifan Wu +6Secure Code Generation LLM Code Generation
Chongqing University · The State Key Laboratory of Blockchain and Data Security, Zhejiang University · Peking University

May 22, 2026J/K move · Enter open · S save

Enhancing Reliability in LLM-Based Secure Code Generation

Mohammed F. Kharma, Mohammad Alkhanafseh, Ahmed Sabbah, David Mohaisen

the Department of Computer Science, Birzeit University, Ramallah, Palestine · the Department of Computer Science, University of Central Florida, Orlando, FL 32816 USA

LLM Evaluation CoT Reasoning LLM Security Secure Code Generation Prompt Engineering

Open paper AskPDF arXiv

Large language models (LLMs) are widely used for code generation, but their security reliability remains inconsistent across languages and prompting strategies. Existing prompt engineering improves functional correctness but rarely ensures consistent security outcomes. We introduce the Mitigation-Aware Chain-of-Thought (MA-CoT) framework, which embeds task-specific CWE mitigation guidance and language-aware safeguards to reduce recurring vulnerabilities in generated code. We evaluate MA-CoT across three LLMs (gpt-5, claude-4.5, gemini-2.5), three programming languages (C, Java, Python), and four prompting strategies (Vanilla, Zero-shot, CoT, MA-CoT) on a 200-task primary dataset, with external validation on LLMSecEval. Using static analysis with expert validation, MA-CoT reduces total security findings from 92 to 39 (57.6%) on the primary dataset and from 73 to 4 (94.5%) on LLMSecEval. High-severity findings (Blocker + Critical) drop from 90 to 39 (56.7%) and from 45 to 2 (95.6%), respectively. Across both datasets, MA-CoT is the only strategy that consistently improves security reliability; Zero-shot and CoT are less reliable and may increase vulnerability, especially in C. We further introduce a strict layered attribution of vulnerability drivers (language-core vs. stack layers) and show that residual risk concentrates in hardening-oriented patterns (e.g., OS- and toolchain-dependent), motivating secure-by-construction primitives alongside prompting.

Abstract

Secure Code Generation

Explore similar work

Enhancing Reliability in LLM-Based Secure Code Generation

Learn from Your Mistakes: Tree-like Self-Play for Secure Code LLMs

DeepGuard: Secure Code Generation via Multi-Layer Semantic Aggregation