Optimas: Optimizing Compound AI Systems with Globally Aligned Local Rewards
Authors: Shirley Wu
Parth Sarthi
Shiyu Zhao
Aaron Lee
Herumb Shandilya
Adrian Mladenic Grobelnik
Nurendra Choudhary
Edward W Huang
Karthik Subbian
Linjun Zhang
Diyi Yang
James Zou
Jure Leskovec
Venue: ICLR 2026 Poster (2026)
Keywords: Compound AI System
Heterogenous Configuration
Optimization
Local Rewards
Primary area: alignment, fairness, safety, privacy, and societal considerations
TL;DR: Optimas optimizes heterogenous configurations in compound AI systems by maximizing globally aligned local rewards to improve overall system performance

Abstract: Compound AI systems integrating multiple components, such as Large Language Models, specialized tools, and traditional machine learning models, are increasingly deployed to solve complex real-world tasks. However, optimizing compound systems remains challenging due to their non-differentiable structures and diverse configuration types across components, including prompts, hyperparameters, and model parameters. To address this challenge, we propose Optimas, a unified framework for effective optimization of compound systems. The core idea of Optimas is to maintain one Local Reward Function (LRF) per component, each satisfying a local–global alignment property, i.e., each component’s local reward correlates with the global system performance. In each iteration, Optimas efficiently adapts the LRFs to maintain this property while simultaneously maximizing each component’s local reward. This approach enables independent updates of heterogeneous configurations using the designated optimization method, while ensuring that local improvements consistently lead to performance gains. We present extensive evaluations across five real-world compound systems to demonstrate that Optimas outperforms strong baselines by an average improvement of 11.92%, offering a general and effective approach for improving compound systems.

Published as a conference paper at ICLR 2026
OPTIMAS: OPTIMIZING COMPOUND AI SYSTEMS WITH GLOBALLY ALIGNED LOCAL REWARDS
Shirley Wu Parth S Shiyu Zhao"!, Aaron Lee', Herumb Shandilya', Adrian Mladenic Grobelnik', Nurendra Choudhary", Eddie Huang", Karthik Subbian", Linjun Zhang', Diyi Yang', James Zou"*! Jure Leskovec*'
*ttanford University Ammzzo "Jožef Stefan Institute *Rutgers University
https://ogg imas. stanford.edu.
ABSTRACT
Compound AI systems integrating multiple components, such as Large Language Models, specialized tools, and traditional machine learning models, are increas- ingly deployed to solve complex real-world tasks. However, optimizing compound systems remains challenging due to their non-differentiable structures and diverse configuration types across components, including prompts, hyperparameters, and model parameters. To address this challenge, we propose OPTIMAS, a unified framework for effective optimization of compound systems. The core idea of OPTIMAS is to maintain one Local Reward Function (LRF) per component, each satisfying a local-global alignment property, i.e., each component's 's local reward correlates with the global system performance. In each iteration, OPTIMAS effi- ciently adapts the LRFs to maintain this property while simultaneously maximizing each component's 's local reward. This approach enables independent updates of heterogeneous configuration using the designated optimization method, while en- suring that local improvements consistently lead to performance gains. We present extensive evaluations across five real-world compound systems to demonstrate that OPTIMAS outperforms strong baselines by average improvement of I 1.92%, offering a general and effective approach for improving compound systems.
Config.
Compound System
OPTIMAS
Local Imput Output
Local LLM
Iter.
LLM API
Parameters
Global Reward
Local Reward Functions (LRFs)
LLM API
Iter.
parameters
Hyper-
LLM API
LLM
API
Culpul
Local
Retrieuer
Iter.
Model Selection
reward
LLM API
Tool
Retrieuer
Adaptive Local-Global Allgnment
Local Optimization
Given compound AI system's 's heterogeneous configurations (e.g., prompts, parameters) across multiple components, OPTIMAS maintains globally aligned Local Reward Func- tions (LRFs) as the system evolves, where each supervises component and assigns higher local rewards outputs with higher system performance (aka. global rewards). It iteratively adapts LRFs and optimizes each component to maximize its local reward for effective system optimization.
Figure I: Overview.
INTRODUCTION
Modern AI systems increasingly employ compound systems that integrate multiple complex compo- nents, such as Large Language Models (LLMs), tool/function calls, and traditional machine learning models like retrievers (Yuksekgonul et al., 2025; Khattab et al., 2023; Du al., 2023). These components collaborate to process heterogeneous data sources and solve complex tasks through specialized subtask allocation (Zaharia al., 2024; Zhou et al., 2025; Leike al., 2018; Kandogan al., 2025; Chen et al., 2025b). While compound AI systems have yielded performance advantages
"Equal contribution; "'Euall supervision. Correspondence: (shi rwu, psarth Feca atanford.edus

---

Published as a conference paper at ICLR 2026
Component:
Local Model
LIM
Model Selector
Config.:
Paraneters
Prompss
Probability lity vectors
Hyperperaneters
User Belussor Expert
Analgat
Text tooks
Text Scorer
Alggmeet
Medical Contect
Actions
Querg
Vector Darabase
Sim
Entlyy
Scores
Options
product
Pred
Query
Ansuer
Knzuradge
Enditlin
Agpresator
Product Expert
Deciden
Matching
scones
Reterion tools
Relation Sconer
(a) Product Rec. (Amazon)
(b) Medical Analysis (PubMed)
(c) Semi-Structured Knowledge Base Retrieual (STaRK)
Probee
Solution
Query
Kaguuord
Retrieuer
Generator
Genorator
Genarator
Code Modifier
(d) General RAG (HolpotQA)
(e) Self-Verified Code Generation (BigCodeBench))
Figure 2: Five real-world and challenging compound Al systems. The goal is to automatically optimize the configuration across heterogeneous set of components and parameters, e.g,, model parameters, prompts, model selection choice, and hyperparameters. See Appendix D for details.
over monolithic models (Du et al., 2023; Huang et al., 2023; Leblond & et al., 2023; Lewis et al., 2020; Liu et al., 2023), they can be highly sensitive the failure of individual components, which leads to cascading failures in the final results (Cemri et al., 2025; tse Huang et al., 2025; Peng et al., 2025; Chen et al., 2024a). For example, if an LLM misinterprets irrelevant or misleading information. This leads to subsequent tool calls operating on incorrect inputs, producing unreliable outputs throughout the system. Therefore, optimizing these systems as a whole is crucial for maintaining reliability and global system performance (i.e., global rewards).
input query, it can retrieve
However, optimizing these compound AI systems end-to-end fundamentally challenging due to their non-differentiable nature. Second, it is hard to jointly optimize heterogeneous configurations (textual or numerical, continuous or discrete) from individual components, including prompts. hyperparameters, model selections, and even model weights. Moreover, running the entire compound Al system during optimization to achieve global reward is costly.
Previous works have largely focused optimizing specific configurations in isolation, such as optimizing prompts through textual feedback (Yuksekgonul et al., 2025; Khattab et al., 2023; Yang al., 2024; Madaan et al., 2023; Wu al., 2024a) or model selection through iterative search (Chen et al., 2025b; 2024c). Yet these approaches can fail to capture critical bottlenecks. For example, a perfectly optimized prompt can struggle to compensate for poorly chosen model. Even when the components are individually well optimized, they may still collaborate suboptimally, as the upstream component might not have visibility into which inputs effective for the downstream components. Consequently, previous methods typically require costly system runs over many configurations to identify the best configuration for the components to work well together, leading to low data efficiency due to the large number of configurations.
Present work. Here we propose OPTIMAS (Figure 1), unified framework for the effective and data-efficient optimization of compound AI systems. The core idea is to learn globally aligned Local Reward Function (LRF) per component, such that independently maximizing component's local reward still reliably improves the global rewards. We show that under mild conditions, our approach converges reliably, providing strong theoretical guarantees. Furthermore, since the learned LRFs can be used optimize components locally, OPTIMAS has higher data efficiency by avoiding extensive runs of the entire compound AI system to achieve high global reward.
Specifically, each LRF estimates the contribution of a component's ss output to the global reward. All LRFs are implemented using a shared LLM backbone, with separate projection heads for each component produce component-specine rewards. We propose lightweight adaptation mechanism using mini-batch preference data to ensure the LRFs remain aligned with the evolving system configuration (Figure 3). Leveraging this decentralized structure, OPTIMAS applies specific
2

---

Published as a conference paper at ICLR 2026
Table A comparison of OPTIMAS with selected methods. OPTIMAS optimizes compound systems with heterogeneous configurations and enables higher data efficiency with local optimization to reduce number of system runs. We prove OPTIMAS's convergence under mild conditions.
Supports compound system
Optimizes heterogeneous config.
Data efficiency
Convergence guarantee
OPRO (Yang et al.. 2024)
X
DSPy (Khattab et al., 2023)
TextGrad (Yuksekgonul al., 2025)
LLMSelector (Chen al., 2025b)
OPTIMAS
optimization method to each component based on its configuration type. For example, reinforcement learning for model parameters (Rafailov et al., 2023; Schulman et al., 2017) metric-guided search for prompts and hyperparameters (Yang et al., 2024; Opsahl-Ong et al., 2024; Liashchynskyi Liashchynskyi, 2019). Overall, OPTIMAS iteratively updates heterogeneous configurations towards a higher global reward by using each adaptive LRF as an objective. By optimizing a component to maximize its local reward, OPTIMAS reduces the entire system runs maintain higher data efficiency.
We conduct extensive experiments to evaluate OPTIMAS across five real-world compound systems (Figure 2), including challenging settings such behavior-driven product recommendation and medical analysis. OPTIMAS consistently outperforms strong baselines, achieving an average relative improvement of 1 11.92% with higher data efficiency, while baseline methods occasionally improve performance. For example, while DSPy improves the performance on the multi-hop QA system, may degrade performance other tasks, such as product recommendation. In contrast, OPTIMAS is the only method that improves performance across all five tasks, consistent with our theoretical guarantee (Section 4.4) that aligning local and global rewards enables effective optimization.
RELATED WORK
Optimizing LLM single-step generation. Prior work extensively optimizes prompts for Large Language Models (LLMs) in single-step generation to improve performance (Yang et al., 2024; Madaan et al., 2023; Guo et al., 2024; Shinn et al., 2023; Yu et al., 2024; Yin et al., 2025; Chen et al., 2025a; Williams, 1992), but these methods are limited in their ability to handle complex, multi-step tasks. For example, addressing complex queries often requires combining multiple components- such as L.LMs, tools, and machine learning predictors-oe to obtain accurate predictions.
Optimizing nulti-coppoeet/muuli-sttt generation. Compound AI systems consisting of multiple components enable more complex planning and specialized processing at each task step (Khattab et al., 2023; Yao et al., 2023; Liu et al., 2023; Du et al., 2023; Zhang et al., 2025). Previous studies typically optimize different components separately, such as optimizing LLM prompts (Khattab et al., 2023; Wu et al., 2024a; Yuksekgonul et al., 2025), fine-tuning model weights using supervised learning (Zhao et al., 2025; Chen et al., 2024d) or reinforeement learning (Lin et al., 2023; Chen al., 2025c), developing model routing (Chen al., 2025b) and layer grouping (Chen al., 2023) strategies, and selecting hyperparameters (Wang et al., 2023; Falkner et al., 2018; Pham et al., 2018; Liu et al., 2019). In contrast, OPTIMAS enables end-to-end optimization across all components.
Reward modeling for multi-step tasks. To provide more fine-grained supervision, recent works break down global rewards (e.g, answer accuracy) into more targeted signals (i.e., dense/process rewards). Representative approaches include leveraging or bootstrapping from human step-wise annotations (Lightman et al., 2024: Uesato et al., 2022: She al., 2025); hierarchical planning that assigns rewards to eтror correction steps ( (Wang et al., 2025); using Monte Carlo Tree Search to assign credit to intermediate reasoning steps (Wang al., 2024: Ma et al., 2025; Jiao et al., 2024; Chen et al., 2024b; Setlur et al., 2025) or actions (Chen et al., 2025d; Choudhury, 2025). Recently, Chen et al. (2024d) leverage Bayesian optimization to decompose global losses into local losses for optimizing model weight. Optimas differs by dynamically aligning local rewards with global rewards through preference- based adaptation. This design is applicable to both differentiable and non-differentiable configurations, without requiring fixed decomposition or extensive retraining. The local optimization avoids extensive system runs and offers higher data efficiency. Moreover, we provide theoretical analysis to prove the convergence of our framework. We highlight our key contributions in Table 1.
3

---

Published as a conference paper at ICLR 2026
3 PROBLEM FORMULATION: OPTIMIZING COMPOUND AI SYSTEMS
Compound Al system. A compound AI system is represented as a directed acyclic graph G = E), where C = [k is a set of K distinct components (task nodes) and & is the set of all possible directed edges between components. A component of the compound system can be an LLM, general machine learning model, a model selector, etc. We denote the input and output to each component C₁ as 3yg and y/k, respectively. The system input is treated as a source node Co-
The system can operate with dynamic planning: for each input instance x. the connections {xx) C E between the components can be adaptive. A directed edge (C,, (,), {xx) indicates that the output of component C; is routed as input to component Cj when processing instance r. By default, we assume the component indices follow the topological order over E, meaning that C2 is the upstream component of C if <.
Component configurations. component Ck (F&: + yk is controlled by a configuration policy Vk- The configuration space V can either be empty (indicating no optimizable configuration for the component), discrete (e.g., textual prompts or model selections), or continuous (e.g., model parameters or hyperparameters). We denote the joint configuration policy by v (Vi,..., Vx).
Forward execution. For given input x and configuration policy v, the system executes components in topological order over the edge set E(r): y/- = C& (ww: | C. pu(Ck)); va), where pa(Cs) denotes all the parents of component C₁ over E(x). For clarity, we define the overall system as f(; v) := y, where y a collection of outputs from one or more components.
Optimization objective. Given a dataset D with initial inputs and a user-defined global reward function R: X' x y + that evaluates the final system output, the optimization goal is find the configuration policy v* (x) that maximizes the expected global reward:
v = arg max ENED R(x, f(z;; v))]-
4 OPTIMAS: GLOBALLY ALIGNED LOCAL REWARDS FOR OPTIMIZATION
Challenges. Directly optimizing the objective in Eq. 1 difficult. As the configuration spaces are typically non-differentiable, gradient-based optimization cannot be used. Moreover, each policy V; may control a different configuration type, so the joint policy v can span heterogeneous spaces. Therefore, previous efforts ( (Yuksekgonul et al., 2025; Zhao et al., 2025; Chen et al., 2025b; Khattab et al., 2023) largely focus on optimizing the policy for single types of configurations, which simplifies the optimization problem; however, this also leads to suboptimal compound systems.
Key intuition. To address the challenges, our approach (Figure 3) learns Local Reward Functions (LRFs) that align with the global reward for individual components, allowing local and independent optimization on heterogeneous components using different optimization approaches.
Such local-global alignments (Section 4.1) encourage that the global reward to increase during local optimizations (Section 4.3). Moreover, as the system configurations change during optimization, the LRFs should be adapted to remain aligned. To ensure alignment, OPTIMAS employs lightweight adaptation mechanism that updates LRFss with minimal data sampled from the system, preserving consistency with the global reward (Section 4.2).
4.1 LEARNING LOCAL REWARD FUNCTIONS
Definition (Local Reward Function (LRF)). An LRF on component C₁ is defined as Fk (Ik, yk) + R. which evaluates the component's 's output yı given the provided context Zk.
Implementation. We implement all LRFs with a LLM backbone & and separate linear heads hy for component Ct. The backbone encodes the concatenated text inputs [ıg, yk] into an embedding, and the corresponding head projects this embedding to scalar reward value. Using such a multitask neural network ensures scalability with large number of components and reduces memory costs. Specifically, each LRF is modeled as:
T& (31,, yk) ha (( vse))); for all is if V is non-empty.
4

---

Published as a conference paper at ICLR 2026
Config. y
The t-th Iteration
Config. vf+l
G
Gx-1
X
Tkk
Initialize
Local Reward Function Tk
Optimize
lIv
Cs Iv
yi
Yk
Train
Ck+1
Cx+1
Mini-batch Preference data
reward
r vx)
(Xk.
Local Improu.
Global improu.
(xx,
R
Global rewards
Local rewards
Vi 2 Local reward maximization
Local-Global alignment
Figure 3: OPTIMAS optimization iteration. At each iteration, OPTIMAS updates a component Cs by first collecting mini-batch of preference data and adapting its Local Reward Function TĄ to remain aligned with the global task metric. This alignment helps ensure that optimizing the component to maximize its local reward also improves the global reward.
Property (Local-global alignment). An LRF rı is said to be aligned with the global reward R if, for every input x and for any two candidate outputs Dj Mr of Cs.
Tk (=k vit) 2 Th (TKK Vi
(3)
Estuwnstream
RR f(: v
E Eslownstream
R(x, f(; [
where V-A denotes the configurations of all downstream components (those that directly indirectly receive information originating from C₁). The expected global reward for each candidate output estimated via Monte Carlo sampling. This involves executing the downstream components with the candidate output and the outputs from the non-downstream components fixed, capturing their stochasticity, and averaging the resulting global rewards from the final system outputs.
Objective of reward functions. To make each LRF conform to the local-global alignment property, we collect (v),), preference dataset under the current system configuration v, and train each using pairwise log-sigmoid ranking loss:
(x(D(vv) E (sx, oi vE ~ D&(v) log o (ra vt rx( U3 1)
(4)
The collection of (v) follows the following steps: execute the compound system up to C and record the partial trajectory (x, (x1, y1 ..., (I y --- 1)); 2) sample two candidate outputs for Cs (e.g., via higper-temperature decoding or alternate hyperparameters); and 3) estimate their expected task metrics according the expectation terms on the right- side of Eq. 3. The output with the higher expected value is labeled as y and the other as V3 in (v(.).
4.2 ADAPTIVE LOCAL REWARD FUNCTIONS
Problem: Misaligned LRFs the evolving system. As the system configuration changes during optimization, LRFs trained under previous configuration v may become inaccurate under the updated configuration v++1, Specifically, (1) after updating Cg, the same outputs from its upstream component C, (i < k) may lead to different global reward, making r; misaligned. (2) Its downstream components C (j > k) now receive inputs generated by the updated Cx. which may fall outside the distribution seen by their LRFs. These shifts accumulate over time, degrading the local-glober alignment property (Eq. 3) that LRFs are designed to satisfy.
However, retraining all LRFs from scratch after every configuration update is expensive. To address this, we develop a lightweight adaptation strategy that incrementally refines the LRFs as the system changes, maintaining alignment without full retraining.
5

---

Published as a conference paper at ICLR 2026
Initial reward modeling. Given the initial system configuration and dataset with initial inputs, we first construct an offline preference dataset for each component and train its LRF to convergence. This offline phase establishes well-aligned LRFs that accurately reflect each component's 's contribution to the global reward.
Stage 1:
Online reward function adaptation. When any configuration changes, we sample small batch of input data and construct a mini-batch of preference data By for each component Cj using the steps described in Section 4.1. We then optimize the LRF on Cs the objective Lx following the definition in Eq. 4. To enable stable optimization and improve data efficiency, we maintain a buffer of previous generated preference data into Bk_ This adaptation helps maintain the local-global alignment property in Eq. 3.
Stage 2:
4.3 OPTIMIZATION WITH GLOBALLY ALIGNED LOCAL REWARD FUNCTIONS
Local Optimization. As each component has its own LRF, OPTIMAS flexibly applies a specialized optimization method for each component. See details in Appendix F. Concisely,
For textual prompts, we use prompt optimization algorithms such as OPRO (Yang et al., 2024) that ranks candidate prompts by their average local reward and select the best-performing one.
For components that trainable models (e.g.. an LLM), apply reinforcement learning such as Proximal Policy Optimization (PPO) (Schulman et al., 2017) using the LRF as the critic.
For discrete or low-dimensional continuous configurations, such as model selection or hyperpa- rameter tuning, we construct a probability distribution over candidate values based on their local rewards and sample from it to update the configuration.
Overall algorithm (Figure 3). Starting with the initial system configuration, OPTIMAS leverages initial reward modeling to learn a of LRFs that are well-aligned with the global reward. At each iteration optimization, OPTIMAS randomly selects component to optimize, conducts local optimization, and if the local configuration change leads to improved global reward, it updates the system and adapts the LRF using minimal amount of data. To prevent potential cascading errors, the new configuration
accepted only if improves the global reward on a small validation set. Since the optimization with LRFs is conducted locally, the number of system runs to achieve high global reward is reduced, as we later show in the experiments. The detailed algorithm is provided in Appendix A.
4.4 THEORETICAL INSIGHTS
We prove that the local-global alignment property holds for the LRFs constructed in Section 4.1.
Theorem 4.1 (Informal). Under regularity conditions, the minimizer of equation 4 satisfies the local-ge global alignment property (equation 3). In addition, maximizing тk (zk. Cx(zk:Va) vk)) over Vk and maximizing R(x, (x; v- Cs( vx)) over Vy will yield the same solution.
We defer the formal statement for this theorem appendix B. Since solving equation I generally challenging, we introduce some regularity conditions to make the convergence analysis tractable.
As the configurations are heterogeneous, where some of the coordinates discrete, and some continuous, without loss of generality, assume the first M configurations y(1) (V1,..,, V.)
continuous, and the last (K M) configurations (?) ((MM+1, ..., Vx} } discrete. We write the objective function Ez~D R[2z, f(x; v)] as I(v) I(V1, V2. vx) == l(v(1) v!2)).
Assumption 4.1. Suppose for any given configuration (2), the initial level set (v(2) ((y(1), ((2)) < 1(v9,(1), (+))) is compact set, where v0(11) is the initialization used in the algorithm for (1), In addition, for every component k and every fixed w. V(,, v has a unique maximizer.
Theorem 4.2. Under Assumption 4.1, the algorithm will converge to the component wise maximum, that the limit point v" satisfies
I(v") 2 l(VR, v=+);
for any le [K] and any Vk--

---

Published as a conference paper at ICLR 2026
Table 2: Performances of each method the compound systems. The best and second-best results in each column are highlighted. Relative improvement is computed with respect to the best baseline.
AMAZON Product Rec. (Acc.)
PUBMEDQA Medical Analysis (Acc.)
STARK-PRIMM
HotPorQA RAG (FI)
BIGCODEBENCH Verified Code Gen.
Complex Retrieval
(MRR)
(Pass Rate)
Single LLM
24.20=1143
54.13-2.73
0.00=0.00 40.73:10.44
21.58:11.24 33.80::1.11
35.47::0:4 36.67:11.35
Unoptimized
23.71
57.46:00775
REINFORCE
21.89:22.65
L.MSSllcttr
67.93++0.09
HBC
21.55+#2.07
58, 80+00888
36,95-+0.59
21.16+4-0.97
27.78++2.08
TextGrad
20.88+33.53
56,96+22.24
41.31=+1667
24.86++1.19
35.71+40.10
DSPy
18.18+0.82
60.26+00.00
41.40++0.04
44.90+-0.32
33.81+-2.75
OPTIMAS
24.24=0082
69.13:10.33
50.5410.700
50,48-1148
38.92:10.36
Rel. Improv. 14.3%
1.8%
22.1%
12.4%
9.0%
Table 3: Number of equivalent runs on the entire systems (in thousands) by TextGrad, DSPy, and
OPTIMAS in Table We control the optimization process each use comparable system runs.
AMAZON
PUBMEDQA
STARK-PRIME
HotPoTQA
BIGCODEBENCH
Average
TextGrad 0.32
0.70
0.70
2.12
0.18
0.80
DSPy 0.24
0.66
0.66
2.09
0.28
0.79
OPTIMAS
0.31
0.52
0.51
0.71
In fact, our theoretical analysis shows that by conducting local optimization, OPTIMAS essen- tially performing coordinate maximization. Therefore, existing convergence results for coordinate maximization directly apply. Also, note that the block-coordinate (round-robin) updates adopted in OPTIMAS do not guarantee global optimality in non-convex problems, but this does not constitute flaw unique to our method; rather, it reflects a standard limitation non-convex optimization broadly, and our global convergence guarantees only hold under additional structural assumptions, such as Dolyak-Lojasiewice. or Burdyka-Lojasiewiis conditions.
EXPERIMENTS
We firstly summarize the datasets, baselines, and metrics, with details provided in the appendix. Benchmarks & Evaluation (Appendix C). We evaluate OPTIMAS on five real-world tasks:
AMAZON (Jin et al., 2024): A behavior-driven recommendation task based on Amazon products, evaluated using accuracy measure if the predicted next item matches the ground-truth item.
PUbMedQA (Jin et al., 2019): A clinical classification dataset derived from PubMed abstracts (Na- tional Center for Biotechnology Information (NCBI), 2024), evaluated by accuracy, defined as the proportion of predictions that exactly match the ground-truti -ruthh labels.
STARK- (Wu et al., 2024b): A retrieval benchmark over semi-structuree biomedical corpora, evaluated using Mean Reciprocal Rank (MRR).
HotrotQa (Yang et al., 2018): 8): A multi-hop question answering dataset, evaluated using the Fl score between predicted and ground-truie truth answers.
BIG CedeB (Zhuo et al., 2024): The instruction split of BigCodeBench for self-verifying code generation, evaluated using pass rate.
Compound Systems (Figure 2, Appendix D). We design a compound system per benchmark with diverse and common patterns for agentic systems. All systems are accessible in our code repository.
Baselines (See Appendix E for details). We compare OPTIMAS against five baselines: Unoptimized, REINFORCE (Williams, 1992), LLMSelector (Chen et al., 2025b), Hierarchical Behavior Cloning (HBC) (Le et al., 2018), TextGrad (Yuksekgonul et al., 2025), DSPy (Khattab et al., 2023; Opsahl-Onu et al., 2024). Moreover, we provide a single LLM reference which prompts an LLM to complete the task directly. This reference is used for justifying our system design in the experimental setup.
7

---

Published as a conference paper at ICLR 2026
Component
contig.
Condigg
Beform
Contiy. Aftec
surmery feaduak, and recomnndaviies Eraar fixt your cloike sellecs the dosici chanclatiiics oxtined in the arrmary while traepposiing sny creenetive feedhack reerhvo
Nest decider
Bext Itan
Pick the sext iem besed on fredback..
hes ollgns with the york
becider
Gobal Reward on Velidetoon
0.24
Condidate Poolier
Profiler
8.23
Analyaer
Session Analyser
Parameters
Model parameter change
Candidate Profilnc
Prranctere
Based provided Beedbook, the lest align with
Aralyse the provieed zercec (ummary
0.38
lsalback (iem-by-lien which candidur
liecly cuszee with
Mest irem Denider
Cancistate Profier
Declider
Pronup
prefennece
chreen from wall aligne
only spusific dragth wakkneess
15
Iteration
45
60
desined chamstersstics outliced in simmary, incompooning any corstructive feedback recsived.
the faxdltack. Priciise complaneni naher then appliarr what user already siewed
(a) AMAZON system
Companent
Configg
Config.
Before
Config
After
0.50
Global Roward on Velldation 55t
Soctur
Aggregator fox
Hyparparaneter
("rolation weight':
text weight:: 0.1)
('relation weight"; 1.0, text, weight: 0.1)
0.25
Gliven question and a list ot entitie with their property information,
You are solcrd with soussing the relevence- of five cnrinies hoard
specitic question For each entity carefully znalyze its properties relation question's requrements Assign complete
0.50
Text Scores
Prompt
Please provide juatifiaiinn
where
Apgregator
based how
each score, highlighring direct matches any relevant implisit comoctions Suive clailty and prrcisioes ctsure that the sores tocuntlly reflect the eniiies! televance to the question.
wudl properics mutch
0.45
the sequirements
dewnribed in the quedioo
Intil Canfiig.
Text Scozez
-(Same "Config. After" in the list zow)
Given question and list of entisies, each with detailed
Prompt
preserty information, assign relevance score between
Teration
and to each enily based un the following criteria:
(b) STARK PRIME system
Coxponent
Config.
Config. Bedoze
Configg Aét
0.40-
Anower Generator
Keyword Extractor
Prompt
retrieve relevant
Extract keywords from the quety
kllaning eaay oflanee seace adrionl Roza de tand ingucial tere fid reproue the taa iics.
Global Reward Velidation
Gonerator
Rint
Prompt
Genenite usefial himts to answer the query.
Genemee relovst and actimaalc hints effectidly addees the give
ensuring cuaikty and cone serves is your repereas.
0.22
Gooe to clorr and relevan
prodide dilect
Answer Generator
and cnecios snover to the query nnso now than to wash. Your sopeencskeel atictly fecm dee spcctios willoot uny odidiicna wodod diberation If
hinns, dirccily answer
0.30
Hii Gemeratar
Gonezator
Prompt
quary with short answer for the query.
inoficicntto Sevea kit ansee cearty indican ii more stverotion is needed th perper respoue.
Keyword
Betriever
(number retrived passages) k-]
0.22
Inbbal Co (fff..
Lives to dea and adeeeen tara, zovice dricc
and zancie susverte eaery nor scels Your respene
Piease provide ile a clow and dinect answer to hee following query, the hims Limit
20
M0
50
Answez Genseator
Prompt
cebeeriiie. Irite hiss an wehig seg weber.o inaffi te Frreiia at arowee deaty sikate mon
walest abficcal
response words any relesaett numerccal cltss. adlitiona explatation cortest.
and include
Ieeation
proper reporse
(c) HOTPOTQA system
Figure 4: Global reward and configuration updates of the three compound AI systems over the optimization iterations. For conciseness, we only show the local optimization steps that lead to an increase in global reward on the validation sets. The annotations show the optimized components.
Table 4: Average pairwise ranking accuracy on validation sets, measuring how often the method assigns a higher score to the output with higher expected global reward.
AMAZON
PUBMEDQA
STARK-PRIME
HotPOtQA
BIGCODEBENCH
Avg.
LLM Judge
51.25%
49.54%
54,37% 76.64%
50.00% 72.40%
42,45% 90.57%
49.52% 77.96%
OPTIMAS
84.93%
65.28%
5.1 PERFORMANCE OF THE SYSTEMS DURING AND AFTER OPTIMIZATION.
Takeaway 1: OPTIMAS leads to consistent and substantial improvements (Table 2). We compare OPTIMAS with the baselines under similar numbers of system runs. Under this controlled data cost (Table 3), OPTIMAS consistently improves global rewards across all compound systems, achieving an average relative improvement of 11. 11.92% compared to the best baseline.
8

---

Published as a conference paper at ICLR 2026
Ciobal Reward on Validatton Set
Goobll Rewwrd on Valddttion
0.24
Fitted line (R 0.982)
0.36
Fitted line (R 0.8231
0.19
0.20
0.30
0:16
0.26
0.16 0.14
0.24
0.23
0.24
0.25
0.24
0.20
0.12
0.18
0.08
Local-Global Alignment Quality 196)
Local Gioall--lbballloo Alignment Quality (9)
(a) Candi date Profiles ler on AMAZON
(b) Answer Generator on HOTPOTQA
Figure 5: Local reward models with varying alignment quality are used to optimize a selected component in each task, where we observe that higher alignment quality yields higher global rewards.
10-
Before optimizing with
1.0
Confidence
Gluen some hints, directly onswer the query with short answer.
3.0
Raaard
(local reward-O.12 17, global reward-0.32)
2.6
0.4-
After:
Local
0.4
clear relevant hints, prouide direct and concise answer to the query in more 10 words
Conficince Intarval
3.0-
0.00
0.25
0.50
0.75
1.60
(local reward-0.22 global reward=O.36)
20
Local Reward
(#Tokenns)
Licall Reward Correlation
(b) Prompts for Answer Generator HotpotQA
(c) Lengthh-Loca Reward Correlation
Figure 6: An interpretability study on what is learned by LRFs. With (a) well-aligned LRF, we find that (b) the optimized prompt explicitly constrains the output length of the component. We attribute this to that (c) the LRF prefers short outputs, which is consistent with the use of F1 as global metric.
For REINFORCE, while it yields performance gain from the unoptimized baseline, OPTIMAS outperforms it by around 3% Moreover, REINFORCE needs collecting the reward signal from downstream Monte Carlo sampling, which requires more than three times the data than OPTIMAS. While the strong baseline DSPy shows notable improvements on some datasets (e.g., 21.6% gain on HOTPOTQA their performance may be inconsistent and can even degrade the system (e.g., a 14.3% drop on the AMAZON dataset). For LLMSelector, it requires 2.8k times of forwarding through the entire system, which is 3z more expensive than OPTIMAS.
Takeaway 2: Local optimization improves global rewards (Figure 4). We study how the configu- rations change in the local optimization. Within small number of iterations, OPTIMAS achieves a substantial average improvement of 41.7% over the initial global reward on the validation sets, using lightweight and data-efficient updates. Interestingly, we observe mixed updates on prompts, model parameters, and hyperparameters, which can lead to improved global reward. For example, updating the prompt of Text Scorer in the 9-th iteration improves global reward from 0.49 0.56. Among these cases, they involve optimizing different components to achieve the highest global reward empirically, showing the importance of being able to optimize different types of configurations.
5.2 WHY AND HOW OPTIMAS WORKS: ALIGNMENT, INTERPRETABILITY, AND EFFICIENCY
We conduct extensive in-depth study to understand the mechanism in OPTIMAS framework.
Takeaway 3: OPTIMAS yields high local-global alignment quality (Table 4). To measure alignment quality of LRFs, we compute pairwise ranking accuracy: the probability that an output with higher global reward receives higher score than an output with lower global reward. This reflects how well the learned LRFs aligns with global rewards. We compare against a LLM Judge, which prompts gpt-4o 4o model to score the outputs of components based on 20 in-context examples. This approach similar to prior methods such as TextGrad, which rely on few-shot reasoning over textual patterns. In Table 4, LLM Judge performs closer to random guessing, due to the diversity and stochasticity of components' outputs that make it difficult reason reliably. In contrast, substantially higher performances. Moreover, LRFs internalize the local-global alignment within their weights without relying on limited in-context examples, enabling more precise alignment.
LRFs achieve

---

Published as a conference paper at ICLR 2026
Alignment Quality (%0
70.65
Backbone size
Alignment quality (%)
69.24
67.22
67.57
65.46
1B
72.58
3B
70.04
8B
71.20
50
80
100
Training Data Size (%)
Figure 7: Impact of training data size and LRF backbone size alignment quality, measured by pairwise ranking accuracy: the percentage preference pairs where local and global rewards agree.
Takeaway Higher alignment quality usually leads to higher global reward (Figure 5). To understand how alignment quality affects global reward, we conduct controlled experiments, where we select key component, apply local reward models with varying alignment quality to optimize and measure the average global reward achieved after updating the component. In Figure 5, we show a strong positive correlation between the LRF's alignment quality and the global reward improvement.
Takeaway 5: LRFs learn interpretable directions to improve global reward (Figure 6). An important aspect of system optimization interpretability, ie., the configuration updates are reliable and understandable. We provide study in Figure 6, where the LRF Iearns to favor concise answers. In fact, it is more feasible to interpret configuration updates with LRFs. Specifically, one can perturb the component outputs in certain ways, and observe the changes in local rewards obtain insights.
Takeaway 6: OPTIMAS is data- and computationally efficient. (Figure 7 & Appendix G). We study the efficiency of OPTIMAS along two data is needed to learn effective LRFs, and (2) whether large LRF backbones are necessary.
the HOTPOTQA system: (1) how much training
We train LRF models using varying percentages of available training data (12.5% and measure alignment quality via pairwise ranking accuracy. Figure 7 (left) shows the performance degrades gracefully: using just 12.5% of the data yields 65.46% accuracy (92.7% of full-data performance), and 25% recovers 95.1 %. This indicates that OPTIMAS is data efficient, where LRF can be learned from relatively few system runs. Moreover, Figure (right) compares IB, 3B, and 8B LRF backbones and reports alignment quality. The results show that OPTIMAS is computationally efficient, where lightweight models are sufficient for learning local-g local-global aligmment.
In Appendix G, we further show that local and global reward landscapes are closely aligned across retriever top-k settings, that modest numbers prompt candidates and adaptation inputs suffice for strong performance, and provide cost comparison on the AMAZON system showing that OPTIMAS achieves the best performance while using comparable number of effective system runs.
CONCLUSION
OPTIMAS a unified framework to optimize compound Al systems with heterogeneous configura- tions. OPTIMAS' way to maintain globally aligned local reward functions allows every component, whether a fine-tunable LLM,, LLM API, tools, or model selector, to be optimized locally while improving the overall system. On five real-world tasks, OPTIMAS outperforms strong baselines, effectively optimizes components with different configurations, and exhibits high alignment quality and reliable interpretations. We believe OPTIMAS will serve general, data-efficient approach for continually optimizing practical systems. For future work, we aim to further apply OPTIMAS on even larger systems, with the goal of understanding complex reward modeling and scalability.
ACKNOWLEDGEMENT
We gratefully acknowledge the support of Amazon, NSF under Nos. CCF-1918940 (Expeditions), DMS- MMS-2327709 (IHBEM), IIS 2403318 (III; NIH under No. 1U24NS 146314-01, Stanford Data Applications Initiative, Wu Tsai Neurosciences Institute, Stanford Institute for Stuman-Centers AI, Chan Zuckerberg Initiative, Genentech, SAP, and SCBX.
10

---

Published as a conference paper at ICLR 2026
ETHICS STATEMENT
The applications of compound AI systems today include, but are not limited to, decision processes in socially and economically important domains, making their optimization an increasingly significant research problem. While our work aims make optimization of such systems robust, and interpretable, we recognize potential ethical concerns. In high-stakes settings such healthcare recommendation platforms, optimization could inadvertently amplify biases or propagate unsafe behaviors from individual components; our framework mitigates these risks by aligning local rewards with global performance, thereby reducing the likelihood of harmful emergent behaviors and creating opportunities to incorporate fairness- and safety-sensitive objectives. Another concern is malicious use of optimized compound systems; however, our approach relies on safety- aligned LLMs and tools, and we observe that aligning local and global objectives does not reduce the effectiveness of existing safeguards, while the decentralized optimization structure may even increase opportunities to detect misuse. Regarding data, our experiments rely solely on publicly available datasets or synthetic user interactions without any personally identifiable information (PII), and no private or sensitive user data were used. We believe this work contributes to the safe and responsible advancement of AI by providing both theoretical and practical insights into optimizing compound AI systems, and to promote further research in this direction release all code, models, and benchmarks described this paper.
data-efficient,
REPRODUCIBILITY STATEMENT
We provide full open access to our implementation code, including the original datasets, compound system implementations, and the OPTIMAS framework. To support reproducibility, we provide Appendix C (Dataset Details), Appendix D (System Details), Appendix E (Baseline Details), and finally, Appendix (OPTIMAS Details).
REFERENCES
Mert Cemri, Melissa Z. Pan, Shuyi Yang, Lakshya A. Agrawal, Bhavya Chopra, Rishabh Tiwari, Kurt Keutzer, Aditya Parameswaran, Dan Klein, Kannan Ramchandran, Matei Zaharia, Joseph E. Gonzalez, and Ion Stoica. Why do multi-agent I1m systems fail?, 2025.
Angelica Chen, Jason Phang, Alicia Parrish, Vishakh Padmakumar, Chen Zhao, Samuel R. Bowman, and Kyunghyun Cho. Two failures of self-consistency in the multi-step reasoning of lims, 2024a.
Changyu Chen, Zichen Liu, Chao Du Tianyu Pang, Qian Liu, Arunesh Sinha, Pradeep Varakantham, and Min Lin. Bootstrapping language models with DPO implicit rewards. International Conference on Learning Representations (ICLR), 2025a.
Guoxin Chen, Minpeng Liao, Chengxi Li, and Kai Fan. Step-level value preference optimization for mathematical reasoning. In Findings of the Association for Computational Linguisties: EMNLP 2024, 2024b.
Jiaao Chen, Aston Zhang, Xingjian Shi, Mu Li, Alex Smola, and Diyi Yang. Parameter-efffcieett fine-tuning design spaces, 2023.
Lingjiao Chen, Matei Zaharia, and James Zou. Frugalgpt: How use large language models while reducing cost and improving performance. Transactions on Machine Learning Research, 2024c.
Lingjiao Chen, Jared Quincy Davis, Boris Hanin, Peter Bailis, Matei Zaharia, James Zou, and Ion Stoica. Optimizing model selection for compound ai systems, 2025b.
Yiqun Chen, Lingyong Yan, Weiwei Sun, Xinyu Ma, Yi Zhang, Shuaiqiang Wang, Dawei Yin, Yiming Yang, and Jiaxin Mao. Improving retrieval-augmennedd generation through multi-agent reinforcement learning. CoRR, 2025c.
Zhenfang Chen, Delin Chen, Rui Sun, Wenjun Liu, and Chuang Gan. Scaling autonomous agents via automatic reward modeling and planning, 2025d.
Zhiliang Chen, Chuan Sheng Foo, and Bryan Kian Hsiang Low. Towards autoai: Optimizing a machine learning system with black-box and differentiable components. In ICMI, 2024d.
11

---

Published as a conference paper at ICLR 2026
Sanjiban Choudhury. Process reward models for Ilm agents: Practical framework and directions, 2025.
Yilun Du, Shuang Li, Antonio Torralba, Joshua B. Tenenbaum, and Igor Mordatch. Improving factuality and reasoning in language models through multiagent debate. arXiv, 2023.
Stefan Falkner, Aaron Klein, and Frank Hutter. BOHB: Robust and efficient hyperparameter opti- mization at scale. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1437-1446. PMLR, 10-15 Jul 2018.
Qingyan Guo, Rui Wang. Junliang Guo, Bei Li, Kaitao Song, Xu Tan, Guoqing Liu, Jiang Bian, and Yujiu Yang. Connecting large language models with evolutionary algorithms yields powerful prompt optimizers, 2024.
Dong Huang, Qingwen Bu, Jie M. Zhang, Michael Luck, and Heming Cui. Agentcoder: Multi- agent-based code generation with iterative testing and optimisation. CoRR, abs/2312.13010, 2023.
Fangkai Jiao, Chengwei Qin, Zhengyuan Liu, Nancy F. Chen, and Shafiq Joty. Learning planning- based reasoning with trajectory collection and process rewards synthesizing. In EMNLP. Associa- tion for Computational Linguistics, 2024.
Qiao Jin, Bhuwan Dhingra, Zhengping Liu, William Cohen, and Xinghua Lu. Pubmedqa: dataset for biomedical research question answering. In EMNLP-IJCNLP, 2019.
Yilun Jin, Zheng Li, Chenwei Zhang, Tianyu Cao, Yifan Gao, Pratik Jayarao, Mao Li, Xin Liu, Ritesh Sarkhel, Xianfeng Tang, et al. Shopping mmlu: A massive multi-task online shopping benchmark for large language models. arXiv preprint 4rXiv:2410.20745, 2024.
Eser Kandogan, Nikita Bhutani, Dan Zhang, Rafael Li Chen, Sairam Gurajada, and Estevam Hruschka. Orchestrating agents and data for enterprise: A blueprint architecture for compound ai, 2025.
Omar Khattab, Keshav Santhanam, Xiang Lisa Li, David Hall, Percy Liang, Christopher Potts, and Matei Zaharia. remmonstaae-ssarrc-preeddee Composing retrieval and language models for knowledge-intensing nip, 2023.
Hoang Minh Le, Nan Jiang, Alekh Agarwal, Miroslav Dudik, Yisong Yue, and Hal Daumé III. Hierarchical imitation and reinforcement learning. In ICML, 2018.
Rémi Leblond and et al. Alphacode 2 technical report. Technical report, DeepMind, 2023.
Jan Leike, David Krueger, Tom Everitt, Miljan Martic, Vishal Maini, and Shane Legg. Scalable agent alignment via reward modeling: research direction. CoRR, abs/181 11118111177711 2018.
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, Sebastian Riedel, and Douwe Kiela. Reureeeel-augmeetee generation for mnowledgeinntensin NLP tasks. In NeurIPS. 2020.
Petro Liashchynskyi and Pavlo Liashchynskyi. Grid search, random search, genetic algorithm: A big comparison for NAS. CORR, 0tbs//912.0605/ 2019.
Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. In ICLR, 2024.
Kunyang Lin, Yufeng Wang, Peihao Chen, Runhao Zeng, Siyuan Zhou, Mingkui Tan, and Chuang Gan. Deir: Dynamic consistency intrinsic reward for multi-agent reinforcement learning, 2023.
Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search, 2019.
Zijun Liu, Yanzhe Zhang, Peng Li, Yang Liu, and Diyi Yang. Dynamic Ilm-agent network: An llm-agent collaboration framework with agent team optimization. arXiv, 2023.
12

---

Published as a conference paper at ICLR 2026
Yiran Ma, Zui Chen, Tianqiao Liu, Mi Tian, Zhuo Liu, Zitao Liu, and Weiqi Luo. What step-level reward models rewarding? counterintuitive findings from buts-boosted mathematical reasoning, 2025.
Aman Madaan, Niket Tandon, Prakhar Gupta, Skyler Hallinan, Luyu Gao, Sarah Wiegreffe, Uri Alon, Nouha Dziri, Shrimai Prabhumoye. Yiming Yang, Shashank Gupta, Bodhisattwa Prasad Majumder, Katherine Hermann, Sean Welleck, Amir Yazdanbakhsh, and Peter Clark. Self-refine: Iterative refinement with self-feedback. In NeurlPS, 2023.
National Center for Biotechnology Information (NCBI). PubMed, 2024. Accessed: 2025-05-07.
Krista Opsahl-Ong. Michael J. Ryan, Josh Purtell, David Broman, Christopher Potts, Matei Zaharia, and Omar Khattab. Optimizing instructions and demonstrations for multi-stage language model programs. In EMNLP. Association for Computational Linguistics, 2024.
Jingyu Peng, Maolin Wang, Xiangyu Zhao, Kai Zhang, Wanyu Wang, Pengyue Jia, Qidong Liu, Ruocheng Guo, and Qi Liu. Stepwise reasoning error disruption attack of llms, 2025.
Hieu Pham, Melody Guan, Barret Zoph, Quoe Le, and Jeff Dean. Efficient neural architecture search via parameters sharing. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 4095-4104. PMLR, 10-15 Jul 2018.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D. Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly reward model. In NeurIPS, 2023.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov, Proximal policy optimization algorithms. CoRR, 0e4/1707.06347, 2017.
Amrith Setlur, Chirag Nagpal, Adam Fisch, Xinyang Geng, Jacob Eisenstein, Rishabh Agarwal, Alekh Agarwal, Jonathan Berant, and Aviral Kumar. Rewarding progress: Scaling automated process verifiers for LLM reasoning. ICLR, 2025.
Shuaijie She, Junxiao Liu, Yifeng Liu, Jiajun Chen, Xin Huang, and Shujian Huang. R-prm: Reasoning- process reward modeling, 2025.
Noah Shinn, Federico Cassano, Ashwin Gopinath, Karthik Narasimhan, and Shunyu Yao. Reflexion: language agents with verbal reinforcement learning. In NeurIPS, 2023.
Jen tse Huang, Jiaxu Zhou, Tailin Jin, Xuhui Zhou, Zixi Chen, Wenxuan Wang. Youliang Yuan, Michael R. Lyu, and Maarten Sap. On the resilience Ilm-based multi-agent collaboration with faulty agents, 2025.
Paul Tseng. Convergence of block coordinate descent method for nondifferentiable minimization. Journal of optimization theory and applications, 109:475-494, 2001.
Jonathan Uesato, Nate Kushman, Ramana Kumar, H. Francis Song, Noah Y. Siegel, Lisa Wang, Antonia Creswell, Geoffrey Irving, and Irina Higgins. Solving math word problems with process- and outcome-based feedback. CoRR, abs/2211 1. 14275, 2022.
Peiyi Wang, Lei Li, Zhihong Shao, Runxin Xu, Damai Dai, Yifei Li, Deli Chen, Yu Wu, and Zhifang Sui. Math-shepherd: Verify and reinforce Ilms step-by-step without human annotations. In ACL, pp. 9426 09369999..466 Association for Computational Linguistics, 2024.
Teng Wang, Zhangyi Jiang, Zhenqi He, Wenhan Yang, Yanan Zheng, Zeyu Li, Zifan He, Shenyang Tong, and Hailei Gong. Towards hierarchical multi-step reward models for enhanced reasoning in large language models, 2025.
Zhaozhi Wang, Kefan Su, Jian Zhang, Huizhu Jia, Qixiang Ye, Xiaodong Xie, and Zongqing Lu. Multi-agent automated machine learning. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2023, Vancouver, BC, Canada, June 17-24, 2023. IEFE, 2023.
13

---

Published as a conference paper at ICLR 2026
Ronald J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Mach. Learn., 1992.
Shirley Wu, Shiyu Zhao, Qian Huang, Kexin Huang, Michihiro Yasunaga, Kaidi Cao, Vassilis N. Ioannidis, Karthik Subbian, Jure Leskove, and James Zou, Avatar: Optimizing Ilm agents for tool usage via contrastive reasoning. NeurIPS, 2024a.
Shirley Wu, Shiyu Zhao, Michihiro Yasunaga, Kexin Huang, Kaidi Cao, Qian Huang, Vassilis N. Ioannidis, Karthik Subbian, James Zou, and Jure Leskovec. Stark: Benchmarking IIm retrieval on textual and relational knowledge bases. In NeurlPS Datasets and Benchmarks Track, 2024b.
Chengrun Yang. Xuezhi Wang, Yifeng Lu, Hanxiao Liu, Quoc V. Le, Denny Zhou, and Xinyun Chen. Large language models as optimizers. In ICLR, 2024.
Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William W Cohen, Ruslan Salakhutdinov, and Christopher D Manning. Hotpotqa: A dataset for diverse, explainable multi-hop question answering. EMNLP, 2018.
Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik R. Narasimhan, and Yuan Cao. React: Synergizing reasoning and acting in language models. In ICIR, 2023.
Xunjian Yin, Xinyi Wang. Liangming Pan, Xiaojun Wan, and William Yang Wang. Gödel agent: A self-referential agent framework for recursive self- eeff-imprveeeet,, 2025.
Yang Yu, Qiyue Yin, Junge Zhang, Pei Xu, and Kaiqi Huang. ADMN: agent-driven modular network for dynamic parameter sharing in cooperative multi-agent reinforcement learning. In Proceedings of the Thirty-Third International Joint Conference on Artificial Intelligence, IJCAI 2024, Jeju, South Korea, August 3-9, 2024. ijcai.org, 2024.
Mert Yuksekgonul, Federico Bianchi, Joseph Boen, Sheng Liu, Pan Lu, Zhi Huang, Carlos Guestrin, and James Zou. Optimizing generative ai by backpropagating language model feedback. Nature, 639:609 61699,6666660666 2025.
Matei Zaharia, Omar Khattab, Lingjiao Chen, Jared Quincy Davis, Heather Miller, Chris Potts, James Zou, Michael Carbin, Jonathan Frankle, Naveen Rao, and Ali Ghodsi. The shift from models to compound ai systems. BAIR Berkeley blog, 2024.
Jaay Zhang, Jinyu Xiang, Zhaoyang Yu, Fengwei Teng, Xionghui Chen, Jiaqi Chen, Mingchen Zhuge, Xin Cheng. Sirui Hong, Jinlin Wang, Bingnan Zheng, Bang Liu, Yuyu Luo, and Chenglin Wu. Aflow: Automating agentic workflow generation, 2025.
Wanjia Zhao, Mert Yuksekgonul, Shirley Wu, and James Zou. Sirius: Self-improving multi-agent systems via bootstrapped reasoning, 2025.
Han Zhou, Xingchen Wan, Ruoxi Sun, Hamid Palangi, Shariq Iqbal, Ivan Vulic, Anna Korhonen, and Sercan Ó. Arik. Multi-agent design: Optimizing agents with better prompts and topologies. CoRR, abs/2502.02533, 2025.
Terry Yue Zhuo, Minh Chien Vu, Jenny Chim, Han Hu, Wenhao Yu, Ratnadira Widyasari, Imam Nur Bani Yusuf, Haolan Zhan, Junda He, Indraneil Paul, et al. Bigcodebench: Benchmarking code generation with diverse function calls and complex instructions. arXiv preprint 24Xv:2406611877 2024.
14

---

Published as a conference paper at ICLR 2026
A
ALGORITHM
Algorithm 1 Component-wise optimizztion with reward-model adaptation
Input
C = {C]
Cr v
components
v (v
initial configuration policy
θ (
parameters of local reward functions
Global reward R(-). preference dataset size k, total iterations T Training dataset and validation dataset D,
1: V
2: O' -
Optimization loop do
3: for = 0.....,,
Scheduler: choose component to optimize
4: de ~ Uniform ((1...., K})
Local optimization for the chosen component Ci,
5: if C2, is an LLM with prompts then
6:
V ++1 +- PROMPTOFTMEZZTIONNNN 0t+1) 7: else if C., has trainable weights then
v.+1 +- PPOTRAIN(VÍ, 0²++)
8:
9:
else C, has a hyperparametter configuration then
10:
v 1+1 HYPERPARAMEETRRSEARRCYY 1(+1)
end if
12:
V + v with vi, replaced by v
Validation
13:
if Ex.ED. R f(z:: v)) > Ex.ED. MI,, ((xxiv)) then
14:
v++l t V f- vf+1
15:
Reward-model adaptation
16:
D. f- COLLECCTPEEEEEECEDAAAT... O',
create &s (x+,x") x) pairs
17:
@ < REWARDMOODELTRAINOO D.)
18:
else
19:
v t r 0 O
21: end if
end for
23:
return v*
B THEORETICAL ANALYSIS
Formal statement of and proof of Theorem 4.1 According to the procedure described in Sec- tion 4, the positive and negative pairs are determined by comparing the expected task metrics, We assume the estimated metrics at Ve and yr are chosen following PPyt is labeled as positive)
Ja Estownstreamm R(z. (x V (x))) vi Egtonnstraa R(x, v- (x))) y% }). where (a(
1+expl -auj is the sigmoid function with parameter a > 0. a +oo corresponds to the case where the pairs are chosen deterministically.
Theorem B.1. Under the conditions specified above, the maximizer of Eq. 4 satisfies the local-global alignment property Eq. 3. In addition, maximizing rı(xx, Ck (Iz; v&)) over uk and maximizing
R(x, v=);v=)) Ce(Ix:Va)) over Ug will yield the same solution.
Proof. We first present lemma.
Lemma B.1. Suppose (æ, y) E IR x 1- follows the distribution Pyy = 1 x) = =1(pp' (x))
for some function p _P (0, 1). Then
arg Elogal (y p(x)))) p²
15

---

Published as a conference paper at ICLR 2026
Proof. We take the derivative of the left-hand side with respect p and set it to 0:
o (y p(x)) ai a- p(x))
y)
01 p(x))
which is equivalent to We then have
Bloi(-v y p(z)) yl = 0.
01 (1(-p(=))o - or (p(x))
0,
and therefore pp
Applying this lemma, we can obtain the solution to Eq. 4 is a Edoonsstremm R(x,, f(; yk for some positive a, and therefore it satisfies the local-global alignment property Eq. 3. In the following, we then prove that maximizing rx (x)) over Uk and maximizing R(x, f(x;; V-x) (x())) over Us will yield the same solution.
Lemma B.2. Assume the gocal-glbbal alignment property Eq. 3 holds for component Ck- Les (ax) and (a:) be Iwo configuration policies that differ only in the policy for component Cs; denote the corresponding local outputs by yk and yk. respectively. If
Ex[r (2xxxx)) Øx)]
then
Ey [(E
z [(o.
Proof. Fix arbitrary input instance Because the two policies differ only at C₁, all other component configurations remain the same, so we can write
(=;v=)(F), yk) and v<)
where V- denotes downstream configurations (independent of the choice at Ct). By assumption on the expected local reward, we have r' ( ui) 2 (IrK, Uh) for almost every Applying the alignment property Eq. 3 pointwise yields
Exlowsstream [R(,, f(z;; v -x)),, ūk:) [ru] 2 Edownstream [R(2, V mx]
Taking the expectation over (law of total expectation) gives
E R(z, f(z 2
Hence increasing the expected local reward for C2 cannot decrease-and may strictly increase the expected global objective.
D
Proof of Theorem 4.2 We first show that the algorithm is essentially performing coordinate maximization on l(v). In fact, given previous configuration v", at time t, the updated configuration vi+1 only changes the configuration of a single component, say, Ck- As the change is solved by maximizing (2, (x vx)). by Theorem 4.1, this is equivalently maximizing l(Vk. v- for the k-th coordinate.
To rule out eycling, first prove that the discrete block (2) stabilizes. Consider the sequence (v)()) (2)(2)), (v2.01, v2.(2)) (1)(()) vf((2), We first show that there exists T > 0, such that for all k > 0, ++,00) T.(2)
As (2) are discrete, there are finitely many different configurations. addition, according to the assumption that the coordinate-wicg maximum is unique, each update will result in strict decrease in the loss function. Therefore, after finite number of iterations, (2) will not change.
Now when we consider all the iterations that are later than the time T, (2) is fixed, and we only need to consider the update regarding (1), this case, we apply Theorem 4.1 1.1 of (Tseng, 2001) and complete the proof.
16

---

Published as a conference paper at ICLR 2026
C DATASET DETAILS
This appendix provides the essential statistics and source information for the five compound-systee benchmarks used in our experiments (Figure 2). For each dataset we specify the train / validation / test split sizes, the task formulation as used in the compound pipeline, and the evaluation metric reported in the main paper
AMAZON (Behavior-Driven Next-Item Recommendation). The corpus derived from Amazon MMLU dataset (Jin et al., 2024). Each instance consists of a user's 's historical behaviour sequence (views, clicks, purchases) and the target "next item" to be recommended. We split into to 335 / / 99 user sequences after filtering malformed entries. Accuracy whether the predicted item number 1 matches the ground truth the evaluation metric.
PUBMEDQA (Medical Analysis based QA). PUBMEDQA (Jin et al., 2019) contains biomedical abstracts paired with yes/no/maybe answers to research questions. We keep the original "expert" split and discard ambiguous samples, resulting in 475 500 muestion-absterd system frames the task as three-way classification; exact-match accuracy is reported.
pairs. Our compound
STARK- -RIME (Semi-Structured Knowledge Base Retrieval). STARK-PRIIM -Prime origi- nates from STARK benchmark introduced by (Wu et al., 2024b). blends free-text passages with relational triples from biomedical knowledge graphs. Queries are natural-language questions; rele- vance labels are automatically propagated from the original STARK annotations. We uses the original dataset split: 495/ / 51 96 queries. Performance is measured by Hit@1 1, which is the rate of ranking the ground truth items in the predicted ranking list.
HOTPOTQA DnetrieallAAugmentee Multi-Hop QA). We adopt the HOTPOTQA (Yang et al., 2018) and keep the official train/dev/test splits: 1000, 250, and 100 questions respectively. Each example in the set contains question and its (human-annotated answer. We report answer-level F1 score.
full B1G CODEBBNCH (Self- Verified Code Generation). We use subset of the full-inssa iillrttistutioon subset of BigCodeBench (Zhuo al., 2024) due to efficiency issue. After proportionally drop the data, we obtain 500 / 25 70 coding tasks. Each sample includes a natural- natural-language specification and reference unit tests. Our metric is pass@ the proportion of generated programs that pass all tests in one try.
D COMPOUND AI SYSTEM DETAILS
Table 5 summarizes each pipeline's modules (columns: System, Module, Model, Config, and Opti- mization). In the table below, we clarify the various configuration spaces and optimization methods used across the five systems.
AMAZON (Behavior-Dricee Next-item Recommendation). Session Analyzer and Candidate Pro- filer both use the Qwen 2.5 1.5B model; we optimize their model parameters with PPO reinforcement learning (Schulman al., 2017). This helps each module better encode task-specific knowledge, i.e., user sessions and product candidates. The final Next Item Decider is GPT-4o-mini module, whose prompt we optimize.
PUR MURMEDQA (Medical Analysis based QA). Two modules (Context Model Selector and Solver Model Selector) each do discrete model selection from a list of possible L.MMs. At inference time, these selectors use reward model to pick the best LLM for each input instance. The Context
3. Model Selection (LLMs)): we search over (gpt-4o, gpt-4o-mini, w.t.::gtttttttooooooe.. 5-turbo-0125, gpt-4-turbo, claude-3-7- --w-ude-----scnee---2000 scnnet-20250211111111--
"Aggregator (STaRK): we search ielation_weiie wtt_eeigghtttttteeeeee text_weight E (0.1. 'Retriever (HotpotQA): we search (1,5, 10,10, 25..
17

---

Published as a conference paper at ICLR 2026
Table 5: Modules, models, and optimization methods. Model Selection (LLMs)) indicates a discrete choice of LLMs. Aggregator is tuned over coefficients relation weight and text_weight, Retriever has a hyperparameter k.
System
Module
Model
Config
Optimization
Session Analyzer
Qwen 2.5 1,5B
Model Params
PPO (RL)
Amazon
Candidate Profiler
Qwen 2.5 1,5B
Model Params
PPO (RI.)
Next Item Decider
GPT-4o-mini
Prompi
Prompl Opl.
Context Model Selector
Model Selection (LLMs))
Hyperparam Search
PubMed
Context Analyst
One of (gpt-4o,
Prompt
Prompi Opt.
Solver Model Selector
Model Selection (LLMs))
Hyperpuram Search
Problem Solver
One of (gpt-4o....
Prompt
Prompt Opt.
Text Scorer
Claude 3 Haiku
Prompt
Prompt Opt.
STaRK
Relation Scorer
Claude 3 Haiku
Prompt
Prompt Opt.
Aggregator
Coelficients
Hyperparam Search
Question Rewriter
GPT-4o-mini GPT-4o-mini
Prompi
Prompt Opt.
Info Extractor Retriever'
Prompl
Prompt Opl.
HotpotQA
#Retrieved passages
Hyperparam Search
Hint Generator
GPT-4o-mini
Prompt
Prompt Opt.
Answer Generator
GPT-4o-mini
Prompi
Prompt Opt.
Code Generator
Claude Haiku
Prompt
Prompt Opt.
BigCodeBench
Unit Test Generator
Claude 3 Haiku
Prompi
Prompt Opt.
Final Code Generator
Claude 3 Haiku
Prompt
Prompt Opl.
Analyst and Problem Solver modules then receive the chosen model and optimize only the prompt for improved medical QA performance.
STARK-PRIME (Semi-Structured KB Retrieval). We have two scoring modules (Text Scorer, Relation Scorer), both using Claude 3 Haiku with prompt optimization. The Aggregator merges these two scores; we tune two numeric weights, relation, seciiin_weigge and text_ text_weigiieeeee whtt,, each set in {0.1, 1.0]. We perform global hyperparameter search across the entire training set for the Aggregator module, tuning gelatia relattion_an tito_wnwwee con__wegghtan ghtand text_weight, After identifying the best fixed combination, use it in for inference.
HOTPOTQA (Retrieval-Augmented Multi-Hop QA). Four GPT-4o-mini modules (Question Rewriter, Info Extractor, Hint Generator, Answer Generator) each rely on prompt optimization to improve multi-hop reasoning. Meanwhile, the Retriever is retriever with key hyperparameter k, the number of passages to pull. We search k {1, 10, 25} which we also tune via global hyperparameter search across training instances.
CODE (Self- Verified Code Generation). All three modules (Code Generator, Unit Test Generator, Final Code Generator) are Claude Haiku LLMs; each uses prompt optimization to iteratively refine code solutions based on test outcomes. The global objective is a higher pass rate on the final code.
Overall, the table illustrates how different modules can require different types of optimization from prompt tuning (textual modifications) and model-paraan -araamee fine-tuning (PPO) to discrete model selection and hyperparameter search. By unifying these heterogeneous updates within our OPTIMAS framework, we effectively coordinate local improvements to achieve consistent global reward gains.
All experiments were run node with 8 NVIDIA A100 GPUs (80 GB memory each); depending on the complexity the compound system and hyperparameters, training and optimization typically finished in 2-8 hours.
Before on-policy optimization, we train a reward model on preference pairs ("chosen" vs. "rejected" system outputs) it can assign higher scores better outputs. Table 6 highlights the main hyperpa- rameters for this stage. We adopt LoRA for memory efficiency, and use an strly-stopping mechanism based on the evaluation loss (patience=512 steps) to reduce overfitting.
After training the reward model, we run iterative on-policy optimization for each module in the compound Al system. Table lists the key hyperparameters. For example, the train size (50) limits
18

---

Published as a conference paper at ICLR 2026
Table 6: Key hyperparameters for the local reward model training.
Parameter
Value
Base model
Llama 8B Instruct
LoRA rank
32
LoRA alpha
16
Maximum sequence length
2048 tokens
Learning rate
2e-6
Number of epochs
25
batch size
how many examples used to train each module's local configuration, while the search size (50) sets how many samples we use when searching for the best local update. When a module is selected, we collect small preference data, retrain (or adapt) the reward model as needed, then locally optimize that module's parameters or prompts to maximize its local reward.
Table 7: Key hyperparameters for on-policy optimization.
Parameter
Value / Description
Train size
50
Search size
50
Prompt candidates
3
Local optimization steps
3
Fresh input size
20
Validation size
20
E BASELINE DETAILS
We provide the details to reproduce the reported baseline results and the results OPTIMAS.
This system uses default settings for all components without any optimization. LLMSelector (Chen et al., 2025b): A lightweight policy selects the best LLM per component model routing, without updating other configurations. Only applicable on PUBMEDDA. We run LLMSelcctor (Chen et al., 2025b) with the LLMDIAGNOSER to estimate per-module performance, following their procedure. We perform two rounds of allocation updates with 100 training examples each round.
Unoptimized:
REINFORCE (Williams, 1992): A policy-gradee gradient baseline that directly updates the parameters of local LLM components to maximize the task reward. This method only applicable on A AMAOON, where we deploy two locally hosted L.MSs with trainable parameters. We optimize each component with REINFORCE using sampled trajectories from the full system and propagate the scalar reward back to the corresponding module-level policy. We use 16 rollouts per step estimate the expected system performance, with learning rate 10 for 100 training steps.
Hierarchical Behavior Cloning (HBC) (Le et al., 2018): A hicrarchical imitation learning method that optimizes components to produce outputs similar to those that lead to high global rewards. Specifically, we collect successful trajectories to approximate ground truth intermediate outputs, and then perform supervised updates on the local components to mimic/clone the ideal behavior.
We run HBC using the collected preference dataset by replacing the original reward model with the embedding similarity score. With the same input in the preference dataset, we use wexxt-embeddiing n 3mmmal to embed the module output and the preferred output in the pref- erence dataset and calculate the embedding similarity score. We further weight the similarity score using the gap of the preferred output score and the rejected output score to get the reward for HBC. TextGrad (Yuksekgonul et al., 2025): A gradient-based prompt tuming method using estimated gradients from black-box L.MSs improve prompt efficacy.
We run TextGrad using T- mini to optimize each component's 's prompt independently in separate epochs. Validation is performed every optimization steps using 20 held-out validation
19

---

Published as a conference paper at ICLR 2026
instances. A batch size of 4 is used across all components and datasets. The best pest-performing prompt, as determined by validation frequency --g-ereretted textual feedback, is selected as the final configuration.
DSPy (Khattab al., 2023; Opsahl-Ong al., 2024): A prompt optimization framework using the MIPRO algorithm that jointly refines module-level instructions and few-shot demonstrations. We conduct optimization DSPy's MIPRO (Opsahl-Ong al., 2024) prompt optimization approach. For fair comparison, we disable the few-shot example and system prompt optimization and only conducts optimization on the user instructions. We dynamically set the number of iterations for the MIPRO optimizer to match the budget of system runs in Table 3.
For TextGrad, DSPy, and OPTIMAS, consistently using the same 20 held-out validation instances on each dataset to select the best configurations.
F OPTIMAS DETAILS
Component selection. each iteration t, we randomly select component optimize.
Local Optimization Steps for Different Configurations, Given globally aligned LRF ræ. we perform local optimization on each component C2 improve its configuration Vg. Specifically, we solve:
arg max Ez, [rx( V=(r): subject to d(vs, vt) S 5, V&EV2
(5)
where vf. is the configuration before the 1-th iteration, d(:;-) is a distance function over configurations, and 8 defines trust region that bounds allowable updates. This constraint ensures that r is used within a region where it is expected to produce reliable evaluations.
In practice, explicitly setting the trust region threshold å can be difficult due to heterogeneous configuration types (e.g., continuous weights or discrete tokens). Instead, we adopt conservative number of update steps restrict the magnitude of change during each iteration.
Prompt tuning. For textual prompts, we apply prompt optimization algorithms (Yang et al., 2024; Wu al., 2024a), using Fa as the evaluation metric. We sample multiple prompts limited by max number of prompt candidates. The prompts are ranked by average reward over validation instances, and the best- best-performing prompt is selected.
Model fine-tuning, When C2 is an LLM or neural model with trainable parameters, reinforcement learning algorithms, such as Proximal Policy Optimization (PPO) (Schulman et al., 2017), using r%K the critic. The model parameters are updated for small and fixed number of steps.
apply
Model selection and hyperparameter tuning. For discrete or low-dimensional continuous configurations. such as model selection, tool routing, or scalar hyperparameters. we formulate the optimization as sampling problem parameterized by a probabilistic distribution. Since these configurations are instance-specife the expectation in Eq. Eq. 5 reduces to single input. For each input x, we evaluate set of candidate configurations using the LRF rk, and compute a probability distribution over candidates proportional to fxp (rx( C2 (2z; vx)) This distribution is then used to sample the configuration update for the current iteration.
Under a conservative update to the configuration of component Ck. the expected global reward is guaranteed maintain or improve, the local local-global alignment property in Eq. Eq. holds.
G MORE EXPERIMENT RESULTS
G.1 LOCAL VS. GLOBAL REWARD LANDSCAPES
To better understand how local objectives reflect global performance, we sweep the retriever's top-k setting on MoOTOOTQA and compare local and global rewards (Table 8). The two landscapes are closely aligned: both are unimodal and peak at nearby values (k=5 for the local reward and k=10 for the global reward). The top-3 configurations (5, 10, 15) coincide, differing only in the order of the 22 This alignment shows that the local reward provides reliable proxy for the global objective, lending empirical support to our theoretical guarantee that local optimization drives system-level gains.
20

---

Published as a conference paper at ICLR 2026
Table 8: On HotPoTQA, sweeping the retriever's top-k reveals closely aligned local and global reward landscapes.
k Local reward
2
3
10
15
25
0,4247
0.5578 0,3493
0.5695 0.3325
0,6124
0.6117
0.5949
0.5123
Global reward
0.3398
0.3598
0.3645
0.3568
0.3465
G.2 PROMPT CANDIDATES PER STEP
Table 9: On HOTPOTQA, we vary the number of candidate prompts per optimization step. We compute the global performance (F1) under each experiments.
# candidates
5
10
Final 0.2822
0.1945
0.2968
0.2405
Another factor is the number of candidate prompts considered at each optimization step. On HOT- POTQA, we vary this number from 10 and measure the resulting global performance (Table 9). The best F1 score is achieved with 7 candidates, although results with fewer candidates remain competitive. These findings indicate that exhaustive candidate pools are unnecessary, and that modest numbers already yield strong performance with lower computational cost.
G.3 NEW INPUTS FOR LRF ADAPTATION
Table 10: On HOTPOTQA, number of new inputs used to collect preference pairs for LRF adaptation vs. final global performance (FI).
new inputs
10
20
30
40
Final
0.2773
0.2822
0.2659
0.2533
Finally, we explore how many new inputs are needed when adapting the LRFs. On HOTPOTQA, we test adaptation with 10 40 new inputs (Table 10). Performance improves up about 20 inputs. after which gains plateau and even slightly decline. This suggests that effective adaptation can be achieved with relatively small amounts of new data, reinforcing the practicality of our approach in scenarios where data collection is limited.
G.4 COST ESTIMATION
We report a detailed breakdown of the cost on the AMAZON system. This the only system where components have trainable local models and thus requires PPO training. For all methods, we measure cost in terms of full system runs, i.e., one invocation of the system and assume that all components contribute equally the per-run cost. We label the three components on Amazon system as A, B.C in topological order.
For OPTIMAS, the total of 0.31k effective full system runs decomposes into three parts: (i) LRF training: we collect 60 initial preference pairs per component; each pair requires 2 full system runs, giving 60 x 120 runs; (ii) LRF adaptation: during local optimization, we update the LRFs times, each time collecting 10 new preference pairs per component, again with 2 runs per pair, for total x 10 x 2 = 100 runs; and (iii) global validation: whenever a local update is predicted to improve the global metric, we evaluate the new configuration held-out validation set. We use 20 validation inputs and observe 7 such updates, but the effective cost discounted by a factor 2/3 due to cached trajectories:
20 x 7 R 93.
The discount factor arises because, for each validation, we do not always need to downstream sample the entire system: roughly speaking, if is updated, we need system per validation input; if B
21

---

Published as a conference paper at ICLR 2026
Table 11: Cost-performance comparison on the AMAZON system. Individual cost components are shown in actual runs, with total runs measured in thousands (k). "LRF training" and "LRF adaptation" only apply to OPTIMAS, which learns trainable local reward functions. DSPy and TextGrad optimize only the prompt of the last component C, and report their cost as effective number of full system runs (see text).
Method
LRF training
LRF adaptation
Validation cost
Total runs (k)
Performance (%)
OPTIMAS
120
100
93
0.31
24.24
TextGrad
320
0.32
20.88
DSPy
240
0.24
18. 18
updated, we need 2/3 system run per validation input: updated, we need 1/3 system run per validation input
of full system run. The resulting global performance is 24.24%.
For DSPy and TextGrad, we follow their standard setup and optimize only the prompt for the final component C. while holding A and B fixed. To obtain fair cost comparison in terms of effective full system runs, we (i) pre-compute a pool of 20 validation inputs by running A and B once, and then (ii) only during optimization. Assuming A, B, and C. have comparable cost, convert the number of C-only calls into an effective number of full system runs by dividing by three. With 48 optimization steps for TextGrad and 36 for DSPy, this yields
20 48/3 320 and 20 x 36/3 240
effective full system runs, corresponding to 0.32k and 0.24k in Table 11. These are conservative lower bounds, since any additional evaluation of upstream components would only increase their cost.
Under this normalzzdd view, OPTIMAS achieves the best global performance while using a comparable (slightly lower) effective number of full system runs than TextGrad and only moderately more than DSPy. This demonstrates that local optimization with learned reward functions can deliver stronger performance without incurring prohibitive cost.
Additional PPO training cost. Beyond system-run cost, OPTIMAS incurs additional compute for training local models with PPO. We measure this cost in GPU-hours. For each local PPO update on the AMAZON system, we train for 3 epochs on single NVIDIA A100-SXM4-80C0 GPU. Averaged over runs of local optimization, each run takes approximately 12 minutes, yielding a total of about 6 GPU-hours (equivalently, 1.5 hours on 4 GPUs). This one-time training overhead is modest relative to the cost of repeated system evaluations and is only required for systems with trainable local components.
22


Reviews:
[{"review_id": "H17MpXa8j9", "reviewer": "Reviewer_FLrM", "summary": "The paper proposes an effective and efficient framework, **OPTIMAS**, for optimizing compound AI systems with heterogeneous configurations (e.g., prompts, model parameters, and hyperparameters) as a whole. The key idea is to maintain a **Local Reward Function (LRF)** for each component. Before optimizing a component, the corresponding LRF is updated using preference data so that answers preferred by the global reward are also preferred locally (achieving **local–global alignment**). The component is then optimized to maximize its LRF reward.\n\nTheoretically, the authors show that under certain regularity conditions, training an LRF to prefer outputs associated with higher global reward aligns it with the global reward, and optimizing a component with the LRF is equivalent to optimizing it using the global reward. They further show that, under additional assumptions, the optimization converges to a component-wise maximum.\n\nEmpirically, the method achieves consistent and substantial improvements across five real-world compound AI systems over five baselines (4 other methods and 1 unoptimized baseline). The authors also validate the mechanism itself: local optimization improves global reward, the trained LRFs exhibit stronger alignment with global rewards than a simple LLM judge, and higher alignments lead to higher global rewards. Additionally, they analyze interpretability and efficiency of the method.", "questions": "In Theorem 4.1, should it be the 'minimizer' of equation 4  instead of the 'maximizer'. And in Lemma B.1 should it be formed as an argmax problem instead of argmin?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The presentation is clear and well-organized.\n2. The addressed problem—optimizing a compound AI system that could have heterogeneous configurations as a whole rather than its parts independently—is timely and practically important.\n3. Theoretical analysis under some assumptions provides partial guarantees for the effectiveness of the method.\n4. Empirical evaluation is comprehensive and convincingly supports the method’s effectiveness and mechanism.", "weaknesses": "There appear to be some problems in Theorem 4.1 and Lemma B.1. Specifically, Equation 4 defines a loss with a leading minus sign, so Theorem 4.1 should refer to the **minimizer** of Equation. 4, not the **maximizer**. Similarly, in Lemma B.1 the optimization should be formulated as an **argmax** of the expected log-likelihood. The proof shall remain conceptually sound after correcting these signs.\n\nMinor typo issues:\n- Line 94: “optimized” → “optimize”\n- Line 119: add comma before “While”\n- Line 249: “Eq. equation 3”\n- Line 256: “r_j” → “r_i”\n\nOtherwise, the paper is technically solid.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kjbfHPkXGh", "reviewer": "Reviewer_WMHN", "summary": "The paper proposes OPTIMAS, a framework to optimize compound AI systems by learning component-wise Local Reward Functions (LRFs) that are trained to be locally–globally aligned with the task’s global reward.", "questions": "1. How often does local improvement (per LRF) fail to translate to global improvement in practice, and what diagnostics or safeguards detect misalignment early?\n\n2. How sensitive is performance to the LRF backbone size and adaptation frequency? Could you show accuracy vs. backbone (1B/3B/8B) and vs. adaptation interval?\n\n3. Are there failure cases where local–global alignment deteriorates over long runs? If so, what mitigation (e.g., trust regions, re-labeling frequency) is most effective?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper clearly motivates why compound AI systems are hard (non-differentiable, heterogeneous knobs) and frames a practical local–global alignment objective.\n\n2. The paper unifies optimization across prompts, hyperparameters, model selection, and model parameters within one iterative loop.\n\n3. The paper offers interpretability hooks—probing LRF preferences (e.g., brevity bias matching F1) to explain why updates help.", "weaknesses": "While OPTIMAS reduces the number of full system runs, it introduces non-trivial computational overhead from (i) training and adapting the shared 8B-backbone Local Reward Functions (via LoRA) and (ii) generating preference labels that require downstream sampling to estimate expected global rewards. In settings that enable parameter fine-tuning (e.g., PPO for some modules), total token and wall-clock costs may exceed prompt-only baselines. The paper would be strengthened by a cost breakdown (tokens, time, and dollars) that separates LRF training/adaptation, downstream sampling per preference pair, and any PPO steps, and by a cost-normalized comparison to DSPy/TextGrad.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UuRLW3ALwj", "reviewer": "Reviewer_azQy", "summary": "The paper presents a method to optimize an agentic system comprising of multiple LLMs and other components. To do so, it first elicits preferences data for each local component and artifically get the preference label that \"aligns\" with the global system output, in the sense that if response A gets better system reward, then response A is preferred over B locally.\n\nThis preference elicitation process is done many times, and each LLM in the system is progressively adjusted/fine-tuned so the system performance improves.", "questions": "See my weakness section above.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The presentation of the paper is good. I could quickly understand what their method is about.\n2. Their method makes sense from an intuitive point of view - since we cannot do end-to-end learning of the system (due to perhaps, non-differentiable components), we can slowly adjust each component so that it improves the overall system (by looking at how each local step affects the system output. Despite so, there are some flaws in this approach (see weaknesses).\n3. Experimental results are comprehensive (I do have some questions regarding some of the set up and validity).", "weaknesses": "1. The method is intuitive. However, it has several flaws. The biggest flaw is that it seems to be performing some form of local parameter updates with respect to a non-differentiable reward signal. As far as I know, this has the same flavour as popular gradient-free training such as REINFORCE. Why didn't the authors mention such approaches or use it as the baselines (I'm quite sure we can use it as a baselines).\n2. In addition, what makes this method slightly worse than prior approaches such as REINFORCE is that it optimizes each component in a round-robin/random fashion. This means the approach cannot be globally optimal because we are adjusting one component/subset of parameters at a time right? It is misleading to claim that their approach guarantees optimality in this case (we need some extra convex assumption of the optimization landscape).\n3. In addition, I feel some of the experiment set up is contrived. For instance, how would a single strong model perform in answering the task directly, instead of creating a complicated system to answer it? This seems like a simple baseline that should be discussed. In addition, it seems like the local rewards have a positive correlation with the global rewards - how does simply improving the local models directly (locally) or using stronger local models fare as a baselines?\n4. Lastly, could the authors clarify what they meant that we are reducing the number of system calls? Everytime you try to elicit a preference of the local component, we need to do a forward pass of the system to check the global rewards, is this correct? So we definitely need to accumulate many system calls eventually right?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)