4Bモデルを訓練してPostgreSQLのクエリ計画を81％高速化する

TL;DR

GPT‑6 Astraからオフポリシー蒸留を行い、複数エポックの強化学習を経た40億パラメータのオープンウェイトモデルは、PostgreSQLのJoin Order Benchmarkを平均で81％高速化するpg_hint_planヒントを生成でき、ワークロード全体の遅延を44.7％削減する。

クエリオプティマイザがまだ遅れている理由

Leisら (2015, 2025) は、PostgreSQLのオプティマイザが、特にNP困難な結合順序に関して、大幅な性能向上の余地があることを示した。
オプティマイザは粗い統計と一様分布の仮定に依存しており、1つの不正確な基数推定が劇的に非最適な計画にまで悪化する可能性がある。
計画の品質を検証するのは簡単である——実行時間は単一の観測可能な指標である——このため、強化学習に非常に適した問題である。

問題の定式化

目的: 小規模なモデルを訓練し、pg_hint_planコメントを出力させ、PostgreSQLのデフォルトのコストベースプランナーよりも高速な計画を実行させる。

重要な洞察: 同じクエリが繰り返し実行される分析ワークロードでは、より良い計画のアモルタイズされた利点は、ヒント生成の初期コストを上回る。

実験プラットフォーム

コンポーネント	詳細
データベース	IMDbデータセットのスライス（ディスク上8.5GB）で動作するPostgreSQL 15。
ハードウェア	開発マシン「FLOPper」——2×RTX 3090、16CPUコア、64GB RAM。
トレーニングGPU	Lambdaからレンタルした2×H100 SXM（各80GB VRAM）で約95時間。
ベンチマーク	Join Order Benchmark（JOB — 113クエリ、33結合グラフトポロジー）。
補助ベンチマーク	Cardinality Estimation Benchmark（CEB — 約13.6kクエリ）はトレーニングデータとして使用。

測定ノイズの低減

キャリブレーション装置: 4つのDocker化されたPostgreSQLコンテナ。各コンテナは4CPUコアと8GB RAMにピン止めされ、共有キューからクエリを取得。
ウォームアップポリシー: shared_buffersのヒットブロックおよびリードブロックカウンターが2回連続で2％以内に安定するまでクエリを実行。
ノイズ指標: シミュレートされた「ノーオペ」報酬（候補＝デフォルト）では、shared_buffers＝128MBのとき平均誤差率5％、最悪ケース（p90）13％～20％を示した。
チューニング: shared_buffersを2GBに引き上げることで大部分のノイズを排除（平均誤差≈1.5％、p90≈0％）、JOB全体の実行時間は95秒から60秒に短縮された。

モデルとハーネス

ベースモデル: empero-ai/Qwen3.8-4B-Distill（Qwen 3.8 4B、2.4Tの教師モデルから蒸留）。
エージェントハーネス（qo-agent）: 6つのツール（inspect_relation、get_column_stats、get_plan、evaluate_candidate、keep_default、finish）を提供。
インタラクション形式: モデルはPlanAction JSONを出力。ハーネスはこれをpg_hint_planコメントに変換し、実行時間を測定。

オフポリシー蒸留（教師あり微調整）

教師トレース: 無作為なCEBクエリ上で120回のGPT‑6 Astraロールアウト（各5候補）、推論要約付き。
レンダリングと損失マスク: OpenAIの応答JSONをQwenトークンストリームに変換；システム／ユーザーのプロンプトおよびツール出力トークンをマスク。
LoRAアダプタ: 42.5MB（21.2Mトレーナブルパラメータ）で4Bベースモデルの重みを固定し、単一のRTX 3090でトレーニング可能に。
トレーニングスケジュール: 100件のAstraトレースで1エポック → 小さな改善（幾何平均0.72×）。2エポック → 1.08×幾何平均の高速化；3エポックで過学習し、性能が低下。
データ拡張: 300件のAstraトレースを追加（「デフォルトを保持」の自明な実行を除外）。さらに2エポックで幾何平均を1.10×、ワークロード全体の高速化を1.05×に向上。

エージェント型強化学習

報酬設計:
- 速度向上 = 中央値（デフォルト）／中央値（候補）。
- [0.1, 10]にクリッピングし、0.05のソフトスイッチを適用してノイズを無視。
- 無効な計画（‑0.1）、デフォルトと同一の計画（‑0.02）、有効な候補がないトレース（‑0.1）にペナルティ。
アドバンテージ推定器: グループ平均を引くだけでなく、各ロールアウトの符号を適用する独自の「アンカーデ」GRPO。これにより、本当に高速な計画のみが正のアドバンテージを得る。
トレーニングハイパーパラメータ: LR = 1e‑5、バッチ = 16、1クエリあたり8ロールアウト、各段階で600回の最適化更新。
並行処理テクニック: PostgreSQLコンテナは測定フェーズのみにリースし、20並列ロールアウトを可能にし、vLLM推論を完全に飽和させる。

結果

ステージ	有効な候補クエリ	スコアされたタスク	幾何平均の高速化	ワークロード全体の高速化	勝利（5％以上高速）	誤差（5％以上遅延）
未訓練の4B	14/113	15/113	0.85×	0.85×	3	1
SFT後（1エポック）	48/113	44/113	0.72×	0.76×	5	16
SFT後（2エポック、300追加トレース）	77/113	101/113	1.10×	1.05×	29	13
RL 600更新	99/113	113/113	1.35×	1.16×	34	0
RL 1,200更新	101/113	112/113	1.41×	1.29×	38	2
RL 1,200＋1クエリあたり3ロールアウト（ベストオブ15）	113/113	339/339	1.81×	1.81×	68	0

最終チェックポイントでは、1クエリあたり最大15個のサンプル計画から最良の候補を選択することで、幾何平均1.81×の高速化（全体の遅延44.7％削減）を達成した。

モデルが実際に学習した内容

頻出アクション: Leadingツリーのヒント（917回使用）、スキャンヒント（1,141回）、Parallelヒント（572回）。
結合方法の好み: ハッシュ結合がデフォルトでも、ネストドループ結合を強制する傾向があった。
スキャンの好み: ビットマップまたはシーケンシャルスキャンよりもインデックススキャンを好む。
設定の微調整: enable_sort=offとrandom_page_cost=1.1を頻繁に設定しており、コストモデルの感度を活用していることを示唆。
成功したパターン: Leadingによる結合の再順序付け、より選択性の高いスキャンの追加、並列処理の有効化の組み合わせが、大部分の高速化をもたらした。

コスト内訳

項目	コスト
Lambda 2×H100レンタル（約95時間）	~$800
OpenAI APIによるAstraトレース（約40万トークン）	~$400
FLOPperの電力（GPU連続使用）	~$9/日（無視できる）
合計	~$1,200

コミュニティの反応（選択されたHNコメント）

"8GBのデータセットでメモリに完全に収まる状況で、クエリ計画が81％高速化…" – refibrillator（過学習とスケーラビリティへの懸念）。 "フロンティアインテリジェンスは非常に強力。Astraトレースからの蒸留は、大規模モデルが消えることはないという証拠だ" – devsda（オープンウェイト蒸留の広範な関連性に言及）。 "Postgresより3倍以上良いのは簡単。モデルは必要ない" – huahaiy（手動でチューニングしたヒントでも大きな向上が可能である点を指摘）。 "モデルは定期的にenable_sort=offとrandom_page_cost=1.1を使用していた——これで性能向上が説明できるかもしれない" – zacmps（可能性のある混同要因を強調）。 "繰り返し実行される分析クエリ用のヒントを自動生成する製品が欲しい" – ashley95（実用的なSaaS方向性を提案）。

まとめ

LLMは効果的なオプティマイザヒントを生成できる。 4Bモデルであっても、より大きな教師モデルから蒸留され、実際の実行フィードバックで強化学習された後は、現実的な分析ベンチマークにおいてPostgreSQLのネイティブプランナーを上回る。
ノイズに配慮した測定が重要。 キャッシュサイズ（shared_buffers＝2GB）の適切な設定と、3ペアの中央値の中央値評価により、誤った報酬信号を5％未満からほぼゼロにまで低減した。
オフポリシー蒸留はSFTを加速する。 数百件の教師トレースから得られるトークンレベルの監督により、モデルがハーネス言語と基本的なヒント構文を学習できる。
報酬設計はアルゴリズムの革新よりも重要。 早期のGRPOバージョンはデフォルトと同一の計画を強化していた。独自のアンカーデGRPOとクリッピングされた速度向上報酬が、安定したRL進展に不可欠だった。
小規模モデルはドメイン特化タスクにコスト効率が高い。 トレーニングコストは約$1.2k、単一のRTX 3090での推論コストは無視できるが、繰り返し実行されるワークロードでは80％以上の高速化を達成している。

今後の課題

構造化ヒントスイープ（例：Bao）とLLM生成ヒントの比較。
サンプル効率を比較するためのオンポリシー蒸留。
より豊かな推論トークン監督を得るためのトレース逆転。
より大きなデータセット（TB規模）およびOLTP/OLAP混合ワークロードへのスケーリング。
本番クエリログを継続的に再トレーニングする自動ヒントキャッシュサービスの開発。

コードと引用

すべてのコードはオープンソースで、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}
}