# Improving Retrieval-Augmented Deep Assertion Generation via Joint Training

arXiv: https://arxiv.org/abs/2502.10696 · HTML: https://mlsanity.com/paper/2502.10696/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in RAG for code, one of 15 directions in Retrieval-augmented generation.
- Builds on 47 of its 84 references that are in the map.
- First posted on arXiv 15 Feb 2025.

## Placed in

- Retrieval-augmented generation / [RAG for code](https://mlsanity.com/rag/rag-for-code/)

## Builds on (most cited first, 30 of 47)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Code Llama: Open Foundation Models for Code](https://mlsanity.com/paper/2308.12950/) (2023)
- [CodeBERT: A Pre-Trained Model for Programming and Natural Languages](https://mlsanity.com/paper/2002.08155/) (2020)
- [CodeT5: Identifier-aware Unified Pre-trained Encoder-Decoder Models for Code Understanding and Generation](https://mlsanity.com/paper/2109.00859/) (2021)
- [StarCoder: may the source be with you!](https://mlsanity.com/paper/2305.06161/) (2023)
- [CodeXGLUE: A Machine Learning Benchmark Dataset for Code Understanding and Generation](https://mlsanity.com/paper/2102.04664/) (2021)
- [GraphCodeBERT: Pre-training Code Representations with Data Flow](https://mlsanity.com/paper/2009.08366/) (2020)
- [InCoder: A Generative Model for Code Infilling and Synthesis](https://mlsanity.com/paper/2204.05999/) (2022)
- [An Empirical Evaluation of Using Large Language Models for Automated Unit Test Generation](https://mlsanity.com/paper/2302.06527/) (2023)
- [Large Language Models for Software Engineering: Survey and Open Problems](https://mlsanity.com/paper/2310.03533/) (2023)
- [Software Testing with Large Language Models: Survey, Landscape, and Vision](https://mlsanity.com/paper/2307.07221/) (2023)
- [Retrieval Augmented Code Generation and Summarization](https://mlsanity.com/paper/2108.11601/) (2021)
- [Deep Learning based Vulnerability Detection: Are We There Yet?](https://mlsanity.com/paper/2009.07235/) (2020)
- [Automating Code Review Activities by Large-Scale Pre-training](https://mlsanity.com/paper/2203.09095/) (2022)
- [Retrieval-Augmented Generation for Code Summarization via Hybrid GNN](https://mlsanity.com/paper/2006.05405/) (2020)
- [DebugBench: Evaluating Debugging Capability of Large Language Models](https://mlsanity.com/paper/2401.04621/) (2024)
- [Studying the Usage of Text-To-Text Transfer Transformer to Support Code-Related Tasks](https://mlsanity.com/paper/2102.02017/) (2021)
- [Using Pre-Trained Models to Boost Code Review Automation](https://mlsanity.com/paper/2201.06850/) (2022)
- [A Survey on Large Language Models for Software Engineering](https://mlsanity.com/paper/2312.15223/) (2023)
- [TOGA: A Neural Method for Test Oracle Generation](https://mlsanity.com/paper/2109.09262/) (2021)
- [A Survey of Learning-based Automated Program Repair](https://mlsanity.com/paper/2301.03270/) (2023)
- [RAP-Gen: Retrieval-Augmented Patch Generation with CodeT5 for Automatic Program Repair](https://mlsanity.com/paper/2309.06057/) (2023)
- [On Learning Meaningful Assert Statements for Unit Test Cases](https://mlsanity.com/paper/2002.05800/) (2020)
- [A Survey on Deep Learning for Software Engineering](https://mlsanity.com/paper/2011.14597/) (2020)
- [HITS: High-coverage LLM-based Unit Test Generation via Method Slicing](https://mlsanity.com/paper/2408.11324/) (2024)
- [A Systematic Literature Review on Large Language Models for Automated Program Repair](https://mlsanity.com/paper/2405.01466/) (2024)
- [Neural Transfer Learning for Repairing Security Vulnerabilities in C Code](https://mlsanity.com/paper/2104.08308/) (2021)

## Built on by (most cited first, 2 of 2)

- [Multi-Large Language Model Orchestrated Severity Assessment of Clinical Records (MOSAIC)](https://mlsanity.com/paper/2607.05032/) (2026)
- [ViTeGate: Visual-Textual Triggered Knowledge Poisoning for Vision-Language Retrieval-Augmented Generation](https://mlsanity.com/paper/2609.14685/) (2026)

## Abstract (arXiv)

Unit testing attempts to validate the correctness of basic units of the software system under test and has a crucial role in software development and testing. Very recent work proposes a retrieve-and-edit approach to generate unit test oracles, i.e., assertions. Despite being promising, it is still far from perfect due to some limitations, such as splitting assertion retrieval and generation into two separate components without benefiting each other. In this paper, we propose AG-RAG, a retrieval-augmented automated assertion generation approach that leverages external codebases and joint training to address various technical limitations of prior work. Inspired by the plastic surgery hypothesis, AG-RAG attempts to combine relevant unit tests and advanced pre-trained language models (PLMs) with retrieval-augmented fine-tuning. AG-RAG builds a dense retriever to search for relevant test-assert pairs (TAPs) with semantic matching and a retrieval-augmented generator to synthesize accurate assertions with the focal-test and retrieved TAPs as input. Besides, AG-RAG leverages a code-aware language model CodeT5 as the cornerstone to facilitate both assertion retrieval and generation tasks. Furthermore, the retriever is optimized in conjunction with the generator as a whole pipeline with a joint training strategy. This unified design fully adapts both components specifically for retrieving more useful TAPs, thereby generating accurate assertions. We extensively evaluate AG-RAG against six state-of-the-art AG approaches on two benchmarks and three metrics. Experimental results show that AG-RAG significantly outperforms previous AG approaches on all benchmarks and metrics, e.g., improving the most recent baseline EditAS by 20.82% and 26.98% in terms of accuracy. AG-RAG also correctly generates 1739 and 2866 unique assertions that all baselines fail to generate, 3.45X and 9.20X more than EditAS.