AdaExplore：失敗駆動の適応と多様性を保つ探索による効率的なカーネル生成

Abstract

最近の大規模言語モデル（LLM）エージェントは、テスト時適応において実行フィードバックを用いる可能性を示している。しかし、堅牢な自己改善はいまだ解決されていない。多くのアプローチは依然として各問題インスタンスを独立に扱い、再利用可能な知識を蓄積しない。この制約は、Triton のようなドメイン特化言語で特に顕著であり、Triton は LLM の事前学習データにおいて表現が乏しい。さらに、Triton の厳格な制約と非線形の最適化状況は、単純な生成や局所的な改良を不確実にする。そこで本研究では、AdaExplore というエージェントフレームワークを提案する。これは、性能が重要なカーネルコード生成のために、蓄積された実行フィードバックを通じて自己改善を可能にするもので、2 つの補完的な段階――失敗駆動の適応と、多様性を保つ探索――によって実現する。これにより、追加の微調整や外部知識なしで、正確性と最適化性能の両方を同時に向上させる。適応段階では、エージェントがタスクを合成し、反復して生じる失敗を再利用可能な妥当性ルールのメモリへと変換することで、その後の生成が実現可能な領域（フェーズブル集合）の範囲内に留まるよう支援する。探索段階では、エージェントが候補となるカーネルを木構造として整理し、小さな局所的改良と、より大きな構造的再生成を交互に行う。これにより、局所最適を超えて最適化状況を探索できるようになる。カーネル実行時間最適化ベンチマークでの実験により、この効果が検証される。AdaExplore は KernelBench Level-2 および Level-3 で、それぞれ 100 ステップ以内に 3.12x と 1.72x の速度向上を達成し、さらに計算を追加することで改善を続ける。