Orply.

LATS, SPRINT, and SWiRL Target Different Bottlenecks in Agent Reasoning

Azalia MirhoseiniStanford OnlineMonday, August 3, 202613 min read

Stanford’s Azalia Mirhoseini argues that improving multi-step AI agents requires more than extending a chain of thought: agents must plan, act, use feedback and revise their course. In this CS329A lecture, she examines three approaches to that problem—LATS, which searches alternative action paths at inference; SPRINT, which identifies work that can run in parallel; and SWiRL, which trains next-action choices from offline tool-use trajectories. Each shifts the trade-off among accuracy, latency, inference cost and the reliability of feedback.

Three ways to improve a multi-step agent

Azalia Mirhoseini frames multi-step agency as more than producing a longer answer. An agent must decide what to do, act in an environment, incorporate what those actions reveal, and revise its path toward a final result.

The three methods intervene at different points in that loop. LATS adds inference-time search: it samples alternative actions, evaluates their consequences, and allocates more computation to promising branches. SPRINT changes how a reasoning process is scheduled: it identifies independent work and executes it concurrently rather than decoding every operation in a single serial chain. SWiRL changes training: it uses offline-generated tool-use trajectories and step-wise reinforcement learning to improve the quality of the next action.

MethodIntervention pointMain gainKey dependencyStated constraint
LATSInference-time searchExplores and revises alternative action trajectoriesA searchable environment with useful state feedbackMultiple iterations and samples add inference cost; irreversible actions are unresolved
SPRINTPlanning and execution scheduleReduces sequential decoding by overlapping independent workTasks with exploitable dependency structureParallelism is task-dependent; stragglers and runtime overhead remain
SWiRLOffline data generation and RL trainingImproves next-action and tool-use behavior across stepsSynthetic trajectories and credible process judgmentsTraining signal depends on the quality of trajectories and model-judge rewards
The three methods improve multi-step agents at different points in the loop.

The shared problem is what happens when one linear plan is not enough. For a trip-planning task, reasoning might establish a budget or possible destinations; acting might involve web searches, travel blogs, or advice from people who have visited; search is the subsequent revision of the plan in light of what those actions return. The methods differ in whether they improve that cycle by searching more at test time, exposing parallel work, or training a model to take better steps.

LATS spends more test-time compute to search before it commits

A conventional language-model agent typically generates a plan and follows it. That approach limits exploration of alternative paths and can leave the model weak at planning several actions ahead. Language Agent Tree Search, or LATS, adapts Monte Carlo Tree Search to language-model agents: the model can branch, act, evaluate resulting states, and return to branches that appear promising.

A trip-planning prompt illustrates the move. An agent asked to plan a trip to Hawaii might propose asking friends who have visited, reading relevant subreddits, or pursuing another information source. Rather than committing immediately, LATS can execute candidate actions, score the observations they yield, and expand a higher-scoring branch. Multiple actions—such as asking two friends—can be run in parallel.

Mirhoseini describes LATS as combining chain-of-thought decomposition, environment interaction in the style of ReAct, and tree search. She distinguishes it from Math-Shepherd, where a verifier guides search by scoring reasoning trajectories. LATS instead evaluates the outcomes of actions: what the agent observes after interacting with its environment, supplemented by the model’s reflection on the trajectory.

If we just go based on the highest value at that point, we might miss out other nodes that right now maybe at the current moment have a lower value, but down the road at the end might have higher value.

Azalia Mirhoseini

LATS has six named operations—selection, expansion, evaluation, simulation, backpropagation, and reflection—but its basic logic is direct. It selects a state, samples candidate actions, executes them, scores the resulting states, and continues expanding branches until reaching success, failure, or a search-budget limit. The terminal result is propagated back through the branch; a model-generated reflection on success or failure is made available to later search.

The lecture’s maze example begins with an agent in a dimly lit room with a left door, a right door, and the option to inspect the room. Opening the left door leads to a dark corridor with paintings; the right door is jammed; inspection discovers a key. LATS gives the resulting states values and can continue expanding the left-door branch if it appears most promising.

The value combines two signals. The first is an LM-as-a-judge assessment: given the action and resulting observation, the model estimates how likely the state is to lead to success. The second is self-consistency: if many samples propose an action of the same kind, that action receives a higher signal. In the maze illustration, opening the left door has a combined value of 0.775, compared with 0.275 for the jammed right door and 0.625 for inspection.

Selection uses Upper Confidence bounds applied to Trees, or UCT. A high-value state is attractive because it appears promising, but a state visited relatively little compared with its parent receives an exploration bonus. The purpose is to avoid repeatedly exploiting an early favorite while ignoring a less-tested path that could prove superior. The exploration weight controls the balance.

When a simulated trajectory reaches a terminal result, its reward updates the values and visit counts of the earlier states in that branch. Reflection adds a separate feedback channel. After a successful maze trajectory, the model might record a pattern such as left door, staircase, exit; after a failed one, it can identify what went wrong and incorporate that feedback into later iterations.

The reported benchmarks show the expected trade-off: more test-time search can improve performance on multi-step tasks. On a 100-question HotpotQA subset in an oracle setup, LATS with chain-of-thought plus ReAct reached 0.71 exact match. HotpotQA requires retrieval over at least two Wikipedia passages, making multi-step behavior necessary by construction.

MethodHotpotQA exact match
ReAct0.32
ReAct (best of k)0.38
Reflexion0.51
RAP (ReAct)0.54
LATS (ReAct)0.63
LATS (n = 10)0.65
LATS (CoT + ReAct)0.71
Reported HotpotQA performance on a 100-question oracle-evaluation subset.

On WebShop, a shopping environment in which an agent navigates a website to identify and purchase an item matching an instruction, LATS with ReAct reached a score of 75.9 and a success rate of 38.0. The listed fine-tuning baseline scored 67.5 with a 45.0 success rate; the expert reference was 82.1 score and 59.6 success rate. The comparison makes clear that the metrics need not move together.

The cost of this approach is also central. Candidate generation, tool or environment execution, state evaluation, simulations, and tree updates all consume inference-time compute. Mirhoseini notes that the paper did not thoroughly analyze the cost-benefit trade-off, which limits what can be concluded about time-sensitive uses.

She also treats consequential actions as an open challenge, not a solved boundary. Search is easier when an agent can explore alternatives without permanently changing its environment. Payments and other transactions with no “undo” operation raise the question of how live exploration could be adapted safely; LATS does not answer it.

In discussion, Mirhoseini noted that UCT is only one way to balance exploration and exploitation. The paper did not compare it with other methods from the multi-armed-bandit literature. Its contribution, in her account, was to establish a framework in which alternatives could be tested. Repeated actions are handled through visit accounting within the tree, which assumes a tree rather than a fully connected graph of interchangeable states.

SPRINT makes a reasoning trace less serial

Longer reasoning traces are associated with higher accuracy on hard problems, Mirhoseini says, pointing to DeepSeek-R1 training curves in which AIME accuracy rose alongside average response length. But a longer trace need not be wholly sequential.

Some reasoning operations can be independent: trying alternative approaches, decomposing a task into subtasks, or verifying earlier work. SPRINT starts from the proposition that a reasoning trace can contain dependency structure rather than one unavoidable sequence. It trains a model to expose that structure so independent execution can overlap.

Mirhoseini describes SPRINT as an orchestration with a planner and a pool of executors. The planner produces plans; independent executors carry them out concurrently; their outputs synchronize into a running context; and the model begins another planning round if needed.

The underlying model remains autoregressive. It still produces next tokens rather than becoming a non-autoregressive model with several native streams. The change is behavioral and operational: fine-tuning teaches the model to emit planning and execution structures. Once two independent plans have been generated, their execution can branch out in parallel—potentially through tools such as Python or a calculator—before the results return to context.

The training data begins with existing reasoning trajectories. DeepSeek-R1 generates a trajectory, and GPT-4o decomposes it into distinct steps while annotating planning and execution portions. A dependency graph identifies which steps require preceding outputs and which do not. The graph is then packed into stages: a prerequisite may run first, followed by independent work in parallel, then a synchronization point before dependent work continues.

A sequence such as plan 1, execute 1, plan 2, execute 2 can therefore be reformatted so that plan 1 and plan 2 appear together, followed by their overlapping executions. Supervised fine-tuning on these packed trajectories teaches the reported model, DeepSeek-R1-Distill-Qwen-7B, to generate the structure itself.

The target is fewer sequential decoding steps, which matters for delay and inference cost. It is not a claim that parallelism removes the underlying work. The model still has to produce plans, coordinate branches, ingest outputs, and synthesize a final answer. And as Mirhoseini emphasizes in discussion, independent branches can take unequal amounts of time: a slow branch can become a straggler. Where parallel work is genuinely independent and execution can overlap effectively, the elapsed time need not equal the sum of all branch times; it is still constrained by synchronization, the longest branch, and implementation overhead.

The reported training recipe began with 6,000 DeepSeek-R1 thinking trajectories from the MATH training set and filtered out examples with low parallelization opportunity. SPRINT improved reported accuracy while reducing sequential tokens. On MATH, Mirhoseini says the fine-tuned 7B model gained roughly 3.5 percentage points over its baseline while using fewer sequential tokens than a 32B comparison model.

MethodMATH500 accuracySequential tokensCountdown accuracySequential tokensGPQA-Diamond accuracySequential tokens
Self-consistency80.559078.52,84545.44,735
SoT-chat47.325680.02,36749.43,526
SoT-reasoning90.83,83682.45,82348.07,560
RFT91.02,88084.94,91750.57,103
SPRINT92.52,44085.92,28451.06,336
Reported accuracy and average sequential-token counts across in-domain and out-of-domain tasks.

SPRINT was trained on MATH but also improved on Countdown and GPQA-Diamond. Mirhoseini interprets that result as suggesting that parallel planning may help the model reason more effectively, not merely offer a systems-level opportunity to execute independent work concurrently. The reported results are consistent with that interpretation, but they do not settle why the accuracy gains occur.

Parallelism was concentrated earlier in the solution process. Harder problems required more planning-and-execution rounds, while early stages showed more exploration and broader parallel work. Later stages converged toward fewer plans and more dependent execution.

The advantage remains task-dependent. Short problems or tasks that are intrinsically serial can offer too little independent work to justify planning and synchronization overhead. In the reported comparison with rejection fine-tuning, some short-token bins showed negative savings. Longer traces were more favorable: the cited longer-sequence ranges included reductions of 39% on MATH500, 65% on Countdown, and 45% on GPQA-Diamond.

Class discussion surfaced a second issue besides load balancing: branches that appear individually sound can conflict when brought together. The final model sees the synchronized context and is expected to reconcile contradictions before responding. Mirhoseini’s point was not that parallel reasoning eliminates this risk; serial and parallel reasoning can both contain inconsistent intermediate work.

The next steps she identifies include using RL methods such as GRPO to discover parallelization strategies beyond supervised examples, overlapping expensive tool calls, and developing hardware-optimized implementations that turn sequential-token reductions into lower wall-clock latency.

SWiRL trains the quality of the next action

SWiRL shifts the intervention from inference-time search and execution scheduling into training. Its target is not only a correct final answer, but a model that takes better next actions through a multi-step, tool-using process.

Mirhoseini identifies a compounded problem. Reasoning and tool use are each difficult on their own. Across multiple steps, a model must decide whether a tool is needed, formulate a useful query, retain accuracy across the trajectory, recover from errors, and know when it has enough information to stop.

Using live tools during reinforcement-learning training makes that process harder. Tools can be slow, expensive, buggy, and failure-prone, while RL is already costly. Outcome-only training creates another limitation: if reward arrives only at the final answer, it provides limited information about whether earlier reasoning and tool actions were good.

SWiRL separates trajectory collection from policy optimization. First, it generates multi-step synthetic trajectories offline. Given a prompt, a model can reason, call a tool such as search or a calculator, receive an environmental response, take a further action, or provide a final answer. An LLM judge evaluates each action using the context that exists before that action.

The collected trajectories can be filtered two ways. Process-filtered data retains trajectories in which every action was judged reasonable. Outcome-filtered data retains trajectories whose final answer matches the gold label, regardless of the quality of intermediate steps.

During RL, the model does not execute tools live. It receives the original prompt and the offline trajectory context through a selected point, then generates a next action. A reward model evaluates whether that action is appropriate given the context. The relevant judgment is about the quality of the proposed action or query—not about a future tool result that has not been executed during RL.

The lecture’s question about who is older, Glenn Hughes or Ross Lynch, illustrates the distinction. A first action might be a query for Glenn Hughes’s age. A judge can assess whether that is a reasonable next move before seeing the search result. The previously collected result becomes context for a next action, such as a query for Ross Lynch’s age. Once both results are available in the trajectory context, another action can produce the answer.

At inference, the model again uses tools live. It is prompted with the available tool interface and tags for tool requests and final answers. It can issue a calculation or search query, receive the result, add it to context, and decide whether more work is necessary.

In the lecture’s calculator example, a model is asked about 20 watermelons bought for $80 and sold at a $20 profit. It requests a calculation for 80 plus 20, receives 100, then requests 100 divided by 20, receives 5, and supplies the answer. The intended behavior is iterative: choose an action, receive an environmental response, update context, and choose again.

The reported setup used Gemma-2-27B to generate synthetic data and as the fine-tuned base model. The trajectories came from HotpotQA and GSM8K, and Gemini 1.5 Pro graded process rewards. The HotpotQA collection contained 50,000 trajectories and supported process-filtered, outcome-filtered, jointly filtered, and random data subsets.

One result Mirhoseini emphasizes is that process-filtered data alone appeared more useful for SWiRL than outcome-filtered data or trajectories that were both process- and outcome-correct. Her explanation is that filtering only for correct final answers may overrepresent problems the generating model already knew how to solve. A trajectory that ends incorrectly can still contain useful process behavior. Step-wise RL gives the trained model a new opportunity to choose a better action from the available context rather than simply reproduce the source trajectory.

That distinction also separates SWiRL from supervised fine-tuning. In the reported comparison, multi-step RL outperformed SFT. SFT benefited more from data whose processes and final outcomes were both correct, which Mirhoseini relates to its imitation-learning character: it is trained to reproduce supplied trajectories. SWiRL instead rewards a newly generated next action.

The broader reported result is transfer across tools and domains. Training SWiRL on GSM8K math with calculator-style support raised HotpotQA accuracy from 0.65 for the base model to 0.71. Training on HotpotQA with a search tool raised HotpotQA to 0.73 and also transferred back to mathematical reasoning.

ModelGSM8KHotPotQACofCABeerQAMuSiQue
Base model0.650.650.540.590.45
SWiRL trained on GSM8K0.790.710.560.680.49
SWiRL trained on HotPotQA0.760.730.620.680.50
Reported generalization performance of SWiRL-trained Gemma-2-27B across math and multi-hop QA tasks.

Scaling process-filtered HotpotQA data from roughly 100 to 10,000 examples improved results not only on HotpotQA but also on other QA datasets and GSM8K. Mirhoseini presents this as evidence that synthetic data generated in one setting can teach broader multi-step behavior, including behavior relevant to different tools and domains.

Process-level measurements support her interpretation. The base model’s mean process label was 82.5% on in-distribution HotpotQA and 87.5% on out-of-distribution GSM8K. SWiRL trained on HotpotQA reached 91.0% and 91.6%, respectively.

91.6%
Reported mean process-label score on out-of-distribution GSM8K after SWiRL training on HotPotQA

The method’s dependence on synthetic trajectories and model-judge rewards remains material. If the offline data fails to capture useful next actions, or if the judge cannot distinguish good actions from poor ones, the RL signal weakens accordingly. SWiRL avoids calling tools during the RL loop; it does not eliminate the need to collect tool-use trajectories offline, nor does it avoid live tool execution when the trained model is used at inference.

The frontier, in your inbox tomorrow at 08:00.

Sign up free. Pick the industry Briefs you want. Tomorrow morning, they land. No credit card.

Sign up free