arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2606.03489cs.CRcs.AI
收藏

从错误中学习:面向安全代码LLM的树状自博弈

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

Wenqi Chen, Ziyan Zhang, Bin Wang, Lin Liu, Hengheng Zhang, Zhengsu Chen

AI总结:

提出树状自博弈(TSP)框架,将安全代码生成建模为细粒度序列决策过程,通过构建决策树探索安全与脆弱路径,使模型在关键决策节点自我纠正,显著提升代码安全性并实现跨语言泛化。

摘要:

尽管大型语言模型(LLM)在代码生成方面表现出色,但它们仍然容易复制训练数据中固有的细微但关键的安全漏洞。当前的校准技术,如监督微调(SFT)和强化学习(RL),通常在序列级别应用粗粒度的优化。这种方法往往无法解决安全缺陷的局部性,即单个错误的token选择可能危及整个程序。为了弥合这一差距,我们引入了树状自博弈(TSP),一个将安全代码生成重新定义为细粒度序列决策过程的框架。与盲目最大化似然的标准方法不同,TSP构建了一个决策树,模型在其中探索分支轨迹——同时生成安全的“黄金路径”和易受攻击的变体。通过将代码生成视为自博弈游戏,模型学会严格区分自身的局部错误。这提供了一个密集的、在策略的学习信号,迫使模型在通常出现漏洞的关键决策节点进行自我纠正。我们的实验表明,TSP从根本上提高了模型的可靠性。在Python安全基准测试中,TSP将CodeLlama-7B的通过率(SPR@1)提升至75.8%,显著优于SFT(57.0%)和非结构化自博弈基线。关键的是,TSP引发了鲁棒的分布外泛化:模型不仅将未见类别(CWE)中的漏洞减少了24.5%,还成功将从C/C++学到的安全原则迁移到多种语言,包括Python、Go和JavaScript。这表明TSP不仅仅是记忆补丁,而是内化了抽象的、与语言无关的安全逻辑。

英文摘要:

While Large Language Models (LLMs) excel in code generation, they remain prone to replicating subtle yet critical vulnerabilities endemic to their training data. Current alignment techniques, such as Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL), typically apply coarse-grained optimization at the sequence level. This approach often fails to address the localized nature of security flaws, where a single incorrect token choice can compromise an entire program. To bridge this gap, we introduce Tree-like Self-Play (TSP), a framework that reframes secure code generation as a fine-grained sequential decision process. Unlike standard methods that blindly maximize likelihood, TSP constructs a decision tree where the model explores branching trajectories--generating both secure "golden paths" and vulnerable variants. By treating code generation as a self-play game, the model learns to strictly discriminate against its own localized errors. This provides a dense, on-policy learning signal that forces self-correction precisely at the critical decision nodes where vulnerabilities typically emerge. Our experiments demonstrate that TSP fundamentally enhances model reliability. In Python security benchmarks, TSP boosts CodeLlama-7B's pass rate (SPR@1) to 75.8%, significantly outperforming SFT (57.0%) and unstructured self-play baselines. Crucially, TSP induces robust out-of-distribution generalization: the model not only reduces vulnerabilities in unseen categories (CWEs) by 24.5% but also successfully transfers security principles learned from C/C++ to diverse languages, including Python, Go, and JavaScript. This suggests that TSP does not merely memorize patches, but internalizes abstract, language-agnostic security logic.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑