Causal language modeling can elicit search and reasoning capabilities on logic puzzles

We thank the reviewer for their review and constructive feedback. We will make a pass over the paper and fix the typos pointed out and any other errors we find. We address your comments below.

• Selecting puzzles solvable by a solver: We emphasize that we did this filtering on both our train and test sets to ensure that the puzzles can be solved in reasonable time. Additionally, although our solver only uses 7 strategies, many of the strategies are highly advanced (a hobby Sudoku solver might not be aware of some of them). We thus argue that our solver is a complex algorithm, albeit one which doesn’t use backtracking. We agree that this rules out puzzles which can be solved in reasonable time with backtracking and hence our results currently do not show that causal language modeling on transformers yields them with planning capabilities in complex scenarios requiring extensive backtracking.

• Difficulty rating of puzzles on Kaggle: We wish to point out that the difficulty rating provided on Kaggle is an imperfect measure of the difficulty of the puzzle. This rating is computed as follows. To rate a puzzle, it consider a solver (different from the one we use to generate our solver-order data) which tries to iteratively make progress on a puzzle using some elimination techniques. When the solver gets stuck, it makes guesses and tries to solve the puzzle. The difficulty rating is the average number of guesses the solver had to make to solve the puzzle. Therefore, even a puzzle rated 0.5 can require complex strategies beyond simple scanning to solve them without guessing. Moreover, even for puzzles with rating 0, the elimination technique employed by the solver includes the hidden single strategy which is quite computationally intensive. That being said, we provide an analysis about how the complete puzzle accuracy changes as we increase the rating of puzzle in the attached PDF of the global response.

• “Polynomial time claim”: To clarify the claim, by polynomial time we mean that polynomial in $n$ if the size of the puzzles is $n \times n$. We say that all our filtered puzzles are solvable in polynomial time because if we consider the general version of Sudoku of size $n \times n$ and also consider a generalized version of the solver that employs the same set of strategies, then all the examples filtered by the solver will be solved by polynomial time in $n$ because each strategy in the solver runs in polynomial time in $n$ and none of the filtered Sudoku requires guessing a value in some empty cell. We will clarify this claim in the revised version of the paper.

In addition to these comments, we have added additional ablations and error analysis in the general response above. We have also shown additional experiments in a different puzzle solving setting - solving Zebra puzzles (also known as Einstein riddles) which shares some similarity with Sudoku puzzles but has many differences. We observe a similar outcome with Zebra puzzles as well where causal language modeling on solver order data is able to teach the model to perform well on these puzzles. Please see the general response for more details on these experiments.

We thank the reviewer for their review and constructive feedback. We address your comments below.

• Lacks theoretical or methodological advancements: We reiterate that our main contribution is not proposing a new theory or a new method to train LLMs. Rather it is an advance in our understanding of the precise capabilities and limitations of causal language modeling. Our work is similar to works such as [1], [2], [3] which show that Transformers can learn useful world models. That being said, we posit that our insight on showing the right style of data to train models on for reasoning puzzles might be helpful in developing better pre-training data collection strategies for general LLMs. Our contribution is novel because on general-purpose LLMs such as GPT-4, there is much evidence [5], [6], [7] that these models do not do well at challenging reasoning/planning tasks. Hence, apriori, it is not obvious that Sudoku puzzles can be solved by causal LM alone.

• Narrow scope - evaluations on other reasoning puzzles/variants:

  • Firstly, we note that Sudoku is a challenging task where prior work had not shown success of the causal LM approach (see point in general response about performance of frontier LLMs). While Sudoku is a specific puzzle setting, this enables us to perform a controlled study where we can quantify the exact amount of generalization we observe.Taking the reviewer’s feedback we added experiments with a second puzzle type: Zebra puzzles, where we observed similar results. Please refer to the general response and the attached PDF for more details. We will also add these to the paper.
  • In addition, we believe that Sudokus involve many sub-reasoning tasks involving deductive logic, mapping to abstract concepts etc and are representative of many challenging deterministic puzzles. Hence, we believe, our main message is robust to changes to the puzzle such as the logical dependency structure while keeping the required components of logical inference, search, and planning the same.
  • Rule-less puzzles require world knowledge to solve and complicate a principled and controlled study of the type we conducted. It is challenging to generate CoT style data for such puzzles on a large scale. We believe this is a very interesting direction for future research.
  • We also like the idea of stochastic puzzles. It would be interesting to see how causal LM performs on them but we believe this is outside the scope of the current paper.

• Comparison to traditional Sudoku solvers and AI approaches: Please see the corresponding point in the general response above.

• Overstated claims of reasoning vs pattern matching: We disagree with this comment. Human reasoning also involves a great deal of pattern matching. However, this happens in an abstract concept space rather than the raw input space. For instance, when presented with a novel coding problem, one might perform an abstract pattern matching to understand whether dynamic programming or divide and conquer are applicable. A lot of frontier research also involves pattern matching. Hence, we argue that learning to pattern match in an abstract concept space consistently over a number of steps is reasoning. While we agree that our setting focuses on a single puzzle type, the model learns to solve unseen puzzles; it learns abstract “strategies” during training and applies them for a test puzzle. This is pattern matching in the abstract concept space. Moreover, some of these abstract strategies can involve O(n^3) computation to find the right square to apply the strategy to. The model is learning to do this search! In addition, it learns to do this consistently for over 50 steps. Hence, our claim of the model learning to reason.

• Computational Efficiency: The computational cost of the different orders of the data in our paper is the same. In particular, giving only the value filled in each cell in a step by step manner is more efficient than giving an entire search trace as is done in [4]. Beam search only adds a constant K^2 factor to the decoding process. While further improving the computational efficiency of LLMs is an important area of research, that is not our primary focus and we leave it for future work. We will add a detailed discussion on the computational efficiency aspects in the paper.

• Error analysis: Thanks for the feedback. We have a preliminary error analysis in the paper. Please refer to the general response for details on this. We will be adding a more extended error analysis in the paper.

• Explaining model’s reasoning process: Our toy model can’t generate English explanations. However, our probing study shows that the model implicitly keeps track of candidate sets (the set of possible values in a given cell). This is a commonly used technique by humans. In addition, we have done an edit distance analysis to study how much the choice of cell locations to decode varies between the model vs the solver. We find that there is a great deal of alignment between these two orders thus providing additional evidence of a human-like reasoning process.

[1] Emergent world representations: Exploring a sequence model trained on a synthetic task. Li et al.

[2] Physics of language models: Part 1, context-free grammar. Allen-Zhu et al.

[3] Physics of Language Models: Part 2.1, Grade-School Math and the Hidden Reasoning Process. Allen-Zhu et al.

[4] Beyond a*: Better planning with transformers via search dynamics bootstrapping. Lehnert et al.

[5] [1] Large language models still can’t plan (a benchmark for llms on planning and reasoning about change). Valmeekam et al.

[6] Travelplanner: A benchmark for real-world planning with language agents. Xie et al.

[7] Limits of transformers on compositionality. Dziri et al.

We thank the reviewer for their review and constructive feedback. We will make a pass over the paper and fix the typos pointed out and any other errors we find. We address your specific comments and questions below.

• No SoTA models are evaluated on the sudoku puzzle data: We reiterate that our main focus is not to compete with SoTA approaches to Sudoku solving. We are studying the capabilities and limitations of causal language modeling (which is accepted as the dominant approach to train language models today) in a controlled setting and not looking to compare with other approaches. Nonetheless in our general response above, we present a comparison with some other approaches to solve Sudoku puzzles and will add this discussion to the paper. Note that in our approach we don’t handcraft any parts of the architecture for Sudoku puzzles.
• How do models like GPT-4, Gemini-1.5 and Llama-3 perform on these Sudoku puzzles (with and without CoT)?: Frontier LLMs like GPT-4 and Gemini-1.5 are expected to perform very poorly on a challenging reasoning task such as Sudoku. The poor performance of these models on planning/reasoning tasks has been reported in [1], [2], [3] and others. For completeness, we performed a small-scale study on Gemini-1.5 and GPT-4o where we queried them with 3000 Sudoku puzzles each (in a 4-shot with CoT manner) and we analyzed how well they can solve the puzzle. Overall, we found that they got 0% of the puzzles completely right and their accuracy on a per-cell basis was around 8-10% which is close to random guessing. We will include these results in the paper.
• A full ablation analysis is not included in the paper: Thank you for this feedback. We will be extending the number of ablation studies we have in our paper. We wish to point out that we did include some ablations such as experiments with CoT data vs random order data vs fixed order data. In addition we have also included ablations for different beam search widths. In addition, we will have now performed ablations with respect to puzzle difficulties. The results are described in the general response.
• Are the plots in Figure 3 in the CoT+Beam search setting?: We apologize for the confusion. These plots are just CoT training without beam search. We will clarify this in the figure’s caption.

In addition to these comments, we have added additional ablations and highlighted some error analysis into where the model struggles in the general response above. We have also shown additional experiments in a different puzzle solving setting - solving Zebra puzzles (also known as Einstein riddles) which shares some similarity with Sudoku puzzles but has many differences. We observe a similar outcome with Zebra puzzles as well where causal language modeling on solver order data is able to teach the model to perform well on these puzzles. Please see the general response for more details on these experiments.

[1] Large language models still can’t plan (a benchmark for llms on planning and reasoning about change). Valmeekam et al. 2022.

[2] Travelplanner: A benchmark for real-world planning with language agents. Xie et al. 2024

[3] Limits of transformers on compositionality. Dziri et al. 2024.

Thank you for your response! Let us know if there are any additional comments/concerns that we can answer.

If we have addressed your concerns, can you reconsider the score? Thank you.

We thank the reviewer for their review and constructive feedback. We will make a pass over the paper and fix the typos pointed out and any other errors we find. We address your comments below.

• Clarifying our probing methodology: We agree that the more common way to perform a probing analysis involves learning a linear/non-linear model which takes in the embedding and outputs a label indicating a concept. However, we use probing in more general sense to refer to understand some of the innerworkings of the model. In addition, we want to clarify some confusion about how we do our probing analysis. We DO NOT only evaluate the top-k candidates on the cell the model chooses to predict. Although during its natural course of decoding the model might wish to decode cell location A, we force it (by conditioning) to decode at every other location and evaluate the top-k candidates.

• Probing potential values of one cell while predicting another: We had considered this approach but felt it was harder to get it working for the following reasons. It is unlikely that the same probe will work across different cells (since the embedding conditioned for predicting a particular cell might suppress information about other cells). This might require us to train separate probes for the values of each cell for which the amount of data available becomes sparse. We did try an alternative approach instead where we take the whole sequence of embeddings the model has produced as input to the probe. However, this concatenated embedding becomes a ~25000 dimensional vector which is very high dimensional and makes it hard to train the probe.

• Takeaway of Section 3.5: This section shows the additional gains we get when we use beam search decoding instead of greedy decoding. Beam search decoding with beam width K, maintains K `plausible’ candidates to continue decoding at all times. Note that we prevent a combinatorial explosion with the decoding length by always truncating the list to what we believe are the K most plausible prefixes so far. Using this approach, we show that we are able to further boost the complete puzzle solving accuracy. The main takeaway from this section is that, even in situations where the model’s most likely next token is incorrect, the correct token is in the top-K most likely tokens.

• Wrong Template: We are unsure of what the reviewer meant by the wrong template. We apologize if a formatting error slipped past us. Please let us know what specifically you are referring to and we will fix it.

We answer the reviewer’s technical questions below:

• Showing example data in the appendix: We thank the reviewer for this feedback. Indeed, we agree that this would help explain the challenges in our work better and we will add examples to this end in the Appendix.

• When training with random data, do you include cell location tokens for loss calculations?: We have tried both including it and excluding it. The results are similar in both settings.

• Additionally, how do you ensure that the model predicts every cell during testing?: We deliberately don’t explicitly ensure this. The model is supposed to learn the basic rules of Sudoku as well purely from data. How well they are able to adhere to the rules during inference is a measure of their generalization ability.

• I assume that the input context always includes previously predicted cells and values: Yes this is correct. Note that the input puzzle context length can vary as different puzzles have a different number of filled cells to begin with.

In addition to these comments, we have added additional ablations and a detailed error analysis into where the model struggles in the general response above. We have also shown additional experiments in a different puzzle solving setting - solving Zebra puzzles (also known as Einstein riddles) which shares some similarity with Sudoku puzzles but has many differences. We observe a similar outcome with Zebra puzzles as well where causal language modeling on solver order data is able to teach the model to perform well on these puzzles. Please see the general response for more details on these experiments.

Thank you for your response and the additional results. My concerns are mostly addressed. I like the additional results on the Zebra puzzles and I encourage the authors to include them in the next revision.

• Regarding the probing methodology: I'm still unsure this is the best way to perform probing. However, I agree that the current setup can provide evidence that the model is solving the puzzle in a way that resembles that of a logical solver.
• Additional notes on the difficulty measurement: in Figure 1 of the general response, I notice that the model performs almost perfectly on the easiest puzzles and performs worse and worse as the difficulty increases. Although this behavior is expected, I think it's worth discussing the difficulty distribution of the training and test data. Especially, how will the training data mixing affect the generalization performance to other difficulties? If the model truly learns the reasoning strategy, should we expect it to generalize from hard to simple questions? I feel like there are many interesting points that can be explored here.

Regardless, I'm happy to increase the rating of the current version.