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

原文作者：公众号“为机器立心”

原文链接:https://mp.weixin.qq.com/s/ZkUNbTfyXY5-zMRpiJxdQg

一、研究背景与动机1. 现状：代码大模型的安全隐患

大语言模型已广泛落地代码生成场景，但模型训练数据包含海量开源漏洞代码，导致其天生会复刻漏洞（如SQL注入、缓冲区溢出、不安全函数调用等）。这类漏洞往往由单个Token/单行代码的局部错误引发，并非整段代码逻辑问题，对工程安全构成严重威胁。

2. 传统对齐方案的核心缺陷

目前主流的代码模型对齐手段（SFT监督微调、RL强化学习）均为序列级粗粒度优化，无法适配代码漏洞“局部触发”的特性，三大硬伤尤为突出：

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

设计一种细粒度、自驱动、高泛化的训练框架：

• 精准定位代码中易触发漏洞的风险节点，在Token层级做针对性优化；
• 利用模型自我博弈自动生成正负样本，摆脱对人工标注的依赖；
• 让模型学习通用安全逻辑，而非记忆特定漏洞修复方案，实现跨语言、跨未知漏洞类型的泛化。

论文提出一系列专属概念，是理解TSP的基础：

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

TSP 是一套迭代式自我博弈+节点级对比学习的训练范式，整体分为三大流程：风险节点标注、树状自我博弈生成、节点级训练更新，同时配套专属损失函数与迭代机制。

3.1 整体流程总览

完整训练循环为迭代闭环，每一轮包含3步：

1. CWE风险节点标注：通过LLM自动识别安全代码中所有潜在漏洞触发点；
2. 树状自我博弈生成：模型（对手玩家）在风险节点处生成漏洞分支，构造「安全路径-漏洞路径」偏好对；
3. TSP损失优化：主模型学习区分正负路径，完成参数更新后，本轮主模型变为下一轮的对手玩家，循环迭代。

基于开源数据集DiverseVul（包含150类CWE、近20万条C/C++代码），筛选出1353条修复后的安全代码作为基准，规避原始漏洞数据集标签噪声问题。

（2）自动化标注管线（核心数据准备）

摒弃人工标注，使用专用标注LLM完成风险节点提取，严格遵循三条规则：

• 因果性：仅标注直接决定代码安全的代码行；
• 唯一性：多行逻辑仅保留最核心的1个风险节点；
• 结构化输出：固定格式输出节点ID、代码行、对应CWE编号、漏洞描述。

1. 随机抽取15%数据由两名安全专家人工交叉标注，科恩卡帕系数κ=0.89，证明风险节点定义统一；
2. LLM自动标注与人工标准比对，κ=0.86，自动化管线精度达到学术研究级标准。

论文先从数学层面论证SFT、RL的不足，再引出TSP优化目标：

1. 代码生成基础建模代码生成是条件语言建模任务，模型逐Token生成序列：$$p_{theta}(y | x)=prod_{t=1}^{T} p_{theta}left(y_{t} | x, y_{
2. SFT损失（全局优化）最大化整段安全代码的对数似然，对所有Token一视同仁，无法聚焦风险节点
3. 传统RL损失（全局奖励）仅在代码生成完毕后计算全局奖励，搭配KL散度约束，无法回溯定位漏洞Token，存在严重信用分配问题。

TSP将模型拆分为两个角色（同源模型），实现纯自驱动对抗学习，无需外部数据：

• 对手玩家（Opponent Player）：上一轮迭代后的模型，固定参数，负责在风险节点生成漏洞分支（自我犯错）；
• 主玩家（Main Player）：当前待优化模型，学习区分安全路径与自我生成的漏洞路径，完成自纠错。

对每一条安全代码（黄金路径），遍历其所有风险节点：

1. 保留风险节点之前的所有代码前缀（保证上下文一致）；
2. 对手玩家从该节点开始生成新序列，得到漏洞分支（自我博弈路径）；
3. 最终得到成对样本：黄金路径自我博弈路径。

优势：正负样本前缀完全一致，仅在风险节点分叉，学习信号高度聚焦，彻底解决传统样本上下文差异大的问题。

（2）TSP损失函数（节点级对比损失）

TSP放弃全局序列优化，仅对风险节点计算损失，分层聚合所有节点误差：

1. 单风险节点损失：采用光滑损失，避免梯度爆炸；
2. 整体TSP损失：对单条代码的所有风险节点损失取平均，实现分层监督：
3. 打分函数：借鉴DPO思路，计算主玩家与对手玩家的似然比值，衡量模型迭代提升幅度，衔接偏好学习与生成任务。

一轮迭代完整流程：

1. 生成阶段：固定对手玩家，批量生成漏洞分支，构建偏好样本集；
2. 训练阶段：主玩家最小化TSP损失，强化风险节点的安全决策能力；
3. 角色更替：训练完成后，主玩家参数赋值给对手玩家，进入下一轮迭代。

该机制让模型持续和“过去的自己”博弈，不断修正越来越隐蔽的漏洞，形成渐进式自进化。

3.5 步骤4：梯度与收敛性分析（理论支撑）

论文从梯度角度证明TSP的优越性：

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

1. 训练框架：基于PyTorch + DeepSpeed ZeRO-2，上下文长度扩展至8192 Token（适配树状长序列）；
2. 超参数：学习率，BF16混合精度，全局批次128；
3. 推理配置：训练阶段自我博弈采用高温度（1.0）保证漏洞样本多样性；评估阶段采用贪心解码（温度=0）保证代码确定性；
4. 推理加速：评估阶段使用vLLM + 4卡张量并行。

• 基座模型：CodeLlama-7B、Qwen2.5-Coder-7B、Qwen2.5-Coder-3B（覆盖主流开源代码模型）；
• 对比基线：原始基座模型、SFT、SafeCoder（专业安全代码模型）、无树结构的普通自我博弈（消融实验核心对照组）。

1. Python安全评测：SecurityEval，指标SPR@1（安全通过率），由CodeQL静态漏洞检测；
2. C/C++安全评测：DiverseVul，指标漏洞总数（数值越低越安全）；
3. 通用代码能力：HumanEval，指标pass@1/pass@10（验证安全加固不会损害基础编码能力）；
4. 泛化测试

• 跨CWE泛化：训练集110类漏洞，测试集40类从未见过的CWE
• 跨语言泛化：仅用C/C++数据训练，测试Python、Go、JavaScript、Ruby。

以CodeLlama-7B为例，核心数据如下：

方法Python SPR@1（安全通过率）C/C++ 漏洞总数HumanEval pass@1（通用编码）原始模型55.0%11534.5SFT57.0%11034.1SafeCoder73.7%-33.9普通自我博弈69.6%10333.3TSP（本文）75.8%9434.0

关键结论：

1. TSP将CodeLlama-7B的安全通过率从SFT的57.0%提升至**75.8%**，超越所有基线；
2. 对比“普通自我博弈”（69.6%），证明树状风险节点结构是核心有效设计，单纯自我博弈效果有限；
3. HumanEval指标几乎无下降，说明安全加固不会造成“能力遗忘”，模型基础编码能力稳定。

该结论在Qwen2.5-Coder系列模型上完全复现，证明TSP具备模型通用性。

RQ2：跨编程语言泛化能力

仅使用C/C++安全数据训练模型，在Python/Go/JS/Ruby多语言场景测试：

• 所有模型中，TSP组的跨语言漏洞数量最低
• 证明TSP并非学习C/C++语法，而是提炼语言无关的通用安全逻辑（如输入校验、函数返回值检查、内存规范等），安全知识可自由迁移。

测试模型对训练集从未出现的漏洞类型的防御能力：

1. CodeLlama-7B：TSP漏洞数27，远低于SFT的50、原始模型的64；
2. 部分基线（如SFT在Qwen2.5-Coder-7B）出现严重过拟合：训练见过的漏洞修复较好，但未知漏洞数量反而暴涨；
3. 漏洞严重度拆解：TSP对高危漏洞的削减效果最显著，有效抵御新型高危安全威胁。

本质区别：SFT是“记忆漏洞补丁”，TSP是“理解安全规则”。

4.3 消融实验补充

1. 风险节点必要性：移除树结构与节点标注，仅保留自我博弈，性能大幅下降，验证细粒度节点监督不可替代；
2. 迭代轮数：多轮自我博弈持续优化，模型逐步修复自身产生的隐蔽漏洞；
3. 节点损失聚合：单节点损失优于全局损失，进一步证明局部优化的价值。

TSP 并非全能，存在明确短板，也是未来研究方向：

1. 擅长局部漏洞，弱于长链路隐式漏洞TSP对局部显式控制流漏洞（CWE-079 XSS、CWE-215敏感信息泄露、缓冲区溢出）效果极佳；但对长距离数据流/内存漏洞（CWE-416野指针、CWE-125越界读取、CWE-690空指针解引用）表现一般。 原因：这类漏洞的“危险决策”和“漏洞爆发点”相隔数十行代码，中间逻辑看似正常，TSP的节点级监督无法追踪跨节点的长距离依赖。
2. 依赖初始风险节点标注质量框架效果上限由标注LLM的精度决定，全新未知漏洞可能无法被识别为风险节点，导致模型漏防。
3. 自我博弈样本难度递减随着模型安全能力提升，自身生成的漏洞样本越来越简单，难以挖掘深层、复杂的漏洞模式，训练后期增益收窄。
4. 模型规模限制实验仅基于3B~7B轻量级模型，TSP在千亿级超大代码模型上的效果、适配性未验证。

总结四大核心创新点：

1. 范式创新：首次将树状决策树+自我博弈引入安全代码对齐，把代码安全问题转化为节点级序列决策问题；
2. 数据创新：模型“向自己的错误学习”，自动生成高质量正负样本，摆脱漏洞数据集稀缺的行业痛点；
3. 优化创新：放弃全局优化，聚焦漏洞风险节点，学习信号精准、训练效率更高；
4. 能力创新：模型习得抽象安全逻辑，而非补丁记忆，实现跨语言、跨未知漏洞的强泛化。

1. 降低安全对齐成本：企业无需大规模人工标注漏洞数据，现有代码模型可通过TSP自迭代完成安全加固，适配中小型研发团队；
2. 提升AI代码助手安全性：集成到GitHub Copilot、IDE代码补全工具中，减少日常开发的漏洞代码输出；
3. 跨语言安全统一加固：一套训练逻辑可服务多编程语言，降低多语言代码模型的安全运维成本。

1. 代码LLM对齐提供新范式：“局部节点监督+自我博弈”可迁移到代码漏洞检测、代码修复、代码规范校验等任务；
2. 拓展自我博弈的应用边界：将AlphaGo类博弈思想从推理任务落地到代码生成领域；
3. 为解决“大模型安全泛化”提供思路：从“记忆样本”转向“学习底层规则”，是通用安全对齐的重要方向。

1. 动态风险节点识别：摆脱静态标注，训练中自动发现新型漏洞节点；
2. 融合数据流/污点分析：结合程序分析技术，追踪长距离隐式漏洞，弥补内存类漏洞短板；
3. 分层自我博弈：设计难度递增的漏洞样本生成策略，避免后期样本难度不足；
4. 超大模型适配：针对千亿级代码LLM优化TSP架构与显存策略。

这篇ICML 2026论文提出的TSP树状自我博弈框架，精准击中了当前代码大模型安全加固的核心痛点：粗粒度优化、数据依赖、泛化薄弱。其核心思想可概括为“让模型在犯错的地方纠正错误”：将代码生成拆解为树状决策节点，定位漏洞触发点，通过模型自我博弈生成漏洞样本，在节点层级做细粒度对比学习。

实验充分证明，TSP在安全通过率、漏洞削减上全面超越SFT、传统RL、普通自我博弈等基线，同时保留模型原有编码能力，且习得可跨语言、跨未知漏洞的通用安全逻辑。尽管在长链路隐式漏洞、动态节点识别上存在局限，但它开辟了“自纠错+细粒度节点对齐”的全新技术路线，是代码大模型安全领域兼具理论创新与工程实用性的优秀工作。

链接：https://arxiv.org/pdf/2606.03489

雷峰网