Training a 4B model to produce 81% faster query plans than Postgres

This describes training a 4-billion-parameter language model to generate Postgres query plans using supervised fine-tuning and an agentic reinforcement-learning loop, with the goal of minimizing query execution time. It frames query optimization as a hard, combinatorial problem - join ordering is NP-hard and a single three-table query can have thousands of physical-plan variants (an example counts 4,608 combinations when join algorithms, orientations, and scan types are considered). The approach treats execution time as a single verifiable reward signal: the model emits candidate plans in multiple rollouts, each plan is executed by Postgres, measured against the DBMS’s default plan, and scalar rewards flow back to update model weights, nudging behavior toward consistently faster plans.

The experiment shows substantive gains: up to 81% faster plans in specific cases and a 44.7% latency reduction across 113 join-heavy queries from a 4B model that initially failed to plan for many queries. Key implementation details include a low-noise Postgres measurement rig, a custom GRPO-style reward design for noisy measurements, splitting RL work between vLLM and a rented 2×H100 trainer while running Postgres containers locally, off-policy distillation from hundreds of GPT-6 Astra trajectories, and practical tricks like LoRA, rendering/loss masking, and unrolling. Benchmarks include the Join Order Benchmark and a cardinality-estimation suite, and the result demonstrates that a small open model can learn optimization policies that outperform Postgres defaults.