一句话:训练小型安全模型与代码 LLM 协同解码,在不访问目标模型内部参数的前提下降低漏洞代码生成率。
代码大语言模型(LLM)因在海量未经严格过滤的数据上预训练,容易生成含潜在漏洞的代码。现有缓解手段主要依赖构造无漏洞数据后重新训练或微调模型,但面对十亿级以上参数时计算与数据开销巨大;更重要的是,越来越多代码 LLM 以服务形式发布,内部表示不可见,只能通过 API 调用,使既有策略失效。为此,作者提出 CoSec——一种基于安全模型引导协同解码的代码 LLM 即时安全加固方法。其核心思想是额外训练一个体量小得多的安全模型,与目标代码 LLM 共同解码:由于安全模型对安全 token 具有更高置信度,它可在解码的每一步调整 token 的概率分布,从而在不访问目标模型内部参数的情况下引导其生成更安全的代码。作者在 CodeGen、StarCoder、DeepSeek-Coder 等不同参数规模的模型上开展实验,六个基座模型的平均安全率提升 5.02%–37.14%,同时保持目标模型的功能正确性。
Large Language Models (LLMs) specialized in code have shown exceptional proficiency across various programming-related tasks, particularly code generation. Nonetheless, due to its nature of pretraining on massive uncritically filtered data, prior studies have shown that code LLMs are prone to generate code with potential vulnerabilities. Existing approaches to mitigate this risk involve crafting data without vulnerability and subsequently retraining or fine-tuning the model. As the number of parameters exceeds a billion, the computation and data demands of the above approaches will be enormous. Moreover, an increasing number of code LLMs tend to be distributed as services, where the internal representation is not accessible, and the API is the only way to reach the LLM, making the prior mitigation strategies non-applicable. To cope with this, we propose CoSec, an on-the-fly Security hardening method of code LLMs based on security model-guided Co-decoding, to reduce the likelihood of code LLMs to generate code containing vulnerabilities. Our key idea is to train a separate but much smaller security model to co-decode with a target code LLM. Since the trained secure model has higher confidence for secure tokens, it guides the generation of the target base model towards more secure code generation. By adjusting the probability distributions of tokens during each step of the decoding process, our approach effectively influences the tendencies of generation without accessing the internal parameters of the target code LLM. We have conducted extensive experiments across various parameters in multiple code LLMs (i.e., CodeGen, StarCoder, and DeepSeek-Coder), and the results show that our approach is effective in security hardening. Specifically, our approach improves the average security ratio of six base models by 5.02%-37.14%, while maintaining the functional correctness of the target model.
@inproceedings{paperbot8513,
title = {CoSec: On-the-Fly Security Hardening of Code LLMs via Supervised Co-decoding},
author = {Dong Li and Meng Yan and Yaosheng Zhang and Zhongxin Liu and Chao Liu and Xiaohong Zhang and Ting Chen and David Lo},
booktitle = {ISSTA},
year = {2024},
doi = {10.1145/3650212.3680371}
}