Prompt optimization for secure generation of functional source code with large language models (LLMS)

Abstract

The Disclosed Invention introduces PromSec, an algorithm for prompt optimization for secure and functioning code generation using LLMs. PromSec combines 1) code vulnerability clearing using a generative adversarial graph neural network, dubbed as gGAN, to fix and reduce security vulnerabilities in generated codes and 2) code generation using an LLM into an interactive loop, such that the outcome of the gGAN drives the LLM with enhanced prompts to generate secure codes while preserving their functionality. Introducing a new contrastive learning approach in gGAN, this disclosure formulates the code-clearing and generation loop as a dual-objective optimization problem, enabling PromSec to notably minimize the number of LLM inferences. As a result, PromSec becomes a cost-effective and practical solution for generating secure and functioning codes.

Description

Attorney Docket No.5600234.01331 TITLE METHOD AND SYSTEM FOR PROMPT OPTIMIZATION FOR SECURE GENERATION OF FUNCTIONAL SOURCE CODE WITH LARGE LANGUAGE MODELS (LLMs) PRIORITY CLAIM AND CROSS-REFERENCE TO RELATED APPLICATIONS

[0001] The present application claims priority to and the benefit of U.S. Provisional Patent Application No. 63 / 561,573 filed March 5, 2024, which is incorporated herein by reference in its entirety and relied upon. BACKGROUND

[0002] Large language models (LLMs) are at the forefront of machine learning (ML), demonstrating exceptional performance in various natural language processing (NLP) tasks. These models, including the generative pre-trained transformer (GPT) series by OpenAI utilize a large-scale, transformer-based architecture trained on diverse NLP datasets. Recently, LLMs have attracted attention for generating high-quality source code, significantly benefiting software development by reducing time and expertise costs. Their growing popularity is evident, with tools like GitHub’s Copilot attracting over a million paid subscribers across 37,000 organizations. The effectiveness of these LLMs stems from their extensive training using vast source code databases. However, this reliance on open-source training data introduces significant security risks. The LLMs often replicate securityflaws present in their training examples, leading to vulnerabilities in the generated code. This issue raises concerns about the reliability of LLMs for secure code generation, also known as program synthesis, and highlights the necessity to address and mitigate these inherent security vulnerabilities.

[0003] The generation of secure and functionally correct code with LLMs faces the fundamental challenge of prompt optimization. Significant research efforts – have been directed towards engineering and optimizing LLM prompts for specific applications, highlighting the vast prompt search space and developing heuristics navigate this space effectively. Similarly, in code generation, a recent state-of-the-art (SoTA) approach shows the diversity in security and functionality of code generated with respect to variations in handcrafted prompts. However, the search space challenge is more severe in the code generation application due to the need for external tools to assess code security and functional fidelity, making prompt optimization more complex. Moreover, the lack of differentiability of LLM prompts with respect to these security and functionality metrics, coupled with the absenceAttorney Docket No.5600234.01331 of intuitive heuristics for maintaining security and functionality, adds to the challenge. The one-to-many nature of function-code implementation further complicates the application of traditional, data-driven security solutions, particularly those based on a supervised learning approach.

[0004] Therefore, a need exists for a method and system for prompt optimization for secure generation of functional source code with large language models (LLMs). SUMMARY

[0005] The present disclosure provides for a method and system for prompt optimization for secure generation of functional source code with large language models (LLMs).

[0006] Aspects of the subject matter described herein may be useful alone or in combination with one or more other aspects described herein. In an exemplary aspect of the present disclosure, a method for prompt optimization for secure generation of functional source code with large language models (LLMs), the method including: obtaining an initial prompt; forwarding the initial prompt to a Large Language Model (LLM); generating an initial code using the initial prompt, wherein the LLM generates the initial code; forwarding the initial code to a Security Analysis module; transforming the initial code into a graph representation; forwarding the graph representation to a graph generative adversarial network; enhancing the graph representation to a secure graph representation; converting the secure graph representation into a secure code; reverse-engineering the secure code a secure prompt; feeding the secure prompt to the LLM; and producing a final output.

[0007] In a second aspect of the present disclosure, which may be combined with any other aspect listed herein unless specified otherwise, the initial prompt is one of a direct prompt from a user, a verbal description characterizing the intent of the user, or a source code.

[0008] In a third aspect of the present disclosure, which may be combined with any other aspect listed herein unless specified otherwise, the Security Analysis module is a static security tool to identify the number of CWEs in the initial code.

[0009] In a fourth aspect of the present disclosure, which may be combined with any other aspect listed herein unless specified otherwise, the initial prompt is forwarded directly to the Security Analysis module if the initial prompt is a source code.

[0010] In a fifth aspect of the present disclosure, which may be combined with any other aspect listed herein unless specified otherwise, the initial code is the final output if the Security Analysis module identifies no CWEs.Attorney Docket No.5600234.01331

[0011] In a sixth aspect of the present disclosure, which may be combined with any other aspect listed herein unless specified otherwise, no CWEs is denoted as K=0, where K represents the count of identified CWEs by the security analysis module.

[0012] In a seventh aspect of the present disclosure, which may be combined with any other aspect listed herein unless specified otherwise, the graph generative adversarial network is configured to enhance the graph representation by reducing CWEs while maintaining intended functionality.

[0013] In an eight aspect of the present disclosure, which may be combined with any other aspect listed herein unless specified otherwise, the secure graph representation is converted into the secure code through a Code reconstruction module.

[0014] In a ninth aspect of the present disclosure, which may be combined with any other aspect listed herein unless specified otherwise, further comprising feeding the secure prompt to the LLM until K=0.

[0015] In a tenth aspect of the present disclosure, which may be combined with any other aspect listed herein unless specified otherwise, the secure prompt is fed to the LLM for a set number of user-selected cycles.

[0016] In an eleventh aspect of the present disclosure, which may be combined with any other aspect listed herein unless specified otherwise, a system for prompt optimization for secure generation of functional source code with large language models (LLMs) including: a server; a processor; and a memory storing instructions, which when executed by the processor, cause the processor to apply an algorithm to optimize prompts for or secure generation of functional source code with large language models (LLMs).

[0017] In a twelfth aspect of the present disclosure, which may be combined with any other aspect listed herein unless specified otherwise, the algorithm to optimize prompts for or secure generation of functional source code with large language models (LLMs) is: where ^^ is a prompt, ^^∗is an optimal prompt, ^^ is a codecode either given as a code starting point or a code generated with an initial prompt ^^^, wherein the term ^^(^^) denotes a CWE count in the codec, and ^^(^^, ^^^)is a metric representing a functional discrepancy between ^^^ and ^^.

[0018] In a thirteenth aspect of the present disclosure, the algorithm to optimize prompts for or secure generation of functional source code with large language models (LLMs) is configured to: obtain an initial prompt; forward the initial prompt to a Large Language ModelAttorney Docket No.5600234.01331 (LLM); generate an initial code using the initial prompt, wherein the LLM generates the initial code; forward the initial code to a Security Analysis module; transform the initial code into a graph representation; forward the graph representation to a graph generative adversarial network; enhance the graph representation to a secure graph representation; convert the secure graph representation into a secure code; reverse-engineer the secure code a secure prompt; feed the secure prompt to the LLM; and produce a final output.

[0019] In a fourteenth aspect of the present disclosure, which may be combined with any other aspect listed herein unless specified otherwise, the initial prompt is one of a direct prompt from a user, a verbal description characterizing the intent of the user, or a source code.

[0020] In a fifteenth aspect of the present disclosure, which may be combined with any other aspect listed herein unless specified otherwise, the Security Analysis module is a static security tool to identify the number of CWEs in the initial code.

[0021] In a sixteenth aspect of the present disclosure, which may be combined with any other aspect listed herein unless specified otherwise, the initial prompt is forwarded directly to the Security Analysis module if the initial prompt is a source code.

[0022] In a seventeenth aspect of the present disclosure, which may be combined with any other aspect listed herein unless specified otherwise, the initial code is the final output if the Security Analysis module identifies no CWEs.

[0023] In an eighteenth aspect of the present disclosure, which may be combined with any other aspect listed herein unless specified otherwise, no CWEs is denoted as K=0, where K represents the count of identified CWEs by the security analysis module.

[0024] In a nineteenth aspect of the present disclosure, which may be combined with any other aspect listed herein unless specified otherwise, further comprising feeding the secure prompt to the LLM until K=0.

[0025] In a twentieth aspect of the present disclosure, which may be combined with any other aspect listed herein unless specified otherwise, which may be combined with any other aspect listed herein unless specified otherwise, an algorithm to optimize prompts for or secure generation of functional source code with large language models (LLMs), wherein thealgorithm is: where ^^ is a prompt, ^^∗is an optimal prompt,^^ is a code to be generated, and ^^^ is an original code either given as a code starting point or acode generated with an initial prompt ^^^, wherein the term ^^(^^) denotes a CWE count in thecode c, and ^^(^^, ^^^)is a metric representing a functional discrepancy between ^^^ and ^^.Attorney Docket No.5600234.01331

[0026] In a twenty-first aspect of the present disclosure, any of the structure, functionality, and alternatives disclosed in connection with any one or more of Figs. 1 to 21 may be combined with any other structure, functionality, and alternatives disclosed in connection with any other one or more of Figs.1 to 21.

[0027] Additional features and advantages are described in, and will be apparent from, the following Detailed Description and the Figures. The features and advantages described herein are not all-inclusive and, in particular, many additional features and advantages will be apparent to one of ordinary skill in the art in view of the figures and description. In addition, any particular embodiment does not have to have all of the advantages listed herein and it is expressly contemplated to claim individual advantageous embodiments separately. Moreover, it should be noted that the language used in the specification has been selected principally for readability and instructional purposes, and not to limit the scope of the inventive subject matter. BRIEF DESCRIPTION OF THE DRAWINGS

[0028] Figure 1 shows a diagram showing without prompts from average users, LLMs tend to generate codes with security vulnerabilities, according to an example embodiment of the present disclosure.

[0029] Figure 2 shows a histogram of CWEs in 100 generated code bases according to the prompts summarized in Table 1 (left) and the histogram of the CWE count per code base (right), for Bard, CodeLlama-13B-Instruct, GPT-3.5 Turbo, and GPT4 row-wise, respectively, according to an example embodiment of the present disclosure.

[0030] Figure 3 shows an overview of present disclosure’s pipeline, according to an example embodiment of the present disclosure.

[0031] Figure 4 shows a comparison of the average inter-version and average intra- version graph edit distances per code base, this result invalidates the notion of a data label in supervised secure code generation, according to an example embodiment of the present disclosure.

[0032] Figure 5 shows code reconstruction from CFG graph edits, according to an example embodiment of the present disclosure.

[0033] Figure 6 shows performance evaluation plots with CFG-type graphs: (a) per- code base CWE count before and after applying PromSec and (b) the histogram CWEs before and after applying PromSec, according to an example embodiment of the present disclosure.Attorney Docket No.5600234.01331

[0034] Figure 7 shows (a) The number of iterations required to cleanse all CWEs and (b) the count of remaining CWEs after the first iteration with CFG graphs, according to an example embodiment of the present disclosure.

[0035] Figure 8 shows (a) Comparing PromSec and BL1 in terms of per-code CWE counts, (b) remaining CWE distribution, and (c) code graph similarity, according to an example embodiment of the present disclosure.

[0036] Figure 9 shows the best template distribution across code bases for BL1, for 33% of the codes, no context is required, according to an example embodiment of the present disclosure.

[0037] Figure 10 shows (a) Comparing PromSec with BL2 in terms of per-code CWE count, (b) remaining CWE distribution, and (c) code graph similarity, according to an example embodiment of the present disclosure.

[0038] Figure 11 shows performance evaluation with AST-type graphs: (a) per code base CWE count before and after applying PromSec and (b) CWE counts before and after applying PromSec, according to an example embodiment of the present disclosure.

[0039] Figure 12 shows performance evaluation with DFG-type graphs: (a) per-code base CWE count before and after applying PromSec and (b) CWE counts before and after applying PromSec, according to an example embodiment of the present disclosure.

[0040] Figure 13 shows the per-code base CWE count before and after applying PromSec, and the histogram CWEs before and after applying PromSec with masking: CWE 259, CWE 89, and CWE 78 in thefirst, second, and third columns, respectively, according to an example embodiment of the present disclosure.

[0041] Figure 14 shows GPT-to-BARD cross-LLM transferability: (a) per-code base CWE count before and after applying PromSec, (b) CWE counts before and after applying PromSec, and (c) code graph similarity with using GPT 3.5 Turbo to optimize the prompts, and Bard to generate code, according to an example embodiment of the present disclosure.

[0042] Figure 15 shows GPT-to-Llama cross-LLM transferability: (a) per-code base CWE count before and after applying PromSec, (b) CWE counts before and after applying PromSec, and (c) code graph similarity with using GPT 3.5 Turbo to optimize the prompts, and CodeLlama-13B-Instruct to generate code, according to an example embodiment of the present disclosure.

[0043] Figure 16 shows for SC1 (training and testing with Java), and SC2 (training with Python and testing with Java): (a) the CWE count before and after applying PromSec, (b)Attorney Docket No.5600234.01331 the CWE ID histograms before and after applying the present disclosure, and (c) graph similarity, according to an example embodiment of the present disclosure.

[0044] Figure 17 shows a visual representation of: (a) a source code, and its corresponding, (b) AST graph, (c) CFG graph, and (d) DFG graph., according to an example embodiment of the present disclosure.

[0045] Figure 18 shows for the Python dataset used, histograms of: (a) the CWE count per code base, (b) the CWEs in the dataset, (c) the line count per code base, and (d) the function count per code base, according to an example embodiment of the present disclosure.

[0046] Figure 19 shows a brief explanation of the CWEs in the Python dataset, according to an example embodiment of the present disclosure.

[0047] Figure 20 shows a comparison between (a) the original insecure code, which exhibits a vulnerability to CWE-78 (Command Injection) and (b) the secured code, which mitigates the CWE-78 vulnerability by separating user input from commands, the difference in security arises from distinct prompts provided for code generation in (c) and (d), according to an example embodiment of the present disclosure.

[0048] Figure 21 shows, with removing the gGAN from the present disclosure’s pipeline, performance evaluation plots of: (a) the per-code base CWE count before and after applying PromSec and (b) the CWEs before and after applying PromSec (With CFG-type code graphs), according to an example embodiment of the present disclosure. DETAILED DESCRIPTION

[0049] The present disclosure generally relates to a method and system for prompt optimization for secure generation of functional source code with large language models (LLMs).

[0050] The present disclosure (“PromSec”) addresses these challenges using a graph generative adversarial network (gGAN) model, trainable on a differentiable contrastive loss that synergizes security reduction while upholding intended functionality. This model effectively guides the LLM in generating better prompts that enhance security without compromising functionality. The model’s effectiveness is demonstrated through an iterative interaction loop with the LLM, progressivelyfine-tuning the generation process to meet these dual objectives.

[0051] Prior Approaches and Limitations. As pictorially depicted in Fig. 1, LLMs are inclined to produce functional and insecure code bases if one depends on regular prompts by average users. This is especially the case whenever the intended code’s application has anyAttorney Docket No.5600234.01331 security notation. Efforts in prompt engineering for secure code generation are exhaustive and time-consuming. Methods like bounded model checking (BMC) are restricted to memory access issues. The only effort attempting prompt optimization is in the literature, but it requires white-box access, limiting its applicability to proprietary models like GPT and Bard. Overall, these approaches fall short of offering a unified framework for optimizing prompts for security and functionality, as depicted in Fig.1.

[0052] PromSec Overview. Based on the previous discussion, the Disclosed Invention generates a prompt optimization in code generation as a two-objective optimization problem aiming at code security and functionality preservation. To meet these objectives despite the non-differentiable nature of prompts with respect to their quantifications, the Disclosed Invention follows an iterative approach where in each iteration; the Disclosed Invention improves the code and then improve the prompt. Code improvement is achieved by utilizing a gGAN model specially designed for this purpose; given an input code, it generates an output code with semantic similarity and reduced vulnerability. The model is trained to do so by minimizing a novel (differentiable) contrastive loss that involves the number of CWEs 1 as defined by MITRE and the graph embedding distance. The LLM is then used to produce an improved prompt based on the improved code. In essence, the key difference this approach has over existing methods is the use of actual codefixes to steer the generation of improved prompts, rather than leaning on diagnostic tools such as what is done in the literature. This approach is anticipated to guide prompts toward a more thorough exploration of the solution space, as evidenced in performance assessments where vulnerabilities are swiftly resolved with minimal iterations at a low cost.

[0053] Introduction of PromSec - an algorithm for automatically optimizing prompts in LLM to generate secure source code while upholding its intended functionality.

[0054] Development of a generative graph neural network model, characterizing semantic-preserving and security-enforcing source codefixing as a two-objective optimization problem. This model, trained by minimizing a novel contrastive loss function, enables the application of semantic-preserving securityfixes to the code graph.

[0055] Demonstrating the efficacy of PromSec through comprehensive empirical validation, which confirms its capability in enhancing security and functionality in code generation using LLMs, and transferability across different LLMs, CWEs, and programming languages.

[0056] Code Generation with LLMs. The idea of code generation traces back to sequence-to-sequence models, such as recurrent neural networks (RNN) and long short-termAttorney Docket No.5600234.01331 memory (LSTM). Here, the model takes a user’s natural language prompt, representing their intent, and produces a sequence of tokens representing the desired code. Subsequently, Transformer models , and later LLMs, have contributed to better code generation. Today, LLMs have gained prominence in thisfield. Trained on extensive datasets encompassing various programming languages and patterns, LLMs excel in producing syntactically correct, logically consistent, and functioning source code in many programming languages. Their use extends beyond code generation to debugging, code translation, and even full software application development. As a result, employing LLMs can boost productivity by 30%to 50% in software development and reduce the costs of Information Technology (IT) by up to 65%.

[0057] However, concerns remain about their ability to generate secure and vulnerability-free code consistently. The existence of such vulnerabilities in code can enable malicious exploitation, potentially leading to cyberattacks, data breaches, and system instability. Code vulnerabilities, including buffer overflow, injection attacks, and authentication bypass, can lead to unauthorized access and data theft.

[0058] Graph Representation of Source Code. Graph representation of code abstracts relationships among different program elements and serves as input for a wide range of software analysis tasks. The three widely used types of graph representations of code are abstract syntax trees (ASTs), controlflow graphs (CFGs), and dataflow graphs (DFGs). Fig. 18 illustrates constructing these graphs from a given code example. ASTs provide a hierarchical and structured representation of the syntactic elements within a program, enabling in-depth examination of code syntax and organization. CFGs, on the other hand, capture the controlflow within a program, representing how control is transferred between different code blocks and aiding in understanding program execution paths. DFGs represent data dependencies and theflow of information within the program, offering insights into data- centric aspects of code behavior.

[0059] Graph-based code repair has emerged as a useful tool mainly applied for detecting malfunctioning portions of source code (patches) and suggesting remedies to them. As leading graph representation models, GNNs can effectively capture the complex relationships and dependencies inherent in programming languages. In the context of automatic code repair, GNNs are mainly used to identify buggy portions of source code thereby hinting at potentialfixes. Thesefixes can range from minor syntax corrections to significant structural changes aimed at closing security gaps.

[0060] GNNs and Generative GNNs. Graph neural networks (GNNs) extend deep learning to graph data, utilizing neural network layers for message passing and aggregation.Attorney Docket No.5600234.01331 Efficiently transforming graph information into node embeddings, GNNs excel in capturing complex relationships within graphs, leading to SoTA performance in various applications. Among generative GNNs, graph generative adversarial networks (gGANs) and graph variational autoen-coders (gVAEs) are particularly noteworthy. A gGAN employing adversarial frameworks consists of a generator and a discriminator, for complex network synthesis. The generator creates graphs, while the discriminator assesses their authenticity. Conversely, a gVAEs uses a variational autoencoder architecture, encoding graph data into a latent space and then decoding it to generate new graphs.

[0061] Contrastive Learning. Contrastive learning is a model training technique that emerged as a remedy for the data-labeling requirement of supervised learning. Instead of using specific data labels, a contrastive learning scheme identifies positive and negative data pairs. A positive pair comprises two data points of relatively strong similarity as opposed to a negative pair. Therefore, the model is trained to minimize the differences between items in a positive pair and maximize it for items in a negative pair. A clear advantage of this learning scheme is alleviating the need for data labels explicit supervision. Contrastive learning has shown notable advantages in tasks like image and text representation, and its applicability extends to the domain of GNNs. In GNNs, specifically generative GNNs, contrastive learning facilitates the learning of node and graph embeddings that accurately capture the underlying structure and features of graph data. This makes it particularly suitable for tasks where understanding the relationships and structures within data is crucial, such as in social network analysis, recommendation systems, and bioinformatics.

[0062] How Secure are LLMs in Code Generation? Since LLMs are commonly perceived to generate codes prone to security vulnerabilities, the initial focus was to assess this concern on modern LLMs. To achieve this objective, the Disclosed Invention conducted the following experiment. Designing a total of 50 prompts, such that each prompt asks to generate a source code that does a certain task. While tasks have some security notions, these prompts do not explicitly ask the LLM to beware of security issues, nor do they represent bad programming practices. They are just subtle verbal descriptions of what an LLM is requested to generate, and they are aimed to represent the skill level of an ordinary LLM user without security expertise. Table 1 summarizes the main topics covered by these prompts. These topics cover a wide range of Python scripting and application development areas, including database access, security, encryption, network and system management, user authentication,file and data processing, web scraping and mathematical and data analysis.Attorney Docket No.5600234.01331

[0063] One feeds the prompts to a set of LLMs and observe their generated codes. Then, one analyzes the security of the generated code bases using tools such as Bandit and SpotBugs. Each prompt is used twice resulting in 100 generated code bases. Fig. 2 shows the outcomes of the security analysis across four prominent LLMs: Google’s Bard, Meta’s CodeLlama-13B-Instruct, OpenAI’s GPT-3.5 Turbo, and OpenAI’s GPT4, presented in rows, respectively. For each LLM (row), onefirst plots a histogram of the identified CWEs in the generated codes (left) and a histogram of the counts of CWEs in each code base (right). The figure supports previous studies – that LLMs generate codes prone to security vulnerabilities as evidenced by the identified CWEs.

[0064] CWE 78 stands out as the most frequently recurring security vulnerability among the identified ones. There is a noticeable overlap in the identified CWEs across the four LLMs. Moreover, the security analysis reveals that no more than 36% of the generated codes are entirely secure, according to the security analysis tools one uses for Python and Java codes. Thesefindings are derived without prompting the LLMs to explicitly or implicitly create CWEs. CWEs surfaced solely because the LLMs were instructed to generate code for security- related tasks. This experiment illustrates that SoTA LLMs often produce insecure code when tackling security-related tasks. Hence, it is crucial not to unquestionably rely on their outputs.

[0065] The key idea of PromSec involves the application of contrastive learning within a graph generative adversarial network alongside static security analysis. This unified method aims to direct the LLM in producing secure and functional code as an alternative to conventional optimization approaches that are hindered by the non-differentiable nature of generating secure and functional code.Attorney Docket No.5600234.01331

[0066] Fig.3 shows the main steps of the proposed PromSec approach. First, an initial prompt is obtained. Such a prompt is taken directly from the user as a verbal description characterizing her / his intent. The initial prompt is forwarded to the LLM and a code is generated. The generated code is then forwarded to the Security Analysis module, which is a static security tool to identify the number of CWEs in the code. Alternatively, if the input to the pipeline is a source code ^^, then this code input is directly passed to the Security Analysis module. If the security analysis identifies no CWEs, denoted as “^^ =0” (where ‘^^’ represents the count of identified CWEs by the security analysis tool), the code is considered thefinal output. Otherwise, the code is transformed into a graph representation g (e.g., AST) using the Graph Construction module, then forwarded to a graph generative adversarial network, named gGAN. The objective of this pre-trained gGAN model is to enhance the input graph representation, striving for a more secure (with reduced CWEs) one, ^^^, while maintaining the intended functionality. Once generated, the graph ^^^ is converted back into code cˆ through a Code Reconstruction module. This reconstructed code is then reverse-engineered by the LLM to generate a prompt corresponding to ^̂^, ^̂^. The new prompt ^̂^ is fed back into the LLM to produce a new version of the code, initiating a loop that continues for a set number of cycles or until the CWE count is ideally reduced to zero.

[0067] Code Fixes are One-to-Many - Function-implementation relationship in codes is well known to be a one-to-many mapping, i.e., the same semantic function can be implemented as a code in many ways. Furthermore, code similarity detection systems like Facebook’s Aroma and Intel’s Machine Inferred Code Semantics Analysis (MISIM) illustrate the ability to recognize different code implementations fulfilling identical functions, supporting the one-to-many mapping.

[0068] It is interesting to discuss the implication of the one-to-many function-to-code relationship on suitable learning of generative models on code graphs. Thus, it is possible to obtain an improved code by utilizing a generative GNN model trained to generate better code graphs. In supervised learning of generative models, model training requires labeled data, where a label of a data point represents the desired outcome that the model needs to generate. In an application like source codefixing, an immediate approach is to simply feed a generative model with an input code graph and train it to generate a secure version of this code graph. However, recalling the aforementioned one-to-many function-code mapping, it is evident that the same code can have many secure versions. In other words, for a given insecure code snippet, there could be multiple, structurally diverse but functionally equivalent secure versions.Attorney Docket No.5600234.01331 Accordingly, there is no notion of a single desired outcome, or equivalently, no notion of a unique data label.

[0069] To showcase this limitation, the following experiment: consider a set of 100 Python code bases, where each code base has multiple CWEs. For each code base, one generates three CWE-free versions. A CWE-free version is a secure version of the original insecure code (i.e., free of CWEs) and functionally similar to the original code. Next, one defined and calculated the following quantities. First, the average inter-version distance which is the mean graph edit distance between the code graph of the original code and that of its corresponding CWE-free versions. Second, the average intra-version distance denotes the average graph edit distance among the graph representations of the CWE-free versions of each code base. One calculated these metrics for the 100 considered code bases.

[0070] Fig. 4 illustrates the results of those two metrics per code base (normalized by the maximum value). It is worth noting that the average inter-version and average intra-version distances are comparable. Furthermore, on many occasions, the inter-version distance is even less than the intra-version distance. This observation indicates that the edit distance between the original code base and its CWE-free versions is less than or at least comparable to the distance among its CWE-free versions.

[0071] The observed diversity in secure and functional code versions suggests the viability of integrating contrastive learning into the gGAN model for code improvement, alongside static security analysis. This combination provides an effective alternative approach for guiding LLMs in generating secure, functional code using optimized prompts. Contrastive learning’s strength lies in distinguishing between various examples, making it suitable for scenarios with multiple valid outcomes. It can train models to recognize a broad range of secure coding patterns beyond replicating a single secure version, expanding their understanding of diverse secure coding practices.

[0072] The Training and Operation Pipeline - The Disclosed Invention formulates prompt optimization for secure and functioning code generation with LLMs as follows:to be generated, and ^^^is theoriginal code either given as a code starting point or the code generated with an initial prompt^^^. The term ^^(^^) denotes the CWE count in the code ^^, and ^^(^^, ^^0) is a metric representingthe functional discrepancy between ^^^and ^^.Attorney Docket No.5600234.01331

[0073] As highlighted previously, the essential challenge in this optimization is that pis not differentiable with respect to ^^(^^) and ^^(^^, ^^0). Even modeling these functions in termsof p is challenging since it requires modeling the LLM with external security analysis tools. To address this challenge, one approximates the optimal prompt ^^∗⇤ in Eq. 1 by using a gGAN model to generate improved code graphs that lead to improved prompts for generating secure and functioning code in an iterative manner, as follows.

[0074] The graph-based generative adversarial network (gGAN) improves code security by modifying its graph representations, which are then used to refine input prompts for large language models. It follows a generator and discriminator framework, where the generator produces secure and functionally accurate graph structures, and the discriminator evaluates whether they meet security and functionality requirements. Since gGAN operates on graph data, it is classified as a GNN-based GAN. Architectures like the graph convolutional network (GCN) can be used for this purpose. The graphs, such as Abstract Syntax Trees (ASTs), Control Flow Graphs (CFGs), and Data Flow Graphs (DFGs), are extracted from the original code and analyzed to enhance security while preserving functionality.

[0075] The learning mechanism for gGAN is described in Algorithm 1. The model is trained with a contrastive loss function that guides the generator to produce secure and functionally correct code graphs. This loss encourages modifications that reduce security vulnerabilities while preserving the intended functionality.:

[0076] The generator loss consists of adversarial and contrastive components, denoted as ℒ^ௗ௩and ℒ^^^௧^^^௧^௩^, respectively. These losses are combined using a weighting factor ^^, which balances security improvements with functional preservation. The generator loss isexpressed as follows:

[0077] the generator ^^ can produce outputs that the discriminator misclassifies as valid. It is defined as follows.is designed to minimize the number of Common Weakness Enumerations (CWEs), denoted as ^^, while preserving codeAttorney Docket No.5600234.01331 functionality through an embedding similarity measure S: original code ^^^, by the generated graph ^^^^. The term ^^^denotes the embedding similarity between ^^^and ^^^, ensuring that structural and functional are preserved. Similarly, ^^^^represents the differencein CWE counts between the codes corresponding to graphs ^^^and ^^^, while ^^^^measures theirsimilarity in the embedding space. The function ^^(^^) = exp(−^^) the weighting ofthese terms, and ^^ is a balancing parameter. The summation is taken all graphs in the training dataset to guide the generator toward producing secure yet functionally consistent graph representations.

[0079] The design of the contrastive loss function aims to encourage the generation of graphs with fewer CWEs and greater similarity in embeddings between the original and generated graphs, thereby maintaining code functionality.

[0080] Discriminator Loss. The discriminator loss ^^ follows the standard binary cross- entropy formulation used in generative adversarial networks. It quantifies how well the discriminator ^^ distinguishes between real graphs ^^^and generated graphs ^^(^^^), ensuring that the generator improves over time. The loss is defined as:the discriminator parameters ^^^and generator parameters ^^ீare updated as followswhere η represents the learning rate, controlling the step size in the gradient descent updates.

[0082] The disclosed system converts source code into structured graph representations, such as Abstract Syntax Trees (ASTs), Control Flow Graphs (CFGs), and Data Flow Graphs (DFGs), using standard parsing tools. Converting an AST back into code is straightforward through an unparser. However, reconstructing code from a CFG is moreAttorney Docket No.5600234.01331 complex due to the nature of control flow dependencies. To address this, the system integrates the initial code, its AST, and CFG with the newly generated CFG to systematically reconstruct an updated version of the code, as illustrated in Fig.5. This process modifies specific elements of the original code by resolving discrepancies between the AST and CFG representations of both the original and updated versions. These modifications ensure that the reconstructed code aligns with the new CFG while preserving its original structure and functional consistency. The approach used for DFG-to-code conversion follows a similar methodology.

[0083] Experiments - The Setup, Dataset, and Baselines - The setup includes an LLM used for code generation. An LLM prompt is a verbal description of the intent for generating the code. This description is either obtained directly from an assumed user or obtained by instructing the LLM to generate an approximated prompt in case of starting with source code. The LLM is by default GPT-3.5 Turbo unless otherwise stated. Two datasets are used. Thefirst is a Python code base set obtained from and comprising source codes from MITRE’s documentation, CodeQL documentation, and handcrafted codes from the literature.Attorney Docket No.5600234.01331

[0084] Comparison of performance of PromSec with the following baselines (BLs): BL1 characterizes the literature’s work accustomed to code generation.7 prompt templates are defined with increasing context. The contexts of templates are summarized in Table 3. This is because security analysis tools typically return the IDs of identified CWEs, the lines of code where they occur, and the confidence in their detection. BL2 applies BL1 iteratively and at each iteration, greedily selecting the solution with the minimal CWEs count.1) Security Enhancement. Quantified by the reduction in the number of CWEs (k) as commonly employed in. This metric is computed based on the outcomes of static security analysis techniques that detect and enumerate CWEs. Preservation of Code Functionality. Quantified by the similarity between code graphs as commonly employed in many works such as the literature. To supplement this, fuzzing tests are conducted on a randomly selected subset of generated code bases, providing an additional layer of evaluation. Operational Cost. Evaluated based on three parameters, including the time taken for code execution (execution time delay), thefinancial and delay implications of utilizing LLMs, and the costs of static security analysis.

[0085] Q1: How successful is PromSec in code security enforcement? The gGAN model is trained with 500 training graphs. Then, one feeds 100 test code bases to PromSec and set the maximum number of iterations to 10. First, with a CFG graph type, Fig.6(a) shows the number of CWEs for each code base before and after processing with PromSec. Next, Fig.6(b)Attorney Docket No.5600234.01331 presents a histogram of the CWEs in the code bases before and after applying PromSec. It is clearly noticed from bothfigures that PromSec resolves all CWEs within 10 iterations.

[0086] Q2: Does PromSec need iterative LLM interaction? For this purpose, one plotted a histogram of the required number of iterations to fully secure each code based on Fig. 7 (a). One also plotted a histogram of the remaining CWEs in each code base after one PromSec iteration in Fig. 7(b). The two histograms confirm that one iteration is not sufficient to secure the majority of code bases (68%). On the other hand, around 49% of the code bases are fully cleaned from CWEs within thefirst two iterations.

[0087] To answer Q3: How does PromSec perform relative to the baselines regarding security, functionality preservation, and cost? The disclosure replicates the experiment of Q1 and compare PromSec’s performance with that of the Baseline 1 (BL1). The results are shown in Fig.8. The number of CWEs per code base is shown in 8(a), the CWE distribution in 8(b), and the comparison of the graph similarity metric (a proxy for functionality preservation) in Fig.8(c).

[0088] Observing how well BL1 does in the best and worst scenarios in terms of reducing the number of CWEs: In the best scenario, one chose the prompt template in BL1 that leads to the fewest CWEs. Conversely, in the worst scenario, one chose the one that leads to the maximum CWE count. The results show that PromSec is significantly better than BL1 in resolving CWEs. BL1 doesfix some CWEs, but the code it produces is still not completely secure. In terms of preserving code functionality, especially in its best situation, BL1 does not do as well as PromSec. This result shows that PromSec has superior performance both infixing security bugs (CWEs) as well as upholding the intended code functionality.

[0089] (BL1) report that adding more context does not always make code completions more secure. To examine this observation in code generation, Fig. 9 shows how often each prompt template is the most effective for a given test code. The templates are indexed according to the increase in the context. The histogram in thisfigure reveals an interesting observation: Template 1, which does not use any security report context, turns out to be the best choice 33% of the time. This result suggests that while including some context is beneficial, adding too much does not always help. In fact, there was a decrease in preference for templates with more context.

[0090] Next, one compares the effectiveness of PromSec against the other baseline, BL2, and present the results in Fig.10. Thisfigure reveals that BL2 outperforms BL1 in terms of security enhancement, as evidenced by a notably lower count of remaining CWEs. However, this improvement in security enforcement comes at the expense of a reduction in codeAttorney Docket No.5600234.01331 functionality preservation as shown in Fig.10(c), where BL2 results in a greater degeneration in code functionality.

[0091] Fuzzing Tests. This experiment addresses Q4: How effective is PromSec in preserving the intended code functionality? In software engineering, fuzzing testing (FT) is a procedure typically used to compare the functionality of two source codes. It involves giving the same inputs to both codes and then measuring differences in their outputs. However, FT usually requires a lot of manual work and can be time-consuming, especially for large numbers of code bases. Also, FT only works well when the code bases have clear and measurable inputs and outputs. Because of these limitations, one uses FT as a secondary method to check how well PromSec preserves the original semantics and function of code while improving its security. A comparison of PromSec with the BL1 is shown in Table 4.

[0092] For conducting FT, one uses 20 specially designed code bases that can be tested with fuzzing. Each code base has a main function with specific input and output parameters and calls several internal functions. This setup allows one to quantify any changes in the outputs caused by PromSec. These tests are repeated for 1,000 trials. In each trial, one samples random inputs and calculate the differences in outputs. Then, one reports the average value of the output differences.

[0093] Table 4 shows that PromSec perfectly preserves the same functionality as quantified by the zero change in outcomes. On the other hand, BL1 exhibits some degree of functionality preservation but lower than that achieved by PromSec. BL2 is excluded from the FT comparison due to the clear evidence that its functionality preservation is significantly degraded compared to BL1, as indicated in the previous experiment.

[0094] Empirical Cost Analysis The third aspect of comparing PromSec to the baselines is the operational cost. One compares the costs of BL1 along with 1 iteration of PromSec denoted by (PromSec(1)), and 10-iteration BL2 with 10-iteration PromSec (denoted by PromSec(10)). The results are shown in Table 5. As seen in the table, PromSec(1) demonstrates improvements over BL1. The overall time cost is notably lower with PromSec(1), with a mean of 9.95 seconds, in contrast to 35.62 seconds. The LLM query cost category shows a reduction in the number of queries, input and output token counts, and LLM model query time, with PromSec(1) requiring 2 queries versus 7 in BL1, and token counts of 359.2 and 256 compared to 1,153.1 and 574.9, respectively. The query time is also reduced to 9.5 seconds from 33.94 seconds. In the security analysis cost, PromSec(1) requires fewer security analysis queries (1 vs. 7) and less time (0.237 seconds vs. 1.652 seconds). These results indicate that PromSec(1) is more efficient in terms of both time and resource management compared to theAttorney Docket No.5600234.01331 existing BL1 baseline. A more obvious advantage is observed for PromSec(10) over BL2, where PromSec(10) has around an order of magnitude reduction compared with BL2.

[0095] Q5: What is the impact of the code graph type on performance? One evaluates the performance of PromSec when it utilizes either AST or DFG graph types instead of CFG. Fig.11 shows the performance analysis based on AST graphs. This performance is acceptable but inferior when compared to the scenario involving CFG graphs. Part of the reason for this is that although ASTs encapsulate the code’s syntax, they lack the depth needed to encompass controlflow and operational semantics, crucial for a comprehensive security analysis. Next, Fig. 12 shows the corresponding performance analysis plots with DFG graphs. Although the performance of AST is slightly better than DFG graphs, both are inferior to CFG graphs. DFGs focus on dataflows and dependencies but do not capture the logic and behavior of codes as done by CFGs.

[0096] PromSec’s Transferability to Unseen CWEs - As shown in the experiments, the ability of the trained gGAN model tofix CWEs is established. Q6: Can PromSec transfer to unforeseen CWEs in its training process? For this purpose, one conducts the following experiment. Excluding a given test CWE from the training set of the gGAN model and refer toAttorney Docket No.5600234.01331 it as a masked CWE. Then, one picks test examples that exhibit this masked CWE. In other words, the model does not see any training examples exhibiting the masked CWE.

[0097] Fig.13 depicts the effects of masking specific CWEs, namely, CWE 259, CWE 89, and CWE 78, on PromSec’s performance. Several observations can be drawn from this figure. First, masking a given CWE results in a slight decrease in PromSec’s effectiveness against it and a minor reduction in handling some other related CWEs. This observation validates PromSec’s ability to transfer its security-enhancing capabilities to CWEs not directly included in its training. This transferability is attributed to the generalization capability of the gGAN model. Furthermore, it is intuitive to assume that resolving a specific CWE may have a beneficial impact on mitigating related CWEs, especially if they originate from similar underlying issues. For example, prompt optimizations directed to prevent CWE 78 (input injection) might implicitly also mitigate CWE 89 (SQL injection) risks. This result is because both vulnerabilities typically stem from inadequate control of user input, whether it be in operating system commands (in CWE 78) or database queries (in CWE 89). This analysis hints at the possibility of mitigating new and salient security vulnerabilities not explicitly included in the training process of the gGAN.

[0098] PromSec’s Transferability across LLMs - Considering the pipeline in Fig. 3, one addresses Q7: Are promptfixes obtained with a given LLM exclusively valid for that LLM or can they be transferable to other LLMs? For this purpose, one assumes GPT-3.5 Turbo as the base LLM used for optimizing the prompts while feeding these optimized prompts to another LLM to generate the code. One considers Google Bard and Meta’s CodeLlama-13B- Instruct as these LLMs and show the corresponding results in Figs.14 and 15. In bothfigures, prompts optimized with GPT-3.5 Turbo work well with both Bard and CodeLlama-13B- Instruct. The performance exhibits negligibly small remaining CWEs. This result validates the ability of PromSec’s prompt optimizations to transfer across LLMs.

[0099] PromSec’s Transferability across Programming Languages - Q8: How successfully can PromSec’s optimized prompts generalize across different programming languages? For this purpose, a dataset of Java prompts from the literature is collected. Next, one considers these prompts as inputs to PromSec’s pipeline illustrated in Fig. 3. One intentionally employs these prompts to examine the operation of this pipeline starting with prompts, rather than source code as done in the tests considered so far. Next, one trains a gGAN model over 200 code bases generated from these prompts. After that, one uses the remaining 100 prompts to test the operation of PromSec (denote this scenario by SC1). To examine the generalization of PromSec across programming languages, one tests this prompt dataset withAttorney Docket No.5600234.01331 the previous gGAN trained on Python code (denote this scenario by SC2). Fig. 16 shows a histogram of the CWEs in the test set before and after applying PromSec for these two scenarios in (a) and (b), respectively. Fig.16(a) clearly shows that PromSec perfectly resolves the CWEs, similar to the case with Python code experiments. Note that one uses the same hyperparameters and number of iterations. Next, Fig. 16(b) shows almost the same performance is attainable with PromSec despite being trained using codes from another language. The result in Fig.16(b) showcases the transferability of PromSec across source code of different programming languages.

[0100] As afinal remark, one conducts an ablation experiment to examine the importance of the role of the gGAN model in the process of code security enforcement. The results show that PromSec exhibits a significant degradation in its security enhancement performance once the gGAN model is removed. This observation validates the important role the gGAN model plays in this pipeline.

[0101] Exemplifying Code Graphs - As an example illustrating the construction of code graphs, Fig.17 shows an example code snippet in (a) with its AST, CFG, and DFG graphs in (b), (c), and (d), respectively. Part (b) constructs the AST showing a hierarchical tree structure that represents its syntactic organization, including nodes for functions, arguments, and operations. Part (c) shows the CFG, which captures the possible paths of execution. It outlines the sequence of operations and decision points, revealing how the controlflows through the program. Part (d) displays the DFG, which focuses on theflow and usage of data within the code. It maps out where data is generated (inputs), how it is manipulated (through operations), and where it ends up (outputs). Together, these graphs provide a multi-faceted view of the source code, each emphasizing a different aspect of its structure and execution logic.

[0102] Python Dataset Information - Fig 18 summarizes key information about the Python dataset used in the majority of the experiments. Fig.18(a) shows a histogram of the CWE count per code base while Fig. 18(b) shows a histogram of the CWEs available in the whole data set. These CWEs are clarified in Fig.19. Also, Fig.18(c) shows a histogram of the number of lines per code base whereas Fig.18(d) shows a histogram of the number of functions per code base.

[0103] A Tracing Example - In addition to the numerical experiments on code security enhancement with PromSec, it is interesting to trace the evolution of a prompt and how that leads to code security. One exemplifies this with the following example code snippet. The code shown in Fig. 20(a) exhibits a security vulnerability characterized by CWE-78Attorney Docket No.5600234.01331 (Command Injection). This is clear as it allows direct execution of user-provided commands without proper validation or separation of user input from the commands. This implementation enables attackers to potentially inject malicious commands, posing a significant security risk. Conversely, Fig. 20(b) presents the secured code, which effectively mitigates the CWE-78 vulnerability. It achieves this by utilizing the shlex module to safely parse user input and separate it from the commands, ensuring that only intended commands are executed. Fig.20(c) shows the original prompt leading to the generation of the original code, and in Fig. 20(d), is the optimized prompt that leads to the generation of the secure code. The optimized prompt differs in asking for separating the user input from the actual commands, thereby leading to the resolution of CWE 78.

[0104] An Ablation Study - Since the pipeline of PromSec is based on iteratively using the gGAN model and interaction with the LLM, it is interesting to investigate whether the gGAN or solely interacting with the LLM is the actual source of code securing enforcement. For this purpose, one conducts the following ablation experiment. One applies PromSec while not using the gGAN and only querying the LLM for 10 iterations. The results of this experiment are shown in Fig.21.

[0105] As can be seen in Fig. 21, it is evident that the exclusion of the gGAN model from PromSec’s pipeline strongly harms the performance. Some CWEs have indeed been removed. However, many others are remaining. This highlights the importance of the involvement of the generative GNN in PromSec.

[0106] It should be understood that various changes and modifications to the presently preferred embodiments described herein will be apparent to those skilled in the art. Such changes and modifications can be made without departing from the spirit and scope of the present subject matter and without diminishing its intended advantages. It is therefore intended that such changes and modifications be covered by the appended claims.

Claims

Attorney Docket No.5600234.01331 CLAIMS The invention is claimed as follows:

1. A method for prompt optimization for secure generation of functional source code with large language models (LLMs), the method comprising: obtaining an initial prompt; forwarding the initial prompt to a Large Language Model (LLM); generating an initial code using the initial prompt, wherein the LLM generates the initial code; forwarding the initial code to a Security Analysis module; transforming the initial code into a graph representation; forwarding the graph representation to a graph generative adversarial network; enhancing the graph representation to a secure graph representation; converting the secure graph representation into a secure code; reverse-engineering the secure code a secure prompt; feeding the secure prompt to the LLM; and producing a final output.

2. The method of claim 1, wherein the initial prompt is one of a direct prompt from a user, a verbal description characterizing the intent of the user, or a source code.

3. The method of claim 1, wherein the Security Analysis module is a static security tool to identify the number of CWEs in the initial code.

4. The method of claim 3, wherein the initial prompt is forwarded directly to the Security Analysis module if the initial prompt is a source code.

5. The method of claim 3, wherein the initial code is the final output if the Security Analysis module identifies no CWEs.

6. The method of claim 5, wherein no CWEs is denoted as K=0, where K represents the count of identified CWEs by the security analysis module.Attorney Docket No.5600234.01331 7. The method of claim 1, wherein the graph generative adversarial network is configured to enhance the graph representation by reducing CWEs while maintaining intended functionality.

8. The method of claim 1, wherein the secure graph representation is converted into the secure code through a Code reconstruction module.

9. The method of claim 6, further comprising feeding the secure prompt to the LLM until K=0.

10. The method of claim 1, wherein the secure prompt is fed to the LLM for a set number of user-selected cycles.

11. A system for prompt optimization for secure generation of functional source code with large language models (LLMs) comprising: a server; a processor; and a memory storing instructions, which when executed by the processor, cause the processor to apply an algorithm to optimize prompts for or secure generation of functional source code with large language models (LLMs).

12. The system of claim 11, wherein the algorithm to optimize prompts for or secure generation of functional source code with large language models (LLMs) is: where ^^ is a prompt, ^^∗to be generated, and ^^^is an original code either given as a code starting point or a code generated with an initial prompt ^^^, wherein the term ^^(^^) denotes a CWE count in the code c, and ^^(^^, ^^^)is ametric representing a functional discrepancy between ^^^and ^^.Attorney Docket No.5600234.01331 13. The system of claim 11, wherein the algorithm to optimize prompts for or secure generation of functional source code with large language models (LLMs) is configured to: obtain an initial prompt; forward the initial prompt to a Large Language Model (LLM); generate an initial code using the initial prompt, wherein the LLM generates the initial code; forward the initial code to a Security Analysis module; transform the initial code into a graph representation; forward the graph representation to a graph generative adversarial network; enhance the graph representation to a secure graph representation; convert the secure graph representation into a secure code; reverse-engineer the secure code a secure prompt; feed the secure prompt to the LLM; and produce a final output.

14. The system of claim 13, wherein the initial prompt is one of a direct prompt from a user, a verbal description characterizing the intent of the user, or a source code.

15. The system of claim 13, wherein the Security Analysis module is a static security tool to identify the number of CWEs in the initial code.

16. The system of claim 15, wherein the initial prompt is forwarded directly to the Security Analysis module if the initial prompt is a source code.

17. The system of claim 15, wherein the initial code is the final output if the Security Analysis module identifies no CWEs.

18. The system of claim 17, wherein no CWEs is denoted as K=0, where K represents the count of identified CWEs by the security analysis module.

19. The system of claim 18, further comprising feeding the secure prompt to the LLM until K=0.Attorney Docket No.5600234.01331 20. An algorithm to optimize prompts for or secure generation of functional source code with large language models (LLMs), wherein the algorithm is: where ^^ is a prompt, to be generated, and ^^^is anoriginal code either as a or a code generated with an initial prompt ^^^, wherein the term ^^(^^) denotes a CWE count in the code c, and ^^(^^, ^^^)is ametric representing a functional discrepancy between ^^^and ^^.

Images