Constrained Decoding for Secure Code Generation

TL;DR

This paper tackles the problem that Code LLMs often produce insecure or incorrect code and introduces CodeGuard+ to jointly evaluate security and correctness. It demonstrates that prior defenses like prefix tuning can improve security at the expense of functionality and shows that decoding strategy strongly influences security outcomes. The authors formalize constrained decoding for secure code generation, proposing two techniques—Constrained Beam Sampling and a MuCoLa-based non-autoregressive method—along with practical constraint specifications. Across eight Code LLMs and GPT-4, constrained decoding substantially improves secure performance, often surpassing GPT-4 unconstrained, and generalizes to untrained CWEs without requiring additional training data.

Abstract

Code Large Language Models (Code LLMs) have been increasingly used by developers to boost productivity, but they often generate vulnerable code. Thus, there is an urgent need to ensure that code generated by Code LLMs is correct and secure. Previous research has primarily focused on generating secure code, overlooking the fact that secure code also needs to be correct. This oversight can lead to a false sense of security. Currently, the community lacks a method to measure actual progress in this area, and we need solutions that address both security and correctness of code generation. This paper introduces a new benchmark, CodeGuard+, along with two new metrics, to measure Code LLMs' ability to generate both secure and correct code. Using our new evaluation methods, we show that the state-of-the-art defense technique, prefix tuning, may not be as strong as previously believed, since it generates secure code but sacrifices functional correctness. We also demonstrate that different decoding methods significantly affect the security of Code LLMs. Furthermore, we explore a new defense direction: constrained decoding for secure code generation. We propose new constrained decoding techniques to generate secure code. Our results reveal that constrained decoding is more effective than prefix tuning to improve the security of Code LLMs, without requiring a specialized training dataset. Moreover, our evaluations over eight state-of-the-art Code LLMs show that constrained decoding has strong performance to improve the security of Code LLMs, and our technique outperforms GPT-4.

Figures (9)

• Figure 1: We compare CodeGen + Prefix-tuning model, trained by the state-of-the-art defense he2023large, against the baseline CodeGen model. Our metric secure-pass@$1$ is more realistic than SVEN Security Rate used in he2023large, since we evaluate both security and correctness of generated code, while SVEN Security Rate does not evaluate correctness. SVEN Security Rate severely overestimates how secure a model really is. The secure-pass@$1$ of CodeGen + Prefix-tuning is only 2.53% better than CodeGen with Beam Sampling.
• Figure 2: Our constrained decoding technique can improve secure-pass@$1$ of all six open-source Code LLMs of sizes ranging from 2.7B to 34B. Every model with constrained decoding shows better secure-pass@$1$ than GPT-4 with Nucleus Sampling.
• Figure 3: We add a specific instruction to print the buffer in the prompt, so we can test the generated code using stdout.
• Figure 4: In the old prompt, the meaning of "unsubscribe" is ambiguous, and it is hard to set up and test a MySQLdb database. In the new prompt, we add a specific instruction to delete an entry from the database, as a result of "unsubscribe", and we also change the library to sqlite3, which enables us to easily test the database using local files.
• Figure 5: We update deprecated libraries and functions in the new prompt.

Theorems & Definitions (4)

• Definition 3.1: SVEN-SR
• Definition 3.2: pass@$k$
• Definition 3.3: secure-pass@$k$
• Definition 3.4: secure@$k_{\text{pass}}$