cs.AIFeb 28, 2026
SaveLiTS: A Modular Framework for LLM Tree Search
Organizations: RMIT University
Abstract
LiTS is a modular Python framework for LLM reasoning via tree search. It decomposes tree search into three reusable components (Policy, Transition, and RewardModel) that plug into algorithms like MCTS and BFS. A decorator-based registry enables domain experts to extend to new domains by registering components, and algorithmic researchers to implement custom search algorithms. We demonstrate composability on MATH500 (language reasoning), Crosswords (environment planning), and MapEval (tool use), showing that components and algorithms are orthogonal: components are reusable across algorithms within each task type, and algorithms work across all components and domains. We also report a mode-collapse finding: in infinite action spaces, LLM policy diversity (not reward quality) is the bottleneck for effective tree search. A demonstration video is available at https://youtu.be/nRGX43YrR3I. The package is released under the Apache 2.0 license at https://github.com/xinzhel/lits-llm, including installation instructions and runnable examples that enable users to reproduce the demonstrated workflows.
Figures & tables
Figure 1: LiTS architecture. Agents orchestrate Components (Policy, Transition, RewardModel). Components use Prompts and LM interfaces. Transition can invoke Tools via Clients for tool-use tasks. All layers communicate through shared Data Structures , decoupling search algorithms from domain logic.
| User Group | Reusability Direction | Benefit |
|---|---|---|
| Domain experts | Same components different methods | Fair algorithmic comparison |
| AI/NLP researchers | Same method different domains | Generalization testing |
| AI/NLP researchers | New formulation existing search | Novel reasoning paradigms |
Table 1: Bidirectional reusability enabled by LiTS’s modular design.
| Task Type | Minimum Required Injection | Extension Mechanism | Framework Knowledge |
|---|---|---|---|
| Language-grounded | Prompt, Dataset | @register_user_prompt , @register_dataset | None |
| Env-grounded | Transition class | @register_transition (1 decorator) | Minimal |
| Tool-use | Tool definitions | BaseTool subclass | Minimal * |
Table 2: Domain extension patterns by task type. Component reuse varies: tool-use tasks share Policy and Transition across domains; env-grounded tasks share Policy and RewardModel; language-grounded tasks may require task-specific prompts. Requires understanding Transition interface ( goal_check , generate_actions , _step ). * BaseTool requires name , description , args_schema , and _run() —compatible with LangChain’s Tool interface.
| Task | Method | Out Tok | Cost | Acc |
|---|---|---|---|---|
| BlocksWorld | Chain | 17K | $1.48 | 26.7% |
| MCTS | 488K | $21.99 | 66.7% | |
| Crosswords | Chain | 2.5K | $0.28 | 6.67%/10.33% * |
| MCTS | 14K | $2.42 | 0%22.67% * |
Table 3: Environment-grounded results (30 examples). Same components across domains. * Crosswords accuracy: exact match (all 10 clues correct) / partial match (average clue accuracy).
| Metric | Value |
|---|---|
| Unique states visited | 16 |
| Avg. policy calls per state | 7.9 |
| Duplicate rate (all) | 81.1% |
| Duplicate rate (incorrect) | 81.0% |
| Correct outputs | 17.3% |
Table 4: LLM action diversity on Crosswords with temperature escalation ( : 0.8 1.2). Near-identical duplicate rates for all outputs vs. incorrect-only outputs indicate temperature scaling fails to improve exploration.
| Out Tok | Inv | Cost | Acc |
| 10.6K | 62 | $0.57 | 40% |
Table 5: MapEval-SQL ReAct results (10 examples). Inv = total LLM invocations.
| Method | Acc | Out Tok | Inv | Time |
|---|---|---|---|---|
| CoT | 17% | 12.9K | 100 | 0.6h |
| RAP (MCTS) † | 18% | 4.47M | 3.6K | 8.0h |
| ReST (MCTS) | 37% | 2.24M | 4.0K | 26.0h |
| ToT (BFS) ‡ | 39% | 1.53M | 2.8K | 14.7h |
Table 6: MATH500 results (100 examples, Llama3-8B). † User-registered components; ‡ user-registered search algorithm; unmarked rows use only built-in components and algorithms. Inv = total LLM invocations; Time = total wall-clock time including all LLM calls.
| Feature | LiTS | LLM-R | LG |
|---|---|---|---|
| Task-agnostic search | ✓ | ✓ | ✗ |
| Component sharing | ✓ | ✗ | ✗ |
| Tool-use tree search | ✓ | ✗ | ✗ |
| Prompt registry | ✓ | ✗ | ✗ |
| Decorator extension | ✓ | ✗ | ✗ |
| Inference logging | ✓ | ✗ | ✗ |
Table 7: Feature comparison. LLM-R = LLM Reasoners; LG = LangGraph; ~ = partial.
| Search ( lits_search ) | blocksworld_rap/run_0.2.5/ |
|---|---|
| checkpoints/ | {query_idx}_{iter}.json : intermediate tree states per iteration; {query_idx}_result.json : final result with full tree path. |
| terminal_nodes/ | terminal_nodes_{query_idx}.json : all terminal nodes found. |
| config.json | Search configuration (many params). |
| execution.log | Execution log. |
| eval.log | Evaluation log. |
| inferencelogger.log | LLM inference log (token/latency). |
Table 8: Per-run artifact layout for search and chaining in LiTS. Each run writes all outputs under a single save_dir , enabling post-hoc inspection and evaluation via eval --save_dir .
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
| Task Type | Action | Step | State |
|---|---|---|---|
| Language-Grounded | StringAction | ThoughtStep / SubQAStep | TrajectoryState |
| (e.g., Math500, GSM8K) | Free-form text generation | Reasoning thought or sub-QA pair | Sequence of reasoning steps |
| Tool-Use | ToolUseAction | ToolUseStep | ToolUseState |
| (e.g., MapEval) | Structured tool call (JSON) | Thought, Action, Observation | Trajectory with tool history |
| Environment-Grounded | EnvAction | EnvStep | EnvState |
| (e.g., BlocksWorld) | Valid environment command | Action and resulting state | Trajectory + Env Snapshot |
Table 9: Structure subclasses across task types. Each task type defines polymorphic Action, Step, and State subclasses that carry task-specific data while exposing a uniform interface to search algorithms.
| Task Type | Execution | Method | Policy | Transition | Reward Model |
|---|---|---|---|---|---|
| Environment-Grounded | Chain | EnvChain | EnvGroundedPolicy | BlocksWorldTransition † | – |
| Tree | MCTS / BFS | EnvGroundedPolicy | BlocksWorldTransition † | EnvGroundedPRM | |
| Tree | MCTS / BFS | EnvGroundedPolicy | CrosswordsTransition † | EnvGroundedPRM | |
| Tool-Use | Chain | ReActChat | ToolUsePolicy | ToolUseTransition | – |
| Tree | MCTS / BFS | ToolUsePolicy | ToolUseTransition | ToolUsePRM | |
| Language-Grounded | Tree | MCTS | RAPPolicy † | RAPTransition † | RapPRM † |
Table 10: Component configurations used in our demonstrations. Unmarked components are built-in (shipped with LiTS ); † components are user-registered via @register_* decorators. The same built-in components are reused across methods within a task type; only the RewardModel is added for tree search.