論文：MAC-SQL: A Multi-Agent Collaborative Framework for Text-to-SQL

????

arXiv:2312.11242, 北航 & Tencent

Code: MAC-SQL | GitHub

一、論文速讀

本文提出了一個使用多個 agents 進行互相合作的框架 MAC-SQL 來解決 Text2SQL 任務。

MAC-SQL 主要由 3 種 agents 組成：

• Selector：通過去除掉 inference 時無關的信息來把一個大的數據庫分解為更小的數據庫
• Decomposer：通過 prompt 的方法來把一個復雜的 question 分解為漸進性的幾個可以被獨立解決的 sub-questions
• Refiner：用于檢測和自動改正 SQL 的錯誤

原論文對三者的描述：

Specifically, the Decomposer disassembles complex questions into simpler sub-questions and addresses them sequentially through chain-of-thought reasoning. If required, the Selector decomposes a large database into smaller sub-databases to minimize interference from irrelevant information. Meanwhile, the Refiner utilizes an external tool for SQL execution, acquires feedback, and refines any incorrect SQL queries.

下面是基于三個 agents 來實現 MAC-SQL 的算法流程：

二、MAC-SQL

2.1 Selector agent

The Selector agent is designed to automatically decompose a large database into smaller sub-databases to minimize interference from irrelevant information.

由于現實世界種的數據庫特別大，包含了許多 tables 和 columns，一次 LLM 的 API Call 可能無法處理這些多的 schemas，因此需要使用 Selector 來去除掉無關信息并得到的一個較小的 schema。

如下是一個使用 Selector 的 prompt 示例：

可以看到，Selector 的 prompt 包含四部分：task desc、instruction、demonstrations 和一個 text example，期望的輸出是一個 JSON，里面枚舉了所有選擇后的 tables，并將它們分成 3 類：

• “keep_all”
• “drop_all”
• 一個相關的 column list

這個 JSON 中的 schema 輸入給 Decomposer agent。

2.2 Decomposer agent

The primary purpose of the Decomposer is to systematically decompose complex questions into progressively refined sub-questions, which can then be solved individually.

當面對復雜問題時，生成的 SQL 經常會有缺陷，因此一個自然的想法就是像 CoT 一樣將其分解為 sub-questions 再解決。

圖示：

可以看到，最后一個 sub-question 就是原來的 user question 了，所以最后一步生成的 SQL 就是 Decomposer 輸出的 SQL。

具體來說，Decomposer 可以用 CoT 或者 least-to-most 兩種 prompt 策略來實現。

2.3 Refiner agent

The primary function of the Refiner is to detect and automatically rectify SQL errors.

對于一個 SQL，Refiner 首先從下面三個角度來診斷：

• 句法正確性
• 可執行性
• DB 檢索后是否為非空結果

如果檢查通過，這個 SQL 被輸出為最終答案，否則，就執行 correction 操作。之后，修正后的 SQL 仍需重新診斷，重復這個過程直至檢查通過或者達到最大重復次數。

具體的 correction 過程包括基于原始 SQL 和錯誤反饋信息或修改引導信號進行推理，生成修改后的結果。

圖示：

最后論文還指出，單靠一個 refiner agent 提供的幫助也是有限的，不可避免還存在一些錯誤的問題，這需要再系統層面進行更多的優化。

三、指令微調的 SQL-Llama

本文基于 Code Llama 7B，使用前面介紹的 3 個 agents 的指令數據進行微調，得到了 SQL-Llama，讓 model 在 database simplification、question decomposition、SQL generation 和 SQL correction 方面的能力得到增強。

用于微調的數據是基于三個 agents 在 BIRD 和 Spider 數據集上得到的。

這個過程的關鍵挑戰是：model 的訓練過程需要平衡它的 complexity 和 performance。也就是需要在維持其較高的 performance 時有效處理 db 相關任務的復雜度。

四、實驗

論文在 BIRD 數據集上做的測試，baseline 只選擇了 LLM-based 的方案，并沒有選擇非 LLM 的 baseline。實驗結果如下圖所示：

論文還做了消融實驗，證明了三個 agents 在提高 acc 方面都發揮了重要的作用。

另外論文還發現，ICL 中增加 demonstrations 的數量可以讓效果更好。

論文最后還給出了 error cases 的統計分析，可以參考原論文。

五、總結

本論文利用多個 agents 合作的思路來解決 Text2SQL 任務，同時提供了一個開源的 SQL-Llama 模型，在 BIRD 數據集上實現了 SOTA 效果。