Training a 4B Model to Produce 81% Faster PostgreSQL Query Plans

TL;DR

A 4‑billion‑parameter open‑weights model, after off‑policy distillation from GPT‑6 Astra and several epochs of reinforcement learning, can generate pg_hint_plan hints that make PostgreSQL execute the Join Order Benchmark 81 % faster on average, cutting total workload latency by 44.7 %.

Why query optimizers still lag

Leis et al. (2015, 2025) showed that PostgreSQL’s optimizer leaves substantial performance on the table, especially for join ordering, which is NP‑hard.
Optimizers rely on coarse statistics and uniform‑distribution assumptions; a single bad cardinality estimate can cascade into dramatically sub‑optimal plans.
Verifying plan quality is trivial – execution time is a single, observable metric – making the problem well‑suited for reinforcement learning.

Problem formulation

Goal: Train a small model to emit pg_hint_plan comments that steer PostgreSQL to faster plans than its default cost‑based planner.

Key insight: For analytic workloads where the same query runs repeatedly, the amortized benefit of a better plan outweighs the one‑off cost of hint generation.

Experimental platform

Component	Details
Database	PostgreSQL 15 on a slice of the IMDb dataset (8.5 GB on disk).
Hardware	Development machine "FLOPper" – 2 × RTX 3090, 16 CPU cores, 64 GB RAM.
Training GPU	2 × H100 SXM (80 GB VRAM each) rented from Lambda for ~95 h.
Benchmark	Join Order Benchmark (JOB – 113 queries, 33 join‑graph topologies).
Auxiliary benchmark	Cardinality Estimation Benchmark (CEB – ~13.6 k queries) used for training data.

Reducing measurement noise

Calibration rig: Four Dockerized PostgreSQL containers, each pinned to 4 CPU cores and 8 GB RAM, pull queries from a shared queue.
Warm‑up policy: Run a query until shared_buffers hit‑block and read‑block counters stabilize within 2 % over two consecutive runs.
Noise metric: Simulated "no‑op" reward (candidate = default) showed a 5 % mean error rate and a 13 % – 20 % worst‑case (p90) error when shared_buffers = 128 MB.
Tuning: Raising shared_buffers to 2 GB eliminated most noise (mean error ≈ 1.5 %, p90 ≈ 0 %) and reduced total JOB runtime from 95 s to 60 s.

Model and harness

Base model: empero-ai/Qwen3.8-4B-Distill (Qwen 3.8 4 B distilled from a 2.4 T teacher).
Agent harness (qo-agent): Provides six tools (inspect_relation, get_column_stats, get_plan, evaluate_candidate, keep_default, finish).
Interaction format: Model outputs PlanAction JSON; the harness converts it to a pg_hint_plan comment and measures execution time.

Off‑policy distillation (Supervised Fine‑Tuning)

Teacher trajectories: 120 GPT‑6 Astra rollouts (5 candidates each) on random CEB queries, with reasoning summaries.
Rendering & loss‑masking: Converted OpenAI response JSON to Qwen token streams; masked out system/user prompts and tool‑output tokens.
LoRA adapter: 42.5 MB (21.2 M trainable parameters) frozen the 4 B base model weights, enabling training on a single RTX 3090.
Training schedule: One epoch on 100 Astra trajectories → modest improvement (geometric mean 0.72×). Two epochs → 1.08× geometric mean speedup; three epochs over‑fit and regressed.
Data expansion: Added 300 Astra trajectories (filtered to exclude trivial "keep default" runs). Two more epochs raised geometric mean to 1.10× and total workload speedup to 1.05×.

Agentic Reinforcement Learning

Reward design:
- Speedup = median(default) / median(candidate).
- Clip ratio to [0.1, 10] and apply a 0.05 soft‑threshold to ignore noise.
- Penalize invalid plans (‑0.1), duplicate‑default plans (‑0.02), and trajectories with no valid candidate (‑0.1).
Advantage estimator: Custom "anchored" GRPO that subtracts the group mean and applies a per‑rollout sign, ensuring only truly faster plans receive positive advantage.
Training hyper‑params: LR = 1e‑5, batch = 16, 8 rollouts per query, 600 optimizer updates per stage.
Concurrency trick: Lease PostgreSQL containers only during measurement phases, allowing 20 concurrent rollouts and keeping vLLM inference fully saturated.

Results

Stage	Valid candidate queries	Tasks scored	Geometric mean speedup	Total workload speedup	Wins (≥ 5 % faster)	Regressions (≥ 5 % slower)
Untrained 4 B	14/113	15/113	0.85×	0.85×	3	1
After SFT (1 epoch)	48/113	44/113	0.72×	0.76×	5	16
After SFT (2 epochs, 300 extra traj.)	77/113	101/113	1.10×	1.05×	29	13
RL 600 updates	99/113	113/113	1.35×	1.16×	34	0
RL 1 200 updates	101/113	112/113	1.41×	1.29×	38	2
RL 1 200 + 3 rollouts per query (best‑of‑15)	113/113	339/339	1.81×	1.81×	68	0

The final checkpoint achieved a 1.81× geometric‑mean speedup (44.7 % total latency reduction) when selecting the best candidate among up to 15 sampled plans per query.

What the model actually learned

Frequent actions: Leading tree hints (917 uses), scan hints (1 141), and Parallel hints (572).
Join method preference: Nested‑loop joins were often forced, even when hash joins were default.
Scan preference: Index scans over bitmap or sequential scans.
Configuration tweaks: Frequently set enable_sort=off and random_page_cost=1.1, indicating the model exploits cost‑model sensitivities.
Successful motifs: Reordering joins via Leading, adding a single more selective scan, and enabling parallelism together accounted for the majority of speedups.

Cost breakdown

Item	Cost
Lambda 2 × H100 rental (≈ 95 h)	~ $800
OpenAI API for Astra trajectories (≈ 400 k tokens)	~ $400
Electricity for FLOPper (continuous GPU use)	~ $9 / day (negligible)
Total	≈ $1,200

Community reaction (selected HN comments)

"81 % faster query plans on an 8 GB dataset that fits entirely in memory…" – refibrillator (concern about over‑fitting and scalability).
"Frontier intelligence is extremely powerful; the distillation I did off Astra trajectories is proof enough that large models are not going anywhere" – devsda (notes the broader relevance of open‑weight distillation).
"It’s not hard to be more than 3× better than Postgres; you don’t need a model" – huahaiy (points out that hand‑tuned hints can also yield large gains).
"The model regularly used enable_sort=off and random_page_cost=1.1 – could explain gains" – zacmps (highlights possible confounding configuration changes).
"Would love a product that auto‑generates hints for recurring analytic queries" – ashley95 (suggests a practical SaaS direction).

Takeaways

LLMs can learn to emit effective optimizer hints. Even a 4 B model, once distilled from a larger teacher and reinforced with real execution feedback, outperforms PostgreSQL’s native planner on a realistic analytic benchmark.
Noise‑aware measurement is critical. Proper cache sizing (shared_buffers = 2 GB) and a three‑pair median‑of‑medians evaluation reduced spurious reward signals from < 5 % to near‑zero.
Off‑policy distillation accelerates SFT. A few hundred teacher trajectories provided enough token‑level supervision to teach the model the harness language and basic hint syntax.
Reward design matters more than algorithmic novelty. Early GRPO variants reinforced duplicate‑default plans; the custom anchored GRPO and clipped speedup reward were essential for stable RL progress.
Small models are cost‑effective for domain‑specific tasks. Training cost ≈ $1.2k, inference cost negligible on a single RTX 3090, yet the model delivers > 80 % speedup on repeatable workloads.

Future work

Structured hint sweeping (e.g., Bao) vs. LLM‑generated hints.
On‑policy distillation to compare sample efficiency.
Trace inversion for richer reasoning token supervision.
Scaling to larger datasets (TB‑scale) and mixed OLTP/OLAP workloads.
Automated hint‑caching service that continuously re‑trains on production query logs.

Code & citation

All code is open‑source at https://github.com/polyphilz/qorl.

@article{bansal2026qorl,
  title = {Training a 4B model to produce 81% faster query plans than Postgres},
  author = {Bansal, Rohan},
  journal = {rohanbansal.com},
  year = {2026},
  month = {September},
  url = {https://rohanbansal.com/qorl}
}