ProSec: Fortifying Code LLMs with Proactive Security Alignment

We appreciate your feedback and respectfully clarify key differences between CodeLMSec and ProSec to illustrate our unique contributions:

Different goals

CodeLMSec and ProSec serve fundamentally different purposes. CodeLMSec is a codeLM security benchmark that evaluates codeLMs with vulnerability-inducing prompts. In contrast, ProSec is an alignment training technique that employs scalable data synthesis to produce a high-quality dataset, effectively securing codeLMs without harming utility.

Each entry in the ProSec dataset includes a vulnerability-inducing prompt, an insecure implementation, and a secure implementation. The alignment training loss requires a model to generate secure implementation with a higher probability than generating the insecure one.

Scalability

Effective alignment requires more samples than a typical benchmark. While CodeLMSec includes 280 prompts, the previous SOTA alignment dataset contained 1.5k prompts. ProSec scales up to more than 10k entries. ProSec automates the data synthesis process by composing diverse vul-inducing scenarios using CWE definitions and existing instruction-tuning datasets, avoiding the labor-intensive manual curation required by CodeLMSec.

Reference code

Benchmarks need only prompts; alignment requires paired insecure and secure code. ProSec uses rejection sampling with static analyzers to curate insecure code snippets, and it instructs LLMs to produce secure fixes, thereby constructing pairwise alignment data entries.

Balancing security and utility

An alignment dataset needs to balance between enhancing security and preserving utility. ProSec proposes a training dynamic-based data selection algorithm to construct a utility-preserving dataset, ensuring models’ utility is not compromised during security alignment.

Technically, we follow standard practices by evaluating model utility on coding benchmarks (i.e., HumanEval and MXEval) and assessing correctness via test cases. E.g., Table 1 shows that for Phi3mini-Inst, ProSec achieves 44% on MXEval versus SafeCoder's 42%, showing better utility preservation.

Empirical comparison

Although the two works are different, we did experiments during the rebuttal to show unique challenges in synthesizing alignment data.

We used the prompts from CodeLMSec to construct an alignment training dataset with the same pipeline as ProSec. (The coding instructions were obtained from CodeLMSec rather than synthesized by ProSec.) The ratios of vulnerable code generated are (lower is better): the original model: 45.6, ProSec: 19.7, CodeLMSec: 38.9. We can see that the model aligned with ProSec is more secure, underscoring the challenge of effective alignment data synthesis.

Clarifications

Q1: What is preferred/win and a less preferred/lose?

Each entry in an alignment dataset consists of three parts: a prompt, a preferred (or “win”) response, and a less-preferred (or “lose”) response [1]. In ProSec, the prompt is a vulnerability-inducing prompt, while the “win” and “lose” responses correspond to secure and insecure implementations, respectively.

Q2: The goal of ProSec?

The goal of ProSec is to prevent the generation of vulnerable code. It trains the model to favor generating the secure code over the insecure ones. We will revise section 2.

Q3: How to measure functionality?

We follow established practices[2, 3, 4] to set up the experiments. The functionality correctness is evaluated by test cases. Code generations that pass all test cases are considered correct.

Q4: Why is data selection necessary?

The data selection is critical for balancing security and utility of the aligned model. Intuitively, including too many utility-preserving data samples weakens the security alignment, while too few impairs the model’s utility. ProSec employs a training dynamic-based algorithm to identify and selectively include those utility-preserving data samples whose distribution is disrupted by the security alignment.

Q5: How to measure utility?

As shown in Table 1, the utility is measured as the performance on two coding benchmarks, HumanEval and MXEval.

Q6: What and how are CWEs selected?

We select 38 CWEs that overlap between PurpleLlama and SafeCoder to set up a fair evaluation. Please see Section 4 (line 269) for details.

Q7: How to measure data diversity?

We use the cosine similarity of semantics embeddings to analyze the diversity of the dataset (line 359, Fig. 4). We use string editing distance (line 255, Section 3.3) to deduplicate code snippets and thus increase dataset diversity.

[1] ​​Rafailov, Rafael, et al. Direct preference optimization: Your language model is secretly a reward model. NeurIPS’23

[2] He, Jingxuan, et al. Instruction tuning for secure code generation. ICML’24.

[3] He, Jingxuan,et al. Large language models for code: Security hardening and adversarial testing. CCS’23

[4] Wei, Yuxiang, et al. Magicoder: Empowering code generation with oss-instruct. ICML’24.

Thank you for your detailed and supportive review.

ProSec’s relevance to ICML

We respectfully contend that our work aligns well with previous contributions recognized at ICML. For example, prior works such as data selection for language model training [Qurating, ICML’24], data synthesis for code language models [MagicCoder, ICML’24] and security-focused SFT for code language models [SafeCoder, ICML’24] introduced innovative practices in LLM training and alignment, thereby establishing a strong precedent for valuable engineering-focused work.

Furthermore, we note that ICML has a tradition of supporting application-driven research. For instance, [Invariant; ICML’23] prompts LLMs to generate loop invariants – addressing a significant challenge in the software engineering domain; [Next; ICML’24] enhances program repair performance by incorporating execution trace information to prompts; and [StackSight; ICML’24] leverages LLMs to translate low-level code into readable C++ code via CoT prompts.

Similarly, ProSec offers a scalable approach to securing code language models during the alignment stage. In particular, our introduction of utility-preserving dataset and our formulation of alignment side effects via training dynamics provide valuable technical insights into balancing security alignment and model utility. We believe these contributions address important practical challenges and pave the way for further advancements in aligning codeLM with domain specific constraints.

Ablation study for normal preference data

We appreciate the suggestions and will include the following discussion in the paper.

During rebuttal, we conducted experiments to illustrate the effectiveness of introducing the utility-preserving dataset (DNorm). Due to time constraints, we evaluate all models on a subset of PurpleLlama and MXEval for security and utility assessment, respectively. The results are summarized below:

We can see that DNorm indeed helps preserve the models’ utility. By adding DNorm to the SafeCoder dataset, the utility of the aligned model improves slightly while maintaining a comparable level of security performance compared to the vanilla SafeCoder dataset. In contrast, when DNorm is removed from the ProSec dataset, the model’s utility performance drops dramatically, indicating that aligning the model solely on the security dataset would significantly compromise its utility.

[Qurating] Wettig, Alexander, et al. "Qurating: Selecting high-quality data for training language models." ICML’24

[MagicCoder] Wei, Yuxiang, et al. Magicoder: Empowering code generation with oss-instruct. ICML’24.

[SafeCoder] He, Jingxuan, et al. Instruction tuning for secure code generation. ICML’24.

[Invariant] Pei, Kexin, et al. "Can large language models reason about program invariants?."ICML’23

[Next] Ni, Ansong, et al. "Next: Teaching large language models to reason about code execution." ICML’24

[StackSight] Fang, Weike, et al. "StackSight: Unveiling webassembly through large language models and neurosymbolic chain-of-thought decompilation." ICML’24

Thank you for the supportive review.

Is the improvement only affected by the fact that the generated dataset is 7x larger than SafeCoder? Can you do a control experiment where both datasets have the same size (and same secure / vulnerable mixture ratio), and test if the data quality of ProSec is also higher (e.g. more diverse)?

The quality of ProSec dataset is higher than SafeCoder because ProSec scales to more diverse scenarios without requiring manual efforts.

SafeCoder constructs alignment data entries by collecting real-world vulnerabilities and the corresponding fixes. However, real-world vulnerabilities and their fixes are sparse. For example, SafeCoder only collects 465 entries from 145 million git commits.

On the other hand, ProSec leverages LLMs to enumerate potentially vulnerable-inducing scenarios, systematically scaling up to more diverse scenarios. Therefore, the overall quality of the ProSec dataset is better than that of the SafeCoder dataset.

Empirical Support

During rebuttal, we did another set of experiments as suggested by the reviewer.

• [Size] We randomly sample a subset of the ProSec dataset to match the size of the SafeCoder dataset and run the alignment algorithm again. The resulting model is noted as “ProSec[Subset]”.
• [Mixture ratio] The original SafeCoder dataset does not contain utility-preserving data for preference optimization. To make sure the ratio of security-focused and utility-preserving data samples are the same between ProSec and SafeCoder, we mix the vanilla SafeCoder dataset with the utility-preserving dataset of ProSec at the same ratio. The resulting model is noted as “SafeCoder+DNorm”.

The empirical results show that the ProSec dataset is better than the SafeCoder dataset when both datasets contain the same number of entries and when both datasets contain the same ratio of security/utility data samples.

Note that we evaluate all models on subsets of PurpleLlama and MXEval for security and utility measurements due to time constraints.

We can see that the model trained on the ProSec subset is safer than that trained on the SafeCoder dataset with the same size. It also outperforms the SafeCoder dataset mixed with the same ratio of benign data samples. That indicates the data quality of ProSec is indeed higher. Moreover, the model trained on the full ProSec dataset is safer than that trained on the subset, as expected. That is because the full ProSec dataset covers more scenarios.

The alignment training on all datasets do not significantly affect coding utilities.

We will include the experiments in the paper.

Thank you for the supportive and detailed review. We will include the discussions below to the paper.

Q1: Explain static analyzers

We use the static analyzers in PurpleLlama. It consists of three tools: regular expressions, semgrep, and weggli. All tools work on the generated source code. Regular expressions simply match insecure patterns in code (e.g., matching deprecated APIs). Semgrep and weggli first build the AST from a code snippet, and then search for problematic patterns (e.g., a potential NULL pointer that is dereferenced without being checked).

Q2: How alignment data is filtered

The alignment dataset contains two parts: the secure practice preference data (DSec) and the utility preservation preference data (DNorm). For both dataset, a data entry contains a prompt, a preferred code snippet, and a less-preferred code snippet.

We control the quality of DSec by heuristics. We first check the syntactic correctness of the code snippets. Then we use static analyzers to make sure the preferred code is secure. After that, we use heuristics to make sure the preferred (secure) code snippet has corresponding functionality with the vulnerable code. Finally, we use string similarity (based on editing distance) to deduplicate data entries.

For DNorm, we develop a data selection algorithm that takes training dynamics into consideration. Intuitively, the algorithm identifies and prioritizes the data entries whose utility are broken during the security alignment to mitigate degradation in utility. Specifically, we first align the target model with DSec only. For each candidate entry in DNorm, the algorithm keeps track of how its generation probabilities change during the alignment training. Then the algorithm selectively includes the entries whose generation probabilities decrease significantly during the security training.

Please refer to section 3.3 (line 240) in the paper for details.

Q3: Comparison with post-processing

Yes, we focus on CWEs that can be detected by static analyzers. There are more complex CWEs that rely on global information or knowledge about program functionality. Detecting such CWEs are open challenges[1,2]. We leave it as future work to improve alignments on complex CWEs.

An agentic workflow incurs higher computational costs and increased latency because it runs a static analyzer for every coding request and may require multiple queries to the code language model.

This design could degrade the user experience in scenarios like code copilots, where swift completions are expected.

In fact, post-processing with an agentic design complements code model alignment techniques: an aligned codeLM may reduce the number of conversational turns needed, while the agent can capture edge cases where the model produces insecure code.

We further use empirical results to show the cost an agentic workflow may introduce. During rebuttal, we additionally implemented an agentic baseline that uses static analyzers to check generated code and iteratively asks the code generation model to fix. For each fix request, we provide the feedback from the static analyzers, the problematic code, and the initial coding instruction to the model. We evaluate the agentic workflow on a randomly sampled subset of PurpeLlama due to time constraints. Here are the statistics:

Moreover, on average, a coding request requires five rounds of fixes to achieve the secure performance of ProSec. This demonstrates that using a security-aligned model is more efficient than simply applying post-processing to a secure code generation agentic workflow.

Note that we choose to use the tested codeLM to fix the code, instead of using black box LLMs. That is because in a realistic use scenario of a smaller code LM, querying a larger black box model might be less preferred (due to the latency and cost) or even not allowed (due to privacy and policy concerns).

Q4: Cost

We use Claude-3.5-haiku to synthesize the instructions. For each CWE, we synthesize ~10k initial instructions, and then cluster them to identify the most diverse 2k instructions. The cost to synthesize instructions for each CWE is around 5 USD.

[1] Ding, Yangruibo, et al. "Vulnerability detection with code language models: How far are we?." arXiv preprint arXiv:2403.18624 (2024).

[2] Google, Project Zero. https://googleprojectzero.blogspot.com/2024/06/project-naptime.html