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

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 models that write software code often accidentally copy hidden security vulnerabilities from their training data. Fixing these is hard because traditional training evaluates the whole program at once, missing the exact moment a single bad choice ruins the code. We created Tree-like Self-Play (TSP), a training method that treats code generation like a branching decision tree. Instead of just memorizing the "right" answer, the AI plays a game against itself, generating both secure paths and its own vulnerable mistakes at critical decision points. By comparing these paths, the AI learns to pinpoint and correct its own specific errors. This method significantly boosts the ability of AI to write secure code across multiple programming languages. More importantly, it helps the AI internalize underlying security principles, allowing it to successfully defend against completely new types of software vulnerabilities it has never seen before.