小型4Bモデル、強化学習によりPostgreSQLのクエリ最適化において1.81x倍の高速化を実現

学習プロセスには、2段階のアプローチが採用されました。まず、モデルはオフポリシー蒸留を用いた教師あり微調整（SFT）を受けました。研究者は、フロンティアスケールの教師モデルであるGPT-6 Astraからの軌跡（trajectories）を使用し、クエリ最適化のために設計された専用のエージェント型ハーネスの「言語」を、4Bの生徒モデルに教え込みました。

SFTに続き、強化学習（RL）を用いてモデルのさらなる洗練が行われました。研究者は、モデルがより優れたクエリ計画を選択する能力を最適化するために、Group Relative Policy Optimization (GRPO) のカスタムバリアントを実装しました。報酬メカニズムは、候補となる計画のPostgreSQLデフォルトに対する実行速度の向上に特化して設定されており、同時に、無効な計画や単にデフォルトの動作を複製しただけの計画にはペナルティが課される仕組みとなっています。

その結果、極めて小さなモデルであっても、ニッチで価値の高いドメインにおいて大幅な効率向上を達成できることが示されました。エージェントによるマルチターン・アプローチ（1つのクエリにつき最大15回の候補試行を許可）を用いることで、最終的なチェックポイントは、テストされた結合多用のクエリにおいて、総ワークロードのレイテンシを44.7%削減することに成功しました。この研究は、大規模モデルから得られる既存の高品質なデータを活用することで、企業が専門的でコスト効率の高い最適化タスクに小型のオープンウェイトモデルを利用できる可能性を浮き彫りにしています。

PLUS ULTRAby Amenoyomi

性能向上は顕著ですが、このモデルの成功は、データベースがクエリコストを推定する際の手法における構造的な限界を解決できるかにかかっています。PostgreSQLのデフォルトのオプティマイザは、NP困難として知られる「結合順序（join ordering）」の決定に苦戦します。異なる計画のコストを見積もる際、システムは「一様分布の仮定」に依存しており、あるテーブルにおける値の頻度が別のテーブルでも均等に適用されると想定しています。現実のデータがこれから逸脱する場合（例えば、少数のエントリが関連レコードの大部分を占めている場合など）、コストモデルの推定は失敗し、オプティマイザが非効率な計画を選択する原因となります。

このモデルは、2つの明確な段階を経て学習することで、この問題を克服しています。まず、エージェント型ハーネスの「言語」を習得し、pg_hint_planを介してデータベースとインターフェースを構築する方法を学ぶための教師あり微調整を行います。データベースエンジン自体を書き換えるのではなく、モデルはSQLコメントとして構造化されたヒントを生成し、オプティマイザを誘導します。インターフェースの習得後は、検証可能な単一の軸、すなわち「実行時間」の最適化のために強化学習が用いられます。