[Paper Review] Large Language Model-Aware In-Context Learning for Code Generation

Jia Li, Ge Li+5 more|arXiv (Cornell University)|Oct 15, 2023

Software Engineering Research7 citations

TL;DR

This paper proposes LAIL, a learning-based in-context example selection method for code generation that leverages large language models (LLMs) themselves to evaluate and label candidate examples based on their generation probability for ground-truth programs. By training a contrastive learning retriever on LLM-provided feedback, LAIL achieves state-of-the-art performance, improving Pass@1 by up to 11.58% over baselines on CodeGen and 4.38% on GPT-3.5 across three benchmark datasets.

ABSTRACT

Large language models (LLMs) have shown impressive in-context learning (ICL) ability in code generation. LLMs take a prompt consisting of requirement-code examples and a new requirement as input, and output new programs. Existing studies have found that ICL is highly dominated by the examples and thus arises research on example selection. However, existing approaches randomly select examples or only consider the textual similarity of requirements to retrieve, leading to sub-optimal performance. In this paper, we propose a novel learning-based selection approach named LAIL (LLM-Aware In-context Learning) for code generation. Given a candidate example, we exploit LLMs themselves to estimate it by considering the generation probabilities of ground-truth programs given a requirement and the example. We then label candidate examples as positive or negative through the probability feedback. Based on the labeled data, we import a contrastive learning objective to train an effective retriever that acquires the preference of LLMs in code generation. We apply LAIL to three LLMs and evaluate it on three representative datasets (e.g., MBJP, MBPP, and MBCPP). LATA outperforms the state-of-the-art baselines by 11.58%, 6.89%, and 5.07% on CodeGen, and 4.38%, 2.85%, and 2.74% on GPT-3.5 in terms of Pass@1, respectively.

Motivation & Objective

• To address the instability of in-context learning (ICL) in code generation, which heavily depends on the quality of selected examples.
• To overcome the limitations of heuristic-based example selection methods that rely only on textual similarity or random sampling.
• To develop a learning-based approach that captures the intrinsic preferences of LLMs for effective example selection.
• To train a retriever using contrastive learning that learns from LLM-generated feedback on example quality.

Proposed method

• Leverage LLMs themselves to estimate the quality of candidate examples by measuring the generation probability of the ground-truth program given a requirement and the example.
• Design a novel metric to quantify this probability feedback and use it to label examples as positive or negative based on their score.
• Train a neural retriever using a contrastive learning objective to learn the preference of LLMs for code generation.
• Use the labeled data from LLM feedback to optimize the retriever, enabling it to retrieve examples that are most helpful for LLMs to generate correct programs.
• Apply the trained retriever to select in-context examples for new requirements, improving ICL performance.

Experimental results

Research questions

• RQ1 Can LLMs themselves provide reliable feedback for evaluating the quality of in-context examples in code generation?
• RQ2 How does LLM-aware example selection compare to traditional heuristic methods like BM25 or random selection in terms of ICL performance?
• RQ3 Can a contrastive learning-based retriever trained on LLM feedback generalize across different LLMs and datasets?
• RQ4 To what extent does the proposed method improve Pass@1 and other code generation metrics?

Key findings

• LAIL improves Pass@1 by 11.58% on CodeGen, 6.89% on MBJP, 5.07% on MBPP, and 4.38% on GPT-3.5, outperforming state-of-the-art baselines.
• The method demonstrates strong transferability across different LLMs and datasets, indicating broad applicability.
• Human evaluation confirms that LAIL-generated examples lead to higher-quality code in terms of correctness, readability, and relevance.
• The approach significantly reduces performance variance compared to random or BM25-based selection, enhancing ICL stability.
• LAIL’s reliance on LLM-generated feedback makes it more aligned with the actual preferences of LLMs than lexical similarity-based methods.
• The contrastive learning objective effectively captures nuanced preferences of LLMs, leading to better example selection than heuristic approaches.