Implicit Chain of Thought Reasoning via Knowledge Distillation

Thank you for your comments.

Re: this method trades interpretability and performance (explicit CoT) for speed, and it doesn't seem like a good trade-off given the extensive literature on making faster inference Transformers

Our method is orthogonal to faster inference methods such as quantization and pruning (although not speculative decoding). Also, efficiency is not our only goal — our deeper aim is to break from the notion that LMs must adhere to human-like reasoning. Instead, we hope to enable LMs to develop their own reasoning pathways, leveraging their capabilities quite different from those we provide them with.

Re: Another weakness of the proposed approach is that it requires significantly more training data. This is potentially due to the much different way of training the student network compared to the pre-trained model as mentioned by the authors.

We suspect that explicit CoT benefits more from pretraining than implicit CoT, since the pretraining data likely contains similar explicit CoT examples (which is supported by the experiment of training models from scratch below). That said, we used the same amount of training data for all experiments, so the comparison is still fair.

Re: Did you try to train a model from scratch

Thanks for the insights! We conducted experiments on $5\times5$ Mult where all models are trained from scratch. We found that the teachers (Explicit CoT) can still reach 100% accuracy. Interestingly, training from scratch improves the performance of implicit CoT: e.g., implicit CoT reaches 87.9% accuracy under no pretraining (using GPT-2 Small), much higher than with pretraining (9.5%). We suspect that pretraining for explicit CoT might position the models in a mode not amenable for implicit CoT reasoning. We have included the full results in the below table:

Re: Literature review is light

We have added more related work on multi-step reasoning and efficient transformer inference methods in the revised draft (Appendix F).

Re: For this assumption to be true, you also need to assume that the attention matrix of the teacher transformer is autoregressive

We have clarified in the updated paper (Section 3.1 - Information Extraction) that we assume that the teacher is autoregressive when we discuss the intuition behind taking the diagonal elements.

Re: it may be the case that doesn’t reach the last token in the sequence

Thanks for attending to this detail! Under the dynamic $\Delta$ setting, the last position taken will always be T, since the last layer is $l=L$, and $\Delta=\frac{T-1}{L-1}$, so $t_l=1+\Delta(l-1)=1+T-1=T$. In your example, L=4, T=3, the positions before rounding would be [1, 1.67, 2.33, 3], so after rounding the last one will be 3 instead of 2.

That said, we do have a typo in the paper: instead of taking the floor in computing $t_l$ (Section 3.1 - Variable Length Chain-of-Thought), we actually rounded it to the nearest integer in our implementation (see https://anonymous.4open.science/r/implicit_chain_of_thought/src/models/teacher.py line 48). We have fixed this typo by adding 0.5 such that taking the floor is the same as rounding.

Re: Did you try to include the original explanation? What was the impact on performance?

In response to your question, we conducted an experiment on GSM8K using GPT-2 Small with implicit CoT. We found that incorporating the full explanation into the model resulted in a reduced validation accuracy of 15.6%, compared to 20.2% when only equations were used. We suspect that equations provide reasoning information in a more consistent and uniform manner. In contrast, combining text with equations introduces a degree of non-uniformity, where the numerical and symbolic components of equations carry more substantive information relative to the accompanying textual words. This discrepancy in information density between text and equations might be a contributing factor to the observed difference in model performance.

Re: Why not “last column”?

Thanks for pointing this out. We have added the last column result in the updated paper. The result is 58.7%, which is worse than taking the diagonal (100%).

Thank you for your comments.

Re: it is hard to tell whether indeed the proposed system is conducting CoT reasoning or is simply learning some reasoning shortcuts

Since the model trained for implicit CoT can do 5-by-5 multiplication, a task for which we don’t know of reasoning shortcuts, it is very likely that the model has learned to conduct reasoning internally.

Re: The proposed method may not generalize compositionally to questions requiring more reasoning steps or just out-of-distribution data

We are not making a claim that implicit CoT can generalize compositionally or to out-of-distribution data. Also, explicit CoT methods face similar issues (https://arxiv.org/abs/2305.18654). As one example, if we trained a transformer model on 4-by-4 multiplication, it is unlikely that it can do 5-by-5 multiplication. As another example, language models are known to not generalize to out-of-distribution data (https://openreview.net/forum?id=73OmmrCfSyy).

Re: It forces the model to conduct CoT with limited computation

We agree, and from empirical evidence it seems that implicit CoT needs a certain number of layers for tasks requiring more reasoning steps (such as GPT-2 Small on $5\times5$ multiplication). However, our method shows that there is untapped potential in large language models like GPT-4, which might have much more number layers than the number of reasoning steps in most tasks.

Re: Does implicit chain-of-thought also work on other types of reasoning tasks

We think it has the potential to do so, though demonstrating that is beyond the scope of this first paper.

Thank you for your comments.

Re: matrix $z$ is often not a square matrix. In that case, what does it mean to extract the elements from $z_{11}$ to $z_{LL}$?

We presented how to deal with the case of (almost certainly) non-square matrices in the paragraph next to “Information Extraction” (page 5 - Variable Length Chain-of-Thought). In short, we introduced a hyperparameter to specify the distance between the columns taken, with two variants, one of which used a variable spacing which is determined by the ratio of reasoning steps to number of layers. In this variant, when there are more tokens than layers ($T>L$), the spacing will be greater than 1; when there are more layers than tokens ($T<L$), the spacing will be less than 1 (the spacing is a float number but the actual position taken is rounded to an integer). In both cases, the last vector is taken from the last column of the teacher's hidden states.

Re: Missing an important baseline

We agree that the baseline you mentioned is more powerful than our current information extraction method. However, we wanted to show that even with the simple method we can realize the concept of implicit chain of thought. We consider our method as one step towards implicit chain-of-thought and we believe there are a lot of improvements that can be made on top of our current method, including the one that the reviewer posits.

That said, we have conducted an experiment on 5X5 multiplication following your suggestions. In particular, we view the teacher states as a sequence of column-wise vectors $c_1, c_2,\cdots c_T$ following your notation. For the encoder, we used a BERT followed by mean pooling and layer normalization (we replaced the embedding layer with a linear projection layer to make shapes match). For the decoder, we used a GPT-2 model (again, we replaced the embedding layer with a linear projection layer to make shapes match). The auto-encoding objective is to predict the teacher hidden hidden states at the next timestamp, conditioned on past teacher hidden states and the encoded teacher states, where the conditioning is implemented via addition to all input hidden states. We then use the encoded teacher states to replace the extracted diagonal elements for both training the mind-reading student and the emulator. We conducted experiments on $5\times5$ Mult and found that GPT-2 Small got 0.2% accuracy and GPT-2 Medium got 0.7% accuracy even using oracle teacher states (code can be found at https://anonymous.4open.science/r/implicit_chain_of_thought/src_autoencoder/train_mind_reading_student.py). That said, this doesn’t mean that this approach is worse — there’s plenty of room for architectural optimization and hyperparameter search which we didn’t have time to exhaust. We consider this approach a potentially improved version of our current implementation, which is only a preliminary step towards realizing implicit chain-of-thought.

Re: Not comparable to explicit CoT

We agree that our method has lower accuracy than explicit CoT. Besides the higher inference throughput, our deeper aim is to break from the notion that LMs must adhere to human-like reasoning. Instead, we hope to enable LMs to develop their own reasoning pathways, leveraging their capabilities quite different from those we provide them with. There could also be interesting follow-up work from ours, such as a hybrid mode of explicit/implicit reasoning, where LMs could use implicit CoT for quick responses when confident, and switch to explicit CoT for challenging queries, or even use the internal reasoning to solve multiple problems at once (similar to the data multiplexing task in https://arxiv.org/abs/2302.12441).