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

Abstract: While 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.

• The paper introduces TSP, a framework that refines secure code generation by using decision trees to target identified CWE risk nodes.
• It employs adversarial branching and contrastive loss optimization to mitigate vulnerabilities and achieve higher security pass rates.
• Experimental results demonstrate improved SPR@1 and robust cross-language generalization, reducing vulnerabilities across diverse CWEs.

Tree-like Self-Play for Secure Code LLM Alignment

Overview of the TSP Framework

The propagation of subtle security vulnerabilities in LLM-generated code arises primarily from the coarse-grained granularity of standard alignment strategies, which often target sequence-level alignment without distinguishing the precise origin of security flaws. To address these shortcomings, the paper introduces the Tree-like Self-Play (TSP) framework, reframing secure code generation as a hierarchical, fine-grained sequential decision process. TSP explicitly identifies and operates at Critical Weakness Enumeration (CWE) risk nodes by dynamically constructing and traversing decision trees within code generation. These risk nodes are annotated either via automated LLMs or expert verification, highlighting the loci where isolated token-level choices can induce vulnerabilities.

The core loop consists of three iterative steps: (1) CWE risk node annotation, (2) adversarial tree-like branching generation, and (3) fine-grained, preference-based optimization grounded in contrastive feedback at each identified node. This process ensures that models receive a dense, on-policy learning signal focused precisely on the vulnerable points within the program logic, combining the efficiency of self-play with the semantic rigor of hierarchical supervision.

Figure 1: The TSP paradigm: LLMs locate critical nodes, branch to generate secure and vulnerable paths, and iteratively refine via local self-play.

Methodology and Theoretical Underpinnings

TSP departs from traditional SFT and RL algorithms by utilizing model-centric negative sampling. Secure code generation is formulated as conditional language modeling where the generation is represented as a path in a decision tree, with each node corresponding to a unique sequence prefix. At each CWE-annotated risk node, the framework samples self-play perturbation paths under the historical model, forming pairwise comparisons against the corresponding secure (golden) path.

The optimization objective is a convex, contrastive loss that compels the main player to strictly prefer the golden path to any self-play-generated vulnerable alternative at each risk node. Unlike policy-gradient RL approaches that incur high variance and suffer from sparse program-level rewards, TSP reduces gradient variance by leveraging closely related prefix pairs, yielding more stable optimization and direct remediation of security-critical decisions. The adversarial self-play curriculum autonomously adapts, forcing the model to target and correct its evolving failure modes without reliance on external annotation or static negative samples.

Dataset Construction and Automated Risk Node Annotation

High-resolution risk node annotation is necessary for TSP’s supervision signal. The authors leverage the DiverseVul corpus, covering 18k+ vulnerable and 330k+ non-vulnerable functions across 150 CWEs. To obtain high-fidelity risk node labels, they employ an automated pipeline where LLMs, using structured prompts, pinpoint critical, causal lines within secure code that, if perturbed, yield the classic vulnerability variant. Human expert cross-validation (Cohen's $\kappa = 0.89$) demonstrates the pipeline’s alignment with expert standards, and the LLM annotator achieves $\kappa = 0.86$ with this ground truth.

Experimental Results

TSP is evaluated across Python and C/C++ tasks using CodeLlama-7B, Qwen2.5-Coder-7B, and Qwen2.5-Coder-3B, benchmarking on SecurityEval, DiverseVul, and HumanEval. Metrics include Security Pass Rate (SPR@1), HumanEval pass@k, and total vulnerabilities detected. Standard SFT and current self-play/dataset-augmented RLHF approaches constitute the main baselines.

Figure 2: TSP consistently outperforms SFT, SafeCoder, and standard self-play on security and general code generation benchmarks—yielding improved SPR@1 and lower vulnerability counts.

• Security Uplift: On CodeLlama-7B, TSP increases SPR@1 to 75.8% (vs. SFT’s 57.0%)—the highest among all methods. Vulnerabilities on DiverseVul are minimized, outperforming ablations and vanilla SFT.
• Generalization: TSP yields robust out-of-distribution transfer for both unseen languages and previously unobserved CWEs. Models fine-tuned only on C/C++ generalize their security heuristics to Python, JavaScript, Go, and Ruby. On tests containing 40 CWEs held out from training, TSP achieves a 24.5% reduction in vulnerabilities relative to SFT.
• Severity-Specific Reduction: TSP demonstrates pronounced mitigation of high-severity vulnerabilities for novel CWEs, reducing category counts where baseline methods fail or even regress due to overfitting to training distributions.

  Figure 3: TSP attains the lowest high-severity vulnerability rates on held-out CWEs, indicating improved model generalization to unseen security risks.

Analysis, Implications, and Limitations

TSP introduces a highly data-efficient, structured mechanism for self-improving LLM alignment without external negative annotation. The results indicate that the TSP curriculum forces models to internalize abstract, language-agnostic security principles rather than simply memorizing training data patches. This is supported by empirical robustness to previously unseen vulnerability types and cross-linguistic generalization.

However, TSP's supervision is bounded by the locality of its annotated risk nodes. It achieves strong results on local or explicit control-flow vulnerabilities (e.g., CWE-215, CWE-079) but is less effective on those with long-range, implicit data/control-flow dependencies, such as out-of-bounds or use-after-free cases, due to the limited capacity of token-level annotation to represent dispersed multi-line invariants. The performance ceiling on models >7B remains an open question, and risk node discovery may need to transition towards more dynamic, runtime-aware annotation.

Future Directions

The TSP paradigm can extend to other code quality criteria beyond security, e.g., efficiency or maintainability, by substituting task-specific node annotations. Integrating on-the-fly annotator networks, hybridized with agile human-in-the-loop correction, may further broaden coverage of complex vulnerability classes. Scaling TSP to frontier models and self-discovered risk node detection are logical next milestones in autonomous code LLM alignment.

Conclusion

Tree-like Self-Play provides a generalizable and scalable solution to the fine-grained alignment of code LLMs for secure generation. By reframing the error correction process around the locus of vulnerability emergence, TSP achieves strong security performance and enhanced generalization, addressing both practical and theoretical limitations of prior SFT and RLHF-based approaches. The method stands as an effective blueprint for targeted, self-improving model alignment in security-critical generative tasks.