論文：DBCopilot: Scaling Natural Language Querying to Massive Databases

????

Code: DBCopilot | GitHub

一、論文速讀

論文認為目前的 Text2SQL 研究大多只關注具有少量 table 的單個數據庫上的查詢，但在面對大規模數據庫和數據倉庫的查詢時時卻力顯不足。本文提出的 DBCopilot 能夠在大規模數據庫上查詢模式不可知的 NL question。

論文指出，實現這個的核心是：從能夠構建各種 NL question 到海量數據庫模型元素的 semantic mapping，從而能夠自動識別目標數據庫并過濾出最少的相關 tables。但目前的基于 LLM 的方法有兩個主要挑戰：

• 由于 token 限制，無法將所有 schema 都輸入給 LLM
• LLM 仍然難以有效利用長上下文中的信息

而在解決可擴展性的問題時，主要有基于 retrieval 的方法和基于 fine-tune 的方法，但是，

• 基于 retrieval 的方法往往是將 doc 視為檢索對象，忽略了 DB 和 DB table 之間的關系；
• fine-tune LLM 來為其注入 schema 的相關知識是資源密集型的方式，且有時候 LLM 是無法微調的

DBCopilot 的做法如下圖所示：

主要分成兩步：

1. Schema Routing：輸入 user question，使用 DSI 技術找到所需要用的 DB 和 DB tables，也就是 DB schema。
2. SQL Generation：輸入 user question、DB schema，通過 prompt LLM 生成 SQL query。

二、問題定義

2.1 Schema-Agnostic NL2SQL

Schema-Agnostic NL2SQL 指的是：只給定 user question 而不給定預期的 SQL query schema（DB 和 DB tables），來生成一個可以在一個數據庫集合中的某個 DB 上執行的 SQL。

像之前 WikiSQL 數據集上，都是指定 question 在哪個 DB 上的。

2.2 Schema Linking VS. Schema Routing

在以往的 NL2SQL 中，Schema Linking 的 input 是 question 和 schema，用于尋找 NL question 中提及到的 schema 元素（比如 tables、columns 或者 database value），可以被視作是一個 NL question 和 DB elements 之間的橋梁。

Schema Routing 的 input 只有不知道 schema 的 question，它的輸出是一個 indexed or memorized schema。

三、方法

3.1 Schema Routing

本文使用一個輕量級的 seq2seq 模型來作為 router，實現將 NL 識別出對應的 DB schema。

由于 space schema 很大（是 table 和 column 的笛卡爾積）、且 DB schema 可以發生變化，因此本文提出了一個 relation-aware、end-to-end joint retrieval 方法來解決 schema routing 問題。

具體做法是，先為 databases 構建一個 schema graph，然后設計一個 schema 序列化算法來將一個 schema 轉化為 token-sequence，利用 graph-based contrained decoding 解碼算法來讓 seq2seq 模型生成 routing 的結果 DB schema。

3.1.1 Schema Graph

schema graph 包含了 databases 的 schema 信息，這個 graph 的 nodes 包含三類：

• $v_s$：一個特殊節點，指代含有所有 databases 的集合
• database
• DB table

graph 的 edge 包含兩類：

• Inclusion relation：表示一個 db 是否是一個 db collection 的一部分；或者一個 table 是否屬于一個 db
• Table relation：包含顯式的 PRIMARY-FOREIGN 關系和隱式的 FOREIGN-FOREIGN 關系

隱式的 FOREIGN-FOREIGN 關系指的是：A 表和 B 表的某個 column 共同連接到另一個 C 表的 key

由此，任何有效的 SQL query schema 都是這個 schema graph 上的一個 trail（或者叫一個 path）。

3.1.2 Schema Serialization

這個序列化算法將一個 SQL query schema 序列化為一個 token seq，當然也可以將一個 token seq 解碼出一個 DB schema。

具體的做法可以參考原論文，這里主要是基于 DFS（深度優先遍歷）的思想。

有了這個序列化算法，當我們訓練 seq2seq 的 schema router 模型時，由于需要監督它的 training data 是 (NL question, DB schema) pair，其中的 DB schema 就是序列化了的 schema。另外，router 的輸出是一個 token seq，也需要反序列化將其轉為結構化的 DB schema。

3.1.3 graph-based 的解碼算法

在讓 schema router 生成 token seq 時，為保證其生成的 schema 的有效性，每一個自回歸生成的 step 中，都受到一個動態前綴樹的約束，這個 tree 包含了解碼后 schema 元素的可能訪問節點的名稱，如下圖所示：

這樣，每個生成 step 的可用 tokens 都可以通過搜索前綴樹來獲得，前綴就是在最后一個元素分隔符之后生成的 token。同時這里使用 diverse beam search 來生成多個候選序列。

3.1.4 schema router 的訓練和推理

我們需要使用 (NL question, DB schema) 這樣的 pairs 來作為 training data 來訓練 router，但是目前缺少這樣的訓練資料。所以，本文提出了使用一個訓練數據合成方法來生成 question-schema pairs。

這個訓練數據合成方法具體來說就是：茨貝格 schema graph 中采樣出一批合法的 schema，然后對每一個 schem 生成一個 pseudo-question，如下圖所示：

具體的這個模型的訓練可以參考原論文。

由此就可以得到用于訓練 schema router 的 question-schema pairs。

之后，我們就可以訓練 Schema Router 了。訓練數據集是 $\{(N_i,S_i)\}$，也就是 quetsion-schema pairs，模型的訓練損失函數如下：

訓練出來之后，就可以使用 graph-based 的解碼算法來做推理了。

3.2 SQL Generation

通過將 NL2SQL 任務解耦為 schema routing 和 SQL generation 兩個部分，DB Copilot 可以與現在的 LLM-advanced NL2SQL 的解決方案進行融合，無論是 in-context prompt engineering 方法或者特定的 NL2SQL LLM。

前面的 schema router 可以為 NL 生成來自多個 db 的多個 schemas，這里探索了 3 種 prompt 策略來為 LLM 選擇和合并這些不同的 DB schema：

1. Best Schema Prompting：從 schema router 種選擇生成的最高概率的 schema 來 instruct LLM
   • 實驗發現這種方式是最優的
2. Multiple Schema Prompting：將 beam search 得到的多個 table schemas 簡單連接起來一起用來 instruct LLM。
3. Multiple Schema COT Prompting：使用多個 candidate schemas 通過 COT 來 instruct LLM

四、實驗

論文在 Spider、Bird、Fiben 數據集上對 schema retrieval 和 NL2SQL 兩個任務上進行實驗對比，DBCopilot 有不錯的表現。

這里 NL2SQL 任務并沒有與其他 SOTA 模型做實驗對比

五、總結

本文提出了 DBCopilot 模型，給出了一種將 NL 查詢擴展到大規模數據庫的思路，通過 LLM 協作來解決模式無關的 NL2SQL 任務。

總之，DBCopilot 突破了 NL2SQL 的界限，使得研究人員能夠更好地執行數據可訪問性的策略。