Learn from Your Mistakes: Tree-like Self-Play for Secure Code LLMs
2026-06-02 • Cryptography and Security
Cryptography and SecurityArtificial Intelligence
AI summaryⓘ
The authors found that large language models making code often repeat small security mistakes from their training data. To fix this, they created a method called Tree-like Self-Play (TSP), where the model explores many possible code paths, learning to spot and avoid its own errors at very specific points. This helps the model improve its code security much better than older training methods. Their tests showed TSP worked well across different programming languages and could handle new kinds of security problems it hadn't seen before.
Large Language ModelsCode GenerationSecurity VulnerabilitiesSupervised Fine-TuningReinforcement LearningSelf-PlayDecision TreeOut-of-Distribution GeneralizationCommon Weakness Enumeration (CWEs)Multi-language Transfer Learning
Authors
Wenqi Chen, Ziyan Zhang, Bing Wang, Lin Liu, Hengheng Zhang, Zhengsu Chen
Abstract
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.