AIME: AI System Optimization via Multiple LLM Evaluators

Pre-trained foundation models, such as Large Language Models (LLMs), have developed rapidly over the recent years (Achiam et al., 2023; Touvron et al., 2023). With these advancements, AI systems have grown in popularity for various tasks such as code generation (Chen et al., 2024; Gulwani, 2010), question-answering (Patel et al., 2024; Wang et al., 2024), mathematical reasoning (Trinh et al., 2024; Song et al., 2024), exploration (Dorbala et al., 2024; 2023; Ren et al., 2024), and information retrieval (Gao et al., 2023) etc. As the application complexity increases, the shift to AI systems containing multiple components such as LLM-based agents and web search (Xiong et al., 2024), will continue (Zaharia et al., 2024; Yuksekgonul et al., 2024). Thus, automatically optimizing these systems, AI system optimization (Yuksekgonul et al., 2024), becomes increasingly necessary.

An emerging paradigm is text-based optimization, also known as prompt optimization (Cheng et al., 2023; Wang et al., 2023; Zhou et al., 2022), whereby the natural language input prompt is tuned to generate an optimal output. This method requires no numerical gradient descent updates typical in optimization for machine learning models (Van Der Malsburg, 1986; Hassoun, 1995; Barto, 1992) and is thus appropriate for optimizing AI systems with fixed LLM components. Recently, there has been a growing class of iterative online methods for text-based optimization (Cheng et al., 2024; Yuksekgonul et al., 2024; Shinn et al., 2024), where a single LLM generates an evaluation based on the current output to help generate the next iteration’s prompt.

While prior art has compared the abilities of a single LLM for evaluations against those of multiple LLMs (Kocmi & Federmann, 2023; Ankner et al., 2024), in AI system optimization literature, there has been a lack of studies questioning the capabilities of using a single LLM evaluator to drive the optimization process. Recently, Yuksekgonul et al. (2024) has viewed the evaluation as a text-based analogy to the objective function for backpropagation (Hinton et al., 2006; Rumelhart et al., 1986) in deep learning optimization. The objective function is a crucial element in optimizing machine learning models (Christiano et al., 2017; Mescheder et al., 2018; Chakraborty et al., 2023; Kingma & Welling, 2014). This importance motivates us to analyze and strengthen the evaluation protocol of state-of-the-art (SoTA) AI system optimization frameworks by addressing a critical research question: What are the failure cases or tasks of utilizing only one LLM-based evaluator for text-based AI system optimization?

For this question, we empirically demonstrate the shortcomings of a single evaluator protocol in judging complex outputs like code based on multiple diverse criteria, such as correctness, readability, and runtime. We emphasize its practical limitations to give optimal evaluation while being instructed to judge based on all criteria simultaneously. Figure 1 illustrates the suboptimality in the practice of an AI system optimization framework with a single-evaluator approach to code generation. Furthermore, by assuming there exists an optimal evaluation policy that in expectation samples the true evaluation between the generated response and ground truth, we also theoretically highlight that the suboptimality gap between a single evaluator and an optimal evaluator is fixed and cannot be reduced given the same output and problem task. With this insight, we then naturally ask the following subsequent query: Can we develop a principled evaluation method for text-based optimization to handle multiple criteria? We address this question by assuming there exists an optimal evaluation policy that in expectation samples the true evaluation between the generated response and ground truth. We then theoretically prove that, under a linear additivity assumption, increasing the number of evaluators can reduce the suboptimality gap. We capitalize on this theoretical insight by proposing AIME: AI system optimization via Multiple Evaluators. AIME generates and combines via concatenation independent natural language evaluations from multiple evaluators based on different evaluation instructions. We demonstrate on code generation tasks with LeetCodeHard and HumanEval benchmarks the superior performance of AIME over a single evaluator in code error detection and the success rate of test cases.

Our main contributions are as follows:

• •

  Novel Evaluation Approach for AI system Optimization: We propose using multiple LLM-based evaluators and introduce our AIME approach for iterative AI system optimization. We concatenate independent diverse samples from multiple LLM-based evaluation policies to better critique system outputs.

• •

  Theoretical Motivation for Multiple Evaluators: We prove that through a linear additivity assumption increasing the number of evaluations can reduce the suboptimatity gap from an optimal evaluation policy while a single evaluator has a fixed gap. This theoretical result helps justify our formulation for a multiple evaluation-based protocol.

• •

  Empirical Performance Over Single-Evaluation Approach: Using popular code generation dataset, LeetCodeHard (Shinn et al., 2024) and HumanEval (Chen et al., 2021), we perform an extensive study showing the superior prowess of AIME with $6$ evaluators over single evaluation to detect errors, with AIME achieve up to $62\%$ higher error detection rate than single evaluation. We then show that AIME-based optimization achieves up to a $16\%$ higher success rate on test cases than optimization with only a single evaluator. We also reveal that the choice of the number of evaluators and the combination of criteria to utilize can affect the success rate by up to $12\%$, emphasizing the design of AIME-based optimization is non-trivial. We provide a code repository. ¹¹1Repository to code: https://github.com/Bridge00/aime

Objective Function. In this section, we now characterize mathematically text-based prompt optimization as a system of LLM-based policies. Let $\pi(\cdot|x)$ be the LLM-based AI system parameterized by fixed LLM-based policy that samples an output response $y\sim\pi(\cdot|x)$ given an input prompt $x\in\mathcal{X}$ from the set of prompts $\mathcal{X}$. We aim to sample a $y\sim\pi(\cdot|x^{*})$ by finding an input prompt $x^{*}$ corresponding to $x$ prompt such that $y$ is closer to the optimal response $y^{*}$. For code generation, $\pi_{\theta}$ would be the LLM generator; $x$ would be the input prompt; $y$ is the generated code; and the $y^{*}$ here would be a code snippet that is a readable, efficient solution to the problem. Mathematically, we can write

where $l$ is a loss function to capture the closeness of sampled response $y$ to the ground truth $y^{*}$.

Iterative text-based optimization. Given an initial prompt $x_{1}$, we perform an iterative text-based optimization method to find $x^{*}$ as follows. For each iteration $t=1$ to $T$, we start by (i) sampling $y_{t}\sim\pi_{\theta}(\cdot|x_{t})$, (ii) evaluate the response $y_{t}$ to obtain evaluation $e_{t}=l(y^{*},y_{t})$, and then finally (iii) generate the next prompt $x_{t+1}\sim\pi(\cdot|y_{t},e_{t},x_{t})$. Recent work by Yuksekgonul et al. (2024) decompose step (iii) into two separate steps and (iii.a) first generate the feedback $f_{t}\sim\pi(\cdot|y_{t},e_{t},x_{t})$, and then (iii.b) generate the next prompt $x_{t+1}\sim\pi(\cdot|y_{t},f_{t},x_{t})$. For simplicity, we use the same variable $\pi$ for all LLM-based policies because the outputs are dependent on the input variables the policy is conditioned on, so the same LLM model can be utilized. In this paper, we use the same model, GPT-4o, for all steps. However, distinct LLM models can be employed at different steps.

Challenges. In an ideal setting, if we had the access to $y^{*}$ as in supervised learning (Tiwari, 2022), then we can achieve the optimal performance with larger data. However, in practice, they are hard to obtain or simply unknown for many tasks such as code generation (Chen et al., 2024). Therefore, a direct comparison to an optimal output $y^{*}$ and the resulting calculation of $e$ in step (ii) are both infeasible. Current SoTA work instead sample an evaluation $e$ from an evaluation policy conditioned by the response output $y$ and prompt $x$ as $e\sim\pi(\cdot|x,y)$. Let us denote $\pi_{e}=\pi(\cdot|x,y)$ for notation simplicitiy. Ideally, we would like the evaluation $e$ of $y$ to be $l(y^{*},y)$. More specifically, let’s assume the existence of an optimal evaluator LLM denoted by $\pi^{*}_{e}$, sampling from which will give us samples of the true loss function $l(y^{*},y)$.

Fixed Gap in Evaluation with Single Evaluation Policy from Prior SOTA. As $\pi_{e}^{*}$ is unavailable as discussed before, current SoTA methods sample the evaluation loss from a single evaluator as $e\sim\pi_{e}$. Now, we know that in the majority of the scenarios $\pi_{e}$ will not be the true evaluator policy $\pi_{e}^{*}$. Thus $e=l(\hat{y},y)$, where $\hat{y}$ is an implicit approximation of $y^{*}$ from $\pi_{e}$. Under this scenario, we define the suboptimality gap in evaluation of prior SOTA as

where we first expand upon the sub-optimality in evaluation and then upper-bound using the total variation distance (Sriperumbudur et al., 2009). We see that the term $d_{\text{TV}}(\pi_{e}^{*}(\cdot|x,y),\pi(\cdot|x,y))$ is fixed and it cannot be improved once we have the evaluator $\pi$. This result shows the hardness of a single evaluator reaching $\pi_{e}^{*}$ due to this constant gap and it will only reduce if our current LLM evaluator is near-optimal which is not true in majority of the scenarios. Empirically, Figure 1 demonstrates a practical observation where a single evaluator lets code errors go undetected, causing a large suboptimality gap from oracle performance in code generation tasks.

Our key idea is to utilize multiple evaluations than single evaluators used in state-of-the-art. The thought that multiple evaluators would work better than one sounds intuitive but a naive introduction of multiple evaluators does not work in practice. We theoretically prove the merit of multiple evaluators and then discuss how to introduce them into the pipeline described in Section 2.

We test the merits of our AIME approach via the code generation task because of its practicalness and its multiple plausible criteria (e.g., correctness, efficiency). Here, the AI system is an LLM generator that is given a code prompt and must produce a code snippet that passes the unit tests for that prompt. This code generation task is a form of instance optimization (Yuksekgonul et al., 2024), whereby the optimization variable, the input prompt, is defined as $x_{t+1}:=(y_{t},f_{t})$. $y_{0}$, $f_{0}$ are empty strings. We provide empirical results showing that AIME is superior to the single-evaluation (Single-Eval) approach in detecting code errors and that AIME-based optimization achieves higher success in test cases than Single-Eval-based optimization. Experiments were run on an Apple M1 Pro and macOS 14.5.

AIME and Single-Eval Implementation Details: We use TextGrad from Yuksekgonul et al. (2024) to implement AIME and Single-Eval. We chose TextGrad because it separates the evaluation and feedback into two separate LLM calls, making it better to analyze the evaluation module in isolation. In TextGrad, the system prompt that generates the initial code, $p_{\text{init}}$, is different from the system prompt that updates the code in the following refinement iterations $p_{\text{update}}$. At $t=0$, $p_{\text{init}}$ specifies to the LLM that it is a code generator while the $p_{\text{update}}$ from $1\leq t\leq T$ specifies that it generates a new version $y_{t+1}$ given the current code $y_{t}$ and the feedback $f_{t}$. The transition from $p_{\text{init}}$ to $p_{\text{update}}$ is explicitly programmed and not caused by the optimization process. Because the scope of this paper lies within the evaluation protocol, our AI system is a single LLM generator. ²²2We repeat the link to the repository: https://github.com/Bridge00/aime

LLM Setup Details: We use GPT-4o for all LLM calls and run $10$ iterations of optimization for each coding problem. Across all trials for both methods, we use the same initial generated code for a given problem so both evaluation protocols can judge the same code in the initial iteration. For Single-Eval, the solitary LLM evaluator call is allowed $3600$ max output tokens. For our AIME approach, each of the $K$ evaluators is allowed $\frac{3600}{K}$ max output tokens. This decision is to model a uniform distribution of weights ${\alpha}$. Note that when $k=1$, Single-Eval and AIME are equivalent. We share the evaluation system prompt for both methods in Appendix A.1. We ablate on the temperature of the evaluation LLM. All other LLM calls in the Textgrad pipeline are given $2000$ max output tokens with call temperature set to $0$ similar to Yuksekgonul et al. (2024). For all experiments, the $\text{top\_p}=0.99$.

Roles for Evaluating Code: The set of evaluation roles $R$ we used for this task are as follows: syntax errors, logic errors, correctness, readability, runtime, and code redundancy. The following results are based on utilizing all these roles. We chose three roles that correlate to maximizing the number of passed test cases: correctness, logic, and syntax. We specifically chose these three to incorporate an overall correctness role with two more specific roles. We will see in Section 4.1 that having overlapping roles can help with the robustness of evaluation in terms of error detection. The three other roles (readability, runtime, redundancy), correlate to criteria such as clarity and efficiency. We will later see in Section 4.3.2 that utilizing only these roles for evaluation decreases the overall performance of the code generation task.

Datasets: We use the following two datasets, LeetCodeHard (Shinn et al., 2024) and HumanEval (Chen et al., 2021), where each dataset contains a set of coding problem prompts and multiple unit tests for each problem to evaluate the generated code. We use the entire LeetCodeHard dataset of $39$ problems with an average of $2.2$ unit tests per problem and the first $20$ problems of HumanEval with an average of $4.4$ unit tests per problem. We withhold giving any of the evaluators of either method any information on unit tests to simulate the scenario where unit tests may be unavailable to help judge (Chen et al., 2024).

AI System Optimization: Many prior works have studied the optimization of complex AI systems. Madaan et al. (2024) was one of the first works to propose a text-based iterative feedback loop for refining LLMs, and Pryzant et al. (2023) established text-based gradients, or Textual Gradients, as feedback to an AI system. DSPy (Khattab et al., 2024; 2022; Singhvi et al., 2023), Trace Cheng et al. (2024), and TextGrad (Yuksekgonul et al., 2024) have formulated LLM and AI-based systems as a network of multiple layers and provided methods to optimize these system analogous to backpropagation and autodifferentiation. Chakraborty et al. (2024a); Ding et al. (2024) used a bi-level optimization formulation to align AI agents and systems. Text-based reinforcement learning has also been used to improve LLM-based systems (Shinn et al., 2024). Decoding and RLHF is an alternative method to optimize or align an LLM with gradient descent (Chakraborty et al., 2024b; Mudgal et al., 2023; Chakraborty et al., 2024c). While these works have shown tremendous results, there has been a gap in the literature we aim to address analyzing the effect of using multiple independent evaluations to optimize the AI system for a complex task, code generation (Chen et al., 2024; Zeng et al., 2024; Zhang et al., 2023; Jha et al., 2010; Shinn et al., 2024; Yuksekgonul et al., 2024; Zan et al., 2022; Jiang et al., 2024; Chen et al., 2021; Gulwani, 2010).

LLM-based Evaluation: LLM-based evaluation, or LLM-as-a-Judge (Zheng et al., 2023), has been growing in interest due to the ability of LLMs to evaluate large outputs like text (Sellam et al., 2020; Kocmi & Federmann, 2023) quickly and to align with human preferences. Verga et al. (2024) showed a panel of smaller LLM judges can provide numeric scores correlating to human judgment than a single larger LLM model can. Prior work has also studied finetuning LLMs to be judges (Zhou et al., 2022). Ankner et al. (2024) used LLM-generated critiques to augment the scalar reward from a reward model. Li et al. (2023) used discussion between multiple LLMs to select a strong LLM-based evaluator for question-answering. Strong LLM judges have been shown to generalize across tasks (Huang et al., 2024). Weak LLM evaluators have been used to judge the debate between two stronger LLMs (Kenton et al., 2024). We are the first to use multiple LLM-based evaluators for iterative AI system optimization.

In this work, we tackle AI system optimization by introducing AIME. AIME utilizes multiple LLM-based evaluators to provide natural language evaluation for the current system output, improving on prior methods that only use a single evaluator. Our key insight is to condition each evaluator with a specific role rather than giving all the roles to a single evaluator. We prove that increasing the number of evaluations reduces the suboptimality evaluation gap, and empirically demonstrate that AIME outperforms Single-Eval in code generation tasks, analyzing success, completion, and error detection rates. Furthermore, we study AIME’s robustness to the adversarial evaluator that generate incorrect evaluations. We also provide ablations such as on the diversity of roles, role combinations, and evaluation temperature, consistently demonstrating AIME’s superior performance and the need for multiple evaluators.

Limitations and Further Work. We only empirically study our approach in code generation. Further work could extend this evaluation approach to other tasks that require multiple criteria like molecule optimization or text generation. In terms of system complexity, we only study multiple evaluators for AI systems comprising a single LLM-based agent, and using a compound system with multiple elements such as a web search agent (Agentic AI system) could be interesting. Another aspect of the work that can be explored further is weighting the different LLM-based evaluations. We gave uniform weighting to all evaluations by giving them the same max output tokens and concatenating them. Future research could investigate methods of weighting and aggregation, possibly using another LLM to summarize or perform best-of-N on the evaluations.