CompoundAISystemsOptimization
Issue Date: 2025-08-15 [Paper Note] Optimas: Optimizing Compound AI Systems with Globally Aligned Local Rewards, Shirley Wu+, arXiv'25 Summary- 複合AIシステムの最適化のために、統一フレームワークOptimasを提案。各コンポーネントにローカル報酬関数を維持し、グローバルパフォーマンスと整合性を保ちながら同時に最大化。これにより、異種構成の独立した更新が可能となり、平均11.92%の性能向上を実現。 Comment
元ポスト:
framework: https://github.com/snap-stanford/optimas
複数のコンポーネントのパイプラインによって構成されるシステムがあったときに、パイプライン全体のパフォーマンスを改善したい。このとき、パイプライン全体のパフォーマンスをユーザが定義したGlobal Reward Functionを最大化するように最適化したい。しかし、多くの場合このような異種のコンポーネントが複雑に連携したパイプラインでは、global rewardsは微分不可能なので、end-to-endで最適化することが難しい。また、個々の異種のコンポーネントのコンフィグ（e.g., textual, numerical, continuous vs. discrete）を同時に最適化することがそもそも難しい。全体のAIシステムを動作させて、global rewardを最適化するのは非常にコストがかかる。先行研究では、特定のコンポーネントを別々に最適化してきた（たとえば、promptをフィードバックに基づいて改善する Large Language Models as Optimizers, Chengrun Yang+, N/A, ICLR'24
, モデル選択をiterative searchで改善するなど）。が、個別のコンポーネントを最適化しても別のコンポーネントの最適化が不十分であれば全体の性能は向上せず、全てのコンポーネントを個別に最適化しても、相互作用が最適ではない場合はglobal rewardが最大化されない可能性がある。
このため、個々のコンポーネントにlocal reward function (LRFs)を定義する。local reward functionは、これらが改善することでglobal reward functionも改善することを保証するような形（local-global alignment properfy）で定義され、これらのlocal reward functionを異なるコンポーネントごとに異なる形で最適化しても、global reward functionが改善されるように学習する。個々のコンポーネントごとにLRFsを最適化することは、全体のシステムの実行回数を削減しながら高いglobal rewardを実現可能となる。加えて、他のコンポーネントのコンフィグが改善されたら、それらに適応してLRFsも改善されていく必要があるので lightweight adaptationと呼ばれる、システムからサンプリングされた最小のデータからLRFsをアップデートする手法も提案する、みたいな話な模様。
<img width="868" height="501" alt="Image" src="
<a href="https://github.com/user-attachments/assets/2d6a1422-c087-455a-813c-97d47da5976d"" target="_blank" rel="noopener noreferrer">https://github.com/user-attachments/assets/2d6a1422-c087-455a-813c-97d47da5976d"</a>
/>
LRFsを定義するときは、共通のLLMをバックボーンとし、個々のコンポーネントに対して別々のヘッドを用意してrewardを出力するようなモデルを定義する。コンポーネントkのinput x, output y が与えられたときに、それらをconcatしてLLMに入力し[x_k, y_k]最終的にヘッドでスカラー値に写像する。また、LRF r_kが *aligned* の定義として、LRF r_kがある共通のinputに対してr_kが高くなるようなoutputをしたときに、downstreamのコンポーネント全体のglobal reward Rが同等以上の性能を達成する場合、alignedであると定義する。このような特性を実現するために、現行のシステムのコンフィグに基づいてそれぞれのコンポーネントを実行し、trajectoryを取得。特定のコンポーネントC_kに対する二つのoutputを（異なるコンフィグに基づいて）サンプリングしてパイプライン全体のmetricを予測し、metricが高い/低いサンプルをchosen/rejectedとし preference dataを用意する。このようなデータを用いて、個々のコンポーネントのLRFsを、chosenなサンプルの場合はrejectedよりもrewardが高くなるようにペアワイズのranking lossを用いて学習する。
(ここまでが4.1節の概要。4.2, 4.3節以後は必要に応じて参照する。4.2ではどのように他コンポーネントが更新された際にLRFsを更新するか、という話と、4.3節では個々のコンポーネントがtext, trainable models, continuous configurationなどの異なるコンポーネントの場合にどのような最適化手法を適用するか、といった話が書かれているように見える。)
評価では5つの実世界のタスクを実現するための複数コンポーネントで構成されるシステムの最適化を試みているようであり、
<img width="1119" height="634" alt="Image" src="
<a href="https://github.com/user-attachments/assets/be6fcb09-a68c-4b4c-998e-9f940cad677f"" target="_blank" rel="noopener noreferrer">https://github.com/user-attachments/assets/be6fcb09-a68c-4b4c-998e-9f940cad677f"</a>
/>
提案手法によって、パイプライン全体の性能がベースラインと比べて改善しシステム全体の実行回数もベースラインと比較して少ない試行回数で済むことが示されている模様。
<img width="1112" height="582" alt="Image" src="
<a href="https://github.com/user-attachments/assets/a4f4c274-920f-4f60-9594-6b19b84f6b34"" target="_blank" rel="noopener noreferrer">https://github.com/user-attachments/assets/a4f4c274-920f-4f60-9594-6b19b84f6b34"</a>
/>