原文作者：公眾號“為機器立心”

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

一、研究背景與動機1. 現狀：代碼大模型的安全隱患

大語言模型已廣泛落地代碼生成場景，但模型訓練數據包含海量開源漏洞代碼，導致其天生會復刻漏洞（如SQL注入、緩沖區(qū)溢出、不安全函數調用等）。這類漏洞往往由單個Token/單行代碼的局部錯誤引發(fā)，并非整段代碼邏輯問題，對工程安全構成嚴重威脅。

2. 傳統(tǒng)對齊方案的核心缺陷

目前主流的代碼模型對齊手段（SFT監(jiān)督微調、RL強化學習）均為序列級粗粒度優(yōu)化，無法適配代碼漏洞“局部觸發(fā)”的特性，三大硬傷尤為突出：

1. 監(jiān)督信號稀釋：SFT對整段代碼做全局似然優(yōu)化，無法區(qū)分“安全關鍵Token”（如輸入校驗、安全函數）和普通代碼，漏洞點得不到重點訓練；
2. 獎勵信號滯后：RL僅在代碼生成完成后給出全局獎勵，存在信用分配問題，模型無法定位具體哪一步Token選擇引發(fā)漏洞；
3. 數據依賴嚴重：高質量漏洞標注數據集稀缺、人工標注成本極高，傳統(tǒng)方法依賴外部正負樣本，難以規(guī)?；?；
4. 泛化能力差：模型容易“死記硬背補丁”，無法將安全邏輯遷移到未知漏洞類型（CWE）跨編程語言場景。

設計一種細粒度、自驅動、高泛化的訓練框架：

• 精準定位代碼中易觸發(fā)漏洞的風險節(jié)點，在Token層級做針對性優(yōu)化；
• 利用模型自我博弈自動生成正負樣本，擺脫對人工標注的依賴；
• 讓模型學習通用安全邏輯，而非記憶特定漏洞修復方案，實現跨語言、跨未知漏洞類型的泛化。

論文提出一系列專屬概念，是理解TSP的基礎：

1. CWE風險節(jié)點（CWE Risk Node）代碼生成決策樹中的關鍵分叉點，對應代碼里極易觸發(fā)特定漏洞（CWE通用漏洞枚舉）的位置。例如選擇strcpy（不安全）而非strncpy（安全）的函數調用行，就是CWE-120（緩沖區(qū)拷貝未校驗長度）的風險節(jié)點。漏洞并非整段代碼問題，而是該節(jié)點的單次決策失誤。
2. 黃金路徑（Golden Path）模型生成無漏洞安全代碼的完整路徑，作為訓練正樣本。
3. 自我博弈路徑（Self-Play Path）模型在風險節(jié)點處自主生成的漏洞代碼分支（負樣本），前綴與黃金路徑完全一致，僅在風險節(jié)點后分叉。
4. 樹狀生成邏輯將代碼自回歸生成建模為決策樹遍歷：根節(jié)點為輸入Prompt，每一個Token對應樹的一個分支，完整代碼對應從根到葉子的一條路徑。

TSP 是一套迭代式自我博弈+節(jié)點級對比學習的訓練范式，整體分為三大流程：風險節(jié)點標注、樹狀自我博弈生成、節(jié)點級訓練更新，同時配套專屬損失函數與迭代機制。

3.1 整體流程總覽

完整訓練循環(huán)為迭代閉環(huán)，每一輪包含3步：

1. CWE風險節(jié)點標注：通過LLM自動識別安全代碼中所有潛在漏洞觸發(fā)點；
2. 樹狀自我博弈生成：模型（對手玩家）在風險節(jié)點處生成漏洞分支，構造「安全路徑-漏洞路徑」偏好對；
3. TSP損失優(yōu)化：主模型學習區(qū)分正負路徑，完成參數更新后，本輪主模型變?yōu)橄乱惠喌膶κ滞婕?，循環(huán)迭代。

基于開源數據集DiverseVul（包含150類CWE、近20萬條C/C++代碼），篩選出1353條修復后的安全代碼作為基準，規(guī)避原始漏洞數據集標簽噪聲問題。

（2）自動化標注管線（核心數據準備）

摒棄人工標注，使用專用標注LLM完成風險節(jié)點提取，嚴格遵循三條規(guī)則：

• 因果性：僅標注直接決定代碼安全的代碼行；
• 唯一性：多行邏輯僅保留最核心的1個風險節(jié)點；
• 結構化輸出：固定格式輸出節(jié)點ID、代碼行、對應CWE編號、漏洞描述。

1. 隨機抽取15%數據由兩名安全專家人工交叉標注，科恩卡帕系數κ=0.89，證明風險節(jié)點定義統(tǒng)一；
2. LLM自動標注與人工標準比對，κ=0.86，自動化管線精度達到學術研究級標準。

論文先從數學層面論證SFT、RL的不足，再引出TSP優(yōu)化目標：

1. 代碼生成基礎建模代碼生成是條件語言建模任務，模型逐Token生成序列：$$p_{theta}(y | x)=prod_{t=1}^{T} p_{theta}left(y_{t} | x, y_{
2. SFT損失（全局優(yōu)化）最大化整段安全代碼的對數似然，對所有Token一視同仁，無法聚焦風險節(jié)點
3. 傳統(tǒng)RL損失（全局獎勵）僅在代碼生成完畢后計算全局獎勵，搭配KL散度約束，無法回溯定位漏洞Token，存在嚴重信用分配問題。

TSP將模型拆分為兩個角色（同源模型），實現純自驅動對抗學習，無需外部數據：

• 對手玩家（Opponent Player）：上一輪迭代后的模型，固定參數，負責在風險節(jié)點生成漏洞分支（自我犯錯）；
• 主玩家（Main Player）：當前待優(yōu)化模型，學習區(qū)分安全路徑與自我生成的漏洞路徑，完成自糾錯。

對每一條安全代碼（黃金路徑），遍歷其所有風險節(jié)點：

1. 保留風險節(jié)點之前的所有代碼前綴（保證上下文一致）；
2. 對手玩家從該節(jié)點開始生成新序列，得到漏洞分支（自我博弈路徑）；
3. 最終得到成對樣本：黃金路徑自我博弈路徑。

優(yōu)勢：正負樣本前綴完全一致，僅在風險節(jié)點分叉，學習信號高度聚焦，徹底解決傳統(tǒng)樣本上下文差異大的問題。

（2）TSP損失函數（節(jié)點級對比損失）

TSP放棄全局序列優(yōu)化，僅對風險節(jié)點計算損失，分層聚合所有節(jié)點誤差：

1. 單風險節(jié)點損失：采用光滑損失，避免梯度爆炸；
2. 整體TSP損失：對單條代碼的所有風險節(jié)點損失取平均，實現分層監(jiān)督：
3. 打分函數：借鑒DPO思路，計算主玩家與對手玩家的似然比值，衡量模型迭代提升幅度，銜接偏好學習與生成任務。

一輪迭代完整流程：

1. 生成階段：固定對手玩家，批量生成漏洞分支，構建偏好樣本集；
2. 訓練階段：主玩家最小化TSP損失，強化風險節(jié)點的安全決策能力；
3. 角色更替：訓練完成后，主玩家參數賦值給對手玩家，進入下一輪迭代。

該機制讓模型持續(xù)和“過去的自己”博弈，不斷修正越來越隱蔽的漏洞，形成漸進式自進化。

3.5 步驟4：梯度與收斂性分析（理論支撐）

論文從梯度角度證明TSP的優(yōu)越性：

1. 梯度方差更低：正負樣本共享長前綴，梯度估計更穩(wěn)定，訓練不易震蕩；
2. 優(yōu)化目標精準：梯度僅來自風險節(jié)點，所有優(yōu)化壓力集中在漏洞觸發(fā)點，不會被普通代碼稀釋；
3. 收斂更穩(wěn)健：模型不僅保證整段代碼語法正確，更在每一個關鍵決策點實現安全約束。

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（專業(yè)安全代碼模型）、無樹結構的普通自我博弈（消融實驗核心對照組）。

1. Python安全評測：SecurityEval，指標SPR@1（安全通過率），由CodeQL靜態(tài)漏洞檢測；
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%），證明樹狀風險節(jié)點結構是核心有效設計，單純自我博弈效果有限；
3. HumanEval指標幾乎無下降，說明安全加固不會造成“能力遺忘”，模型基礎編碼能力穩(wěn)定。

該結論在Qwen2.5-Coder系列模型上完全復現，證明TSP具備模型通用性。

RQ2：跨編程語言泛化能力

僅使用C/C++安全數據訓練模型，在Python/Go/JS/Ruby多語言場景測試：

• 所有模型中，TSP組的跨語言漏洞數量最低
• 證明TSP并非學習C/C++語法，而是提煉語言無關的通用安全邏輯（如輸入校驗、函數返回值檢查、內存規(guī)范等），安全知識可自由遷移。

測試模型對訓練集從未出現的漏洞類型的防御能力：

1. CodeLlama-7B：TSP漏洞數27，遠低于SFT的50、原始模型的64；
2. 部分基線（如SFT在Qwen2.5-Coder-7B）出現嚴重過擬合：訓練見過的漏洞修復較好，但未知漏洞數量反而暴漲；
3. 漏洞嚴重度拆解：TSP對高危漏洞的削減效果最顯著，有效抵御新型高危安全威脅。

本質區(qū)別：SFT是“記憶漏洞補丁”，TSP是“理解安全規(guī)則”。

4.3 消融實驗補充

1. 風險節(jié)點必要性：移除樹結構與節(jié)點標注，僅保留自我博弈，性能大幅下降，驗證細粒度節(jié)點監(jiān)督不可替代；
2. 迭代輪數：多輪自我博弈持續(xù)優(yōu)化，模型逐步修復自身產生的隱蔽漏洞；
3. 節(jié)點損失聚合：單節(jié)點損失優(yōu)于全局損失，進一步證明局部優(yōu)化的價值。

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

1. 擅長局部漏洞，弱于長鏈路隱式漏洞TSP對局部顯式控制流漏洞（CWE-079 XSS、CWE-215敏感信息泄露、緩沖區(qū)溢出）效果極佳；但對長距離數據流/內存漏洞（CWE-416野指針、CWE-125越界讀取、CWE-690空指針解引用）表現一般。 原因：這類漏洞的“危險決策”和“漏洞爆發(fā)點”相隔數十行代碼，中間邏輯看似正常，TSP的節(jié)點級監(jiān)督無法追蹤跨節(jié)點的長距離依賴。
2. 依賴初始風險節(jié)點標注質量框架效果上限由標注LLM的精度決定，全新未知漏洞可能無法被識別為風險節(jié)點，導致模型漏防。
3. 自我博弈樣本難度遞減隨著模型安全能力提升，自身生成的漏洞樣本越來越簡單，難以挖掘深層、復雜的漏洞模式，訓練后期增益收窄。
4. 模型規(guī)模限制實驗僅基于3B~7B輕量級模型，TSP在千億級超大代碼模型上的效果、適配性未驗證。

總結四大核心創(chuàng)新點：

1. 范式創(chuàng)新：首次將樹狀決策樹+自我博弈引入安全代碼對齊，把代碼安全問題轉化為節(jié)點級序列決策問題；
2. 數據創(chuàng)新：模型“向自己的錯誤學習”，自動生成高質量正負樣本，擺脫漏洞數據集稀缺的行業(yè)痛點；
3. 優(yōu)化創(chuàng)新：放棄全局優(yōu)化，聚焦漏洞風險節(jié)點，學習信號精準、訓練效率更高；
4. 能力創(chuàng)新：模型習得抽象安全邏輯，而非補丁記憶，實現跨語言、跨未知漏洞的強泛化。

1. 降低安全對齊成本：企業(yè)無需大規(guī)模人工標注漏洞數據，現有代碼模型可通過TSP自迭代完成安全加固，適配中小型研發(fā)團隊；
2. 提升AI代碼助手安全性：集成到GitHub Copilot、IDE代碼補全工具中，減少日常開發(fā)的漏洞代碼輸出；
3. 跨語言安全統(tǒng)一加固：一套訓練邏輯可服務多編程語言，降低多語言代碼模型的安全運維成本。

1. 代碼LLM對齊提供新范式：“局部節(jié)點監(jiān)督+自我博弈”可遷移到代碼漏洞檢測、代碼修復、代碼規(guī)范校驗等任務；
2. 拓展自我博弈的應用邊界：將AlphaGo類博弈思想從推理任務落地到代碼生成領域；
3. 為解決“大模型安全泛化”提供思路：從“記憶樣本”轉向“學習底層規(guī)則”，是通用安全對齊的重要方向。

1. 動態(tài)風險節(jié)點識別：擺脫靜態(tài)標注，訓練中自動發(fā)現新型漏洞節(jié)點；
2. 融合數據流/污點分析：結合程序分析技術，追蹤長距離隱式漏洞，彌補內存類漏洞短板；
3. 分層自我博弈：設計難度遞增的漏洞樣本生成策略，避免后期樣本難度不足；
4. 超大模型適配：針對千億級代碼LLM優(yōu)化TSP架構與顯存策略。

這篇ICML 2026論文提出的TSP樹狀自我博弈框架，精準擊中了當前代碼大模型安全加固的核心痛點：粗粒度優(yōu)化、數據依賴、泛化薄弱。其核心思想可概括為“讓模型在犯錯的地方糾正錯誤”：將代碼生成拆解為樹狀決策節(jié)點，定位漏洞觸發(fā)點，通過模型自我博弈生成漏洞樣本，在節(jié)點層級做細粒度對比學習。

實驗充分證明，TSP在安全通過率、漏洞削減上全面超越SFT、傳統(tǒng)RL、普通自我博弈等基線，同時保留模型原有編碼能力，且習得可跨語言、跨未知漏洞的通用安全邏輯。盡管在長鏈路隱式漏洞、動態(tài)節(jié)點識別上存在局限，但它開辟了“自糾錯+細粒度節(jié)點對齊”的全新技術路線，是代碼大模型安全領域兼具理論創(chuàng)新與工程實用性的優(yōu)秀工作。

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

雷峰網