SQL-R1: Treinando Modelos de Raciocínio de Linguagem Natural para SQL por Aprendizado por Reforço

SQL-R1: Training Natural Language to SQL Reasoning Model By Reinforcement Learning

April 11, 2025

Autores: Peixian Ma, Xialie Zhuang, Chengjin Xu, Xuhui Jiang, Ran Chen, Jian Guo

cs.AI

Resumo

Natural Language to SQL (NL2SQL) permite interações intuitivas com bancos de dados ao transformar consultas em linguagem natural em instruções SQL estruturadas. Apesar dos avanços recentes na melhoria da interação humano-computador em aplicações de banco de dados, desafios significativos persistem, especialmente em relação ao desempenho de inferência em cenários complexos que envolvem junções de múltiplas tabelas e consultas aninhadas. As metodologias atuais utilizam principalmente o ajuste fino supervisionado (SFT) para treinar o modelo NL2SQL, o que pode limitar a adaptabilidade e a interpretabilidade em novos ambientes (por exemplo, finanças e saúde). Para melhorar o desempenho de raciocínio do modelo NL2SQL nas situações complexas mencionadas, introduzimos o SQL-R1, um novo modelo de raciocínio NL2SQL treinado por algoritmos de aprendizado por reforço (RL). Projetamos uma função de recompensa baseada em RL especializada para tarefas NL2SQL e discutimos o impacto do problema de inicialização a frio na eficácia do treinamento intensivo. Além disso, alcançamos precisão competitiva utilizando apenas uma pequena quantidade de dados sintéticos NL2SQL para treinamento aumentado e exploramos ainda mais a engenharia de dados para RL. Nos experimentos existentes, o SQL-R1 alcançou precisão de execução de 88,6% e 66,6% nos benchmarks Spider e BIRD, respectivamente, utilizando apenas o modelo base de 7B.

English

Natural Language to SQL (NL2SQL) enables intuitive interactions with databases by transforming natural language queries into structured SQL statements. Despite recent advancements in enhancing human-computer interaction within database applications, significant challenges persist, particularly regarding the inference performance in complex scenarios involving multi-table joins and nested queries. Current methodologies primarily utilize supervised fine-tuning (SFT) to train the NL2SQL model, which may limit adaptability and interpretability in new environments (e.g., finance and healthcare). In order to enhance the reasoning performance of the NL2SQL model in the above complex situations, we introduce SQL-R1, a novel NL2SQL reasoning model trained by the reinforcement learning (RL) algorithms. We design a specialized RL-based reward function tailored for NL2SQL tasks and discussed the impact of cold start on the effectiveness of intensive training. In addition, we achieve competitive accuracy using only a tiny amount of synthetic NL2SQL data for augmented training and further explore data engineering for RL. In existing experiments, SQL-R1 achieves execution accuracy of 88.6% and 66.6% on the benchmark Spider and BIRD, respectively, only using the 7B base model.