Self-Reasoning Tokens, teaching models to think ahead.

🌈 Abstract

The article discusses the concept of "Reasoning Tokens" - a way to teach language models like ChatGPT to think ahead and plan for future tokens, rather than just predicting the next token. The author shares the results of experiments that explore this idea and its potential benefits.

🙋 Q&A

[01] Introduction

1. What is the key insight from the "Interpretability in the Wild" paper that the author builds upon?

• The author notes that the computation of the next token in a language model includes information processed in previous steps, indicating that the model expends internal "cognitive power" to process and store information that will be useful for predicting tokens beyond just the very next one.

2. What did the "Do Language Models Plan for Future Tokens?" paper find?

• The paper found a small performance gap when the model was "myopic" or incapable of planning for future tokens, suggesting that while language models do plan ahead, most of their power is used to predict only the next word in the sequence.

[02] Reasoning Tokens

1. What is the key idea behind the "Reasoning Tokens" experiment?

• The author introduces "Reasoning Tokens" - where the model produces two tokens for each token in the original sequence. The first token is used to predict the next token, while the second token duplicates the input of the first one but does not receive a gradient "answer" from the very next token, only from future tokens. This incentivizes the model to "pre-cache" or store information that is useful for the future.

2. What were the results of the "Reasoning Tokens" experiment?

• The results showed a 35% reduction in loss, from 0.621 to 0.401, for a GPT-2 Small model trained on the Open Web Text Corpus. This validates the hypothesis that models can be taught to plan for the future.

[03] Next Steps

1. What is the author's hypothesis for the next experiment with Reasoning Tokens?

• The author is experimenting with Reasoning Tokens in fine-tuned instruction following models, where the model can choose when to start the internal reasoning process before producing the next word in the sequence. The hypothesis is that Reasoning Tokens can substitute and outperform models where a "step by step" explanation is included in the training phase.

2. How does the author see Reasoning Tokens fitting into Mixture of Experts (MoE) models?

• The author suggests that Reasoning Tokens would be a great fit for MoE models, where there could be an expert just for the reasoning phase, allowing the model to leverage the internal mathematical dimensions in a way that may not necessarily make sense to humans.