Paper: Implicit Chain of Thought Reasoning via Knowledge Distillation (ICLR-2024)

Reviewer comment (Experiments):
Given that the method requires a teacher model that does explicit CoT reasoning to distill into the emulator, it should be evaluated against these models. Unfortunately, the proposed method is weaker than explicit CoT models, although faster. Overall, 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.

Author response:
### 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.

Reviewer comment (Experiments):
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. Could a non-pretrained model be better at this task?

Author response:
### 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:

|                     	| GPT-2 Small | GPT-2 Medium |
|-------------------------|------------:|--------------|
| **With Pretraining**	|         	|          	|
| Explicit CoT        	| 100%    	| 100%     	|
| Implicit CoT        	| 9.5%    	| 96.4%    	|
| **Without Pretraining** |         	|          	|
| Explicit CoT        	| 100%    	| 100%     	|
| Implicit CoT        	| 87.9%   	| 98.2%    	|

Reviewer comment (Novelty):
Eventually, the literature review section is very light. GSM8k has been a popular benchmark for many reasoning methods. The paper could benefit from further discussion on related methods to do multi-step reasoning. In addition, fast inference transformers is also an active research domain. Some discussion about the field should be added.

Author response:
### 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).

Reviewer comment (Theory):
In Section 3.1, the $L \times T$ matrix of the teacher model comes from a Transformer architecture, which (by default) has an attention matrix over **all** previous layer tokens. So the assumption that “_Progressing diagonally, from z11 to zLL, we gradually add more intermediate tokens and layers_” is not always true. For this assumption to be true, you also need to assume that the attention matrix of the teacher transformer is autoregressive, i.e., conditional from left-to-right. Such clarification should be added. It is only clear that this is the case by looking at the architecture chosen (GPT-2) which is indeed auto-regressive.

Author response:
### 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.

Reviewer comment (Experiments):
The proposed method may not generalize compositionally to questions requiring more reasoning steps or just out-of-distribution data. It forces the model to conduct CoT with limited computation.

Author response:
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).

Reviewer comment (Experiments):
Besides maths problems, does implicit chain-of-thought also work on other types of reasoning tasks? If not, what are the barriers that implicit CoT faces?

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

Reviewer comment (Experiments):
**Missing an important baseline.** The proposed method can be regarded as **compressing CoTs into a vector**. There could be multiple ways to compress CoTs into a single vector and then train another model to generate the compressed representations for CoTs during the inference process. Why it is the best choice to extract the intermediate states of the teacher model as the compression results? From my perspective, a more elegant way to compress CoTs should be to train an auto-encoder to map the original CoTs into a single vector. The auto-encoder ensures that the compressed vector contains all the necessary information to recover the CoTs and can be used to predict the answer directly. Specifically, denoting the teacher model as $f_{T}(\cdot; \theta_{T})$ with parameter $\theta_{T}$, the auto-encoder model as $g_{enc}(\cdot;\phi_{enc})$ and $g_{dec}(\cdot;\phi_{dec})$ with parameter $\phi$, the student model as $f_{s}(\cdot; \theta_{s})$. We extract the CoTs from the teacher model given a question $q$ as $c = f_{T}(q;\theta_{T})$ where $c = (c_{1}, \dots, c_{N})$ is the CoT. Then the auto-encoder is trained with the objective 

   $$\max _{\phi} P(\hat{c _ i}|z, \hat{c} _ {<i}), \text{where} z=g(c;\phi _ {enc}) \text{and} P(\hat{c _ i}|z, \hat{c} _ {<i}) = g _ {dec}(z, \hat{c} _ {<i};\phi _ {dnc})$$

   Then the student model is trained to maximize the generation probability of the target answer conditioned on the question and $z$. The proposed method is exactly a special case of the above framework, where the auto-encoder is replaced by the proposed information extraction method. The author should demonstrate why the proposed method outperforms the general framework above, or why we should choose the proposed design. Otherwise, the method is too intuitive.

Author response:
### 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.

Reviewer comment (Experiments):
**Insufficient empirical evaluation.** The authors only verify that implicit CoT can boost performance. However, it is still not comparable with standard CoTs. Also, the efficiency improvement is less significant (or necessary) considering the poor performance of implicit CoTs. Finally, I would like to see further discussions on why should we consider implicit CoTs besides the reason for efficiency, especially considering the significantly increased training cost and data collection cost.

Author response:
### 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).