A Mixture of Linear Corrections Generates Secure Code
This addresses the issue of code vulnerabilities in AI-generated code for developers and security practitioners, offering an incremental improvement through inference-time steering.
The paper tackled the problem of large language models (LLMs) generating vulnerable code by investigating whether they internally encode concepts for identifying vulnerabilities, and developed a mixture of corrections (MoC) technique to steer LLMs toward producing less vulnerable code, resulting in an 8.9% improvement in security ratio and a 2.1% boost in functionality on benchmarks.
Large language models (LLMs) have become proficient at sophisticated code-generation tasks, yet remain ineffective at reliably detecting or avoiding code vulnerabilities. Does this deficiency stem from insufficient learning about code vulnerabilities, or is it merely a result of ineffective prompting? Using representation engineering techniques, we investigate whether LLMs internally encode the concepts necessary to identify code vulnerabilities. We find that current LLMs encode precise internal representations that distinguish vulnerable from secure code--achieving greater accuracy than standard prompting approaches. Leveraging these vulnerability-sensitive representations, we develop an inference-time steering technique that subtly modulates the model's token-generation probabilities through a mixture of corrections (MoC). Our method effectively guides LLMs to produce less vulnerable code without compromising functionality, demonstrating a practical approach to controlled vulnerability management in generated code. Notably, MoC enhances the security ratio of Qwen2.5-Coder-7B by 8.9\%, while simultaneously improving functionality on HumanEval pass@1 by 2.1\%.