ICML 2026 | 電子科大：樹狀自我博弈 TSP，面向安全代碼大模型的細粒度自糾錯框架

原文作者：公眾號“為機器立心”原文鏈接:https://mp.weixin.qq.com/s/ZkUNbTfyXY5-zMRpiJxdQg 一、研究背景與動機1. 現狀：代碼大模型的安全隱患大語言模型已廣泛落地代碼生成場景，但模型訓練數據包含海量開源漏洞代碼，導致其天生會復刻漏洞（如SQL注入、緩衝區溢出、不安全函數調用等）。這類漏洞往往由單個Token/單行代碼的局部錯誤引發，並非整段代碼邏輯問題，對工程安全構成嚴重威脅。2. 傳統對齊方案的核心缺陷目前主流的代碼模型對齊手段（SFT監督微調、RL強化學習）均為序列級粗粒度優化，無法適配代碼漏洞“局部觸發”的特性，三大硬傷尤為突出：監督信號稀釋：SFT對整段代碼做全局似然優化，無法區分“安全關鍵Token”（如輸入校驗、安全函數）和普通代碼，漏洞點得不到重點訓練；獎勵信號滯後：RL僅在代碼生成完成後給出全局獎勵，存在信用分配問題，模型無法定位具體哪一步Token選擇引發漏洞；數據依賴嚴重：高質量漏洞標註數據集稀缺、人工標註成本極高，傳統方法依賴外部正負樣本，難以規模化；泛化能力差：模型容易“死記硬背補丁”，無法將安全邏輯遷移到未知漏洞類型（CWE） 或跨編程語言場景。3. 核心研究目標設計一種細粒度、自驅動、高泛化的訓練框架：精準定位代碼中易觸發漏洞的風險節點，在Token層級做針對性優化；利用模型自我博弈自動生成正負樣本，擺脫對人工標註的依賴；讓模型學習通用安全邏輯，而非記憶特定漏洞修復方案，實現跨語言、跨未知漏洞類型的泛化。二、核心概念定義論文提出一系列專屬概念，是理解TSP的基礎：CWE風險節點（CWE Risk Node）代碼生成決策樹中的關鍵分叉點，對應代碼裡極易觸發特定漏洞（CWE通用漏洞枚舉）的位置。例如選擇strcpy（不安全）而非strncpy（安全）的函數調用行，就是CWE-120（緩衝區拷貝未校驗長度）的風險節點。漏洞並非整段代碼問題，而是該節點的單次決策失誤。黃金路徑（Golden Path）模型生成無漏洞安全代碼的完整路徑，作為訓練正樣本。自我博弈路徑（Self-Play Path）模型在風險節點處自主生成的漏洞代碼分支（負樣本），前綴與黃金路徑完全一致，僅在風險節點後分叉。樹狀生成邏輯將代碼自迴歸生成建模為決策樹遍歷：根節點為輸入Prompt，每一個Token對應樹的一個分支，完整代碼對應從根到葉子的一條路徑。三、TSP 整體技術框架（核心創新）TSP 是一套迭代式自我博弈+節點級對比學習的訓練範式，整體分為三大流程：風險節點標註、樹狀自我博弈生成、節點級訓練更新，同時配套專屬損失函數與迭代機制。3.1 整體流程總覽完整訓練循環為迭代閉環，每一輪包含3步：CWE風險節點標註：通過LLM自動識別安全代碼中所有潛在漏洞觸發點；樹狀自我博弈生成：模型（對手玩家）在風險節點處生成漏洞分支，構造「安全路徑-漏洞路徑」偏好對；TSP損失優化：主模型學習區分正負路徑，完成參數更新後，本輪主模型變為下一輪的對手玩家，循環迭代。3.2 步驟1：數據集構建與風險節點自動標註（1）數據源基於開源數據集DiverseVul（包含150類CWE、近20萬條C/C++代碼），篩選出1353條修復後的安全代碼作為基準，規避原始漏洞數據集標籤噪聲問題。（2）自動化標註管線（核心數據準備）摒棄人工標註，使用專用標註LLM完成風險節點提取，嚴格遵循三條規則：因果性：僅標註直接決定代碼安全的代碼行；唯一性：多行邏輯僅保留最核心的1個風險節點；結構化輸出：固定格式輸出節點ID、代碼行、對應CWE編號、漏洞描述。（3）標註質量驗證隨機抽取15%數據由兩名安全專家人工交叉標註，科恩卡帕係數κ=0.89，證明風險節點定義統一；LLM自動標註與人工標準比對，κ=0.86，自動化管線精度達到學術研究級標準。3.3 步驟2：問題建模與傳統方法數學缺陷論文先從數學層面論證SFT、RL的不足，再引出TSP優化目標：代碼生成基礎建模代碼生成是條件語言建模任務，模型逐Token生成序列：$$p_{\theta}(y | x)=\prod_{t=1}^{T} p_{\theta}\left(y_{t} | x, y_{<t}\right) $$="" $y_{<t}$="" 代表當前token之前的代碼前綴，對應決策樹的路徑前綴。<="" section="">SFT損失（全局優化）最大化整段安全代碼的對數似然，對所有Token一視同仁，無法聚焦風險節點：傳統RL損失（全局獎勵）僅在代碼生成完畢後計算全局獎勵，搭配KL散度約束，無法回溯定位漏洞Token，存在嚴重信用分配問題。3.4 步驟3：樹狀自我博弈（TSP核心機制）TSP將模型拆分為兩個角色（同源模型），實現純自驅動對抗學習，無需外部數據：對手玩家（Opponent Player）：上一輪迭代後的模型，固定參數，負責在風險節點生成漏洞分支（自我犯錯）；主玩家（Main Player）：當前待優化模型，學習區分安全路徑與自我生成的漏洞路徑，完成自糾錯。（1）樣本構造邏輯對每一條安全代碼（黃金路徑），遍歷其所有風險節點：保留風險節點之前的所有代碼前綴（保證上下文一致）；對手玩家從該節點開始生成新序列，得到漏洞分支（自我博弈路徑）；最終得到成對樣本：黃金路徑自我博弈路徑。優勢：正負樣本前綴完全一致，僅在風險節點分叉，學習信號高度聚焦，徹底解決傳統樣本上下文差異大的問題。（2）TSP損失函數（節點級對比損失）TSP放棄全局序列優化，僅對風險節點計算損失，分層聚合所有節點誤差：單風險節點損失：採用光滑損失，避免梯度爆炸；整體TSP損失：對單條代碼的所有風險節點損失取平均，實現分層監督：打分函數：借鑑DPO思路，計算主玩家與對手玩家的似然比值，衡量模型迭代提升幅度，銜接偏好學習與生成任務。（3）迭代更新規則（閉環自進化）一輪迭代完整流程：生成階段：固定對手玩家，批量生成漏洞分支，構建偏好樣本集；訓練階段：主玩家最小化TSP損失，強化風險節點的安全決策能力；角色更替：訓練完成後，主玩家參數賦值給對手玩家，進入下一輪迭代。該機制讓模型持續和“過去的自己”博弈，不斷修正越來越隱蔽的漏洞，形成漸進式自進化。3.5 步驟4：梯度與收斂性分析（理論支撐）論文從梯度角度證明TSP的優越性：梯度方差更低：正負樣本共享長前綴，梯度估計更穩定，訓練不易震盪；優化目標精準：梯度僅來自風險節點，所有優化壓力集中在漏洞觸發點，不會被普通代碼稀釋；收斂更穩健：模型不僅保證整段代碼語法正確，更在每一個關鍵決策點實現安全約束。3.6 實現細節（工程配置）訓練框架：基於PyTorch + DeepSpeed ZeRO-2，上下文長度擴展至8192 Token（適配樹狀長序列）；超參數：學習率，BF16混合精度，全局批次128；推理配置：訓練階段自我博弈採用高溫度（1.0）保證漏洞樣本多樣性；評估階段採用貪心解碼（溫度=0）保證代碼確定性；推理加速：評估階段使用vLLM + 4卡張量並行。四、實驗設計與核心結果4.1 實驗基礎設置（1）基線模型與對比方法基座模型：CodeLlama-7B、Qwen2.5-Coder-7B、Qwen2.5-Coder-3B（覆蓋主流開源代碼模型）；對比基線：原始基座模型、SFT、SafeCoder（專業安全代碼模型）、無樹結構的普通自我博弈（消融實驗核心對照組）。（2）評測數據集與指標Python安全評測：SecurityEval，指標SPR@1（安全通過率），由CodeQL靜態漏洞檢測；C/C++安全評測：DiverseVul，指標漏洞總數（數值越低越安全）；通用代碼能力：HumanEval，指標pass@1/pass@10（驗證安全加固不會損害基礎編碼能力）；泛化測試：跨CWE泛化：訓練集110類漏洞，測試集40類從未見過的CWE；跨語言泛化：僅用C/C++數據訓練，測試Python、Go、JavaScript、Ruby。4.2 核心實驗結論（分三大研究問題）RQ1：TSP是否顯著提升代碼安全性？（主實驗）以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關鍵結論：TSP將CodeLlama-7B的安全通過率從SFT的57.0%提升至**75.8%**，超越所有基線；對比“普通自我博弈”（69.6%），證明樹狀風險節點結構是核心有效設計，單純自我博弈效果有限；HumanEval指標幾乎無下降，說明安全加固不會造成“能力遺忘”，模型基礎編碼能力穩定。該結論在Qwen2.5-Coder系列模型上完全復現，證明TSP具備模型通用性。RQ2：跨編程語言泛化能力僅使用C/C++安全數據訓練模型，在Python/Go/JS/Ruby多語言場景測試：所有模型中，TSP組的跨語言漏洞數量最低；證明TSP並非學習C/C++語法，而是提煉語言無關的通用安全邏輯（如輸入校驗、函數返回值檢查、內存規範等），安全知識可自由遷移。RQ3：跨未知CWE漏洞泛化能力（最具價值結論）測試模型對訓練集從未出現的漏洞類型的防禦能力：CodeLlama-7B：TSP漏洞數27，遠低於SFT的50、原始模型的64；部分基線（如SFT在Qwen2.5-Coder-7B）出現嚴重過擬合：訓練見過的漏洞修復較好，但未知漏洞數量反而暴漲；漏洞嚴重度拆解：TSP對高危漏洞的削減效果最顯著，有效抵禦新型高危安全威脅。本質區別：SFT是“記憶漏洞補丁”，TSP是“理解安全規則”。4.3 消融實驗補充風險節點必要性：移除樹結構與節點標註，僅保留自我博弈，性能大幅下降，驗證細粒度節點監督不可替代；迭代輪數：多輪自我博弈持續優化，模型逐步修復自身產生的隱蔽漏洞；節點損失聚合：單節點損失優於全局損失，進一步證明局部優化的價值。五、方法侷限性（論文客觀總結）TSP 並非全能，存在明確短板，也是未來研究方向：擅長局部漏洞，弱於長鏈路隱式漏洞TSP對局部顯式控制流漏洞（CWE-079 XSS、CWE-215敏感信息洩露、緩衝區溢出）效果極佳；但對長距離數據流/內存漏洞（CWE-416野指針、CWE-125越界讀取、CWE-690空指針解引用）表現一般。 原因：這類漏洞的“危險決策”和“漏洞爆發點”相隔數十行代碼，中間邏輯看似正常，TSP的節點級監督無法追蹤跨節點的長距離依賴。依賴初始風險節點標註質量框架效果上限由標註LLM的精度決定，全新未知漏洞可能無法被識別為風險節點，導致模型漏防。自我博弈樣本難度遞減隨著模型安全能力提升，自身生成的漏洞樣本越來越簡單，難以挖掘深層、複雜的漏洞模式，訓練後期增益收窄。模型規模限制實驗僅基於3B~7B輕量級模型，TSP在千億級超大代碼模型上的效果、適配性未驗證。六、對比現有方案：TSP的核心優勢技術方案優化粒度數據依賴泛化能力定位漏洞能力SFT整段序列（粗粒度）依賴大量人工標註安全代碼弱，易過擬合無法定位RL（傳統）整段序列（粗粒度）依賴程序級獎勵一般信用分配問題，無法定位普通自我博弈整段序列中等（自生成樣本）一般無法精準定位SafeCoder序列+指令微調依賴專業安全數據集中等較弱TSPToken/風險節點（細粒度）極低（全自生成樣本）強（跨語言/跨CWE）精準定位漏洞觸發節點總結四大核心創新點：範式創新：首次將樹狀決策樹+自我博弈引入安全代碼對齊，把代碼安全問題轉化為節點級序列決策問題；數據創新：模型“向自己的錯誤學習”，自動生成高質量正負樣本，擺脫漏洞數據集稀缺的行業痛點；優化創新：放棄全局優化，聚焦漏洞風險節點，學習信號精準、訓練效率更高；能力創新：模型習得抽象安全邏輯，而非補丁記憶，實現跨語言、跨未知漏洞的強泛化。七、行業價值與落地場景1. 工業落地價值降低安全對齊成本：企業無需大規模人工標註漏洞數據，現有代碼模型可通過TSP自迭代完成安全加固，適配中小型研發團隊；提升AI代碼助手安全性：集成到GitHub Copilot、IDE代碼補全工具中，減少日常開發的漏洞代碼輸出；跨語言安全統一加固：一套訓練邏輯可服務多編程語言，降低多語言代碼模型的安全運維成本。2. 學術與技術啟發為代碼LLM對齊提供新範式：“局部節點監督+自我博弈”可遷移到代碼漏洞檢測、代碼修復、代碼規範校驗等任務；拓展自我博弈的應用邊界：將AlphaGo類博弈思想從推理任務落地到代碼生成領域；為解決“大模型安全泛化”提供思路：從“記憶樣本”轉向“學習底層規則”，是通用安全對齊的重要方向。3. 未來研究方向（基於侷限性延伸）動態風險節點識別：擺脫靜態標註，訓練中自動發現新型漏洞節點；融合數據流/汙點分析：結合程序分析技術，追蹤長距離隱式漏洞，彌補內存類漏洞短板；分層自我博弈：設計難度遞增的漏洞樣本生成策略，避免後期樣本難度不足；超大模型適配：針對千億級代碼LLM優化TSP架構與顯存策略。八、全文總結這篇ICML 2026論文提出的TSP樹狀自我博弈框架，精準擊中了當前代碼大模型安全加固的核心痛點：粗粒度優化、數據依賴、泛化薄弱。其核心思想可概括為“讓模型在犯錯的地方糾正錯誤”：將代碼生成拆解為樹狀決策節點，定位漏洞觸發點，通過模型自我博弈生成漏洞樣本，在節點層級做細粒度對比學習。實驗充分證明，TSP在安全通過率、漏洞削減上全面超越SFT、傳統RL、普通自我博弈等基線，同時保留模型原有編碼能力，且習得可跨語言、跨未知漏洞的通用安全邏輯。儘管在長鏈路隱式漏洞、動態節點識別上存在侷限，但它開闢了“自糾錯+細粒度節點對齊”的全新技術路線，是代碼大模型安全領域兼具理論創新與工程實用性的優秀工作。鏈接：https://arxiv.org/pdf/2606.03489雷峰網