[RFC][Docs] Design for JAX-MCTS (Tree of Thoughts) Inference Strategy

[yipkingster]

Description

This PR introduces a comprehensive design document for the implementation of JAX-MCTS (Monte Carlo Tree Search) for the "Tree of Thoughts" (ToT) decoding strategy. This represents an evolution from MaxText’s current linear decoding strategies (greedy, sampling, and beam search) toward a non-linear, search-based reasoning framework.

Context & Problem Solved:
Standard decoding strategies often struggle with "high-regret" tasks like mathematical proofs, scientific reasoning, and complex code synthesis, where a single early error can invalidate a long sequence. JAX-MCTS addresses this by enabling the model to explore multiple branching "thoughts," evaluate their potential, and backtrack as necessary to find the most robust logical path.

Proposed Solution:
The design leverages MaxText’s existing high-performance architecture to ensure the search is both scalable and hardware-efficient:

• Prefix Sharing: Highly efficient branch forking by utilizing the Paged Attention PageManager to share KV caches across search nodes.
• Vectorized Search: Implementation of the MCTS logic (Selection, Expansion, Evaluation, Backpropagation) as a functional, JITTED loop using JAX Pytrees.
• Flexible Boundaries: Support for dynamic "Thought Boundaries" (e.g., newlines) through modifications to the MaxEngine.generate loop.

Implementation Highlights:

• jax_mcts.md: Outlines the architectural changes to MaxEngine, PageManager, and the math behind the UCT selection policy.
• Inference Fallback: Describes the transition from search-guided reasoning to greedy sampling for structural completion.

Shortcomings & Future Improvements:

• Currently requires an external Reward Model or self-evaluation prompt; future iterations could include integrated Value Heads.
• Initial implementation will focus on depth-first search; subsequent optimizations will introduce Batched MCTS for higher TPU utilization.

Tests

This is a documentation-only PR. I have:

• Performed a self-review of the mathematical formulas (UCT and Reward update logic).
• Validated that the new file reflects the existing MaxText coding standards (License headers, WIP warnings).

Checklist

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests (N/A - Documentation only).
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive).