Learning to Generate Secure Code via Token-Level Rewards
This work is significant for developers and organizations that rely on LLMs for code generation, as it aims to reduce the prevalence of security vulnerabilities in automatically generated code, which is a critical concern for software security.
This paper addresses the problem of large language models generating insecure code by proposing Vul2Safe, a framework that uses LLM self-reflection to create high-confidence repair pairs and a new dataset called PrimeVul+. It also introduces SRCode, a training framework that uses token-level rewards in reinforcement learning to improve the precision of security pattern optimization, leading to a substantial reduction in security vulnerabilities and improved code quality.
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.