The Pulse
A 4B Model Cut PostgreSQL Query Latency by 44.7%
Rohan Bansal reports that a trained 4B model improved query plans across 113 join-heavy PostgreSQL queries, reducing latency by 44.7% in a benchmark against PostgreSQL’s default plans.
AI.info Team ·
A trained 4B language model reduced query latency by 44.7% across 113 join-heavy PostgreSQL queries, according to a technical write-up by Rohan Bansal. The experiment used supervised fine-tuning and agentic reinforcement learning to teach the model to generate alternative query plans and test them against PostgreSQL’s defaults.
“What follows is a breakdown of an experiment I ran to explore the question: can a small, open-weights model be post-trained via supervised fine-tuning (SFT) and agentic reinforcement learning (RL) to produce Postgres query plans that beat Postgres’s default plans?”
Rohan Bansal, the write-up’s author
Bansal frames the work as an attempt to address a longstanding weakness in query optimization. PostgreSQL must choose plans using estimates rather than executing every possible join during planning. Exact measurement would defeat the purpose of a fast optimizer, while join ordering becomes difficult as the number of tables grows.
The experiment therefore treats query optimization as a search problem. The model generates a candidate strategy, sends it to PostgreSQL for measurement, and receives feedback based on whether the candidate runs faster than the database’s default plan. The goal is not to replace PostgreSQL’s optimizer during interactive execution, but to find better plans for workloads in which the same analytical queries run repeatedly.
An initially unusable baseline
The benchmark used 113 join-heavy queries from the Join Order Benchmark, based on an IMDb database. Before training, the 4B model performed poorly when connected to the query-optimization harness. Bansal reports that the model could not produce a valid query plan for 99 of the 113 queries.
Only a small number of the remaining attempts produced candidates that could be compared with PostgreSQL’s default plans. Some candidates duplicated the default plan, while other trajectories failed during selection or timed out. The results left only a limited set of usable trajectories for the initial evaluation.
That baseline gave Bansal a way to measure the effect of training. Supervised fine-tuning first improved the model’s ability to communicate with the agent system. Agentic reinforcement learning then adjusted the model toward strategies that produced faster measured plans. The training also used off-policy demonstrations generated by a larger model, giving the smaller model examples of how to inspect a query and propose changes.
What the trained model changed
After training, the model regularly explored changes to scan methods, join order and parallel execution. It also favored nested-loop joins in some cases, chose index scans over bitmap or sequential scans, and adjusted planner settings. Bansal identifies three recurring sources of improvement: rewriting the join order, correcting an individual scan without changing the join tree, and enabling parallel execution.
One example involved a query with a costly sequential scan and a lossy filter. The model identified the scan as a problem and forced a bitmap scan instead. Bansal reports that this individual query ran 90 times faster under the alternative plan.
The broader benchmark result was less extreme but still substantial. When the system evaluated multiple candidates for each query and selected the best measured option, it achieved a 1.81x geometric-mean speedup across the 113-query workload. Bansal also reports a 44.7% reduction in latency across the workload.
A workload-tuning system
The results apply to repeated analytical workloads rather than one-off requests. Generating and measuring alternative plans requires additional computation, so the approach makes the most sense when a plan can be reused across many executions of the same query.
Bansal describes the project as evidence that a small open-weights model can learn a specialized optimization task when its actions are connected to a measurable outcome. The model does not need to replace PostgreSQL’s planner in every situation. Instead, it can search for alternatives offline, retain plans that perform better, and apply those results to a recurring workload.
The write-up presents the experiment as a research system rather than a drop-in PostgreSQL feature. Its performance depends on the benchmark, the measurement process and the strategy used to select among candidate plans. Even so, the reported results show that a 4B model can improve on PostgreSQL’s default choices for a substantial set of join-heavy queries after targeted training.