M_SQaLe: a 2B text-to-SQL agent trained on SQaLe

Project page · Dataset · All three models · Citation

M_SQaLe is the text-to-SQL model from the paper SQaLe: a large realistic dataset to empower small specialised text-to-SQL models. It is Qwen3.5-2B trained with GRPO directly on SQaLe, a semi-synthetic text-to-SQL dataset of 1,408,056 questions and 176,761 SQL queries over 9,259 real-world schemas from SchemaPile. SQaLe's schemas are large, with a median of 113 tables and 538 columns, so the model learns to work as an agent. It explores the database through tools, looks up the values it needs, tests its query and then submits it. The schema is never part of its prompt.

The model is trained with reinforcement learning only. There is no supervised warm start, no reasoning traces distilled from a larger teacher, and no test-time scaffolding at evaluation.

Quickstart

The model is an agent, so it needs a loop that executes its tool calls against a database and feeds the results back. sqale_agent.py in this repository is a self-contained reference loop that reproduces the training environment: the same prompt, tools, turn budget and final-answer fallback. It opens your database read-only.

pip install vllm huggingface_hub
hf download trl-lab/qwen3.5-2b-grpo-sqale sqale_agent.py --local-dir .

python sqale_agent.py --model trl-lab/qwen3.5-2b-grpo-sqale \
  --db path/to/database.sqlite \
  --question "Which supplier shipped the most back-ordered items last quarter?" \
  --trace

From Python, one agent can answer many questions with a single loaded engine:

from sqale_agent import SQaLeAgent, vllm_generate

agent = SQaLeAgent(*vllm_generate("trl-lab/qwen3.5-2b-grpo-sqale"))
episode = agent.run("How many customers placed more than three orders in 2024?", "shop.sqlite")
print(episode["sql"])       # the submitted query
print(episode["ended_by"])  # submit_sql, final_turn, fallback_to_last_working_query, ...

To build your own loop, follow the protocol in The agent loop below.

Highlights

• +33.0 points on BIRD dev over the untrained base model (19.3% → 52.3%), and +27.6 points on the SQaLe test set (38.7% → 66.3%).
• Strong for its size. On 100 moderate SQaLe test questions it reaches 52% at 224 TFLOPs per question. Every general-purpose model in the 9–12B range except Qwen3.5-9B scores lower, and Qwen3.5-27B is 11 points ahead at 1,710 TFLOPs per question.
• Finds the right tables in large schemas. On full-size SQaLe test schemas it opens every gold table in 89.7% of episodes.
• Reads values from the data. Before submitting, it has looked up 88% of the string literals the gold answer depends on. The untrained base model looks up 35%.

_{Accuracy on 100 moderate questions from the SQaLe test set against model size, with the average compute per question as colour. MSQaLe, MBIRD and MSynSQL are Qwen3.5-2B trained with the same recipe on SQaLe, BIRD train and SynSQL-2.5M; the red arrow is the gain of MSQaLe over the untrained Qwen3.5-2B. Figure from the paper's appendix.}

The SQaLe model family

The paper trains Qwen3.5-2B three times with one recipe, once per training corpus, to measure how much of a small model's ability comes from the data it is trained on. Base model, agentic environment, tool set, reward, GRPO objective, step count, batch size and evaluation protocol are identical across the three. Only the training corpus and the strength of the length curriculum differ, so the gaps between the models come from the data.

Evaluation

Execution accuracy (%) on 300 questions per benchmark, as reported in the paper. Every model runs in the agentic environment used for training, with the schema withheld and no test-time scaffolding. The SQaLe test split shares no schema with the training split and is evaluated over the full schema. EHRSQL is a single-domain benchmark over electronic health records.

On the SQaLe test set, training on SQaLe adds 27.6 points to the base model. The same recipe adds 15.3 points when trained on BIRD and 12.0 points when trained on SynSQL.

Schema size

The SQaLe test questions are also evaluated at three database sizes: the gold tables only, the gold tables plus 32 distractor tables, and the full schema. Brackets give the share of episodes in which the model opened every gold table.

Accuracy falls for every model as the schema grows, and M_SQaLe leads at every size. When only the gold tables are present, all three models find them. At full size, M_SQaLe opens all gold tables in 89.7% of episodes, against 85.0% for M_BIRD and 83.3% for M_SynSQL. BIRD and SynSQL schemas hold 5–10 tables, so the models trained on them never had to learn to search for the relevant tables. M_SQaLe learns this search during training.

Looking up values

Many questions depend on literals that exist only in the table values. Before submitting, M_SQaLe has looked up 88% of the string literals the gold answer depends on, against 55% for M_SynSQL, 40% for M_BIRD and 35% for the base model. The difference is largest for literals the question does not state. M_SQaLe searches the data for 61% of them and uses the correct value in 55% of cases, against 41% for M_BIRD and 39% for M_SynSQL.

The agent loop

What an episode looks like

Question   Which supplier shipped the most back-ordered items last quarter?
Turn 1     {"tool": "list_tables"}
           → suppliers, shipments, order_items, … (113 tables)
Turn 2     {"tool": "join_path", "args": {"table_a": "suppliers", "table_b": "order_items"}}
           → suppliers → shipments → order_items
   ⋮
Turn 7     {"tool": "submit_sql", "args": {"sql": "SELECT s.name …"}}

Each turn opens with a <think> block in which the model plans its next call, and a turn may hold several calls.

Using your own harness

The following details from training matter for accuracy if you build your own loop.

• Use the training prompt verbatim. SYSTEM_PROMPT and USER_TEMPLATE in sqale_agent.py are the exact strings used in training. The system prompt holds the task, the tool list and an execution-plan scaffold. The user turn states that the schema is not shown and gives the question.
• Tool calls are bare JSON. After its <think> block the model writes one or more {"tool": "<name>", "args": {...}} objects with no prose and no code fences. Keep thinking enabled.
• Return results as the next user turn, formatted as Tool result: followed by the result and the remaining budget, for example [4 tool turns left before you must answer].
• Keep earlier reasoning in context. During training every earlier <think> block stayed in the context. The Qwen chat template drops those blocks when a conversation is rebuilt from a message list, so append each completion to the prompt string instead, as sqale_agent.py does.
• Six tool turns, then a final turn. After the sixth tool turn, or after four failed run_query calls in a row, send the exploration-limit message and take the next answer as final. If that answer is empty or does not execute, use the last query that ran without error.
• Sampling. Use temperature 0.6, top-p 0.95 and top-k 20, which are the defaults in generation_config.json, and stop on <|im_end|> and <|endoftext|>. An episode is budgeted at 12,288 tokens, so a context of 16,384 tokens leaves headroom.

Training

Environment. Each training question becomes a tool-use episode over a live SQLite database built from its schema and rows, and the schema is never placed in the prompt. A graph-aware sampler decides which tables the episode's database holds. It draws a target size between 4 and 64 tables, starts from the gold tables and grows along foreign keys until the target is reached, then shuffles the table order. The model therefore sees both small and large databases, and list_tables does not reveal which tables matter. Early in training, a submit gate refuses submit_sql until the same query has been run with run_query. The gate's probability decays from 1.0 to 0 over the first 80 steps, and it is never active at evaluation.

Reward. A query whose result set matches the gold result scores 3.0. A query that executes with a wrong result scores 1.0 + 0.9·F₁ + b, where F₁ is the multiset F₁ against the gold result and b adds 0.1 times the share of joins along declared foreign keys and 0.1 for agreeing with the most common result in the group. A query that parses but does not execute scores 0.5, and anything else scores 0. The executable tier tops out at 2.1, so no partial credit ranks a wrong query above a correct one. Correctness is decided by comparing result sets, with no LLM judge.

Length curriculum. Half of SQaLe's gold queries are shorter than 150 characters, so a uniformly sampled batch mostly holds short questions. The curriculum buckets training questions by the length of their masked gold SQL at the 50th, 80th and 95th percentiles. The first 400 steps draw from the short bucket only. The mixture then moves linearly from (0.40, 0.30, 0.20, 0.10) to (0.15, 0.20, 0.25, 0.40) over 1,000 steps and holds there for the remaining 400. After the warm-up no bucket falls below 0.15, so short questions keep contributing gradient for the whole run.

Optimisation. GRPO from the Qwen3.5-2B base checkpoint, with 18 questions × 8 candidates = 144 episodes per update for 1,800 steps. This repository holds the last checkpoint. BIRD dev was monitored during training and not used to select the checkpoint.

All hyperparameters (identical for the three models except the curriculum)

Intended use and limitations

• M_SQaLe answers natural-language questions over SQLite databases through tool access, including databases with hundreds of tables. It is a research model released with the paper.
• It expects the agent loop described above. Putting the DDL in a single prompt and asking for SQL is a different setting from the one it was trained in.
• It was trained and evaluated on SQLite only.
• At 2B parameters it still trails much larger models. The untrained Qwen3.6-27B scores 76.0 / 69.3 / 55.0 on the three benchmarks above, and EHRSQL's clinical questions remain hard for every 2B model in the paper.
• The model's SQL runs against your data, so give it read-only access, as sqale_agent.py does.
• The checkpoint is a full Qwen3_5ForConditionalGeneration export and carries Qwen3.5-2B's vision tower unchanged. Only the language model was trained, and only text was evaluated.

Citation

If you use this model or the SQaLe dataset, please cite:

@misc{wolff2026sqale,
  title  = {{SQaLe}: A Large Realistic Dataset to Empower Small Specialised Text-to-{SQL} Models},
  author = {Wolff, Cornelius and Gomm, Daniel and Hulsebos, Madelon},
  year   = {2026}
}

Authors: Cornelius Wolff and Daniel Gomm (University of Amsterdam, Centrum Wiskunde & Informatica), Madelon Hulsebos (Centrum Wiskunde & Informatica). Questions and feedback are welcome in the Community tab of this repository.

The model builds on Qwen3.5-2B. SQaLe's schemas come from SchemaPile, and the evaluation uses BIRD and EHRSQL. We thank their authors for making them available.