A 4B model trained with RL produces 81% faster query plans than Postgres

A small model against Postgres's planner

Rohan Bansal has published an experiment showing that a small open-weights language model can be post-trained to produce PostgreSQL query plans that beat the database's own. On the Join Order Benchmark (JOB) — 113 join-heavy SQL queries over the IMDb dataset — the trained 4B model reached a 1.81x geometric mean speedup once the best candidate was selected, and summed latency over the whole workload dropped by 44.7%. The same checkpoint recorded 68 wins and no regressions.

The starting point was discouraging. Untrained, the model could not produce a usable plan for 99 of the 113 queries, and only 14 of its attempts yielded a valid candidate. The author's framing is that query optimizers are hard — join ordering is known to be NP-hard, and planners estimate cardinalities from statistics instead of counting rows — but verifying a plan is easy, because the only axis that matters is execution time. That asymmetry is what makes the problem a good fit for reinforcement learning.

Hints, a harness, and distillation

The model does not replace the planner; it steers it. Bansal used pg_hint_plan, a third-party extension that accepts structured hints in SQL comments, and built a six-tool agent harness called qo-agent. The agent inspects relations and column statistics, reads the default plan and submits candidate plan actions, which are compiled into hints, executed and timed against the default.

Training came in two stages. First, supervised fine-tuning by off-policy distillation from 500 GPT-6 Astra agent trajectories, updating a 42.5 MB LoRA adapter with 21.2 million trainable parameters on top of a 4.66-billion-parameter Qwen derivative. Two epochs improved results and a third regressed them, even though validation loss had stopped moving.

Rewards in a noisy environment

The second stage was agentic reinforcement learning with a custom GRPO variant. Speedup was scored as the median of three default-plan measurements divided by the median of three candidate measurements, and the measurement rig had to be built carefully: four PostgreSQL containers on one machine create page-cache contention that pollutes timings. Training was split across two locations — vLLM and the trainer on a rented 2x H100 node, the database containers on the author's desk. Early reward shaping pushed the model to keep returning Postgres's default plan; the fix was to compare rollouts against each other instead of scoring them in absolute terms.

What the model learned

Analysis of the final checkpoint, after 1,200 updates, shows 295 of 337 searches first inspected a relation, column statistics or the default plan. The project cost about $1,200: roughly 95 hours of a 2x H100 node from Lambda plus $400 in OpenAI API fees for the demonstration trajectories. Bansal's conclusion is not that large models are obsolete — the distillation from Astra is why the small model works at all — but that small models deserve serious attention on narrow, easily verifiable tasks.