KAISTの研究チームが、コーディングエージェントの速度を押し上げるフレームワークAgSpecを発表した。やり方は明快だ。「投機的デコーディング」という既存の加速手法を、コーディングエージェントの動作構造に合わせて削り直した。手を入れたのはたった二箇所。コピーして使うテキストをどこに蓄えておくか、そして一度に何文字を先に書いてみるか。検索エンジンを新しく作るのではなく、すでに動いているエンジンの上に乗せる方式なので、既存のパイプラインにそのまま組み込める。
まず範囲をはっきりさせよう。ここで速くなるのは「モデルが応答を書き下ろす区間」だけだ。テストを走らせツールを呼び出す時間まで含めた作業全体の時間ではなく、そのうちモデルが文字を生み出す時間である。しかもこの加速は、出力を一文字も変えずに成し遂げられる。より速く、同じ答えを返すだけだ。
伸び幅はこうだ。従来もっとも速かった手法より平均18.0%速く、トークン(単語の断片)を一文字ずつ順番に作る基本方式と比べると最大4.76倍だった。ピンと来なければ時間に置き換えてみよう。応答生成に2分かかっていた作業が25秒ほどで終わる計算になる — 最大倍率を時間に換算した一例だ。
ところで、この論文で本当に目を引くのは倍率の大きさではない。同じAgSpecをどのエンジンの上に乗せるかで効果が劇的に分かれたこと、しかも元々さほど賢くないエンジンほど大きく跳ねたという点だ。この話は結果を見ながら続けよう。
大規模言語モデルはトークンを一度に一つずつ順番に作る。トークンを一つ取り出すたびに、モデル全体を一回ずつ回さなければならない。だから答えが長いほど、この逐次ステップが積み重なって遅くなる。
投機的デコーディングは、ここに分業を差し込む。安価な「下書き係」が次に来るトークンを複数あらかじめ当ててみて、本モデルがそれを一括で検証し、当たったものだけを確定する。当てたトークン一つにつき、本モデルが別途回すべき逐次ステップを一つ節約する計算だ。検証は、本モデルが直接取り出したときの確率分布をそのまま保つように設計されており、通過したトークンだけが残る。だからこの論文のように毎回確率がもっとも高いトークンを選ぶ方式(greedy)で回すと、最終出力は本モデルが一トークンずつ作ったものと完全に同じになる。品質の損失がゼロというのは、この条件で成り立つ話だ。
AgSpecが使うのは、その中でも「検索ベース」の投機的デコーディングだ。下書きを別のモデルからではなく、すでに手元にあるテキストから引っ張ってくる。いまのプロンプト、先に生成した応答、あらかじめ貯めておいた保管庫から、いまの文脈に続きそうな箇所を探してそのままコピーする。コードではこのコピーが特によく効く。識別子やコード断片、リポジトリの慣習が繰り返され、エージェントのセッションでは同じサブタスクと修正ループが長く定型的なシーケンスを延々と吐き出すからだ。
問題は、コーディングエージェントの形がチャットボットとは違うところにある。チャットボットは質問一つに答え一つで終わる。コーディングエージェントは要求一つを複数のターンにわたって処理する。一ターンとは、モデルを一回呼び、ツールを一回回すサイクルだ。ファイルを直すコーダー、テストを走らせて結果を確認する実行役のように役割を分け、モデル呼び出しとツール実行を交互に行う。一度で成功することはまれだ。失敗すれば次のターンで前回の試みの大部分を書き直しながら直していく。まさにこの「複数ターンにまたがる再利用」が、従来の加速方式が見落としていた箇所だ。保管庫を何で満たすか、下書きを何文字出すかを決める規則がすべて「要求一つ」の単位にとどまっていたため、ターンとターンをまたぐ再利用を取り込めなかった。
下書きが実際に役立つには、二つが噛み合わなければならない。一つ、コピーしてくるテキストが保管庫になければならない。二つ、その形がいま作ろうとしている成果物と合致しなければならない。従来の方式はこの両方で滑っていた。
材料にたとえてみよう。使えるテキストは三箇所にある。いま食卓に並べた材料(今回のセッションでやり取りしたプロンプトと応答)、冷蔵庫から出してまな板に載せた材料(今回開いたリポジトリのファイル)、そしていつも棚にある調味料(あらかじめ貯めた大容量の参照テキスト)だ。従来の方式は食卓の上の材料だけを見るか、そもそも他人の家の倉庫にあたる調味料だけを漁っていた。肝心の冷蔵庫の中の材料 — 前のターンで自分が作った出力と、たった今開いたファイル — は取り出して使えなかった。ファイルがプロンプトの長さ制限で切り落とされた瞬間に手が届かなくなり、再試行は前回の試みの呼び出しが終わってしまうとその中身を引っ張ってこられなかった。
形も問題だった。保管庫はエージェントが「読んだそのまま」のテキストを収めるのに、コーダーが出すのは「どこをどう直すか記した修正指示書(パッチ)」だ。同じコードでも、読むときと直すときで形式が違う。だから読んだ形のまま保存しておいたテキストは、肝心のコーダーが書く行とずれてマッチングが壊れた。
AgSpecの判断はこうだ。問題は「テキストをどうマッチングするか」というアルゴリズムではなく、そのマッチングを取り巻く方針 — 保管庫に何を収め、何文字を下書きとして出すか — にある。だからエンジンには触れず、その上に方針を二組乗せる。
一つ目の手入れは、保管庫を三つに分けたことだ。先ほどの材料のたとえをそのまま持ってくると、こう対応する。
前の二つはセッションとともに生まれて消える「ライブ」保管庫で、グローバルだけがセッションをまたいで固定だ。ようやく冷蔵庫を開けたこと — 前のターンの自分の出力と、たった今開いたファイル — が検索対象に入ったわけだ。
ここに形の問題を解く仕掛けがもう一つ付く。ファイルを初めて開くとき、読んだ原本だけを収めるのではなく、「エージェントが直すときの形」にあらかじめ変形したコピーも一緒に入れておく。よくある修正フォーマットでは、そのまま残す行の前に空白を、消す行の前にマイナスを付ける。そこでAgSpecは、すべての行の前に空白を付けたコピーと、マイナスを付けたコピーをあらかじめ作っておく。コーダーがパッチにどの行をコピーして入れても、二つのうちどちらかと合う。読んだ形と直す形がずれてマッチングが壊れていた問題を、変形版をあらかじめ敷いておくことで防いだわけだ。
二つ目の手入れは下書きの長さだ。下書きトークンは受容されればステップを節約するが、拒否されれば検証だけ無駄になる。この無駄は、要求が一度に押し寄せて検証が混み合うほど大きくなる。そこで「何文字まで差し出すか」を二つの軸で決める。
一つは、あらかじめ定めておいた上限だ。下書きは後ろに行くほど当たる確率が下がる。一文字を追加で検証するコストを、その文字が埋め合わせてくれる地点までだけ差し出すのが得だ。その限界線を、過去の作業記録を再生してあらかじめ求めておく。役割ごとに有用な長さが違うので、この上限はエージェントの役割別に別々に取る。
もう一つは、随時調整する倍率だ。同じ役割でも有用な長さはターンごとに、ターンの中でも揺れる。そこで直近数回の試みで実際に受容された長さを見て、次の下書きをもっと長く出すか短く出すかをリアルタイムで引いたり縮めたりする。あらかじめ測った上限が天井を押さえ、この倍率がその中で現場に合わせて動く構造だ。
さて、先ほど後回しにした話だ。AgSpecはエンジンの上に乗せるフレームワークなので、同じエンジンを「AgSpecあり/なし」できれいに比較できる。ここで現れた構図が、この論文の白眉だ。
すでに過去の応答を再利用し、マッチング長に合わせて下書きを伸ばすことができるエンジン(SuffixDecoding)の上では、AgSpecを乗せてもスループットが16.5%伸びるにとどまった。スループットとは1秒間に確定するトークン数、つまりどれだけ速く書き下ろせるかだ。一方、その二つの仕掛けがそもそもないエンジン(SAM-Decoding)の上ではスループットが64.6%跳ね、一度の投機で稼ぐトークン数はほぼ二倍になった。
メッセージははっきりしている。エンジンが元々さほど賢くないほど、AgSpecが引き出す潜在力は大きかった。賢いエンジンはすでに似たことを自前でやっているので、さらに乗せる余地が少なかったのだ。
劇的なのは、順位がひっくり返ったという事実だ。SAM-Decodingは単独では12個の設定すべてでSuffixDecodingに後れを取っていた。ところがAgSpecを乗せると、そのうち9個の設定でSuffixDecodingを追い抜いた。元々より遅かったエンジンが、方針を二組入れ替えただけでより速かったエンジンを逆転したのだ。ある事例では速度向上が2.72倍から3.61倍に上がった。
リポジトリがそもそもなく、冷蔵庫(ワークスペース保管庫)が空っぽになる環境でも効果は残った。その条件でも、より速いAgSpecの変種は、もっとも速い従来方式を35〜145%上回った。食卓(セッション)に積み上がるテキストから検索を続け、長さ制御が検証フィードバックで休みなく適応したおかげだ。
評価モデル三つ、リポジトリ単位のベンチマーク二つ(複数のエージェントが役割を分けて複数ターンを回す状況)と、汎化用のベンチマーク二つ(リポジトリのない環境、一つのエージェントが全役割を兼ねる環境)を用いた。測定はプロンプトを読み込む時間を除いた、純粋な応答生成区間のスループットと受容長だ。出力品質は、この評価が用いたgreedy条件では基本方式と同じ出力を出すため、測定の対象ではない。比較群は基本方式、検索ベース五種、学習済み下書き方式一つ。漏れを防ぐため、グローバル保管庫は評価と重ならない別のタスクログで作り、評価タスクは保管庫の構成から除外した。
速くしてくれるが、タダではない。いくつかの前提が敷かれている。
送り手の形にあらかじめ変形しておく仕掛けは、エージェントがコードを読んで直すように包む作業の枠組み(スキャフォールド)がどんな修正フォーマットを使うかに依存する。フォーマットがセッションの前に分かっていなければ変形版を作れないし、使う枠組みが変わればその変形も手直しが要る。役割別の上限も過去の作業記録を再生して取るので、その記録を集められる環境を前提とする。
効果の幅が設定によって大きく揺れる点も念頭に置くべきだ。要求が一度に大きく押し寄せる状況では検証そのものが混み合い、利得が1.08倍まで薄くなる設定もあった。大きなバッチでは期待ほど出ないことがあるという意味だ。ライブ保管庫はセッションが長くなるほど大きくなるが、きわめて長いセッションでそのメモリと検索の負担がどう変わるかは、さらに突き詰めるべき点だ。
最後は限界というより、性格の再確認だ。投機的デコーディングが守るのは本モデルの出力確率分布だ。この論文のようにgreedyで回せば出力はそのままだが、ランダムサンプリングでは同じ分布に従うだけで、個々の生成文まで同一である保証ではない。どちらにせよAgSpecはより速く答えを返すだけで、より良い答えを返すわけではない。エージェントの成功率そのものを上げたいなら、別のレバーが要る。
この論文の実用的な教訓は単純だ。速度を押し上げる出発点は「より賢いエンジンに乗り換えること」ではなく「自分のパイプラインが何を繰り返すかを測ること」だ、ということ。何を保管庫に食わせ、何文字を受け取るかを作業構造に合わせて変えるだけでも、エンジンをそのままにしたまま相当な速度が残っていた。
この点は、ONTLEのように複数の顧客プロジェクトを同時に運用する立場では特に重い。一場面を追ってみよう。顧客が「このリポジトリのバグを直してほしい」と要求を入れる。コーダーがファイルを読んでパッチを作り、実行役がテストを走らせる。失敗すれば次のターンで前回のパッチをほぼ書き直す。顧客の画面にはこの区間が「進行中」と表示され、ターンが積み重なるほど顧客は待つ。まさにこの繰り返しこそ、検索ベースの加速がもっともよく効く場所だ — 前のターンの自分の出力と、たった今開いたファイルを保管庫が覆い、コーダーの修正フォーマットに合わせた変形版が敷かれていれば。
問題は、その手ざわりがプロジェクトごとに違うという点だ。どの役割がどんなフォーマットで何を繰り返すかがまちまちだ。だから一度は測らなければならない。自分たちのパイプラインでどのターンが何を書き直すかを測り、そのパターンを役割別の上限に反映しておいて初めて、この加速は本領を発揮する。顧客の「進行中」の時間を減らす仕事は、逆説的にもエンジンを替えるところからは始まらない。
いいえ。AgSpecが在庫を減らすのは、モデルが応答を書き下ろす区間だけです。初期プロンプトを読み込む区間は、測定からそもそも除いて見ています。長い入力を一度に消化する段階がボトルネックなら、それはこのフレームワークが狙う箇所ではありません。
本モデルに小さな予測部を付け足して学習させる方式は、下書きを「作り」出します。AgSpecは逆に、すでに持っているテキストから「写して」きます。だから追加の学習も、下書き専用のGPUも要りません。その代わり、写してくる材料が保管庫に実際に積み上がり、形が合っていなければ利得は出ません。繰り返しの多いコーディングエージェントが、この条件によく合います。
あらかじめ測った上限は天井の役割なので比較的安定していますが、リアルタイムの倍率がその中で現場に合わせて動きます。ただし役割構成や修正フォーマット、繰り返しのパターンが変われば天井そのものがずれます。パイプラインに手を入れた後は、上限を測り直すほうが安全です。
特定のモデル専用の技法ではありません。異なる複数のモデルと、異なる複数の検索エンジンの上ですべて実装し、効果を確認しました。ただし幅はモデル・バッチ・ワークロードによって揺れるので、自分の設定で小規模にまず測ってから判断するのが正解です。