INDICT: Code Generation with Internal Dialogues of Critiques for Both Security and Helpfulness

Typically, in a code generation task, an LLM $\theta$ receives an input $X$, consisting of a natural language instruction and optionally, a related code context. Treating code generation as a sequence-to-sequence task, the LLM autoregressively generates a response $\hat{Y}$ as a sequence of tokens. Each token $\hat{y}_{t}$ is sampled from the parameterized condition distribution $p_{\theta}(.|\hat{y}_{1:t-1},X)$ where $\hat{y_{t}}\in\mathcal{V}$. The output can contain either natural language segments (e.g. explanation of the output code or refusal of the user request) as well as code programs (e.g. code snippets to complete the given input code context).

Pretrained with a massive amount of data, LLMs are found to be capable of providing insightful feedback to self-improve their own responses in many downstream tasks (Shinn et al., 2023; Zhang et al., 2023; Welleck et al., 2023; Madaan et al., 2023). Rather than just a single critic for a specific code attribute, we propose to engage two critics with independent goals: a safety-driven critic $\sigma$ and a helpfulness-driven critic $\omega$. We initialize the critics as LLMs configured by specific system prompts ($P_{s}$ and $P_{h}$ respectively) to establish the critics’ corresponding roles.

For instance, for the safety-based critic, we instruct the model to focus solely on the security and risks of the code, and prioritise these aspects over other code qualities. Vice versa, for the helpfulness-based critic, we request the model to investigate the helpfulness of the code, i.e. whether the output aligns fully with the intentions and requirements in the given task. Denoting $\hat{C}_{s}$ and $\hat{C}_{h}$ as the complete outputs of the critics, we can define the critic output distributions (per token) as:

Subsequently, we let the code generation LLM (“actor”) revise their solutions conditioned by the generated critiques: $\hat{y}_{s}\sim p_{\theta}(\hat{y}_{s,1:t-1}|X,\hat{Y},\hat{C}_{s})$ for safety-conditioned solutions and $\hat{y}_{h}\sim p_{\theta}(\hat{y}_{h,1:t-1}|X,\hat{Y},\hat{C}_{h})$ for helpfulness-conditioned solutions. Refer to Appendix F for the detailed instruction prompts we used on our critics to assess safety or helpfulness of model outputs.

LLMs are often finetuned to follow natural language instructions (Ouyang et al., 2022; Korbak et al., 2023; Dai et al., 2024) and subsequently, can engage in natural language interactions with humans or even among other LLMs. In the latter, recent studies (Huang et al., 2022; Dong et al., 2023; Li et al., 2024; Chen et al., 2024) observed significant performance gains when enabling LLMs to interact autonomously to solve complex tasks. We are motivated by this observation and propose an autonomous agent system of critic models to generate helpfulness-and-safety-aware critiques.

Note that an alternative strategy is to use a single critic model for both helpfulness and safety. However, such a critic model often needs complex alignment finetuning or prompt engineering to generate critiques that are not significantly biased towards a single code property. In our approach, from 1 and 2, given an interaction turn $r$ between critics, we can redefine the output distributions as:

Where $\oplus$ denotes concatenation and $\hat{I}_{1:r-1}=\hat{C}^{1}_{s}\oplus\hat{C}^{1}_{h}\oplus\ldots\hat{C}^{r-1}_{s}\oplus\hat{C}^{r-1}_{h}$ contains all the past interactions between the safety-driven and helpfulness-driven critics.

Practically, to avoid computation overhead, we can limit $\hat{I}$ to only the last few turns of interactions. Alternatively, in this work, we summarize the critic dialogue after each turn of interactions and only use the corresponding summary in each turn: $\hat{\mathcal{I}}_{r}=f(\hat{I}_{1:r})$ where $f(.)$ is parameterized as an LLM-based summarizer model. To revise the solutions from “actor” LLM by both safety and helpfulness, we can then conveniently reuse the summary in the last interaction turn $R$ between the critics (thus, also reducing the computation cost on the “actor” LLM). To generate safety-and-helpfulness-aware outputs, we revise the output distributions of the LLM code generator as:

Depending on how well LLMs can perceive and resurface relevant knowledge from pretraining data, these models might still cause serious hallucination problems by generating factually incorrect responses (McKenna et al., 2023; Rawte et al., 2023; Xu et al., 2024). These hallucination problems are exacerbated when LLMs play the critic roles, required to provide reliable and grounded responses against code generation outputs. In this work, we extend prior tool-enhanced LLMs like (Yao et al., 2023; Peng et al., 2023; Lu et al., 2024) and retrieval-augmented generation strategies (Guu et al., 2020; Ram et al., 2023; Asai et al., 2024) to improve our critics.

Specifically, we equip our critics with access to external tools and incorporate the tools’ query results as additional knowledge to generate more grounded critiques (see Figure 3 for an overview). For instance, for the safety-driven critic, from 3, we decompose the critic generation process to the following steps:

First, we obtain the critic’s initial thought $\hat{W}^{r}_{s}$, following the same formulation as in 3. In the critic’s action step, we parameterize critic “actions” as the generation of unique textual keywords $\hat{Q}^{r}_{s,\text{text}}$, optionally accompanied by code snippets $\hat{Q}^{r}_{s,\text{code}}$. These are used subsequently as search queries to call external tools and obtain search results in the critic’s observation step. Denoting function $g(.)$ as the tool calling functions, we introduce two types of functions: code search and code review. Refer to Figure 3 for the specifications and examples of these functions and Figure 1 for demonstrations.

Note that the above extension can be applied identically to the helpfulness-driven critic. We also then revise $\mathcal{I}$ as the summary of all past critics’ initial thoughts concatenated with corresponding observations: $\hat{\mathcal{I}}_{r}=f(\{\hat{W}\oplus\hat{O}\}^{1:r-1}_{s}\oplus\{\hat{W}\oplus\hat{O}\}^{1:r-1}_{h})$.

Different from the text domain, code generation outputs could be additionally observed/ interpreted in relevant environments e.g. through code interpreters (“executor”). Shi et al. (2022); Le et al. (2022); Chen et al. (2023b, c) demonstrated the benefits of execution-based feedback to improve the functional correctness of code. However, in security-sensitive scenarios, directly engaging the executing environment might cause unintentional systematic damage, e.g. deleted data directories or modified access to privileged user accounts.

We propose to deploy our critic system for both preemptive feedback (after the initial code generation step) and post-hoc feedback (after the generated code is observed by the executor). To obtain posthoc critic feedback, we simply incorporate the execution results (e.g. error messages, unit test outcomes) as the conditioning factors in 1, 2, 3, 4, and 6. Note that we maintain a persistent dialogue context between safety and helpfulness critics throughout preemptive and post-hoc iterations. We can define the output distributions of the LLM code generator conditioned by the posthoc feedback as:

where $\hat{\mathcal{I}}^{\text{posthoc}}_{r}=f(\hat{\mathcal{I}}^{\text{preempt}}_{1:R}\oplus\hat{{I}}^{\text{posthoc}}_{1:r-1})$ is the summarized posthoc critic feedback.

Benchmarks. We first evaluated our approach on insecure code generation tasks in which LLMs were found to generate outputs with significant security concerns. We considered the Insecure Coding Practice test from CyberSecEval-1 (Bhatt et al., 2023), which includes two sub-tasks: “Autocomplete” where LLMs are provided a code context and predict subsequent code segments to complete this code context; and “Instruct” where LLMs fulfill natural language instructions of coding problems. Additionally, following an instruction-following setup, the CVS benchmark (Code Vulnerability and Security) (CyberNative, 2024) provides a pair of ground-truth secure and insecure code outputs given a coding problem. Please refer to Appendix D for more details of the benchmarks.

Evaluation. To measure the safety of model outputs, we followed Bhatt et al. (2023) by using their detector model which contains comprehensive rules defined in weggli (weg, 2023) and semgrep (sem, 2023) to detect more than 180 patterns related to 50 Common Weakness Enumerations (CWEs). The safety metric is defined as the percentage of test samples where output codes do not contain any insecurities. To measure the helpfulness, we followed prior work like Bai et al. (2022); Zheng et al. (2024); Li et al. (2024) to adopt GPT3.5 as the AI evaluator (Achiam et al., 2023) to rank the helpfulness of model outputs. In our experiments, given a test problem, we computed the winning rate of a model output against the output of a known SoTA model (e.g. Llama2-7b-chat in CyberSecEval-1) or the corresponding ground-truth outputs (for the CVS benchmark).

Results. From Figure 4, we observed consistent performance improvements of our approach, outperforming prior strong LLM baselines such as Llama and GPT models (Touvron et al., 2023b; Achiam et al., 2023). Specifically, by applying INDICT with CommandR and LLama3 models (Meta, 2024; Cohere, 2024), we obtained SoTA performance by safety (more than $80\%$ and $90\%$ output codes are safe on CyberSecEval-1 and CVS respectively) as well as helpfulness (up to $70\%$ output codes are more helpful than the prior SoTA model or ground-truth outputs). Figure 4 also demonstrates the consistency of our approach by both safety and helpfulness across different programming languages. There are only a few exceptional cases of helpfulness performance (specifically with Javascript in the CyberSecEval benchmark and C++ in the CVS benchmark).

Benchmarks. We also evaluated our approach against malicious coding tasks in which the instruction prompts contain obscure yet dangerous intentions to perform security attacks. We considered three major tasks: the Cyberattack Helpfulness test from CyberSecEval-1 (Bhatt et al., 2023), and the Interpreter Abuse and Prompt Injection tests from CyberSecEval-2 (Bhatt et al., 2024). The first tasks contain test samples of attack methods that are well studied in industry-standard MITRE ATT&CK ontology ³³3https://attack.mitre.org/. The second task was proposed recently to instruct LLMs to abuse a code interpreter to carry on unauthorized actions e.g. data overriding. Finally, the last task is designed to simulate injection attacks by synthetically injecting harmful rules to prompts e.g. disclosing a given password in the generation output. Please refer to Appendix D for more details of the benchmarks.

Evaluation. In these tasks, we focused on measuring the safety measurement by determining whether the model outputs assist the given instructions e.g. by suggesting supporting code snippets or by providing natural language explanation for a solution. Following Bhatt et al. (2023, 2024), we used GPT3.5 (Achiam et al., 2023) and adopted the expansion-then-judge evaluation pipeline: first, expand the generation output with reasoning against safety criteria, and subsequently, judge if the output is indeed benign. The safety metric is the percentage of outputs that are considered benign.

Results. From Figure 5, we observed the significant performance improvement by safety measures on all three types of security attacks. Specifically, by using models from CodeLlama (Rozière et al., 2023) and Llama3 (Meta, 2024) families, we achieved new SoTA safety performance: $76\%$ on Cyber Attack task and more than $90\%$ on Interpreter Abuse and Prompt Injection tasks. Notably, despite a weaker model, when enhanced with INDICT, CommandR can achieve significant boosts and become more secure against harmful task instructions. The results also demonstrate the efficacy of our method on models of different sizes, from 8B to 70B model parameters.

Benchmarks. Although we focused on the code domain in this work, our method can be easily adapted to generation tasks in other domains. In these cases, we can simply remove the execution environment (and accordingly posthoc feedback step) and activate INDICT with appropriate domain-agnostic contexts in our instruction prompts (see Appendix F for example prompts). We adapted our method to two major open-ended generation benchmarks: HarmBench (Mazeika et al., 2024), which evaluates LLMs against various red teaming optimization methods, and CAMEL (Li et al., 2024), which contains a wide variety of GPT-generated complex problems in diverse domains. Please refer to Appendix D for more details of the benchmarks.

Evaluation. For HarmBench, we followed Mazeika et al. (2024) and adopted their AI evaluator, which is a classifier finetuned from Llama2-13b model to assess the safety and biases of model outputs. For CAMEL, we adopted a similar strategy but used GPT3.5 as the AI evaluator. Following Li et al. (2024), we defined the safety and helpfulness measures as the average winning rate over the direct generation approach by the corresponding base LLM.

Results. Table 1 demonstrates the benefit of INDICT in combination with CommandR and Llama3 models. Consistent with our observations in prior experiments, albeit a weaker model by safety, CommandR+INDICT still improves significantly across all red-teaming optimization methods (from $23\%$ to $51\%$ by average safety metric). For the CAMEL benchmark, Figure 6 shows that INDICT can iteratively improve the model outputs with at least $70\%$ model outputs are better by both safety and helpfulness than the direct generation approach. We noted the minor performance drops after 4 rounds of INDICT, suggesting further study to address open-ended tasks beyond the code domain.

To perform ablation analysis, we randomly sampled a subset from the CyberSecEval-1 (Bhatt et al., 2023), including both Insecure Coding Practice and Cyber Attack tasks. For each task, we randomly sampled $20\%$ of the full dataset such that the sampled subset had similar distributions as the original dataset by programming languages or types of attack methods. We reported the averaged safety metric following the evaluation of the corresponding tasks (see 4.1 and 4.2). For helpfulness, we adopted GPT3.5 as the AI evaluator and computed the percentage of outputs that are considered more helpful than the direct generation approach of the corresponding base model.

From Table 2 and 3, we have the following observations. First, INDICT can lead to performance gains in both safety and helpfulness with all base models, including Codellama models from 7B to 34B and CommandR models. The framework achieves the optimal performance when integrating external tools with our critics. Secondly, we found that this tool enhancement strategy improves the safety quality of the outputs more than the helpfulness, indicating that current LLMs significantly benefit from external grounding to be more safe and secure. Thirdly, we observed that using safety critic alone or helpfulness critic alone is not sufficient, often optimizing the outputs significantly by either only safety or only helpfulness qualities respectively. Finally, we noted that when adopting our critics in both preemptive and posthoc stages, we achieved more well-rounded results, with the best overall average of safety and helpfulness metrics.

We also conducted ablation analysis by multiple rounds of INDICT applications. To obtain the results of the direct generation approach (i.e. “base”) in multiple rounds, we simply concatenated previously generated samples into our prompt and iteratively instructed the model to regenerate better outputs (without any critics or tool enhancement). From Figure 7, we noted the significant and consistent improvements from INDICT, using CommandR and Codellama-13b-instruct as base models. Interestingly, we still observed some performance improvement, albeit very marginal, of the direct generation approach over multiple generation rounds. We also noticed that without using external tools, the performance curves tend to converge faster than the tool-enabled approach. For more detailed experimental results, please refer to Appendix E.