发布于：2024 年 10 月 29 日
#RAG #Text2 SQL #NL2 SQL
语言模型在将自然语言问题转换为 SQL 查询（文本到 SQL ）的任务中显示出良好的性能。然而，大多数最先进的 （SOTA） 方法都依赖于强大但闭源的大型语言模型 （LLM），例如 ChatGPT 和 GPT-4，这些模型可能存在模型架构不明确、数据隐私风险和昂贵的推理开销等局限性。为了解决这些限制，我们引入了 CodeS ，这是一系列预先训练的语言模型，参数范围从 1B 到 15B，专为文本到 SQL 任务而设计。 CodeS 是一种完全开源的语言模型，它以小得多的参数大小实现了卓越的准确性。实验结果表明，我们的 CodeS 在几乎所有具有挑战性的文本到 SQL 基准测试中都实现了新的 SOTA 准确性和稳健性。
欢迎关注
一、引言
1.1 研究背景与现状
Text-to- SQL 任务旨在将用户的自然语言问题转换为相应的 SQL 查询，以便在数据库上执行并获取答案。传统方法多采用监督微调（SFT），但随着大型语言模型（LLMs）的兴起，如GPT-4、GPT3.5和PaLM-2等，提示学习（prompt learning）或上下文学习（in-context learning）成为新趋势。然而，当前大多数先进的Text-to- SQL 方法依赖于闭源的LLMs，这带来了诸多局限性，如模型架构不透明、数据隐私风险以及高昂的推理开销等。
1.2 论文核心创新点
开源语言模型：本文介绍了 CodeS ，这是一系列完全开源的语言模型，参数规模从1B到15B不等，专为Text-to- SQL 任务设计，解决了闭源模型的诸多问题。
增量预训练：通过收集包含 SQL 相关、自然语言相关和自然语言到代码数据的语料库，对 StarCoder 进行增量预训练，增强了模型在 SQL 生成和自然语言理解方面的能力。
数据库提示构建：提出了一种综合的数据库提示构建策略，包括Schema 过滤器和值检索器，并融入关键的数据库元数据，有效提升了模型生成 SQL 查询的准确性。
双向数据增强：针对新领域适应问题，采用双向数据增强技术，自动生成大量真实且通用的（问题， SQL ）对，减少了标注成本，提高了模型在新领域的适应性。
1.3 CodeS 模型优势
全开源：与依赖闭源LLMs的现有方法不同， CodeS 基于 StarCoder 构建，完全开源，用户可根据自身需求选择合适参数规模的模型，便于进行定制化开发和研究。
高性能：在多个具有挑战性的Text-to- SQL 基准测试中， CodeS 取得了新的最优性能，展现出卓越的准确性和鲁棒性。
小尺寸高效： CodeS 模型规模相较于现有先进的LLMs（如ChatGPT和GPT-4）小10到100倍，但在 SQL 生成能力上表现出色，推理速度更快，所需GPU内存更低，更适合实际部署。
二、相关工作
2.1 监督微调与提示学习
在Text-to- SQL 任务中，监督微调曾是主流方法，通过“编码器-解码器”神经网络模型，利用数据库相关数据进行训练。随着语言模型发展，Text-to- SQL 逐渐被视为序列到序列任务，T5、BART等模型通过微调实现 SQL 查询生成。而提示学习则借助少量示例作为提示，让LLMs在不进行参数微调的情况下完成任务，但在实际应用中，由于使用闭源模型API存在成本和隐私问题，面临一定挑战。
2.2 代码语言模型
现有代码语言模型通常在多种编程语言混合数据上预训练，对于特定编程语言（如 SQL ）的小规模模型，其性能可能受限。 CodeS 则专注于以 SQL 为中心的数据进行优化，以提升在Text-to- SQL 任务中的表现。
2.3 模式链接与数据增强
模式链接在Text-to- SQL 过程中起着关键作用，用于识别自然语言问题中的数据库模式和值，主要有基于字符串匹配和神经网络的策略。数据增强方面，现有方法多采用 SQL 到问题的合成管道，但合成问题与实际用户问题存在差异。本文提出的双向数据增强策略，结合了问题到 SQL 和 SQL 到问题的合成，能更准确地生成多样化的用户问题。
三、预备知识
3.1 Text-to- SQL 任务
Text-to- SQL 任务的目标是根据自然语言问题和数据库生成可执行的 SQL 查询。数据库包含模式（表和列）以及元数据（列类型、注释、值、主键和外键关系）。例如，对于查询银行金融数据库中特定条件的银行名称的自然语言问题，需要生成准确的 SQL 查询来获取答案（见图2）。
3.2 预训练语言模型
基于Transformer架构的语言模型，通过无监督学习目标在大规模文本数据上预训练，主要有“语言建模”和“掩码语言建模”两种范式。预训练后的模型可通过监督微调或上下文学习适应特定任务。
四、方案概述
4.1 增量预训练
为提升模型能力，收集了来自多源的语料库，包括 SQL 相关数据（11GB）、自然语言相关数据（4.5GB）和自然语言到代码数据（6GB）。基于 StarCoder 模型，在该语料库上进行增量预训练，得到不同参数规模（1B、3B、7B、15B）的 CodeS 模型，优化目标为广泛使用的语言建模目标。
4.2 数据库提示构建
通过 Schema 过滤器和值检索器构建高质量数据库提示，并融入数据库元数据（列数据类型、注释、代表性值、主键和外键），为模型提供更丰富的上下文信息，有助于生成准确的 SQL 查询。
4.3 新领域适应
提出双向数据增强方法，包括问题到 SQL 和 SQL 到问题的增强。通过收集少量真实用户问题并手动标注 SQL 查询，利用GPT-3.5扩展数据集；同时从现有基准中提取模板，填充新数据库模式并优化问题，生成大量（问题， SQL ）对，以适应新领域数据库。
4.4 CodeS 的使用方式
CodeS 可通过监督微调（SFT）和少样本上下文学习两种方式使用。在有大量训练数据时，SFT可使模型快速适应数据分布；数据有限时，可利用少样本上下文学习能力，通过高效检索器选择有用示例辅助模型预测。
五、增量预训练
5.1 预训练语料库
为强化Text-to- SQL 模型在 SQL 生成与自然语言理解方面的能力，本研究从三个关键维度收集数据集。在 SQL 相关数据 方面，运用 StarCoder 预训练语料库中的 SQL 片段，数据量达11GB，以此提升模型生成 SQL 查询的能力。针对 自然语言理解能力 的提升，从Alpaca-cleaned、Unnatural-instructions和UltraChat三个来源采样4.5GB的高质量对话数据，这些数据有助于模型更好地理解用户的自然语言问题。在 自然语言到代码 数据维度，整合了CoNaLa、StaQC、CodeAlpaca 20k和新构建的NL- SQL -458K等数据集，总量为6GB，旨在缩小自然语言问题与 SQL 查询之间的差距，特别是NL- SQL -458K数据集，通过从开源语料库提取 SQL 查询并借助GPT-3.5生成对应自然语言问题的方式构建，为模型提供了丰富的学习素材。
5.2 预训练细节
CodeS 模型基于 StarCoder 构建， StarCoder 预先在多种编程语言、Jupyter笔记本、GitHub问题和Git提交等混合数据上进行了预训练。在开发 CodeS 时，针对 SQL 相关数据进行两阶段的增量预训练，而对自然语言相关和自然语言到代码数据各进行一个阶段的训练，这种混合训练方式有利于模型在自然语言和代码相关任务中发挥出色性能。具体而言， CodeS -15B基于 StarCoder -15B，而 CodeS -(1B、3B、7B)则分别源自相应的 StarCoder Base-(1B、3B、7B)。预训练过程中，优化广泛使用的语言建模目标，采用AdamW优化器，其参数设置为 β1 = 0.9、β2 = 0.95、ε = 10-8，学习率设定为 5e-5，同时应用权重衰减（0.1）和余弦退火学习率调度（无热身步骤，最终学习率为初始值的十分之一）。训练采用大批次（4M tokens）并进行梯度裁剪（裁剪值为1.0），借助 DeepSpeed Zero-3 框架及 BF16 混合精度优化 GPU 内存消耗。由于GPU内存限制， CodeS -(1B、3B、7B) 的最大上下文长度为 8,192，而 CodeS -15B 为 6,144。不同参数规模的 CodeS 模型预训练耗时各异， CodeS -(1B、3B、7B、15B) 分别约需 1.5、3、8、16 天。
六、数据库 Prompt 构建
除了模型改进之外，构建有效的数据库提示对于文本到 SQL 任务也至关重要。高质量的提示为语言模型提供了有价值的见解，使其能够更高效地生成精确的 SQL 查询。为了制作这些卓越的数据库提示，我们采用了两个关键策略：Schema 过滤器 和 值检索器，同时还整合了关键的数据库元数据。算法 1 中介绍了详细说明我们的提示构建过程的伪代码。
6.1 Schema 过滤器
在实际场景中，数据库常包含众多表和列，这可能导致数据库提示过长，超出模型上下文长度，进而影响模型性能。为解决此问题，本研究采用 Schema 过滤器，根据用户问题筛选相关表和列，减少提示长度，同时保留关键信息。具体操作是，开发一个模式项分类器，依据用户问题预测表和列的相关性得分，保留得分最高的top k1个表和每个表中得分最高的top k2个列。对于训练集中的样本，依据真实 SQL 查询确定使用的表和列，若使用的表数量少于top k1，则随机选择未使用的表填充；对于列也采取类似操作，确保每个保留的表包含top k2个列。此策略不仅能缩短数据库提示长度，还能减轻模型的模式链接压力，提升模型性能。
6.2 值检索器
从数据库中检索与问题相关的值，有助于模型在生成 SQL 查询谓词时更好地理解数据库内容，提高查询准确性。以往研究多使用最长公共子串（LCS）算法计算数据库值与问题的匹配度，但该算法在处理大数据集时时间复杂度高，效率低下。为改进此问题，本研究提出“粗到细”的匹配方法。首先利用 Lucene 构建 BM25 索引，根据用户问题快速从数据库中提取数百个潜在相关值，然后使用 LCS 算法计算这些潜在值与问题的匹配度，确定最相关的值。这种方法显著减少了 LCS 算法的调用次数，大大提高了值检索速度，尤其适用于包含大量值的数据库，如 BIRD 基准中的 Donor 数据库。
6.3 数据库元数据
本研究在数据库提示中融入多种关键元数据，以提升模型生成 SQL 查询的准确性。列数据类型 决定了数据的操作规则，例如数值类型支持算术运算，而字符串类型则需转换后才能进行算术操作，尽管模型可从列名和注释中推断部分类型信息，但明确列数据类型仍有助于模型生成更准确的 SQL 查询。数据库中的注释 可解决模式中的歧义问题，特别是在实际数据库中，歧义模式可能导致模型选择错误的表或列，将注释纳入输入可辅助模型进行准确的模式链接。代表性数据库值能为模型提供列值的格式信息，例如告知模型“gender”列的取值范围，有助于生成精确的谓词。主键和外键信息 则指导模型正确推导表之间的连接路径，确保生成准确的JOIN ON子句，本研究使用唯一标识符表示主键列，以特定格式表示外键关系，为模型提供关键的结构信息。
七、新领域适应
在现实世界的场景中，人们通常拥有自己的数据库，这些数据库来自金融、基因、生物学、学术界等各种新领域。但是，由于缺乏标记的训练数据，在这些数据库上开发文本到 SQL 模型具有挑战性。在本节中，我们提出了一种 双向数据增强技术，以最小的注释成本自动生成大量真实和通用（问题、 SQL ）对。
7.1 问题到 SQL 增强
针对新领域数据库缺乏标注训练数据的问题，本研究提出问题到 SQL 的增强方法，旨在生成与用户偏好一致的真实（问题， SQL ）对。首先，从实际用户中收集少量具有代表性的自然语言问题，这些问题反映了用户对特定数据库的常见查询需求，然后手动为其标注对应的 SQL 查询，获得少量高质量的（问题， SQL ）对。由于收集到的问题数量有限，为扩充数据集，引入两阶段提示方法，利用GPT-3.5根据收集到的真实问题生成更多潜在问题，捕捉用户意图，再为这些合成问题生成相应的 SQL 查询。为保证生成问题的多样性，在生成过程中打乱用户提供问题的顺序，并采用高温度超参数。最终生成的新数据对能够更好地反映用户偏好，增强了模型对新领域数据的适应性。
7.2 SQL 到问题增强
受Syntax SQL Net启发，本研究探索了 SQL 到问题的增强方法，通过通用模板生成（问题， SQL ）对。具体而言，使用从 Spider 基准中提取的75个常见 SQL 模板，这些模板涵盖多种常见查询模式，如“Return the lowest {COLUMN} of {TABLE}”及其对应的 SQL 模板“SELECT {COLUMN} FROM {TABLE} GROUP BY {COLUMN} ORDER BY COUNT (*) ASC LIMIT 1”。由于单个 SQL 模板可对应多个自然语言问题模板，将新数据库的模式信息填充到模板中，生成大量模板（问题， SQL ）对。然而，这些模板问题因直接插入表和列名而显得生硬，为使其更自然，利用GPT-3.5根据手动编写的优化示例对模板问题进行改写，最终生成的新问题与典型Text-to- SQL 数据集更契合，提高了模型对新领域数据分布的适应性。
八、CodeS 的使用方式
8.1 监督微调
在有大量训练数据可用时，监督微调是优化 CodeS 模型的有效方法。首先，将数据库提示与问题组合作为模型输入， CodeS 在训练过程中学习根据此输入预测准确的 SQL 查询，其学习目标是最小化预测 SQL 查询与真实 SQL 查询之间的差异。经过监督微调后，模型能够快速适应特定的数据分布，在面对测试样本时，根据输入的数据库提示和问题准确生成 SQL 查询，有效提高了模型在实际应用中的性能。
8.2 少样本上下文学习
当训练数据有限时，少样本上下文学习使 CodeS 能够利用其预训练知识处理用户问题，无需进行参数微调。在这种方式下，模型的性能不仅取决于其自身的内在能力，还依赖于提供的示例质量。为提高示例选择的有效性，本研究采用“问题模式感知演示检索器”，综合考虑测试问题与训练问题之间的语义相关性和核心结构相似性，选择最有用的K个示例。具体而言，通过去除问题中的实体获取其核心结构（问题模式），使用nltk工具和SimCSE模型计算相似度，避免仅基于实体匹配选择示例导致的偏差。选择示例后，将其与测试样本合并为统一输入序列，送入模型预测 SQL 查询，从而在数据有限的情况下实现较好的性能表现。
九、实验设置与结果分析
9.1 实验设置
9.1.1 数据集
主要实验在 Spider 和 BIRD 两个基准数据集上进行，同时评估模型在 Spider 的四个变体（ Spider -DK、 Spider -Syn、 Spider -Realistic和Dr. Spider ）上的鲁棒性。此外，还使用了手动创建的金融（Bank-Financials）和学术（Aminer-Simplified）领域的数据库，以测试模型在实际场景中的应用能力。 Spider 数据集包含丰富的训练、开发和测试样本，涵盖多种领域数据库； BIRD 数据集具有更复杂的问题、模糊模式和大量数据值，且提供外部知识（EK）用于特定样本； Spider 变体用于模拟实际场景中的不同挑战；新领域数据库则针对特定领域特点设计任务和评估指标。
9.1.2 评估指标
对于 Spider 家族基准，采用执行准确率（EX）和测试套件准确率（TS）评估，其中EX评估预测与真实 SQL 查询执行结果是否一致，TS通过在多个扩充数据库实例上评估减少EX的误报情况； BIRD 基准主要依赖EX，并引入有效效率分数（VES）评估 SQL 查询执行效率，但VES受硬件、软件和系统状态影响较大，实际应用中EX更稳定可靠。
9.1.3 环境
我们的实验是在运行 CentOS Linux 7 操作系统的计算机上使用 PyTorch 13.1.7 进行的，该计算机配备了 8 个 NVIDIA A800 80GB GPU、一个 Intel（R） Xeon（R） Platinum 8358 CPU 和 1024GB RAM。
9.2 实验结果
9.2.1 少样本上下文学习评估
在 Spider 和 BIRD 基准的少样本上下文学习评估中（见表4），增量预训练显著提升了 StarCoder 的 SQL 生成能力。对比不同版本的 StarCoder 和 CodeS ， CodeS 表现更优， CodeS -15B 成为领先的开源预训练语言模型。较小模型在增量预训练后提升更明显，证明了增量预训练对小规模模型在 SQL 相关任务上的重要性。
9.2.2 监督微调评估
在 Spider 开发集上（见表5），SFT CodeS-3B 超越了基于GPT-4的领先方法，SFT CodeS -7B和SFT CodeS -15B达到新的最优性能，但SFT CodeS -15B相比SFT CodeS -7B优势不明显，可能存在过拟合问题，表明更大模型不一定保证更好的微调结果。在 BIRD 数据集上（见表6），由于其复杂性更高，本方法相比基线方法提升显著，SFT CodeS -15B在测试集上的执行准确率相比ChatGPT + COT提升23.20%，且模型生成的 SQL 查询执行效率更高。
9.2.3 鲁棒性基准评估
在 Spider 变体上（见表7），SFT CodeS -7B 表现出色，在不同数据集上相比最佳基线均有显著提升，表明模型在分布变化场景下具有良好的泛化能力。在Dr. Spider 上（见表8），尽管在数据库扰动方面略逊于ChatGPT + ZeroNL2 SQL ，但在自然语言问题扰动方面表现更好，模型对问题语义的理解更准确，整体上SFT CodeS -7B和SFT CodeS -15B的全局平均性能略优于专为文本到 SQL 鲁棒性设计的ChatGPT + ZeroNL2 SQL 。
9.2.4 消融实验
在 Spider 和 BIRD 数据集上的消融实验（见表9）表明，各关键组件对模型性能均有重要影响。仅使用问题相似度检索示例（-w/o pattern similarity）在 Spider 上会导致性能下降，而随机选择示例（-w/o demonstration retriever）在多数情况下会降低性能。Schema 过滤器和值检索器缺失会分别影响性能和生成速度，以及 SQL 查询谓词生成能力。元数据方面，列数据类型影响较小，而注释、代表性数据库值和主键/外键对模型在两个数据集上的性能至关重要。
9.2.5 实际场景评估
在 Bank-Financials 和 Aminer-Simplified 两个新领域数据集上的评估（见表10）显示， CodeS -7B 在不同使用方式下均展现出一定的领域转移能力。通过双向数据增强策略微调后的 CodeS -7B（SFT CodeS -7B using aug. data）在两个数据集上的准确率显著提升，合并多个数据集训练的统一模型（SFT CodeS -7B using merged data）不仅避免了性能下降，还在 Aminer-Simplified 数据集上进一步提高了性能。
9.3 模型效率与部署
与基于提示的先前最优方法 DIN- SQL +GPT-4 相比， CodeS 模型在推理速度上有显著提升（1B、3B、7B、15B变体在 Spider 数据集上的推理时间分别为0.6、0.9、1.1和1.5秒）。在部署方面， CodeS 适合在本地机器上运行，不同参数规模的 SFT CodeS 模型在 Float16 精度下所需 GPU 内存容量分别为 10GB、13GB、20GB 和 35GB，避免了对昂贵的GPT-4 API的依赖。
十、研究结论与展望
10.1 主要研究成果
本研究成功开发了 CodeS 系列模型，这是专为 Text-to- SQL 任务设计的开源语言模型。通过增量预训练、数据库提示构建和双向数据增强等创新技术， CodeS 在多个基准测试中取得了优异成绩，展现出卓越的 SQL 生成能力、准确性和鲁棒性。同时，模型在实际场景中也表现出良好的适应性和效率，为开源 Text-to- SQL 模型的发展做出了重要贡献。
10.2 未来研究方向
未来研究可从以下几个方面展开：进一步优化模型架构和训练方法，提高模型性能和效率；探索更有效的数据增强策略，以更好地适应不同领域和复杂场景；加强对模型可解释性的研究，提升用户对模型决策的信任度；将Text-to- SQL 技术与其他人工智能技术（如知识图谱、强化学习）相结合，拓展其应用范围和功能。
本研究通过对 CodeS 模型的深入探索，为Text-to- SQL 领域提供了一种强大的开源解决方案，有望推动相关技术在实际应用中的广泛部署和进一步发展。希望本研究成果能为人工智能领域的研究人员和开发者提供有价值的参考，共同推动自然语言处理与数据库技术的融合创新。
代码地址：https://github.com/RUCKBReasoning/ CodeS
https://github.com/RUCKBReasoning/codes