Learning to Generate Secure Code via Token-Level Rewards

📄 arXiv: 2602.23407v1 📥 PDF

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

分类: cs.CR, cs.AI, cs.SE

发布日期: 2026-02-26

备注: 18 pages, 3 figures

💡 一句话要点

提出Vul2Safe框架，通过token级奖励学习生成安全代码，解决安全数据稀缺和奖励信号粗糙问题。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 安全代码生成 大型语言模型 强化学习 Token级别奖励 代码漏洞修复

📋 核心要点

  1. 现有代码生成模型易产生安全漏洞，且缺乏高质量安全数据和细粒度奖励信号。
  2. Vul2Safe框架利用LLM自我反思构建高质量修复对，并引入token级奖励的SRCode训练框架。
  3. 实验表明，PrimeVul+和SRCode能显著减少代码漏洞，并提升代码质量。

📝 摘要（中文）

大型语言模型（LLMs）在代码生成方面表现出强大的能力，但仍然容易产生安全漏洞。现有方法通常存在两个关键限制：高质量安全数据的稀缺和粗粒度的强化学习奖励信号。为了解决这些挑战，我们提出了Vul2Safe，一种新的安全代码生成框架，它利用LLM的自我反思，从真实世界的漏洞中构建高置信度的修复对，并进一步生成多样化的隐式提示来构建PrimeVul+数据集。同时，我们引入了SRCode，一种新颖的训练框架，它率先在代码安全的强化学习中使用token级奖励，这使得模型能够在训练过程中持续关注和加强关键的细粒度安全模式。与传统的实例级奖励方案相比，我们的方法可以更精确地优化局部安全实现。大量的实验表明，PrimeVul+和SRCode在减少生成代码中的安全漏洞的同时，提高了跨多个基准的整体代码质量。

🔬 方法详解

问题定义：现有代码生成模型虽然强大，但生成的代码常常包含安全漏洞。主要痛点在于：一是缺乏足够的高质量安全代码数据用于训练；二是传统的强化学习方法通常使用实例级别的奖励信号，无法精确指导模型学习细粒度的安全模式，导致模型难以有效避免漏洞。

核心思路：论文的核心思路是利用大型语言模型（LLM）的自我反思能力来生成高质量的安全代码修复数据，并设计一种基于token级别的强化学习奖励机制，从而更精确地指导模型学习安全的代码生成模式。通过这种方式，模型可以更好地理解和避免代码中的潜在漏洞。

技术框架：Vul2Safe框架包含两个主要组成部分：PrimeVul+数据集构建和SRCode训练框架。PrimeVul+数据集通过LLM自我反思，从真实漏洞中生成修复对，并生成隐式提示来增加数据多样性。SRCode训练框架则使用token级别的奖励信号，在强化学习过程中，模型生成的每个token都会根据其安全性获得奖励，从而引导模型生成更安全的代码。

关键创新：该论文的关键创新在于：一是利用LLM的自我反思能力来构建高质量的安全代码数据集，解决了数据稀缺的问题；二是提出了token级别的强化学习奖励机制，相比于传统的实例级别奖励，能够更精确地指导模型学习细粒度的安全模式。这种细粒度的奖励机制是该方法的核心优势。

关键设计：PrimeVul+数据集构建过程中，使用了LLM的自我反思能力来验证修复的正确性，并生成多样化的隐式提示。SRCode训练框架中，token级别的奖励函数设计至关重要，需要根据具体的安全漏洞类型进行调整。此外，强化学习算法的选择（例如，Policy Gradient）以及超参数的设置也会影响最终的性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，PrimeVul+数据集和SRCode训练框架能够显著减少生成代码中的安全漏洞。具体而言，在多个代码生成基准测试中，该方法在降低漏洞数量的同时，还提升了代码的整体质量。相较于现有方法，该方法在安全性方面取得了显著的提升。

🎯 应用场景

该研究成果可应用于自动化代码安全审计、安全的代码生成工具、以及软件开发生命周期的早期安全风险识别。通过集成到IDE或CI/CD流程中，可以帮助开发者编写更安全的代码，降低软件安全漏洞的风险，提升软件系统的整体安全性。

📄 摘要（原文）

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.