基于T5的两阶段的多任务Text-to-SQL预训练模型MIGA

越来越多的工作证明了预训练语言模型(PLM)中蕴含着丰富的知识,针对不同的任务,用合适的训练方式来撬动 PLM,能更好地提升模型的能力。在 Text-to-SQL 任务中,目前主流的生成器是基于语法树的,需要针对 SQL 语法进行设计。

近期,网易互娱 AI Lab 联合广东外语外贸大学和哥伦比亚大学基于预训练语言模型 T5 的预训练方式,提出了两阶段的多任务预训练模型 MIGA。MIGA 在预训练阶段引入三个辅助任务,并将他们组织成统一的生成任务范式,可以将所有的 Text-to-SQL 数据集统一进行训练;同时在微调阶段,MIGA 针对多轮对话中的错误传递问题进行 SQL 扰动,提升了模型生成的鲁棒性。

目前对于 Text-to-SQL 的研究,主流的方法主要是基于 SQL 语法树的 encoder-decoder 模型,可以确保生成的结果一定符合 SQL 语法,但是需要针对 SQL 语法进行特殊设计。最近也有一些关于 Text-to-SQL 的研究是基于生成式语言模型,可以很方便地继承预训练语言模型的知识和能力。

为了降低对基于语法树的依赖,更好地挖掘预训练语言模型的能力,该研究在预训练 T5 模型的框架下,提出了一个两阶段的多任务 Text-to-SQL 预训练模型 MIGA (MultI-task Generation frAmework)。

MIGA 分为两阶段的训练过程:

  • 在预训练阶段,MIGA 使用与 T5 相同的预训练范式,额外提出了三个与 Text-to-SQL 相关的辅助任务,从而更好地激发预训练语言模型中的知识。该训练方式可以将所有的 Text-to-SQL 的数据集进行统一,扩充了训练数据的规模;而且也可以灵活地去设计更多有效的辅助任务,进一步发掘预训练语言模型的潜在知识。
  • 在微调阶段,MIGA 针对多轮对话和 SQL 中容易存在的错误传递问题,在训练过程中对历史 SQL 进行扰动,使得生成当前轮次的 SQL 效果更加稳定。

MIGA 模型在两个多轮对话 Text-to-SQL 公开数据集上表现优于目前最好的基于语法树的模型,相关研究已经被 AAAI 2023 录用。

基于T5的两阶段的多任务Text-to-SQL预训练模型MIGA

论文地址:https://arxiv.org/abs/2212.09278

MIGA 模型细节

基于T5的两阶段的多任务Text-to-SQL预训练模型MIGA

图 1 MIGA 模型图。

多任务预训练阶段

该研究主要参考 T5 的预训练方式,基于已经训练好的 T5 模型,设计了四个预训练任务:

  • Text-to-SQL 主任务:上图中的黄色部分,通过设计 Prompt 为”translate dialogue to system query”,然后以一些 special tokens 来将历史对话、数据库信息和 SQL 语句拼接输入到 T5-encoder 中,让 decoder 直接输出对应的 SQL 语句;
  • 相关信息预测:上图中的绿色部分,设计 Prompt 为”translate dialogue to relevant column”,T5-encoder 的输入也是与主任务一致,decoder 则需要输出与当前问题相关的数据表和列,目的是加强模型对 Text-to-SQL 的理解;
  • 当前轮次的操作预测:上图中的灰色部分,设计 Prompt 为”translate dialogue to turn switch”,这个任务主要是针对多轮对话中的上下文理解进行设计,对比上一轮的对话和 SQL,decoder 需要输出当前对话的目的做了哪些变化,比如图中例子是 where 条件进行了改动;
  • 最终对话预测:上图中的蓝色部分,设计 Prompt 为”translate dialogue to final utterance”,目的是为了让模型去更好的理解上下文对话,需要 decoder 输出整个多轮对话下来,最后时刻的 SQL 对应的一句完整问题描述。

通过这样的一个统一的训练方式设计,MIGA 可以通用而又灵活地去处理更多与任务相关的额外任务,而且还有一下优点:

  1. 参考人类编写 SQL 的步骤,对话文本到 SQL 任务被分解为多个子任务,允许主任务从中学习;
  2. 训练样本的构造格式与 T5 一致,可以最大限度地发挥预训练 T5 模型对目标任务的潜力;
  3. 统一的框架允许多个辅助任务的灵活调度。当应用于特定任务时,仅需在特定任务的标记数据中使用相同的训练目标对上面的预训练模型进行微调即可。

在预训练阶段,该研究整合了 Text-to-SQL 数据集 Spider 和对话 Text-to-SQL 数据集 SparC 和 CoSQL 的数据来训练 T5 模型。

微调阶段

在预训练阶段之后,该研究针对目标任务的标注数据,单纯使用 Text-to-SQL 任务来进一步微调模型。该研究在预测当前轮次的 SQL 时,会拼接之前轮次的预测 SQL,在此过程中,为了尽量克服多轮对话和生成中所带来的错误传递问题,该研究提出了 SQL 扰动的方案,对输入数据中的历史轮次 SQL,以 α 概率来进行扰动。SQL 语句的扰动主要以 β 的概率采样相应的 token,然后进行以下其一的扰动:

  • 用相同数据表中的列,来随机修改或者新增 SELECT 部分中的列;
  • 随机修改 JOIN 条件中的结构,比如交换两个表的位置;
  • 修改”*” 所有列为一些其他的列;
  • 交换”asc” 和”desc”。

上述的扰动是该研究在实验中统计发现最常见的一些错误传递导致的 SQL 生成错误情况,因此针对这些情况来进行扰动,降低模型关于这方面的依赖。

实验评估

评估数据集为多轮对话 Text-to-SQL:SparC 和 CoSQL。

评估指标为:

  • QM:Question Match,表示单轮问题中的生成的 SQL 与标注输出完全匹配的比例;
  • IM:Interaction Match,表示多轮对话中整个完整轮次所有生成的 SQL 与标注输出完全匹配的比例。

在表 1 的对比实验中,MIGA 在两个数据集上的 IM 分数,以及 CoSQL 的 QM 分数上,超过了目前最好的多轮对话 Text-to-SQL 模型。而且对比同类型的基于 T5 的方案,MIGA 分别在 IM 上提升了 7.0% 和 QM 上提升了 5.8%。

基于T5的两阶段的多任务Text-to-SQL预训练模型MIGA

表 1 对比实验分析,第一部分为树模型,第二部分为基于预训练生成模型。

在表 2 的消融实验中,该研究对 MIGA 的两阶段训练过程中的几个任务进行了探索,同时证明了这些任务分别都会对目标任务有不同程度的提升。

基于T5的两阶段的多任务Text-to-SQL预训练模型MIGA

表 2 针对 SparC 任务,分别去除各项任务或数据,在指标上都有所降低。

在实际的案例分析结果,MIGA 在生成的稳定性和正确性,对比基于 T5-3B 训练模型更好,可以看到 MIGA 在多表连接操作和列和表格的映射上要优于其他模型。在 Case#1 的 Question#2 中,T5-3B 模型不能为相对复杂的 JOIN 结构(两表连接)生成有效的 SQL,进而导致 Question#3 中更复杂的 JOIN 结构(三表连接)也预测错误。而 MIGA 准确地预测了 JOIN 结构,并较好地维持了之前的条件 t1.sex=”f”。在 Case#2 中,T5- 3B 混淆了不同表中的多个列,并将 earnings 误认为是 people 表的列,而 MIGA 可以正确识别该列属于 poker_player 表,并将其链接至 t1。

基于T5的两阶段的多任务Text-to-SQL预训练模型MIGA

表 3 案例分析。

结语

网易互娱 AI Lab 针对 Text-to-SQL 提出了一个基于 T5 的两阶段的多任务预训练模型:MIGA。在预训练阶段,MIGA 将 Text-to-SQL 任务分解出了额外三个子任务,并将其统一为序列到序列的生成范式,从而更好地激发预训练 T5 模型。并且在微调阶段引入了 SQL 扰动机制,降低多轮 Text-to-SQL 生成场景下的错误传递带来的影响。

未来,研究团队会进一步探索更多有效的策略来撬动超大语言模型的能力,并且探索更优雅更有效的方式来进一步克服因为错误传递而导致的效果降低问题。

© 版权声明

相关文章