Source-linked AI summary
RESCUE: Retrieval Augmented Secure Code Generation
Jiahao Shi, Tianyi Zhang
TL;DR
LLMs can generate vulnerable code, while conventional RAG struggles with noisy security documents and misses security semantics implicit in coding tasks. RESCUE constructs a distilled hierarchical knowledge base and uses multi-faceted retrieval; across its evaluation, it improves SecurePass@1 by an average of 4.8% absolute over the second-best baseline.
Problem
Conventional RAG for secure code generation struggles with irrelevant content in security documents and retrieval methods that miss implicit security requirements in task descriptions.
Source: §1 · p. 1Method
RESCUE combines LLM-assisted summarization and program slicing to build a hierarchical security knowledge base, then retrieves relevant knowledge using multiple security-focused task facets.
Source: §1 · p. 2, §2.2 · p. 4Results
Across four benchmarks and six LLMs, RESCUE improves SecurePass@1 by 4.8% absolute on average over the second-best baseline.
Source: §1 · p. 2, §3.2 · p. 6Takeaways & Limitations
The reported analyses identify refined knowledge-base construction as essential and support the effectiveness of hierarchical multi-faceted retrieval.
Source: §3.3.1 · p. 8, §5 · p. 10RESCUE adds inference latency through proactive analysis and additional LLM calls, though concurrent calls and selective invocation are proposed as mitigations.
Source: F LIMITATIONS · p. 27
Abstract
from arXiv · showhide
Despite recent advances, Large Language Models (LLMs) still generate vulnerable code. Retrieval-Augmented Generation (RAG) has the potential to enhance LLMs for secure code generation by incorporating external security knowledge. However, the conventional RAG design struggles with the noise of raw security-related documents, and existing retrieval methods overlook the significant security semantics implicitly embedded in task descriptions. To address these issues, we propose \textsc{Rescue}, a new RAG framework for secure code generation with two key innovations. First, we propose a hybrid knowledge base construction method that combines LLM-assisted cluster-then-summarize distillation with program slicing, producing both high-level security guidelines and concise, security-focused code examples. Second, we design a hierarchical multi-faceted retrieval that traverses the constructed knowledge base from top to bottom and integrates multiple security-critical facts at each hierarchical level, ensuring comprehensive and accurate retrieval. We evaluated \textsc{Rescue} on four benchmarks and compared it with five state-of-the-art secure code generation methods on six LLMs. The results demonstrate that \textsc{Rescue} improves the SecurePass@1 metric by an average of 4.8 points, establishing a new state-of-the-art performance for security. Furthermore, we performed in-depth analysis and ablation studies to rigorously validate the effectiveness of individual components in \textsc{Rescue}. Our code is available at https://github.com/steven1518/RESCUE.
1 INTRODUCTION
RESCUE addresses noise in raw security documents and retrieval methods that miss security semantics implicit in coding tasks. It combines a distilled security knowledge base with hierarchical, multi-faceted retrieval and reports improved SecurePass@1 across its evaluation.
1 INTRODUCTION
- Prior secure-generation approaches include fine-tuning, constrained decoding, and security-tool feedback, with stated data-curation, oracle, flexibility, and data-leakage concerns.Fine-tuning requires substantial curation and training; constrained decoding requires trained security models or human-crafted rules, while tool-feedback methods raise a data-leakage concern when tools are also used in evaluation.Source: §1 · p. 1
- Existing RAG approaches face noisy security documents and text-similarity retrieval that misses security requirements implicit in task descriptions.Irrelevant example logic can distract generation, while missed security semantics can lead to inaccurate retrieval.Source: §1 · pp. 1–2
- RESCUE distills high-level security guidelines and slices concise security-focused code examples, then retrieves knowledge hierarchically using API patterns, vulnerability causes, and code.The framework’s hybrid knowledge-base construction combines LLM-assisted cluster-then-summarize distillation with static program analysis.Source: §1 · p. 2
- 4.8% absolute average improvement in SecurePass@1 over the second-best baseline establishes RESCUE’s state-of-the-art security performance across four benchmarks and six LLMs.RESCUE also retains 98.7% of the original models’ capability for functionally correct code generation.Source: §1 · p. 2
2 METHOD
RESCUE builds a hierarchical security knowledge base from vulnerability fixes and retrieves tailored security context for coding tasks. It combines summarization and program slicing with task analysis and staged retrieval.
2 METHOD
- The offline knowledge-base construction and online retrieval stages assemble tailored security context to augment the coding prompt.Retrieved knowledge is intended to steer the LLM toward secure code generation.Source: §2 · p. 2
2.1 HIERARCHICAL SECURITY KNOWLEDGE BASE CONSTRUCTION
- The knowledge base organizes generalizable security knowledge above concise code examples, using CWE categories to structure vulnerabilities.Its raw dataset contains 704 vulnerabilities and fixes from CVE and GitHub projects.Source: §2.1 · p. 3
2.1.1 CLUSTER-THEN-SUMMARIZE
- The cluster-then-summarize pipeline groups vulnerability-fix instances by CWE and recursively summarizes them into security guidelines and vulnerability-cause summaries.The guidelines express actionable prevention practices, while cause summaries capture root conditions and failure patterns used during retrieval.Source: §2.1.1 · p. 3
2.1.2 SECURITY-FOCUSED STATIC PROGRAM SLICING
- Static program slicing removes unrelated example logic by tracing statements dependent on patched code in both vulnerable and secure versions.The method builds a program dependence graph, slices backward and forward around changed statements, then compares and complements the resulting variants.Source: §2.1.2 · pp. 3–4
2.2 HIERARCHICAL MULTI-FACETED RETRIEVAL
- Retrieval mirrors the knowledge-base hierarchy, first identifying relevant vulnerability types and then narrowing to secure code examples through multi-faceted queries.The method proactively analyzes security-relevant aspects of each coding task.Source: §2.2 · p. 4
2.2.1 PROACTIVE MULTI-FACETED ANALYSIS
- Task analysis produces a vulnerability-cause explanation, a zero-shot draft program, and API calls extracted from its abstract syntax tree.These facets target implicit security requirements, implementation details, and potential API misuse.Source: §2.2.1 · p. 4
2.2.2 HIERARCHICAL RETRIEVAL PROCESS.
- CWE-level retrieval ranks vulnerability types using dense cause similarity and sparse BM25 API matching, then fuses facet scores with threshold and rank filtering.The modified RRF uses a confidence threshold and keeps facet results ranked within the top 10.Source: §2.2.2 · p. 4
- Code-level retrieval adds dense similarity between draft code and sliced secure examples, then fuses three facets to select examples and assemble generation prompts.Prompts pair selected secure examples with the security guidelines for their CWE type.Source: §2.2.2 · pp. 4–5
3 EXPERIMENTS
RESCUE is evaluated across four benchmarks, six LLMs, and five baseline methods, with SecurePass@1 capturing security and functional correctness together. It outperforms existing methods on that metric, while ablations and additional analyses examine its knowledge construction and retrieval components.
3.1 EXPERIMENT SETUP
- Evaluation spans six LLMs, five baselines, and four benchmarks, using SecurePass@1 to jointly assess security and functional correctness.CodeGuard+ includes unit tests and CodeQL checks; the other three benchmarks are used to assess functional correctness.Source: §3.1 · pp. 5–6
3.2 MAIN RESULTS
- RESCUE achieves a 4.8% absolute SecurePass@1 improvement over the second-best baseline on average, while maintaining comparable functional-correctness performance.The comparison covers six LLMs and four benchmarks; INDICT has higher SecureRate than RESCUE but the lowest Pass@1.Source: §3.2 · p. 6, Table 1 · p. 7
3.3.1 ABLATION STUDY ON SECURITY KNOWLEDGE BASE CONSTRUCTION
- Using raw security data instead of the constructed knowledge base significantly reduces performance across all reported metrics, as irrelevant code logic distracts generation and retrieval.Table 2 reports five construction variants on CodeGuard+ across three models and four metrics.Source: Table 2 · p. 7, §3.3.1 · p. 8
- Removing security guidelines drops Llama3.1-8B-Instruct SP@1 by 4.8 points, while slicing benefits retrieval and generation and reduces input-token costs.The construction ablations compare five variants across three models on CodeGuard+.Source: Table 2 · p. 7, §3.3.1 · p. 8
3.3.2 IMPACT ANALYSIS OF HIERARCHICAL RETRIEVAL
- Hierarchical retrieval’s peak SP@1 exceeds the non-hierarchical setting across all three LLMs; performance follows an inverted U-shaped trend as k increases, motivating k = 4.Figure 3 compares both retrieval settings across varying numbers of selected CWE types.Source: §3.3.2 · p. 8, Figure 3 · p. 8
3.3.3 ABLATION STUDY ON MULTI-FACETED RETRIEVAL
- Across seven facet combinations, multi-faceted retrieval mostly beats single-facet approaches; all three facets usually perform best, while API and vulnerability-cause facets outperform code similarity.The comparison covers three models on CodeGuard+.Source: §3.3.3 · pp. 8–9, Table 3 · p. 9
3.4 IMPACT ANALYSIS OF SUMMARIZATION LLMS
- Performance gains persist when GPT-4o is replaced by three open-source summarization models, indicating that improvements do not depend on GPT-4o.The summarization model distills security knowledge, while the generation model is used for online code generation.Source: §3.4 · p. 9, Table 4 · p. 9
3.5 ADDITIONAL ANALYSIS
- Additional analyses report acceptable overhead, gains on dynamic testing, transfer to unseen CWEs and languages, validated RRF settings, and stronger gains for non-functionality-specific vulnerabilities.Source: §3.5 · p. 9
4 RELATED WORK
Prior code-generation RAG retrieves repositories, API documentation, or functionally similar examples, while secure-code methods span fine-tuning and training-free approaches. RESCUE instead constructs a hierarchical security knowledge base and targets security knowledge retrieval.
Retrieval-Augmented Code Generation
- Code-generation RAG has supported repository-level generation, unfamiliar-API use, and functional correctness through similar examples, while RESCUE targets security without compromising functional correctness.The cited work includes repository retrieval [Wu et al., 2024; Zhang et al., 2023], API documentation [Zan et al., 2022; Zhou et al., 2023; Liu et al., 2023b; Gu et al., 2025], and functionally similar examples [Parvez et al., 2021; Su et al., 2024].Source: §4 · p. 10
Secure Code Generation
- Secure-code approaches include fine-tuning and training-free methods such as prompting, security tools, agents, and RAG; RESCUE automatically builds a hierarchical security knowledge base from raw data and designs specialized retrieval.SecCoder retrieves secure-code examples with a dense retriever, while SOSecure retrieves StackOverflow content with BM25.Source: §4 · p. 10
5 CONCLUSION
RESCUE combines hybrid distillation for a hierarchical security knowledge base with hierarchical multi-faceted retrieval. Across five methods, four benchmarks, and six models, it improves security without compromising functional correctness.
5 CONCLUSION
- Across five state-of-the-art methods, four benchmarks, and six models, RESCUE improves security without compromising functional correctness.In-depth analyses support the necessity of knowledge-base construction and the hybrid distillation and hierarchical multi-faceted retrieval designs.Source: §5 · p. 10
- RESCUE constructs a hierarchical security knowledge base through hybrid distillation and designs a hierarchical multi-faceted retrieval method.The conclusion reports that analyses highlight the necessity of knowledge-base construction and validate both design choices.Source: §5 · p. 10
- Further analyses highlight the necessity of knowledge-base construction and validate the hybrid distillation and hierarchical multi-faceted retrieval designs.Source: §5 · p. 10
ETHICS STATEMENT
The authors state that the work adheres to the ICLR Code of Ethics and aims to reduce vulnerabilities in LLM-generated code. They describe its data sources and defensive focus.
ETHICS STATEMENT
- The authors state that their work adheres to the ICLR Code of Ethics and aims to enhance security in LLM-generated code.They identify reducing software vulnerabilities as the primary goal and describe a positive ethical impact on secure and reliable software development.Source: ETHICS STATEMENT · p. 10
- The knowledge-base dataset uses publicly available known vulnerabilities and fixes from CVEs and public GitHub projects, with no human subjects or personal or private data.Source: ETHICS STATEMENT · p. 10
- RESCUE is designed for defensive use, abstracting and applying secure coding patterns rather than discovering new exploits.The authors note that security tools can have dual-use potential.Source: ETHICS STATEMENT · p. 10
REPRODUCIBILITY STATEMENT
The authors state that they provide resources and detailed descriptions to support reproducibility, and that the full RESCUE implementation is available in the repository linked in the abstract.
REPRODUCIBILITY STATEMENT
- The authors report providing comprehensive resources and detailed descriptions, with the full RESCUE implementation available through the repository linked in the abstract.Source: REPRODUCIBILITY STATEMENT · p. 11
A DATA COLLECTION
The security dataset is derived from SafeCoder and contains vulnerable code, secure code, and CWE labels for Python, C, and C++ instances.
A DATA COLLECTION
- The cleaned dataset contains 372 Python instances and 332 C/C++ instances, each with vulnerable code, secure code, and a CWE type.Empty and duplicate instances were removed.Source: A DATA COLLECTION · p. 15
B METHOD DETAILS
The method details describe hierarchical guideline distillation and bidirectional program slicing, alongside analyses of average method improvements and the reduction in code length after slicing.
B.1 CLUSTER-THEN-SUMMARIZE
- RESCUE clusters similar raw instances, partitions clusters into batches of at most b instances, and recursively summarizes them until one compact security snippet remains per cluster.The default batch size is 10.Source: §B.1 · pp. 15–16
B.2 SECURITY-FOCUSED STATIC PROGRAM SLICING
- RESCUE builds a program dependence graph, marks deleted patch statements as vulnerable points and added statements as secure points, then slices bidirectionally around each point.Backward slicing captures influencing statements, forward slicing captures affected statements, and the hop limit bounds path length.Source: §B.2 · pp. 16–17
C.1 ANALYSIS OF AVERAGE METHOD IMPROVEMENTS
- Table 5 reports methods’ mean percentage-point improvements over their respective LLM Alone baselines across CodeGuard+, HumanEval+, BigCodeBench, and LiveCodeBench.The averages are computed across applicable LLMs; the supplied passages do not give the table’s values.Source: Table 5 · p. 18, §C.1 · p. 18
C.2 IMPACT OF PROGRAM SLICING ON CODE LENGTH
- 81.5% of code lines were removed on average by slicing, reducing vulnerable samples from 89.5 to 16.1 lines and secure samples from 91.3 to 17.3.The authors characterize the removed lines as unrelated to the security aspects under consideration.Source: Table 6 · p. 18, §C.2 · p. 18
C.3 OVERHEAD AND TIME ANALYSIS
The analyses report gains on unseen CWE types and across languages, alongside benchmark, vulnerability-category, retrieval-fusion, and overhead findings. The appendix also describes the guideline-summarization and code-generation prompts.
C.3 OVERHEAD AND TIME ANALYSIS
- Most added overhead comes from draft generation and vulnerability-cause analysis, which use multiple LLM calls; CWE-level and code-level retrieval add only marginal time.The authors describe the overhead as acceptable and say concurrency and engineering optimizations could reduce it.Source: Table 7 · p. 18, §C.3 · pp. 18–19
C.4 ADDITIONAL EVALUATION ON CWEVAL BENCHMARK
- Across three models, RESCUE outperforms zero-shot and SecCoder on CWEval’s dynamically tested Python and C/C++ tasks, extending evidence beyond static-analysis benchmarks.CWEval uses manually crafted test oracles to check functional correctness and security against adversarial inputs.Source: Table 8 · p. 19, §C.4 · p. 19
C.5.2 BREAKDOWN AND ANALYSIS OF PERFORMANCE ON SEEN AND UNSEEN CWE TYPES
- 3.9 average absolute improvement on unseen CWE types persists across models, despite 15 Python and 7 C/C++ CodeGuard+ types being absent from the knowledge base.The authors attribute transfer to shared vulnerability patterns, general guidelines, and retrieved context activating pretrained model knowledge.Source: §C.5.1 · p. 19, Table 10 · p. 20, §C.5.2 · p. 20
C.5.3 CROSS-LANGUAGE GENERALIZABILITY ANALYSIS
- RESCUE consistently improves over zero-shot generation on Go and JavaScript tasks across three models, despite having no language-specific security knowledge for those languages.The evaluation uses CWEval tasks in languages not covered during knowledge construction.Source: Table 11 · p. 20, §C.5.3 · p. 20
C.6 ANALYSIS OF MODIFIED RECIPROCAL RANK FUSION (RRF)
- Modified RRF thresholds filter low-relevance results and perform better than standard RRF; sensitivity tests on 20 CodeGuard+ instances show a stable overall trend.The tests fixed two thresholds while varying the third.Source: Table 12 · p. 21, §C.6 · p. 21
C.7 ANALYSIS OF FUNCTIONALITY-SPECIFIC AND NON-FUNCTIONALITY-SPECIFIC VULNERABILITIES
- RESCUE gains were +7.88 on non-functionality-specific versus +6.02 on functionality-specific vulnerabilities; the authors link this gap to easier retrieval of local constraints than API threat models.The former category includes memory-related flaws, while the latter includes unsafe API use and library-dependent logic.Source: §C.7 · p. 22, Table 16 · p. 23
D.1.1 INITIAL GUIDELINE SUMMARIZATION
- Initial guideline prompts extract common security knowledge from cases, while intermediate prompts merge guidelines with identical meanings into concise sentences.The pipeline uses these stages to summarize raw vulnerability-fix instances.Source: §D.1.1 · p. 23, §D.1.2 · p. 24
D.3 SECURITY KNOWLEDGE AUGMENTED CODE GENERATION
- The augmented-generation prompt supplies security guidelines and secure code examples, instructing the model to preserve the given code snippet and adapt rather than copy the knowledge.The zero-shot completion prompt instead presents the target coding task directly.Source: §D.2.1 · p. 24, §D.3 · pp. 25–26
D.2 PROACTIVE ANALYSIS
The appendix describes a proactive vulnerability-cause prompt and illustrates a vulnerability-fix instance in which a YAML loading call is replaced with a safer alternative.
D.2.2 VULNERABILITY CAUSE ANALYSIS
- The vulnerability-cause prompt asks for the vulnerability’s name and underlying cause before code generation, while explicitly excluding fixing information.Its example describes a potential SQL injection risk from dynamic SQL based on user-supplied input.Source: §D.2.2 · pp. 24–25
E VULNERABILITY-FIX INSTANCE EXAMPLE
- The example vulnerability-fix instance is about 50 lines long, but only one line is security-related: replacing yaml.load(classes) with yaml.safe_load(classes).The example’s remaining function logic largely implements a database query.Source: E VULNERABILITY-FIX INSTANCE EXAMPLE · p. 26
F LIMITATIONS
RESCUE’s evaluation relies primarily on static security analysis, and its current scope is function-level code generation. It also adds inference overhead, mainly from proactive analysis requiring additional LLM calls.
Static Analysis-Based Evaluation
- Static security analysis may produce false positives and negatives, and may not fully capture complex inter-procedural analyses.The authors note that broader adoption of dynamic evaluation remains future work.Source: F LIMITATIONS · p. 27
Function-Level Coding Tasks
- RESCUE focuses on function-level code generation; repository-level deployment requires adaptation for cross-file dependencies and longer code contexts.Source: F LIMITATIONS · p. 27
Inference Overhead
- RESCUE adds inference overhead compared with zero-shot generation, mainly because proactive analysis requires additional LLM calls.The authors identify parallelizing independent calls and selectively invoking RESCUE on potentially vulnerable patterns as possible mitigations.Source: F LIMITATIONS · p. 27
Related papers
This paper cites
Cited by
- Secure and Functional Code Generation with Self-Evaluation2026
- Applying Security by Design at the Point of Execution: How Governed Security Requirements Affect the Security of AI-Generated Code2026
- CodePoisonRAG: Knowledge Poisoning Attacks on Retrieval-Augmented Code Generation2026
- Unsaid, Unsafe? Implicit Security Obligations in LLM-Based RTL Code Generation2026
- MACGen: Toward Functionally Correct and Secure Code Generation via Multi-Agent Collaboration2026