Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

DeTriever: Decoder-representation-based Retriever for Improving NL2SQL In-Context Learning

Yuxi Feng, Raymond Li, Zhenan Fan, Giuseppe Carenini, Mohammadreza Pourreza, Weiwei Zhang, Yong Zhang

TL;DR

The paper tackles the challenge of selecting beneficial in-context demonstrations for NL2SQL by introducing DeTriever, which learns a weighted fusion of LLM hidden states across layers to form retrieval embeddings. The model is trained with a proxy similarity target via a supervised contrastive loss, aligning retrieval with the relative usefulness of demonstrations without requiring expensive pairwise inferences. Experiments on Spider and BIRD show DeTriever substantially outperforms random prompts and external retrievers, with strong gains in both out-of-domain and in-domain settings, and an ablation analysis highlighting the advantages of multi-layer representations and carefully chosen positives. The approach offers a practical, scalable way to enhance NL2SQL ICL and has potential applicability to other open-ended generation tasks, such as code completion and QA.

Abstract

While in-context Learning (ICL) has proven to be an effective technique to improve the performance of Large Language Models (LLMs) in a variety of complex tasks, notably in translating natural language questions into Structured Query Language (NL2SQL), the question of how to select the most beneficial demonstration examples remains an open research problem. While prior works often adapted off-the-shelf encoders to retrieve examples dynamically, an inherent discrepancy exists in the representational capacities between the external retrievers and the LLMs. Further, optimizing the selection of examples is a non-trivial task, since there are no straightforward methods to assess the relative benefits of examples without performing pairwise inference. To address these shortcomings, we propose DeTriever, a novel demonstration retrieval framework that learns a weighted combination of LLM hidden states, where rich semantic information is encoded. To train the model, we propose a proxy score that estimates the relative benefits of examples based on the similarities between output queries. Experiments on two popular NL2SQL benchmarks demonstrate that our method significantly outperforms the state-of-the-art baselines on one-shot NL2SQL tasks.

DeTriever: Decoder-representation-based Retriever for Improving NL2SQL In-Context Learning

TL;DR

The paper tackles the challenge of selecting beneficial in-context demonstrations for NL2SQL by introducing DeTriever, which learns a weighted fusion of LLM hidden states across layers to form retrieval embeddings. The model is trained with a proxy similarity target via a supervised contrastive loss, aligning retrieval with the relative usefulness of demonstrations without requiring expensive pairwise inferences. Experiments on Spider and BIRD show DeTriever substantially outperforms random prompts and external retrievers, with strong gains in both out-of-domain and in-domain settings, and an ablation analysis highlighting the advantages of multi-layer representations and carefully chosen positives. The approach offers a practical, scalable way to enhance NL2SQL ICL and has potential applicability to other open-ended generation tasks, such as code completion and QA.

Abstract

While in-context Learning (ICL) has proven to be an effective technique to improve the performance of Large Language Models (LLMs) in a variety of complex tasks, notably in translating natural language questions into Structured Query Language (NL2SQL), the question of how to select the most beneficial demonstration examples remains an open research problem. While prior works often adapted off-the-shelf encoders to retrieve examples dynamically, an inherent discrepancy exists in the representational capacities between the external retrievers and the LLMs. Further, optimizing the selection of examples is a non-trivial task, since there are no straightforward methods to assess the relative benefits of examples without performing pairwise inference. To address these shortcomings, we propose DeTriever, a novel demonstration retrieval framework that learns a weighted combination of LLM hidden states, where rich semantic information is encoded. To train the model, we propose a proxy score that estimates the relative benefits of examples based on the similarities between output queries. Experiments on two popular NL2SQL benchmarks demonstrate that our method significantly outperforms the state-of-the-art baselines on one-shot NL2SQL tasks.
Paper Structure (20 sections, 5 equations, 5 figures, 4 tables)

This paper contains 20 sections, 5 equations, 5 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. NL2SQL
  4. LLMs and Retrieval-augmented ICL
  5. Approach
  6. Problem Formalization
  7. LLM Hidden States
  8. Answer similarity-based training objective
  9. Experiments
  10. Datasets
  11. Experimental Settings
  12. Baselines
  13. Results
  14. Analysis
  15. Training Target
  16. ...and 5 more sections

Figures (5)

  • Figure 1: Examples of in-context learning (ICL) for NL2SQL.
  • Figure 2: Overview of the training process of our proposed DeTriever method.
  • Figure 3: One-shot ICL results (execution accuracy on development set of Spider dataset) of CodeLlama-13b where prompt is retrieved by representations computed in each decoder layer. Here Mean represents average hidden state of a specific layer, where EOS represents the hidden state of EOS token in a specific layer. The one-shot Prompt is retrieved by the one with biggest cosine similarity.
  • Figure 4: One-shot ICL results (execution accuracy on development set of Spider dataset) of CodeLlama-13b where the prompt is retrieved by our trained retriever with a different number of positive examples. The one-shot Prompt is retrieved by the one with the largest cosine similarity.
  • Figure 5: One-shot ICL execution accuracy on Spider where DeTriever is trained with different batch sizes.