Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Inference-Time Safety For Code LLMs Via Retrieval-Augmented Revision

Manisha Mukherjee, Vincent J. Hellendoorn

TL;DR

This work presents a principled approach to trustworthy code generation by design that operates as an inference-time safety mechanism that employs retrieval-augmented generation to surface relevant security risks in generated code and retrieve related security discussions from a curated Stack Overflow knowledge base.

Abstract

Large Language Models (LLMs) are increasingly deployed for code generation in high-stakes software development, yet their limited transparency in security reasoning and brittleness to evolving vulnerability patterns raise critical trustworthiness concerns. Models trained on static datasets cannot readily adapt to newly discovered vulnerabilities or changing security standards without retraining, leading to the repeated generation of unsafe code. We present a principled approach to trustworthy code generation by design that operates as an inference-time safety mechanism. Our approach employs retrieval-augmented generation to surface relevant security risks in generated code and retrieve related security discussions from a curated Stack Overflow knowledge base, which are then used to guide an LLM during code revision. This design emphasizes three aspects relevant to trustworthiness: (1) interpretability, through transparent safety interventions grounded in expert community explanations; (2) robustness, by allowing adaptation to evolving security practices without model retraining; and (3) safety alignment, through real-time intervention before unsafe code reaches deployment. Across real-world and benchmark datasets, our approach improves the security of LLM-generated code compared to prompting alone, while introducing no new vulnerabilities as measured by static analysis. These results suggest that principled, retrieval-augmented inference-time interventions can serve as a complementary mechanism for improving the safety of LLM-based code generation, and highlight the ongoing value of community knowledge in supporting trustworthy AI deployment.

Inference-Time Safety For Code LLMs Via Retrieval-Augmented Revision

TL;DR

This work presents a principled approach to trustworthy code generation by design that operates as an inference-time safety mechanism that employs retrieval-augmented generation to surface relevant security risks in generated code and retrieve related security discussions from a curated Stack Overflow knowledge base.

Abstract

Large Language Models (LLMs) are increasingly deployed for code generation in high-stakes software development, yet their limited transparency in security reasoning and brittleness to evolving vulnerability patterns raise critical trustworthiness concerns. Models trained on static datasets cannot readily adapt to newly discovered vulnerabilities or changing security standards without retraining, leading to the repeated generation of unsafe code. We present a principled approach to trustworthy code generation by design that operates as an inference-time safety mechanism. Our approach employs retrieval-augmented generation to surface relevant security risks in generated code and retrieve related security discussions from a curated Stack Overflow knowledge base, which are then used to guide an LLM during code revision. This design emphasizes three aspects relevant to trustworthiness: (1) interpretability, through transparent safety interventions grounded in expert community explanations; (2) robustness, by allowing adaptation to evolving security practices without model retraining; and (3) safety alignment, through real-time intervention before unsafe code reaches deployment. Across real-world and benchmark datasets, our approach improves the security of LLM-generated code compared to prompting alone, while introducing no new vulnerabilities as measured by static analysis. These results suggest that principled, retrieval-augmented inference-time interventions can serve as a complementary mechanism for improving the safety of LLM-based code generation, and highlight the ongoing value of community knowledge in supporting trustworthy AI deployment.
Paper Structure (27 sections, 3 figures, 3 tables)

This paper contains 27 sections, 3 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Motivating Example
  3. Method: Inference-Time Safety via Community-Driven Retrieval
  4. Design Overview
  5. Security-Oriented Knowledge Base
  6. Retrieval of Community Discussions
  7. Inference-Time Revision Prompt
  8. Background
  9. Security in LLM-Generated Code.
  10. Vulnerability Detection and Evaluation Datasets.
  11. Security Knowledge in Stack Overflow.
  12. Retrieval-Augmented Generation for Security.
  13. Evaluation
  14. Experimental Setup
  15. Datasets.
  16. ...and 12 more sections

Figures (3)

  • Figure 1: AnswerID: 61307412, which includes community comments providing security insights. This content was used as context to enhance the generated code in SOSecure.
  • Figure 2: Inference-time revision guided by community knowledge. Left: original GPT-4 output containing a usage pattern commonly associated with CWE-078. Right: one possible revised output produced after the model is provided with a relevant Stack Overflow discussion (Figure \ref{['fig:hero']}) as contextual guidance.
  • Figure 3: Conceptual overview of inference-time revision with SOSecure. Community security discussions from Stack Overflow are retrieved based on similarity to LLM-generated code and provided as contextual guidance during revision.