/note/tech

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

要約:

■ 1. 問題設定と実験の概要

  • クエリオプティマイザの現状:
    • Leis らは2015年と10年後の2025年に同じ問いを立て、今なお改善の余地が大きいと結論づけた
    • 結合順序の決定だけでもNP困難であり、最適化は極めて難しい
  • 検証の容易さ:
    • 良いプランは速く、悪いプランは遅いため、プランの良否は実行時間という単一軸で検証できる
    • 言語モデルは検証しやすい出力を持つタスクの学習に強く、この問題は強化学習に適している
  • 実験の問い:
    • 小規模なオープンウェイトモデルをSFTとエージェント型RLで事後学習し、Postgresのデフォルトプランより速いプランを出せるか
    • 答えは明確にイエス
  • 主な成果:
    • 当初113クエリ中99件でプランを出せなかった4Bモデルが、結合の多い113クエリで44.7%のレイテンシ削減を達成
    • 並行コンテナ間のLinuxページキャッシュ競合ノイズを最小化する計測環境を構築
    • ノイズの多い環境でRLロールアウトを評価するため、独自のGRPO変種を設計
    • vLLMとトレーナーは2x H100のレンタルノード、Postgresコンテナ4台は手元のマシンという2拠点構成でRLを実施
    • GPT-6 Astraのエージェント軌跡約500件でオフポリシー蒸留を実施

■ 2. クエリオプティマイザの仕組み

  • 選択的述語の影響:
    • WHERE句の絞り込み条件の有無で、最適な結合順序が変わる
    • 述語がなければ、どの順序で結合しても2つ目の結合に渡る行数は同じ200万行
  • 結合順序による作業量の差:
    • 日本企業5%、2000年代のタイトル20%と仮定すると、cn'⋈mcを先に行えば約10万行となる
    • t'⋈mcを先に行うと約40万行となり、作業量は4倍
  • 組み合わせ爆発:
    • 結合木、内外の向き、結合アルゴリズム(hash、merge、nested loop)、スキャン方式の組み合わせで候補が増える
    • スキャン方式は逐次、インデックス、インデックスオンリー、ビットマップの4種を考慮する
    • 3テーブルの結合で4,608通り、12テーブルでは約7.3×10^22通りに達する
    • 実際には並列実行や集約方式もあるため過小評価であり、Postgresは動的計画法(12結合以上は遺伝的アルゴリズム)で探索空間を枝刈りする
  • 推定であり計数ではない:
    • 実際に結合を実行しなければ行数は数えられず、それでは高速なオプティマイザの意味がない
    • オプティマイザの目的は厳密な最小化ではなく、多様なクエリで十分良い結果を出すこと
    • Postgresはpg_statisticの頻出値と頻度、ヒストグラムからカーディナリティを推定する
  • 一様分布仮定の脆さ:
    • 結合時には、あるテーブルでの値の頻度が相手テーブルにもそのまま当てはまると仮定する
    • 日本企業5%が映画の50%を担っていれば、最初の結合は100万行となり、コストモデルの選択が誤りになる
    • 初期の結合での誤推定は、結合木全体の推定に連鎖的に波及する

■ 3. pg_hint_planによるプランの誘導

  • コストモデルの制約:
    • Postgresは常に最小コストのプランを選び、ソースを改変しない限りコストモデルは変えられない
  • pg_hint_plan:
    • SQL文の上にコメントとして構造化ヒントを付け、指定した結合方式やスキャン方式にプランを誘導するサードパーティ拡張
    • 例としてHashJoinとSeqScanを指定すると、実際のプランもその通りになる

■ 4. 取り組む価値のある問題設定

  • 問い:
    • 言語モデルは、より良いクエリプランにつながるヒントを生成できるようになるか
  • 却下した案:
    • プランナーと同じ情報をモデルに与える案は、より良いカーディナリティ推定器を作る試みに等しい
    • 数十年の研究と競うことになり、推論レイテンシだけで高速なオプティマイザに対する利得を上回る
  • 採用した案:
    • 同じクエリを何千回も実行する重い分析ワークロードに焦点を当てる
    • 学習時に同じクエリを数十〜数百回実行しても、全実行で償却すればコストは大幅に低い
    • 単発クエリではなく、繰り返し実行されるクエリでPostgresを上回ることが目標

■ 5. モデルとエージェントハーネス

  • モデル選定:
    • 自宅の2x RTX 3090機(FLOPper)で学習、推論しやすいよう4Bモデルから始めた
    • Qwen 3.8には4B版がなく、ドイツのEmperoによる蒸留モデルempero-ai/Qwen3.8-4B-Distillを採用した
    • このモデルはQwen 3.8 2.4Tを教師にQwen 3.5 4Bへ蒸留したもので、MMLUでは上回りGSM8Kではやや劣る
    • 本タスクにどちらが有利かは不明だが、蒸留モデルで進めた
  • qo-agentハーネスの6ツール:
    • inspect_relation: 列の型、インデックス定義、推定行数とバイト数を返す
    • get_column_stats: 1〜8列分のプランナー統計を返す
    • get_plan: デフォルトプランの推定値、または提出済み候補のプランを返す
    • evaluate_candidate: 提案プランを検証、実行し、時間とプラン診断を返す
    • keep_default: デフォルトプランを候補として返し、探索を終了する
    • finish: 選択した候補IDまたはデフォルトを受け取り、探索を終了する
  • 構造化出力:
    • エージェントはPlanAction JSONを出力し、evaluate_candidateがそれをヒントに変換して元のクエリの先頭に付ける
    • エージェントは複数の候補を試し、最も速い候補を選んで終了する

■ 6. ベンチマーク

  • Join Order Benchmark(JOB):
    • Leis らが導入し、IMDbデータセットで結合順序最適化を評価する
    • 33テンプレートにまたがる113クエリで構成される
    • テンプレートはテーブルと結合グラフが共通で、選択述語だけが異なるクエリの構造的な族
  • Cardinality Estimation Benchmark(CEB):
    • Flow-lossの論文で導入され、同じIMDbを使う約13.6kの合成クエリ、16テンプレートで構成される
    • CEBのテンプレート定義はJOBより緩く、結合グラフが同じでも述語だけ異なる別テンプレートがある
  • 学習と評価の分担:
    • 規模の大きいCEBを学習に、JOBを評価に使う
  • 同一DBでの学習と評価の妥当性:
    • モデルにIMDbをよく学ばせることがまさに狙い
    • 特定企業の特定DBの分析ワークロードに使う前提なら、全DBへの汎化は不要
    • 真の懸念はJOBのテンプレートへの過学習で、JOBと同形のCEBクエリは除外する必要がある
  • クエリトポロジーの照合:
    • 別名を解決したテーブルをノード、結合をエッジとする結合グラフをトポロジーと定義し、選択述語は除外する
    • JOBは33トポロジー、CEBは12トポロジーで、重複はなくフィルタリングは不要だった

■ 7. Postgresの計測ノイズ除去

  • ノイズ除去の必要性:
    • 同じクエリを20回実行しても実行時間は毎回異なる
    • 学習全体がデフォルトより速いかの計測に依存するため、ノイズを最大限除去する必要がある
  • 環境:
    • FLOPperは16物理コア、64GB RAM、2TB NVMe SSD
    • 使うIMDbのスライスはディスク上で8.5GB
    • Postgresはバッファにページをキャッシュし、OSのファイルシステムキャッシュも一段下で同じことをする
  • キャリブレーションの仕組み:
    • 各4コア、メモリ上限8GBのPostgresコンテナを4台起動した
    • 4台未満では学習が遅すぎ、4台超ではCPU競合によるノイズが増える
    • 共有キューから113クエリを取り出し、ウォームアップ後に20回計測する
  • ウォームアップの判定:
    • Postgresはshared_buffersでヒットするとshared hit blocks(SHB)、Linuxに要求するとshared read blocks(SRB)を加算する
    • EXPLAINのBUFFERSオプションで両カウンタを取得できる
    • 前回との差が両方2%以内でプランも同じならウォームと判定し、最低2回、最大5回でウォームアップを打ち切る
  • 128MB設定での問題:
    • 半数のクエリは2回でウォームと判定されたが、SRBはゼロにならず大きな値で安定していた
    • 8.5GBのDBに対し128MBでは毎回キャッシュミスが起き、安定は常駐を意味しなかった
  • 二峰性のノイズ:
    • job-13bの20回中14回は186〜204ms、6回は227〜253msで、14〜26%遅かった
    • ノイズは一様ではなく、2つの速度が時により現れた
  • 変動係数(CV)の限界:
    • job-13bのCVは10.3%だった
    • CVは平均に基づくため外れ値に左右されやすい
    • 重要なのはばらつきの大きさではなく、学習時の判定がどれだけ欺かれるか

■ 8. 報酬が欺かれる率

  • 計測方式:
    • 候補とデフォルトのペアを交互に3組実行し、それぞれの中央値の比で高速化か低速化かを判定する
    • 差が5%未満なら引き分けとする
  • 幻の報酬:
    • デフォルトと同一の候補を3回ずつ抽出すると、job-13bでは約20%の確率で片側の中央値が遅い塊に入る
    • その結果、存在しない14〜26%の高速化や低速化を学習信号としてモデルに見せてしまう
    • 全抽出で見ると、報酬は約40%の確率で欺かれる
  • 最小化すべき指標:
    • 同じCVでも、一様なばらつきか5%以上離れた2つの塊かで欺かれる率は異なる
    • 最小化すべきはCVではなく、欺かれる率そのもの
  • no-op誤り率の算出:
    • 20回の計測に6連続の窓をスライドさせ、隣接ペアで候補とデフォルトの役割を入れ替える
    • 1窓あたり8通り、15窓で計120通りを模擬する
    • 5%を超える差が出た割合を113クエリで平均したものを平均no-op誤り率とする
    • 各クエリの誤り率を昇順に並べた90%地点をp90クエリとし、最悪のクエリ群の指標とする
  • 128MB設定の結果:
    • 平均no-op誤り率は5.0%と5.4%、p90は13%と20%だった
    • no-opプランの20件に1件が報酬を得て、約10件に1件のクエリが13%超の確率で欺かれる

■ 9. Postgresのメモリ設定の調整

  • 調整対象:
    • shared_buffers: Postgres自身のキャッシュに保持できるDBの量
    • work_mem: ソートやハッシュ1回がディスクに溢れる前に使えるメモリ量
  • 結果:
    • work_memはノイズに全く影響せず、shared_buffersがすべてを決めた
    • 2GBでは中央値クエリのSRBがウォームアップ後にゼロとなり、作業セットが完全に常駐した
    • no-op誤り率は約4分の1に下がり、p90クエリはほぼ欺かれなくなった
    • job-13bのCVは10.3%から0.9%になり、20回すべてが7ms以内に収まった
  • 副次的な効果:
    • デフォルトプラン自体が速くなり、113クエリの合計実行時間が95秒から60秒に減った
    • 計測が速くなり、学習時間も短縮される
  • 最終設定:
    • shared_buffersを2GB、work_memを4MBに固定した

■ 10. 評価指標

  • 幾何平均高速化率:
    • 全クエリを均等に重み付けし、2倍と0.5倍は相殺して1倍となる
    • 絶対時間を問わないため、バッチ全体で実際に学習できたかの指標になる
  • ワークロード全体の高速化率:
    • デフォルト時間の合計を候補時間の合計で割る
    • 分析者にとっての実用性の指標
    • ただし偶然見つけた1件のプランに支配され、他が破綻していても高く出ることがある

■ 11. フロンティアモデルによる対照実験

  • 目的:
    • GPT-6 AstraやQwen 3.8 2.4Tでデフォルトを改善できなければ、4Bモデルにも期待できない
  • 結果:
    • JOBの10クエリで評価し、Astraは候補5件で幾何平均2.54倍、ワークロード2.12倍、退行ゼロだった
    • Qwen 3.8 2.4Tは候補5件で幾何平均2.26倍だった
    • Astraは候補1件では幾何平均0.85倍にとどまった
    • Astraは推論サマリーを有効にすると成績が下がったが、10クエリで1回だけの評価のため偶然のばらつきとみなした
  • 示唆:
    • 候補1件と5件の差から、エージェントは連続する候補実行の間で文脈内学習ができる
    • そのため4Bモデルに一発でプランを出させる方針ではなく、エージェント型の複数ターン方式を維持した
  • 評価手順:
    • 各候補を1回ウォームアップして1回計測し、候補予算を使い切るとfinishのみ呼べるようにした
    • 選択後に交互3ペアを実行し、中央値の比を0.1〜10倍にクリップして極端な値の影響を抑えた
  • フロンティアモデルはエージェント型のクエリ最適化を実行できる

■ 12. 未学習4Bモデルのベースライン

  • 結果:
    • 113クエリ中、有効な候補なしが81件、選択失敗が16件、タイムアウトが1件だった
    • スコア対象は15件のみで、そのうち9件は定義上1.00倍だった
    • 残る6件の新規プランのうち5件は1.02〜1.30倍、1件は0.05倍だった
  • 主な失敗:
    • PlanActionが有効なオブジェクトになっていない
    • 結合木に全リレーションがちょうど1回ずつ含まれていない
    • アクションが余計なactionキーで包まれている
    • Leading木の部分木が、クエリの結合グラフ上で実際にはつながっていない
    • ツールを誤ったタイミングで呼ぶ、または存在しないツールを呼ぶ
  • モデルはタスクが苦手なだけでなく、ハーネス自体も理解できていなかった

■ 13. SFTによるオフポリシー蒸留

  • 方針:
    • まずモデルにハーネスの「言語」を話させ、クエリ最適化の能力はその後で高める
  • オフポリシー蒸留の仕組み:
    • 生徒モデルが教師モデルの軌跡を模倣する手法で、学習データは生徒自身が生成したものではない
    • 各トークンの予測確率からトークン単位の損失を求め、平均して逆伝播で重みを更新する
    • 1つのデモが数千〜数万のトークン予測を含むため、少量のデータでも重みが素早く適応する
  • デモの入手方法:
    • ランダムに妥当な軌跡を生成するツールも試し、ハーネスは学習できた
    • ただし新規の候補プランを出しにくくなる問題が生じた
    • より大きく賢いモデルをハーネスで実行し、その軌跡を蒸留する方法を採った
  • レンダリング:
    • AstraはResponses API形式で軌跡を出力するため、Qwenが理解するトークン列へ変換する必要がある
  • 損失マスキング:
    • システムプロンプト、ユーザープロンプト、ツール結果は文脈として見せるが、損失は計算しない
  • アンロール:
    • 軌跡を(文脈、応答)のペアに分解し、文脈には以前のモデル応答も累積的に含める
    • 応答が3つある軌跡は3つの学習例になる
  • LoRA:
    • 46.6億パラメータの全更新には64GB以上のVRAMが必要で、24GBのRTX 3090では不可能
    • 元の重みを凍結し、小さな行列の積で重みを調整するLoRAを使った
    • 学習したLoRAは42.5MB、学習可能パラメータは2,120万だった

■ 14. 教師モデルの選定

  • 比較:
    • Astraは平均ターン数7.6、最終文脈は平均22,009トークンで、10クエリの所要時間は3分13秒だった
    • Qwen 3.8 2.4Tは平均ターン数12.2、最終文脈は平均45,144トークンで、所要時間は15分51秒だった
    • Qwenはターンあたりの出力トークンが1,795と、Astraの161を大きく上回った
  • Astraの欠点:
    • APIでは推論トークンが提供されず、短い推論サマリーしか得られない
    • 推論サマリーで学習すると生徒の性能が下がるという研究がある
    • その研究は、サマリーを生の推論風に合成展開するトレースインバージョンで改善できると示した
  • 文脈長の制約:
    • RTX 3090でOOMを避けるには約5万トークン以下に収める必要があった
    • Astraの軌跡はすべて収まったが、Qwenの一部は超えた
  • 判断:
    • Astraの軌跡でSFTを試し、性能が落ちればトレースインバージョン、それもだめならより大きなGPUとQwenの軌跡に切り替える

■ 15. SFTの経過と結果

  • 最初の学習:
    • CEBからAstra軌跡を120件生成し、100件を学習、20件を検証に使った
    • Prime Intellectのrenderersで変換し、prime-rlでLoRAを学習した
    • 100軌跡はアンロールとパッキングで382行になり、バッチサイズ1で1エポック学習した
    • 有効な候補は14件から48件に増え、ハーネスを学習できた
  • 追加エポック:
    • 新しい軌跡を生成するより安価なため、既存データでエポックを重ねた
    • 学習を速めるため、Lambdaで2x H100ノードを借りた
    • RTX 3090で4時間かかった1エポックが、H100では45分で済んだ
    • 2エポックで幾何平均1.08倍に改善したが、3エポックでは0.82倍に退行した
    • 2エポック目は検証損失がほぼ変わらなかったにもかかわらず改善した
    • 検証損失が平坦でも、有用な振る舞いの学習が止まったとは限らない
  • データの追加:
    • 学習データを絞り込まず能力の欠落も監査していなかったため、特に有効なLeading木の構築能力が不足していた
    • Astra軌跡を320件追加生成し、候補を1つも試さずデフォルトを選んだ6件だけを除外した
    • 追加で2エポック学習し、幾何平均1.16倍、ワークロード1.06倍、勝ち20件、退行5件に達した
  • 結果:
    • モデルはハーネスを学んだうえで、JOBクエリで実際に良い判断を下すようになった
    • 推論サマリーによる学習は性能を損なわなかった

■ 16. エージェント型強化学習の報酬設計

  • エージェント型RLの仕組み:
    • 現在の方策をハーネス内で訓練クエリごとにN回実行し、各ロールアウトの出力を検証可能な基準で採点する
    • 同一クエリのロールアウト間で相対化したアドバンテージが正なら強化し、負なら抑制する
  • 当初の報酬:
    • 有効な候補があれば、高速化率の自然対数から無効候補数×0.1と重複候補数×0.05を引く
    • 有効な候補がなければ一律で-3とする
  • 当初の報酬の問題:
    • -3の罰が厳しすぎ、モデルは無効プランを恐れてデフォルトプランを返し続けた
  • GRPOによる悪化:
    • prime-rlのGRPOは、各報酬からグループの平均報酬を引くだけでアドバンテージを求める
    • どのロールアウトもPostgresに勝てなくても、平均より良ければ正のアドバンテージを得る
    • デフォルトと同じプランを出すことが最も強く強化されてしまった
  • 改訂した報酬:
    • 高速化率を0.1〜10倍にクリップしてから対数をとる
    • ノイズを考慮して0.05でソフト閾値処理し、デフォルトとの差が5%以内なら0とする
    • keep_defaultやデフォルトでの終了は0とし、タイムアウト時は寛大な扱いの代わりに0.1を差し引く
    • 有効な候補なしは3.0に代えて0.1の罰とし、デフォルトと同じ指紋の候補には0.02の罰を科す
  • アンカー付きGRPO:
    • ロールアウトの基準値を、0と他のロールアウトの平均品質のうち大きい方とする
    • 最もましだという理由だけでは加点されない
    • 同じ4ロールアウトの例では、すべてが負のアドバンテージとなり、何も強化されない

■ 17. RL学習の実施

  • 2拠点構成:
    • LambdaのノードでPostgresを計測すると、非GPU資源を他者と共有していたと見られ、ノイズが桁違いに大きかった
    • FLOPperとLambdaノードをTailscaleで接続した
    • 推論は1台目のH100上のvLLM、計測はFLOPperのPostgres、アドバンテージ計算と重み更新は2台目のH100で行った
  • 最初の試行:
    • SFTアダプタをベースモデルに統合し、新たなアダプタでRLを開始した
    • 学習率1e-06、バッチ8、クエリあたり4ロールアウト、120更新で試したが改善しなかった
  • 積極化:
    • 学習率を1e-05に上げ、バッチを16、ロールアウトを8に増やし、600更新とした
  • 非同期化:
    • ロールアウト時間の92%がvLLM推論に費やされていた
    • 計測時だけPostgresワーカーを借りる方式に変え、4ワーカーで20ロールアウトを並行実行できるようにした
  • 学習の推移:
    • 600更新のRLを2回続けて実行し、正のアドバンテージを得るロールアウトの割合が着実に増えた
    • Postgresより速いプランを出すロールアウトの割合も同様に増えた
  • 結果:
    • 600更新で有効候補99件、幾何平均1.35倍、ワークロード1.16倍、退行ゼロに達した
    • 1,200更新で有効候補101件、幾何平均1.41倍、ワークロード1.29倍、勝ち38件に達した
  • 3ロールアウトでの評価:
    • 実運用では多数のロールアウトから最良の候補を選ぶと想定し、クエリごとに3回実行した
    • 3回の中から最良を選ぶと、最大15候補から選ぶことになり、幾何平均とワークロードの双方が1.81倍となった

■ 18. モデルが学習した内容

  • 探索行動:
    • 候補を提出した337件中295件が、先にリレーション、列統計、デフォルトプランを調べた
    • 339件中235件が5回の候補枠をすべて使った
    • job-01dでは、損失の多い逐次スキャンをビットマップにすれば速くなると推論し、実際に強制して90倍の高速化を得た
  • 好む戦略:
    • 1,347件のアクション中、スキャンヒントが1,141回、Leading木が917回、Parallelヒントが572回だった
    • 行数補正(Rows)の使用は意外にも146回にとどまった
    • ハッシュ結合よりネステッドループ、ビットマップや逐次スキャンよりインデックススキャンを好んだ
    • enable_sort=offとrandom_page_cost=1.1を頻繁に使った
  • 勝ちにつながる型:
    • Leadingによる結合順序の書き換え、結合順序を変えない単一スキャンの修正、Parallelの使用が利得の多くを生んだ

■ 19. コスト

  • 内訳:
    • 手元のマシンだけなら電気代(1日約9ドル)を除き無料で済んだ
    • 学習を速めたいがために、Lambdaの2x H100ノードを約95時間で約800ドル借りた
    • Astraの軌跡生成にOpenAI APIで約400ドルかかった
    • 合計1,200ドル
  • 蒸留なしの可能性:
    • DeepSeek-R1-Zeroが示したように、4BモデルはRLだけでもハーネスを学べた可能性がある
    • ただし、より多くのロールアウトが必要だったと考えられる

■ 20. 結論

  • 成果:
    • オフポリシー蒸留とRLにより、ハーネスを理解できなかった4Bモデルが、最良15候補の条件で幾何平均1.81倍の高速化を達成した
    • 結合の多いSQLワークロード全体で、合計レイテンシが44.7%減少した
  • 小規模モデルの価値:
    • 5T超の巨大モデルがある中でも、小規模モデルを軽視すべきではない
    • Astraからの蒸留が示すように、フロンティアモデルの価値も失われない
  • 企業への示唆:
    • 企業はデータを持っており、RL環境を構築して少額でオープンウェイトモデルを特定領域のタスク向けに学習できる
    • 学習も推論も速く安い小規模モデルが、今後この用途でますます使われる
    • 企業が自社モデルを容易に学習できる既製のソリューションも増えている

■ 21. 今後の課題

  • Baoのような構造化ヒントの網羅探索が、4Bモデルより有効かを検証する
  • オンポリシー蒸留を試し、オフポリシー蒸留と性能を比較する
  • トレースインバージョンがオフポリシー蒸留に与える影響を調べる
  • 専用EC2環境で報酬が欺かれるノイズを測り、規模拡大の可能性を探る