Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Code Security Vulnerability Repair Using Reinforcement Learning with Large Language Models

Nafis Tanveer Islam, Mohammad Bahrami Karkevandi, Peyman Najafirad

TL;DR

This work tackles the security shortcomings of code generated by large language models, showing that purely functional optimization via cross-entropy can overlook security needs. It introduces SecureCode, a reinforcement learning framework using a causal decoder (CodeGen2) to repair vulnerable code by jointly optimizing syntactic (CodeBLEU) and semantic (BERTScore) rewards, ensuring added security lines while preserving functionality. Evaluated on the VulDeeLocator dataset, the RL approach outperforms supervised fine-tuning across BLEU and Rouge-L, with three real-world case studies illustrating the model’s ability to insert security measures without compromising behavior. The study demonstrates that RL-based secure code repair can significantly enhance the practicality of AI-assisted programming in security-critical contexts.

Abstract

With the recent advancement of Large Language Models (LLMs), generating functionally correct code has become less complicated for a wide array of developers. While using LLMs has sped up the functional development process, it poses a heavy risk to code security. Code generation with proper security measures using LLM is a significantly more challenging task than functional code generation. Security measures may include adding a pair of lines of code with the original code, consisting of null pointer checking or prepared statements for SQL injection prevention. Currently, available code repair LLMs generate code repair by supervised fine-tuning, where the model looks at cross-entropy loss. However, the original and repaired codes are mostly similar in functionality and syntactically, except for a few (1-2) lines, which act as security measures. This imbalance between the lines needed for security measures and the functional code enforces the supervised fine-tuned model to prioritize generating functional code without adding proper security measures, which also benefits the model by resulting in minimal loss. Therefore, in this work, for security hardening and strengthening of generated code from LLMs, we propose a reinforcement learning-based method for program-specific repair with the combination of semantic and syntactic reward mechanisms that focus heavily on adding security and functional measures in the code, respectively.

Code Security Vulnerability Repair Using Reinforcement Learning with Large Language Models

TL;DR

This work tackles the security shortcomings of code generated by large language models, showing that purely functional optimization via cross-entropy can overlook security needs. It introduces SecureCode, a reinforcement learning framework using a causal decoder (CodeGen2) to repair vulnerable code by jointly optimizing syntactic (CodeBLEU) and semantic (BERTScore) rewards, ensuring added security lines while preserving functionality. Evaluated on the VulDeeLocator dataset, the RL approach outperforms supervised fine-tuning across BLEU and Rouge-L, with three real-world case studies illustrating the model’s ability to insert security measures without compromising behavior. The study demonstrates that RL-based secure code repair can significantly enhance the practicality of AI-assisted programming in security-critical contexts.

Abstract

With the recent advancement of Large Language Models (LLMs), generating functionally correct code has become less complicated for a wide array of developers. While using LLMs has sped up the functional development process, it poses a heavy risk to code security. Code generation with proper security measures using LLM is a significantly more challenging task than functional code generation. Security measures may include adding a pair of lines of code with the original code, consisting of null pointer checking or prepared statements for SQL injection prevention. Currently, available code repair LLMs generate code repair by supervised fine-tuning, where the model looks at cross-entropy loss. However, the original and repaired codes are mostly similar in functionality and syntactically, except for a few (1-2) lines, which act as security measures. This imbalance between the lines needed for security measures and the functional code enforces the supervised fine-tuned model to prioritize generating functional code without adding proper security measures, which also benefits the model by resulting in minimal loss. Therefore, in this work, for security hardening and strengthening of generated code from LLMs, we propose a reinforcement learning-based method for program-specific repair with the combination of semantic and syntactic reward mechanisms that focus heavily on adding security and functional measures in the code, respectively.
Paper Structure (22 sections, 4 equations, 2 figures, 1 table)

This paper contains 22 sections, 4 equations, 2 figures, 1 table.

Table of Contents

  1. Introduction
  2. Related Works
  3. Code Vulnerability with LLMs:
  4. Vulnerability Repair:
  5. Fine-Tuning of LLMs with RL:
  6. Methodology
  7. Reinforcement Learning for Functional Repair
  8. Policy Optimization:
  9. Reward:
  10. Experiments
  11. Dataset
  12. Evaluation Metrics
  13. BLEU Score:
  14. Rouge-L:
  15. Experimental Results and Discussion
  16. ...and 7 more sections

Figures (2)

  • Figure 1: An illustrative example showcases how the Cross-Entropy (CE) emphasizes only functionality and neglects security.
  • Figure 2: A high-level overview of our proposed CodeRepair System with Reinforcement learning with semantic and syntactic loss.