A Multi-agent Text2SQL Framework using Small Language Models and Execution Feedback

TL;DR

MATS is a multi-agent framework that enhances small language models for Text2SQL tasks by leveraging execution feedback and reinforcement learning, achieving performance comparable to larger models while being more cost-effective.

Contribution

The paper introduces MATS, a novel multi-agent training scheme for small language models that improves Text2SQL performance through execution feedback and reinforcement learning.

Findings

MATS achieves accuracy comparable to large LLMs on benchmark datasets.

It operates efficiently on a single GPU with fewer parameters.

The framework effectively enhances small LLM capabilities for complex tasks.

Abstract

Text2SQL, the task of generating SQL queries from natural language text, is a critical challenge in data engineering. Recently, Large Language Models (LLMs) have demonstrated superior performance for this task due to their advanced comprehension and generation capabilities. However, privacy and cost considerations prevent companies from using Text2SQL solutions based on external LLMs offered as a service. Rather, small LLMs (SLMs) that are openly available and can hosted in-house are adopted. These SLMs, in turn, lack the generalization capabilities of larger LLMs, which impairs their effectiveness for complex tasks such as Text2SQL. To address these limitations, we propose MATS, a novel Text2SQL framework designed specifically for SLMs. MATS uses a multi-agent mechanism that assigns specialized roles to auxiliary agents, reducing individual workloads and fostering interaction. A training scheme based on reinforcement learning aligns these agents using feedback obtained during execution, thereby maintaining competitive performance despite a limited LLM size. Evaluation results using on benchmark datasets show that MATS, deployed on a single- GPU server, yields accuracy that are on-par with large-scale LLMs when using significantly fewer parameters. Our source code and data are available at https://github.com/thanhdath/mats-sql.