ICML 2026 | 电子科大：树状自我博弈 TSP，面向安全代码大模型的细粒度自纠错框架

背景

大语言模型（LLM）在代码生成场景中的应用已广泛落地，但其训练数据中蕴含的海量开源代码往往包含安全漏洞（如 SQL 注入、缓冲区溢出、不安全函数调用等）。这导致模型在生成代码时，容易复刻这些漏洞。此类漏洞通常由单个 Token 或单行代码的局部错误引发，而非整段代码的逻辑缺陷，对软件工程安全构成严重威胁。

目前主流的代码模型对齐手段，如监督微调（SFT）和强化学习（RL），均属于序列级的粗粒度优化，无法适配代码漏洞“局部触发”的特性，存在以下四大硬伤：

1. 监督信号稀释：SFT 对整段代码进行全局似然优化，无法区分“安全关键 Token”（如输入校验、安全函数调用）与普通代码，导致漏洞点得不到重点训练。
2. 奖励信号滞后：RL 仅在代码生成完成后给出全局奖励，存在严重的信用分配问题，模型无法定位具体是哪一步 Token 的选择引发了漏洞。
3. 数据依赖严重：高质量、细粒度的漏洞标注数据集稀缺且人工标注成本极高，传统方法依赖外部正负样本，难以规模化。
4. 泛化能力差：模型容易“死记硬背”特定漏洞的修复补丁，难以将安全逻辑迁移到未知的漏洞类型（CWE）或跨编程语言场景中。

针对上述痛点，电子科技大学的研究团队提出了一种细粒度、自驱动、高泛化的训练框架——树状自我博弈（Tree-based Self-Play, TSP）。该框架旨在精准定位代码中易触发漏洞的风险节点，利用模型自我博弈自动生成正负样本，并让模型学习通用的安全逻辑，而非记忆特定的修复方案。

核心内容

TSP 是一套迭代式的自我博弈与节点级对比学习训练范式，其核心在于将代码自回归生成建模为决策树遍历过程，并通过“对手玩家”与“主玩家”的角色交替，实现模型的自纠错与自进化。

1. 基础概念定义

理解 TSP 需掌握以下四个专属概念：

• CWE 风险节点（CWE Risk Node）：代码生成决策树中的关键分叉点，对应代码中极易触发特定漏洞（CWE 通用漏洞枚举）的位置。例如，选择 strcpy（不安全）而非 strncpy（安全）的函数调用行，即为 CWE-120（缓冲区拷贝未校验长度）的风险节点。漏洞并非整段代码问题，而是该节点的单次决策失误。
• 黄金路径（Golden Path）：模型生成的无漏洞安全代码的完整路径，作为训练的正样本。
• 自我博弈路径（Self-Play Path）：模型在风险节点处自主生成的漏洞代码分支（负样本）。该路径的前缀与黄金路径完全一致，仅在风险节点之后发生分叉。
• 树状生成逻辑：将代码生成视为决策树遍历。根节点为输入 Prompt，每一个 Token 对应树的一个分支，完整代码对应从根到叶子的一条路径。

2. 训练流程与机制

TSP 的训练循环是一个迭代闭环，每一轮包含三个主要步骤：

1. CWE 风险节点标注：

   • 摒弃人工标注，使用专用标注 LLM 自动识别安全代码中所有潜在漏洞触发点。
   • 标注遵循三条规则：因果性（仅标注直接决定代码安全的代码行）、唯一性（多行逻辑仅保留最核心的 1 个风险节点）、结构化输出（固定格式输出节点 ID、代码行、对应 CWE 编号、漏洞描述）。
   • 实验验证：基于 DiverseVul 数据集，随机抽取 15% 数据由两名安全专家交叉标注，科恩卡帕系数 $\kappa=0.89$；LLM 自动标注与人工标准比对，$\kappa=0.86$，证明自动化管线精度达到学术研究级标准。

2. 树状自我博弈生成：

   • 模型被拆分为两个角色（同源模型）：对手玩家（Opponent Player）和主玩家（Main Player）。
   • 对手玩家（上一轮迭代后的模型，参数固定）在风险节点处生成漏洞分支，构造「安全路径-漏洞路径」偏好对。
   • 优势：正负样本前缀完全一致，仅在风险节点分叉，学习信号高度聚焦，彻底解决传统样本上下文差异大的问题。

3. TSP 损失优化：

   • TSP 放弃全局序列优化，仅对风险节点计算损失，分层聚合所有节点误差。
   • 单风险节点损失：采用光滑损失，避免梯度爆炸。
   • 整体 TSP 损失：对单条代码的所有风险节点损失取平均，实现分层监督。
   • 打分函数：借鉴 DPO（Direct Preference Optimization）思路，计算主玩家与对手玩家的似然比值，衡量模型迭代提升幅度，衔接偏好学习与生成任务。
   • 角色更替：训练完成后，主玩家参数赋值给对手玩家，进入下一轮迭代。这种机制让模型持续与“过去的自己”博弈，不断修正越来越隐蔽的漏洞，形成渐进式自进化。

3. 数学论证与梯度优势

论文从梯度角度论证了 TSP 相对于 SFT 和传统 RL 的优越性：

• 梯度方差更低：正负样本共享长前缀，梯度估计更稳定，训练不易震荡。
• 优化目标精准：梯度仅来自风险节点，所有优化压力集中在漏洞触发点，不会被普通代码稀释。
• 收敛更稳健：模型不仅保证整段代码语法正确，更在每一个关键决策点实现安全约束。

4. 实验设置与结果

• 基座模型：CodeLlama-7B、Qwen2.5-Coder-7B、Qwen2.5-Coder-3B。
• 对比基线：原始基座模型、SFT、SafeCoder、无树结构的普通自我博弈。
• 评测指标：
  • Python 安全：SecurityEval（SPR@1，基于 CodeQL 静态检测）。
  • C/C++ 安全：DiverseVul（漏洞总数）。
  • 通用能力：HumanEval（pass@1/pass@10）。
  • 泛化测试：跨 CWE 泛化（训练 110 类，测试 40 类未知 CWE）、跨语言泛化（C/C++ 训练，测试 Python/Go/JS/Ruby）。

核心实验结论（以 CodeLlama-7B 为例）：

1. 安全性显著提升：TSP 将 CodeLlama-7B 的安全通过率从 SFT 的 57.0% 提升至 75.8%，超越所有基线。对比“普通自我博弈”（69.6%），证明了树状风险节点结构是核心有效设计。
2. 能力未遗忘：HumanEval 指标几乎无下降，说明安全加固不会造成基础编码能力的遗忘。
3. 跨语言泛化：仅使用 C/C++ 数据训练，TSP 组在 Python/Go/JS/Ruby 多语言场景下的漏洞数量最低，证明其习得了语言无关的通用安全逻辑（如输入校验、内存规范等）。
4. 跨 CWE 泛化：在测试集 40 类从未见过的 CWE 中，CodeLlama-7B 的 TSP 漏洞数为 27，远低于 SFT 的 50 和原始模型的 64。部分基线（如 SFT 在 Qwen2.5-Coder-7B）甚至出现严重过拟合，未知漏洞数量暴涨。

关键要点

• 范式创新：首次将树状决策树与自我博弈引入安全代码对齐，将代码安全问题转化为节点级序列决策问题。