Reasoning with Latent Thoughts: On the Power of Looped Transformers

We thank the reviewer for their valuable comments and feedback, and insightful questions. We address them below.

W1: “Lack a comparison with other categories of models that use looping in the computation (recurrent Transformers, chain-of-thought reasoning)”

A: We would like to clarify that the motivation and goal of this paper is not to come up with the best kind of looping mechanism. Instead we wanted to show evidence and highlight that even the simplest of looping mechanisms can be quite performant in the era of language modeling, if we focus on the right evaluations like reasoning. Having said that, indeed other approaches like recurrent Transformers and CoT reasoning have a similar flavor to looping, and will probably also be effective for reasoning. This is certainly an interesting question, however more detailed evaluation of different methods will require a lot of care and is not the primary motivation of this work. Please also refer to our response to Q3 below for connections with CoT reasoning.

W2: “The results in Figure 1 and Table 4 are not thoroughly explained and are hard to interpret. …. In Figure 1, can you explain how each of the data points are generated?”

A: Thank you for the feedback on this. We had included a brief description for Figure 1 in Section 3.2, however we will make this part clearer in the revision. In Figure 1, for each model we evaluate the log perplexity on x-axis and downstream metrics on y-axis. This is evaluated every 20k steps of training, starting from 120k steps, and each point represents one of the checkpoints during pretraining. The downstream metric on the y-axis is an average over accuracies for a task category, e.g. for closed-book QA the y-axis is the average accuracy for TriviaQA TydiQA-NoContext, Natural Questions and Web Questions. We plot these values in a scatter plot as training proceeds. For visualization we plot the best linear fit between log perplexity and the corresponding downstream metric. Plotting this as training proceeds provides insights into how improvements in perplexity translate to improvements in downstream tasks.

In Table 4, we pretrain a 24-layer 1.5B-parameter model with regularization from Section 4. $k$ denotes the block size used for regularization in Eq 4 and $\lambda_{reg}$ is the strength of regularization. $\lambda_{reg}=0$ should recover baseline training.

Q1: “In Table 4, .. why performance increases when setting $k=4$ but then is reduced when $k$ increases?”

A: This is a good question. While there is no definitive answer for this, our hypothesis is that the inductive bias of looping is stronger where the number of loops is larger. A smaller value of $k$ leads to a higher number of loops ($24/k$). This is most evident from the results on reasoning primitives in Tables 3 and 4 which consistently gets better for smaller values of $k$. Another potential reason could be that hyperparameter choices (learning rate, $\lambda_{reg}$) may not be optimal for the regularization experiments, since they were chosen using the baseline model. It is possible that the optimal learning rate is different for different $k$. We will add this discussion in the paper as well.

Q2: “… the choices for the values of $k$ and $\lambda_{reg}$ are not well justified”

A: The value of $k$ is picked to be divisors of the total number of layers (i.e. 24). We only tried values larger than 4 to save on compute. We also tried various values of $\lambda_{reg}$ in the range of [0.1, 1, 10, 100]. Smaller values were not leading to any inductive bias, while very large values were understandably causing the training to diverge. In the limit of a very large coefficient, we would end up with fully looped models, which we know would suffer on memorization. Hence we just reported 1 and 10 since these provided a useful signal. Since $\lambda_{reg}=10$ worked best for $k=4$, we picked that for all values of $k$. We will include this discussion in the paper.

We thank the reviewer for their comments and for valuing the contributions. We address the main concern about real world datasets below by highlighting the extensive evaluations with a 1.5B parameter language model.

“There is limited evidence on the performance of looped models in real-world applications or datasets… Authors could scale up the experiments, and validate the effectiveness of the looped model in real-world tasks.”

A: We would like to direct the reviewer’s attention the results in sections 3 and 4 wherein we trained a 1.5B parameter language model on a real-world language pretraining dataset (Pile), and performed evaluations on 15 real-world benchmark datasets like TriviQA, Squad, and mathematical reasoning datasets like SVAMP and also 4 additional reasoning primitive datasets. Perhaps we misunderstood what the reviewer meant by “real-world applications or datasets” and we are happy to provide a more detailed answer. Regarding scaling up experiments further beyond 1.5B, we would like to point out that pre-training experiments are heavily compute hungry (unlike fine-tuning setups which can be more easily performed with 7B sized models) and other leading publications such as “Mamba: Linear Time Sequence Modeling with Selective State Spaces” indeed perform pre-training experiments with similar sized models (up to 2.8B parameters) and similar sized pre-training datasets. We strongly believe that our qualitative insights will extend to larger scale models as well but that would require a significantly larger amount of resources.

Thanks for your response. My concerns have been addressed, I will increase my score.

We thank the reviewer for their thoughtful comments and for appreciating the various contributions in the paper. Below we address the main concerns about implementation details for looped models. We hope this convinces the reviewers about the setup and its simplicity, and we are happy to engage more.

W1: “The paper lacks detailed explanations of how the transformer architecture implements the looping mechanism”

A: Thank you for the comment. We will elucidate more details on the looping mechanism in revision. Actually, the looping mechanism is (intentionally) extremely simple and we can simply be viewed as a weight shared model. As described at the start of Section 2, the final model is just a function $f$ composed $L$ times, i.e. $f\circ … \circ f$. Technically $f$ can be any function of choice, with or without residual connection (i.e. $f(x)$ could look like $x + g(x)$) or may or may not have layer norms. The looping mechanism we consider is oblivious to such architectural choices. In our experiments we pick $f$ to be a standard $k$-layer Transformer architecture with residual connection and layer norms. As per your suggestion, we will include some visual description of the looping mechanism to make it easier for the reader to understand.

“.. how the KV cache is managed when layers are looped multiple times”

A: During decoding, we again treat a looped model as a deeper but weight shared model and use different KV caches for each distinct call to the same layer. So in each loop iteration, for a given layer, we use a fresh KV cache entry so that the previous iterations’ caches are not overwritten.

W2: “.. limited discussion on the computational efficiency and scalability of looped models, …. Does looping result in longer training times … or does it impact memory consumption during training and inference?”

A: Training and inference cost for a looped model of effective depth D should be the same as a non-looped model of depth D in terms of FLOPs. In practice, looped models seem to be slightly faster to train, probably because there are fewer parameters to load from memory and also because the backward pass needs to update fewer parameters. This makes looped models especially well suited for extracting higher quality gains in low-memory environments. We will add a discussion of these aspects to the paper. Thank you for your feedback.

Q1: “Given that looped models may have higher perplexity but better reasoning performance, how do you suggest balancing these aspects in practical applications where both are important”

A: The looping-inspired regularization presented in Section 4 provides a path towards getting better reasoning performance without affecting other evals like perplexity and memorization. A more ambitious approach is to design new architectures that can separate the memory component (that can help with perplexity and memorization) from the reasoning component (through looped models) such that they can work in unison. We believe this is an important next direction but is out of the scope of the current work.

Q2: “Can the looping approach be extended to other domains? (vision transformer, other attention mechanisms)”

A: This is a great question and we think this deserves more exploration in future work. Technically the looping mechanism we consider, and its benefits, should be independent of architecture or dataset choices. In this work we restricted all our analyses to text and language models based on Transformer architecture primarily because it is the most common setting used in practice in the context of reasoning.

We thank the reviewer for their useful feedback and appreciation of contributions. The main concerns are addressed below.

“The fact that looped models can improve performance on algorithmic tasks is known and the extension to the tasks tested in the paper does not seem substantial.” “Work done in the group of Tom Goldstein related to reasoning abilities of recurrent models. … Daniel Selsam's paper which also showed the importance of recurrence”

We thank the reviewer for pointing out the above works. Indeed, they are related to our paper and we will add a discussion on them. In section 5, we also discuss other works on looped models that are in a similar vein. We still believe that our results are novel/surprising and not subsumed by these works. We highlight some key high-level and specific differences below.

• High-level: For algorithmic/reasoning tasks in section 2, the focus is not only on solving them well, but also on the depth-optimality of looped models for those tasks. For instance, a 1-layer model looped 8 times can do as well as an 8-layer model on the i-GSM task. This phenomenon is particularly surprising and novel, and is supported by theory and experiments.

• Moreover the major contribution of this work, that is not covered by any of the prior works, is the role of looped models for language modeling in sections 3 and 4. Most of the focus in the literature has been on perplexity and not so much on downstream task performance. However we uncover an inductive bias that looped models have towards reasoning heavy tasks which we believe is also an interesting and novel finding.

• Comparison to the papers from Tom Goldstein’s group: The two papers by Bansal et al. and Schwarzchild et al. study three synthetic tasks: prefix sum, maze solving and chess puzzles and how well looped ResNet model architectures can generalize to harder problems at test time than those seen at train time on these tasks. While they do demonstrate that looping helps improve the performance of these models on the synthetic reasoning tasks, (a) their main focus is showing the surprising result that looping helps you generalize to out-of-distribution task distributions better, (b) they only consider synthetic examples, (c) they consider non-generative models. In our work we identify an inductive bias of looping towards reasoning tasks in generative language models which are based on the Transformer architecture.

• Comparison to Learning a SAT-Solver from Single-Bit Supervision (Selsam et al.): While this paper resonates deeply with our message on the benefits of looped models, it is a highly stylistic setting in that the architecture is specially tuned for the problem of SAT. And the message passing MPNN used in their paper is quite different from a transformer in how a forward pass works and in particular is not a generative model. Given the generality of the Transformer architecture (not tuned for any specific problem) and the generative nature of our setting we believe our results are novel and different from those in this paper.

Q1: “An obvious goal for the looped models would be to figure out how to make the looping adaptive … Did you consider experiments which would test this ability?”

A: Adaptivity in looped models is an important topic, and some prior works like Universal Transformers [Dehghani et al. 2018] have touched upon it for algorithmic tasks. For language modeling, our very preliminary experiments do suggest that some level of adaptivity can be achieved. However a more systematic and thorough study would be required for this and it is beyond the scope of this work. We do believe it’s an important future direction to pursue.

Q2: Why do the plots in Figure 1 contain datapoints for different perplexities for the two compared models?

A: In Figure 1, for each model we evaluate the log perplexity on x-axis and downstream metrics on y-axis. This is evaluated every 20k steps of training, starting from 120k steps, and each point represents one of the checkpoints during pretraining. The downstream metric on the y-axis is an average over accuracies for a task category, e.g. for closed-book QA the y-axis is the average accuracy for TriviaQA TydiQA-NoContext, Natural Questions and Web Questions. We plot these values in a scatter plot as training proceeds. For visualization we plot the best linear fit between log perplexity and the corresponding downstream metric. Plotting this as training proceeds provides insights into how improvements in perplexity translate to improvements in downstream tasks.