Improving Best Of N With Budget Aware Execution For Swe Agents

Improving Best-of-N with Budget-Aware Execution for SWE Agents | AI21

Skip to Main Menu

Skip to Main Content

Skip to Footer

Back to Blog

-->

Back to Blog

In brief

As part of our broader focus on agent optimization, we’ve released a study demonstrating the efficacy of both horizontal and vertical scaling strategies for SWE agents. However, in our experiments, we have traditionally applied a uniform compute budget across problems, as if every task shares the same difficulty.

But outside of a controlled lab setting, that assumption doesn’t hold up.

In reality, we know that task difficulty follows a heavy-tailed distribution: a meaningful fraction of tasks are solved on low effort, while others require a large budget. For example, in a previous post detailing our efforts reaching a state-of-the-art result on SWE-rebench , we found that roughly 50% of tasks sampled from that benchmark are resolved by a single rollout of GPT-5.2, yet our original policy was spending five rollouts on every single one of them. That’s money being left on the table.

At a moment when so many companies are examining their token consumption and looking for ways to eliminate cost, we present our latest research: budget-aware execution strategies that can be adapted according to task difficulty, so you only spend more when you need to. Continuing our work on SWE-rebench, we show how cascading and parallel execution strategies, both paired with early stopping, let you optimize for cost or speed, while keeping quality still; cascading saves you the cost of the last rollout, while parallel execution spares you the wait time for the last rollout, as visualized in the illustrative diagram below. In this post, we cover how we designed these execution strategies and the results we got back.

Overview: Our existing SWE multi-agent architecture

Before digging into those execution strategies, let’s first recap the agent architecture we’re working with (see diagram below). A full overview can be found in our previous piece on how choosing the right SWE agent execution strategy helped us reach state-of-the-art on SWE-rebench , yet here are the basics:

Parallel generation: Per query, our agent generates N candidates patches (rollouts) in parallel, with the final step of each rollout including a self-confidence score. At the same time, a test agent writes and collects tests that a correct patch should satisfy.

Filter: Any candidates that fail the collected tests are eliminated.

Extract: Repository source code relevant to the remaining candidates is extracted.

Reduce: Agent selects a single, final candidate from the filtered and context-enriched pool.

Since rollout generation is both the dominant cost and parallelizable, it is the primary lever we experiment with in this work by controlling the number of rollouts and the order in which they run.

Adapting best-of-N execution strategies for SWE agent tasks

Working on benchmarks such as BrowseComp-Plus and DeepResearch-Bench, we’ve already explored different execution strategies that utilize the agent’s self-confidence to decide how much compute we invest in each task, including:

Cascading with early stopping (cost-optimized): Start with cheap models and only escalate to more expensive ones if they fail or express low confidence; the cost savings from cheaply handling easy queries offsets the cost of the expensive model on the challenging long tail.

Parallel execution with early stopping (latency-optimized): Fire all rollouts at once; once an acceptable candidate is identified, any rollouts still in progress can be terminated immediately. This avoids paying for completions that will never be used, and on queries that are resolved confidently early.

In the above execution strategies, we assumed that each candidate solution can be evaluated in isolation. However, as we found in our last publication on SWE-rebench , selecting a candidate based on individual self-confidence scores alone is noticeably inferior to using an LLM Judge (otherwise known as a selector) . Instead of analyzing a single candidate, the selector reviews the entire pool of generated candidates alongside the relevant source code before making a final choice.

Because the selector relies on this holistic view, we had to adjust our approach. We can no longer simply ask, “Is this individual solution good enough to stop early?” Instead, our acceptance criteria must shift to a group dynamic: “ Do we have a good-enough pool of candidates for the selector to successfully work with? “

Predicting the selector’s performance

This shift changes how we validate our pipeline. Because the selector makes the final choice or synthesis from a collective pool, we need a way to confidently predict whether that pool contains enough signal for the selector to produce a good result (the only exception is if the list contains a single candidate, in which case a selector isn’t needed anyway).

We found that two predictions were needed to inform the overall prediction of selector aptitude:

Prediction #1: That running a selector on the list of candidates will solve the task. We called this prediction Resolve-Now (RN), calculated by a dedicated classifier – with a calibrated threshold p_target (p for precision) – which we trained for this purpose.

Prediction #2: That generating more candidates, and then running a selector, will solve the task. We called this prediction Resolve-Later (RL). We calculate it by training another classifier to predict the likelihood of a selector running on a larger batch of candidates to solve the task, and subtracting Resolve-Now from it, giving us an estimate for the marginal gain of generating another batch of rollouts. We learn a calibrated threshold g ( g for gain ) for it as well.

If either p_target or g pass our threshold, we take that as a signal that there’s a good enough list of candidates at that moment. If that’s the case, we stop generating more rollouts and proceed in our pipeline. Importantly, p_target and g are parameters we can play with based on our own criteria; for instance, raising p_target ’s threshold or lowering g ’s threshold makes early stops less frequent, but safer (more expensive, yet more accurate).

Generating more signal

However, we found that self-confidence is not a strong enough signal to decide when to stop an SWE task early. We would need to generate some more features for our classifiers to more confidently make...