The raw output of a language neural net model is a set of probabilities over the possible next word (or token). That is then converted into text output by a decoding step. Standard decoding — greedy, beam search, or independent sampling — commits to each token or reasoning step without any lookahead about whether it’s heading toward a correct final answer. For multi-step reasoning tasks that require occasional backtracking (arithmetic puzzles requiring “wrong turn, try a different combination first”), that’s a real limitation. Monte Carlo Tree Seardch (MCTS) offers a principled way to search over reasoning trajectories rather than commit to them greedily, using the model’s own predictions as a proposal distribution and a value estimate to guide which branches are worth exploring further.
Outline work plan
- Pick a task where the output is verifiable e.g. the Game of 24 (given four numbers, find an arithmetic expression using +, −, ×, ÷ and parentheses that evaluates to 24).
- Train a model. I suggest you take an existing pre-trained model and fine-tune it on synthetically generated step-by-step solution demonstrations for Game of 24 puzzles. This gives the model a reasonable policy prior for the step format before MCTS search is layered on top.
- Augment the model with MCTS search. The search selects the next action: choose two remaining numbers and an operator to combine them, based on reward feedback of whether or not this choice is likely to lead to a sequence that solves the puzzle.
- Evaluate the performance of the model using standard decoders vs using MCTS decoding e.g. success rate on a held-out puzzle set, reported as a curve over increasing compute budget.
If time permits, you can extend this to include “self-play” i.e. use the sequences output by the model+MCTS as training data to further fine-tune the model. If the model+MCTS picks good sequences, the model can then learn from them and become better over time.
What you will need to implement
- Puzzle generator and brute-force solver: this is the foundation everything else depends on, it generates the synthetic fine-tuning demonstrations, the held-out evaluation set, and doubles as the ground-truth verifier for any candidate solution.
- Supervised fine-tuning of the base model: use step-by-step demonstrations, so it has a usable prior over “which move looks promising” before any search is added.
- Enumerable-action MCTS decoder: at each state, enumerate all valid (pair, operator) actions, score each by the fine-tuned model’s likelihood of generating that step as text (this is the policy), select via UCT, and expand accordingly. This will require a choice of reward to guide the MCTS search. Two possible choices are: (i) a simple rule e.g. how close the remaining numbers are to a known-combinable set, (ii) a learned reward function (run the fine-tuned model, with or without search, from many intermediate states, record whether they eventually reached 24, and train an estimator of success/failure)
- Baseline decoders: greedy single-pass generation, beam search, and best-of-N/self-consistency sampling (generate N independent full attempts, succeed if any hits 24). For a fair comparison the various decoders need to compared on a level playing field. One way might be to control the #model forward passes to be the same, another might be wall-clock time.