How does LATS differ from previous methods like Math Shepherd or React?

Unlike Math Shepherd which uses a verifier for scoring reasoning choices, LATS scores based on the outcomes of actions. It also improves upon React by incorporating more explicit planning into the process.

What are the main components of the LATS method?

LATS consists of six stages: selection, expansion, evaluation, simulation, backpropagation, and reflection. It uses UCT (Upper Confidence Bounds applied to Trees) to balance exploration and exploitation during the search process.

What is the motivation behind the Sprint framework?

Sprint is motivated by the observation that harder problems require longer thinking from LLMs, and that many reasoning steps within a longer thought process can be independent and executed in parallel, allowing for acceleration.

How does Sprint enable parallel execution in language models?

Sprint fine-tunes a reasoning model to identify parallelization opportunities. It trains the model to output plans and executions that can be run concurrently, interleaving planning and execution phases to speed up the overall process.

What are the benefits of the Sprint framework?

Sprint aims to reduce the cost and latency of LLM inference by enabling parallel execution. It has also shown improvements in accuracy and better out-of-domain generalization compared to sequential models.

What challenges does the Swirl framework address?

Swirl addresses challenges in multi-step reasoning and tool use, particularly the compounding of errors and the difficulty of using tools live during training. It aims to train models to know when to call tools, generate correct queries, maintain accuracy, and recover from errors.

How does Swirl generate training data without live tool execution?

Swirl creates synthetic multi-step data through iterative prompting. It uses an LLM as a judge to score the quality of reasoning steps and tool queries offline, separating data annotation from the RL training process.

What is the role of the LLM judge in Swirl?

The LLM judge evaluates the quality of the generated queries for tool invocation, not the tool's output. This allows for reward generation during training without needing to execute tools, which can be slow or prone to failure.

How does Swirl compare to supervised fine-tuning for multi-step tasks?

Multi-step RL optimization with Swirl performs better than supervised fine-tuning, especially when training data is not perfectly correct. RL allows the model more flexibility to explore and learn from new actions, breaking out of imitation loops.

What is the key generalization capability of Swirl?

Swirl demonstrates significant generalization across datasets and tools. Models trained on one domain (e.g., math problems with a calculator) can perform well on entirely different tasks (e.g., question answering with a search tool), learning general reasoning and tool-invocation strategies.

How does Swirl improve a model's 'process correctness'?

After Swirl fine-tuning, models show improved process correctness, meaning they become better at executing each step of multi-step reasoning correctly, both for in-distribution and out-of-distribution data.

Key Moments

Stanford CS329A Self-Improving AI Agents | Part 5 | Planning and Multi-Step Reasoning

Stanford Online

Education6 min read75 min video

Aug 3, 2026|307 views|38|2

Stanford Stanford Online

Save to Pod

Want to know something specific about what's covered?

We've already dissected every moment. Ask and we will deliver (with timestamps).

Key Moments

TL;DR

AI agents can now plan and execute complex tasks using tree-based search and parallel processing, but doing so incurs significant computational cost and faces challenges with irreversible actions.

Key Insights

The Language Agent Tree Search (LATS) framework integrates reasoning, acting, and planning by employing Monte Carlo Tree Search (MCTS) principles, scoring actions based on outcomes rather than just reasoning choices.

Sprint proposes a post-training/fine-tuning approach to enable language models to identify and execute parallelizable reasoning steps, aiming to accelerate inference and reduce computational cost.

The Swirl framework uses reinforcement learning with LLM-generated synthetic data to train models for multi-step reasoning and tool use, demonstrating generalization across different tools and tasks without explicit tool execution during training.

LATS uses a UCT (Upper Confidence bounds applied to Trees) score to balance exploration and exploitation during search, aiming to discover optimal trajectories.

Sprint's fine-tuning process involves annotating existing reasoning traces to identify parallelizable plans and executions, then supervised fine-tuning the model on this structured data.

Swirl's RL training avoids live tool execution to mitigate failures and slow-downs, instead relying on an LLM judge to score the quality of tool queries and reasoning steps.

Language Agent Tree Search (LATS) unifies reasoning, acting, and planning

The LATS framework addresses multi-step tasks by integrating reasoning (deciding what to do), acting (gathering information), and search (refining plans). Unlike previous methods that might score reasoning choices, LATS scores based on the outcomes of actions taken. It incorporates techniques from reinforcement learning and multi-step planning, including Monte Carlo Tree Search (MCTS), to encourage diversification of solutions and exploration of different paths. A key aspect is its ability to take feedback from actions and refine future plans. The process involves six stages: selection (choosing a node to expand using UCT), expansion (sampling actions from a node), evaluation (scoring new states using an LLM judge and self-consistency), simulation (greedily expanding a path to an end state), backpropagation (updating node values based on trajectory outcomes), and reflection (model's self-assessment of the trajectory). This allows models to refine plans based on real-world feedback and observations.

LATS employs Monte Carlo Tree Search for balanced exploration

At the core of LATS's planning mechanism is the Upper Confidence bounds applied to Trees (UCT) algorithm, borrowed from MCTS. UCT is designed to balance exploration (trying new, less-visited paths) and exploitation (following promising paths). It calculates a score that combines the current value of a node with a term that encourages visiting less explored children nodes. The formula V(S) + C * sqrt(ln(N_p) / N_s) is used, where V(S) is the value of the state, C is a hyperparameter, N_p is the number of times the parent node was visited, and N_s is the number of times the current node was visited. This balance is crucial for discovering optimal solutions that might not be immediately apparent. The backpropagation stage uses the outcome of a completed trajectory (success or failure) to update the values of states visited along that path, influencing future selection decisions.

Sprint enables parallel execution for faster reasoning

The Sprint framework aims to leverage the observation that many reasoning steps in large language models can be performed independently and in parallel. Models often exhibit longer thinking times correlating with higher accuracy, but much of this computation can be redundant or parallelizable. Sprint works by fine-tuning a reasoning model to identify these opportunities. This involves creating a dataset where LLMs (like GPT-4o) annotate existing reasoning traces to distinguish planning steps from execution steps, and identify parallelizable sub-tasks. A Directed Acyclic Graph (DAG) is formed to represent dependencies between these steps. The model is then supervised fine-tuned on this data, learning to output plans and parallel executions simultaneously. At inference, the model generates multiple independent plans, which can then be executed concurrently, significantly reducing wall-clock time and computational cost, while also potentially improving accuracy and generalization.

Sprint fine-tuning creates parallel thinking data

The creation of fine-tuning data for Sprint involves a structured annotation process. An LLM is used to decompose a model's reasoning trajectory into steps, distinguishing between planning and execution. It also identifies which steps are independent and can be run in parallel. This information is used to construct a DAG representing the reasoning process. The model is then fine-tuned using this annotated data, learning to output tagged plans and their parallel executions. This approach teaches the model to think more in parallel, outputting multiple actionable plans at once rather than one sequential step after another. The goal is to reduce the sequential token generation bottleneck, leading to efficiency gains and potentially better performance by enabling the model to explore more options concurrently.

Swirl trains models for multi-step reasoning and tool use with synthetic data

Swirl is a framework designed to improve multi-step reasoning and tool use in LLMs without requiring live tool execution during training. The core challenge is that errors in multi-step processes can compound, and live tool execution during training is slow and prone to failure. Swirl addresses this by generating synthetic multi-step data through iterative prompting. An LLM is prompted with access to tools, and at each step, it can reason, call a tool, or propose a final answer. An LLM judge then scores the quality of the reasoning steps and tool queries offline. This synthetic dataset, labeled with rewards, is used for reinforcement learning. The RL process optimizes the expected reward of single actions given the context, effectively training the model to reason and use tools without actually executing them during training, thereby avoiding slow-downs and failures.

Swirl's RL approach generalizes across tasks and tools

A key finding from Swirl is its strong generalization capabilities. By training on synthetic data generated using one set of tools and tasks (e.g., using a calculator for math problems), the model shows improved performance on different tasks and with different tools (e.g., a search tool for question answering). This suggests that the model learns a more fundamental skill of multi-step reasoning and tool invocation, rather than just memorizing specific tool usage patterns. The framework also demonstrates that training with process-filtered data (where reasoning steps are good, even if the final outcome isn't perfect) can be more effective than solely using outcome-filtered data, as it helps the model tackle problems it couldn't solve previously. This approach, utilizing RL on synthetic data, significantly outperforms supervised fine-tuning, particularly for complex, multi-step reasoning tasks.

Challenges and future directions in AI planning

Despite advancements, challenges remain. The computational cost of advanced planning techniques like LATS and Sprint can be substantial. Moreover, scenarios involving irreversible actions, such as financial transactions, present a significant hurdle for current planning frameworks, as a wrong action cannot be easily undone. The effectiveness of these methods is also task-dependent, with harder problems requiring more iterative planning and execution. Future work could explore more sophisticated RL algorithms, better load balancing for parallel execution, and methods to handle reversible and irreversible actions more robustly, further enhancing the capabilities of AI agents in complex, real-world scenarios.

Mentioned in This Episode

●Software & Apps

●Studies Cited

●Concepts

Common Questions

LATS unifies reasoning, acting, and planning in language models by building upon concepts like Monte Carlo Tree Search (MCTS). It focuses on scoring actions based on their outcomes and incorporating model reflection for improved planning and search processes.

Topics

Tool Use AI & Machine Learning Technology & Innovation Language Models Synthetic Data Generation LLM Inference Multi-step Reasoning Parallel Execution

Mentioned in this video

Concepts

Reinforcement Learning

A machine learning paradigm where agents learn to make sequences of decisions by trying to maximize a reward signal.

Chain of Thought

A reasoning process decomposed into a number of steps, forming the basis for LATS's tree expansion.

Companies

Sprint

A framework designed to enable reasoning models to execute plans in parallel, accelerating the reasoning process.

Software & Apps

Gemini

A system card mentioned to encourage tool use, highlighting the trend towards complex problem-solving with LLMs.

M-GMO 227B

Used to create multi-step synthetic data for the Swirl framework.

Gemini 2.5 Pro

A large language model that exhibits longer thinking for harder problems, motivating the Sprint framework.

Monte Carlo Tree Search

A heuristic search algorithm for decision processes, particularly used in game trees and planning.

GPT-4o

Used to annotate reasoning traces into planning and execution steps for creating the Sprint fine-tuning data.

Swirl

A framework that uses reinforcement learning to train language models for multi-step reasoning and tool use, avoiding live tool execution during training.

LATS

Language Agent Tree Research, a framework that unifies reasoning, acting, and planning in language models, inspired by MCTS.

DeepSeek R1

A reasoning model whose training process shows a correlation between longer thinking and higher accuracy, used as a base for Sprint.

React

A framework for reasoning and acting, which LATS builds upon by incorporating more planning into the process.

Studies & Research

GSM8K

A dataset of math problems used in the Swirl experiments to train and test multi-step reasoning and tool use capabilities.

Hotspot QA

A dataset used to test LATS and Swirl, requiring retrieval from multiple Wikipedia pages to answer questions.

WebShop

A dataset for practical applications used to test LATS, involving multi-step processes for product search.

Ask anything from this episode.

Save it, chat with it, and connect it to Claude or ChatGPT. Get cited answers from the actual content — and build your own knowledge base of every podcast and video you care about.

Get Started Free