TreeQuest: tree search as an inference time scaling layer for LLMs

The search is the product, the model is a component

TreeQuest is a flexible answer tree search library featuring AB-MCTS, useful for, but not limited to, LLM inference-time scaling.

That phrase, inference-time scaling, is the idea worth understanding before any code. A language model produces one answer per call, and if that answer is wrong, the usual remedy is a bigger model or a better prompt. TreeQuest takes a different route: keep the same model, but search. Expand candidate answers into a tree, score them, expand the promising branches, and return the best leaf. The extra compute goes into the search rather than into the weights.

The quick start shows how little of the library you have to write. Each node is associated with a user-definable state, and in the example the state is just a string:

import random
from pathlib import Path

import treequest as tq

# Each node is associated with a user-definable `state`.
State = str


# 1. Define a function to be used for node generation.
def generate(parent_state: State | None) -> tuple[State, float]:
    """Generates new states and scores based on the parent state."""
    if parent_state is None:  # None represents the expansion from root.
        new_state = "Initial state"
    else:
        new_state = f"State after {parent_state}"

    # A score for the new state; It should be normalized to the [0, 1] range.
    score = random.random()
    return new_state, score

The generate function takes the parent state and returns a new state and a score, and the score should be normalized to the range from zero to one. The None parent represents the expansion from the root.

Then you instantiate the algorithm and a search tree, run the search with a generation budget, and pull out the best result:

# 2. Instantiate the algorithm and a search tree object.
algo = tq.ABMCTSA()
search_tree = algo.init_tree()

# 3. Run the search with a generation budget (10 in this case).
for _ in range(10):
    search_tree = algo.step(search_tree, {"Action A": generate})

# 4. Extract the best score and state.
best_state, best_node_score = tq.top_k(search_tree, algo, k=1)[0]
print(f"Best state: {best_state}, Score: {best_node_score}")

# 5. Visualize the search tree.
output_file_basename = Path("ab_mcts_a_search_tree")
# Generates `ab_mcts_a_search_tree.html`
tq.render(search_tree, output_file_basename, format="html")

The actions are a dictionary, so "Action A" is a key and the generate function is the value. That makes the branching factor something you choose rather than something the algorithm assumes, which is the flexibility the project description points at.

Ask and tell, and why batching is the faster path

There is a second interface, and the readme is clear that it is not a stylistic alternative but a performance one.

The alternative is an ask-tell interface with batched AB-MCTS sampling steps:

import random
import treequest as tq

State = str


def generate(parent_state: State | None) -> tuple[State, float]: ...


generate_fns = {"Action A": generate}
actions = list(generate_fns.keys())

# We use batch_size=5 here
batch_size = 5

# It runs AB-MCTS sampling step with 5 processes in parallel
algo = tq.ABMCTSM(max_process_workers=batch_size)
search_tree = algo.init_tree()

total_budget = 50
num_steps = total_budget // batch_size
for _ in range(num_steps):
    # ask_batch returns a list of `Trial` object, which has action, parent_state and trial_id attrs
    search_tree, trials = algo.ask_batch(search_tree, batch_size, actions)

    for trial in trials:
        result = generate_fns[trial.action](trial.parent_state)
        # Call tell method with trial_id to update search_tree
        search_tree = algo.tell(search_tree, trial.trial_id, result)

best_state, best_node_score = tq.top_k(search_tree, algo, k=1)[0]

The mechanics are worth spelling out. ask_batch returns a list of Trial objects, each of which has action, parent_state and trial_id attributes. The caller then evaluates each trial, and calls tell with the trial_id to update the search tree. The Trial object is a handle: it names which action to take and from which parent, and the trial_id is how the result gets attached to the right node later. That indirection is what makes the parallelism possible, because the evaluation of a batch can happen anywhere, in any order, concurrently, and the tree is updated afterwards.

The reason to care is stated directly. In particular for AB-MCTS-M, each step call can be slow, and if you encounter slow execution you should prefer ask_batch over step. That is a strong statement, and it makes sense once you see the mechanism: a step call evaluates candidates one at a time inside the search loop, while ask_batch hands out a batch of trials so the evaluation can run in parallel across processes, with max_process_workers controlling the parallelism.

The tree shape is the catch, and it is called out in the same section. Using a large batch_size can skew the search tree shape, meaning the tree may become too wide, so it is best to avoid overly large batch_size, and the readme points at PROFILING.md in the docs for example trees. The recommendation is batch_size of five or less as a starting point.

That is a real trade-off stated by the authors rather than hidden: batching buys you throughput and costs you search quality, because a wide tree explores breadth at the expense of depth.

Two algorithm variants and what the dependencies tell you

The package offers AB-MCTS-A and AB-MCTS-M, as well as Multi-LLM AB-MCTS support, with the paper linked from the readme. The naming in the quick start is ABMCTSA and ABMCTSM, and the A and M variants are the two you instantiate directly.

The dependency list in the project manifest is the most informative document about what these variants involve, because the optional extras are not cosmetic.

The base install pulls in PyYAML, joblib, scipy, ulid-py and nothing else, so the core library is light: YAML for configuration, joblib for parallel work, scipy for the numeric parts of the search, and ulid-py for the unique identifiers the trial handles use.

The abmcts-m extra is where the weight is. It brings jax at 0.5.2 or newer, numpy, numpyro, pandas, pymc, xarray, packaging and numpy. A probabilistic programming stack is a strong hint about what the M variant does, since pymc and numpyro are the two standard probabilistic programming backends, and the combination with jax and xarray points at fitting a model to the observed search outcomes rather than evaluating a fixed score function.

That distinction matters for a user. In the A variant, you supply a score, and the search is a best-first traversal of a scored tree. In the M variant, the model appears to be learned from the search data, and the heavy numeric stack is the cost of that. The Multi-LLM support sits alongside, and the paper referenced in the readme is the place to find the detail.

The vis extra is the visualisation path, bringing Pillow, graphviz and jinja2. Graphviz for the tree graph, Jinja2 for the template, Pillow for the image output, and the render function in the quick start takes a format argument of html, which is the tie between the two. The all extra combines abmcts-m and vis.

The build backend is hatchling with the package laid out under src/treequest, and the project requires Python 3.11 or newer, which is what lets the example use the State | None union syntax directly.

When tree search is the wrong tool for the problem

The library is built around one assumption, and it is worth stating what happens when that assumption does not hold.

The assumption is that you can score a candidate answer. Every node in the tree carries a score, the search is a best-first traversal driven by those scores, and the return value is the best node found. The quick start makes this explicit by saying the score should be normalized to the range from zero to one, which is a hint that the score is a probability-like quantity rather than an arbitrary ranking.

So the library needs a scoring function that separates good answers from bad ones, and for a task with a verifiable answer that is straightforward. A maths problem can be checked. A code snippet can be run. A question with a known answer can be scored against it. A judgement call about whether a piece of text is a good support reply cannot, at least not without a model to do the judging.

The Multi-LLM AB-MCTS support named in the features list is aimed at exactly that gap, and the M variant with its probabilistic programming stack appears to be the path where the scoring is learned from the search outcomes rather than written by hand. That is a much heavier thing to run, and the readme's warning that each step call can be slow for AB-MCTS-M is the practical warning label on it.

There is a second condition worth naming. Search costs compute linearly in the budget you give it, and the budget you give it is the number of model calls. If your model is cheap and your queries are rare, generating one answer and moving on is simply the better trade, and the extra tree buys you nothing. Tree search pays off when the answer space is wide enough that a single sample has a meaningful chance of being wrong, and when you have a way to tell when a candidate is on the right path.

That last point cuts both ways. A tree that grows wide rather than deep, which is what a large batch size produces, explores breadth at the cost of depth, and a shallow tree is very close to sampling many independent answers. The readme's advice to start at a batch size of five is really advice about keeping the search a search.

The project is young, and the tooling says how they work

The version history and the development configuration together describe a project in its first serious year of public use.

The releases are v0.3.3 from 2026-10-16, v0.3.2 from 2026-02-05 and v0.3.1 from 2025-12-09. A patch version three times in a year, with the last push on the same day as the newest tag, is the pattern of a project that is still finding its API. The authors are Kou Misaki and Takuya Akiba, both at Sakana AI, and the project is Apache 2.0 licensed with a Python 3.11 badge and separate CI workflows for checks and tests.

The development dependency group is where the engineering culture shows. It pins mypy at 2.4.0 and ruff at 0.17.1 exactly, with a comment in the manifest explaining why: so that their new releases do not change CI results, and to keep ruff in sync with the pre-commit configuration. That is a small discipline and a revealing one, because a project that pins its linters to avoid CI churn is a project that runs CI on pull requests from outside.

The rest of the dev group is a normal Python test stack, pytest with xdist for parallel test runs, line-profiler and tqdm, and the type stubs for tqdm and PyYAML. The ruff configuration is deliberately narrow, selecting only the rule groups that were enabled by default before version 0.16, with a comment noting that the newer default rule set was much larger. A codebase that keeps the lint surface small is a codebase where the style rules are ones people actually agreed to.

There is a PROFILING.md in the docs directory that the readme points to for example search trees, and a pre-commit configuration at the root, so the formatting and linting are enforced before a commit rather than in CI alone.

For a user, the practical reading is that the API is settled enough to build against, the quick start is short and complete, and the reference to the paper is the place to go for what the algorithms actually do, since the readme itself is a getting started guide rather than a specification.

Editorial conclusion

Adopt TreeQuest if you are running an LLM that can be scored programmatically and you want inference time scaling, meaning spending more search budget on hard questions rather than training a bigger model. Do not adopt it if your task has no reliable way to score a candidate answer, because the whole search depends on the score you return. Before you commit, write the scoring function first and check it separates good answers from bad ones, then start with a small batch size, since the readme warns that a large one can skew the tree shape.

Frequently asked questions

Official sources

1. Issues
2. License: Apache-2.0
3. README
4. Releases
5. SakanaAI/treequest on GitHub