Stutter makes large language models smarter

Dear Reviewer TpjX,

Thank you for your valuable feedback and for taking the time to review our paper. We appreciate your insightful comments and suggestions, which have helped us improve the quality of our work. Below, we address each of your concerns in detail.

Weakness

1. Comparison with State-of-the-Art LLMs: We understand the importance of evaluating our approach against prominent models such as the LLAMA series. We are currently experimenting with the stutter mechanism on LLAMA 1B. We will provide an update with detailed results and comparisons as soon as we have them.

2. Scalability Analysis on performance and efficiency:

   • Performance: As the model size increases, the stutter mechanism continues to leverage the hidden states from the first pass to enhance the model's ability to utilize information effectively. Larger models typically have more capacity to learn and represent complex patterns, which can lead to improved performance on various tasks. Our preliminary experiments indicate that the stutter mechanism provides performance improvements across different model sizes, including larger models.

   • Efficiency: The stutter mechanism is designed to maintain the same time complexity as the base model by employing a two-pass process. This ensures that the computational complexity does not increase significantly with larger models. The stutter mechanism's efficiency, in terms of computational overhead, remains manageable even as model size increases, as it only involves an additional linear attention mechanism in the second pass.

     To implement the stutter mechanism, we increase the model size by approximately 10% by adding another attention mechanism. Regarding scalability, the time complexity of our method scales linearly with model size. In our paper, we experimented with models of sizes 160M, 410M, and 1B parameters. The stutter mechanism can be applied to larger models, such as 7B, 14B, or even bigger models, due to its efficient time complexity. This linear scaling ensures that as the model size increases, the performance and efficiency remain manageable, making our approach suitable for a wide range of model sizes.

3. Long-Term Training Stability: When fine-tuning, it is common to use around 1 billion tokens. In our approach, we freeze the original model weights and only train the token retrospect part, which consists of about 10% of the original parameter. During our experiments, we observed that increasing the number of tokens to 2 billion provided a slight improvement in performance. However, these performance gains were marginal and eventually plateaued, with some indications of diminishing returns.

   This suggests that while the stutter mechanism can benefit from additional training data, the improvements are not substantial beyond a certain point. The training remains stable over extended periods and larger datasets, but the efficiency of additional training diminishes, indicating that our method is robust and does not require excessive amounts of data to achieve optimal performance.

Thank you again for your valuable feedback. We hope that the additional clarifications and improvements we have made to our paper address your concerns. We kindly request that you consider these enhancements when re-evaluating our paper and feel free to let us know if you have further questions.

4. Extra time cost:

   • Training time complexity: As mentioned above, during training, if the base model has time complexity of O(n), the time complexity of the stutter model is 2 x O(n), remaining the same time complexity as the base model.

   • Inference time complexity: In the benchmark dataset we use for inference, the only token we use for stutter is the choice of the samples. In the Lambada-openai dataset, it only predicts the last token. In the multiple-choice dataset, the only tokens we stutter are the choice tokens.

     For example, in the Lambada dataset, given the context:

     • Given: "He heard Rihanna speak 'The Queen wants you in her carriage.' Tom spoke 'No, I’m not going in some asylum.' Ran was seen standing next to him spoke 'It’s just for a private talk with you that’s all.' Tom groaned and went inside the carriage to sit down next to the"

     • The goal is to predict "Queen"

     • So we actually just stutter at the last token.

     For another example from multiple choice dataset:

     • Question: "George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat?"
     • Choice: ["dry palms", "wet palms", "palms covered with oil", "palms covered with lotion"]

     The test was to concatenate the question and each answer, so there will be four inputs, and the model calculation is given by the sum of the lowest log probability of the choice part:

     1. "George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat? dry palms"

     2. "George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat? wet palms"

     3. "George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat? palms covered with oil"

     4. "George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat? palms covered with lotion"

     For inference, we will only stutter the choice part (which is "dry palms", "wet palms", "palms covered with oil", and "palms covered with lotion") or the last token. Therefore, the time complexity for inference with the stutter mechanism is  O(n) for the first pass and  O(1) for each stuttered token in the second pass. The overall time complexity remains O(n), similar to the base model. This ensures that the stutter mechanism does not significantly increase the computational complexity during inference.

Questions:

1. How does the stutter handle with the sequence length problem? see weakness #3

2. As the paper mentioned, each token stutters once, and the strategy of stutter token selection is out of the scope of the paper, why stuttering/repeating each token once can work? The reason why stuttering/repeating each token once can work is that it allows the model to re-evaluate and refine its predictions by considering the context more thoroughly. This prototype demonstrates that the "think again" approach can be effective. Future work will focus on optimizing the token selection strategy to further enhance the performance and efficiency of the stutter mechanism. This initial prototype serves as a proof of concept, showing that the stutter mechanism has the potential to improve model predictions by allowing for an additional layer of attention and refinement.

3. How does the stutter methods impact the time complexity in both training and inference? see weakness #4

Thank you again for your valuable feedback. We hope that the additional clarifications and improvements we have made to our paper address your concerns. We kindly request that you consider these enhancements when re-evaluating our paper and feel free to let us know if you have further questions.

Dear Reviewer ujpr,

Thank you for your valuable feedback and for taking the time to review our paper. We appreciate your insightful comments and suggestions, which have helped us improve the quality of our work. Below, we address each of your concerns in detail.

Weakness:

1. The reported performance improvement is not convincing: We appreciate the reviewer's feedback regarding the performance improvements reported in our paper. While some differences in performance metrics, such as those in Table 2 for LogiQA (0.230 vs 0.215) and SciQ (0.892 vs 0.894), may appear marginal, we would like to highlight the following points:

   • Proof of Concept: Our primary goal is to introduce and validate the stutter mechanism. The current results demonstrate its potential to enhance model performance.

   • Notable Improvements: In specific benchmarks, such as the WSC and WinoGrande datasets, Pythia-410M-Stutter achieved performance close to or even better than Pythia-1B. For example, Pythia-410M-Stutter achieved an accuracy of 0.670 on the WSC benchmark compared to Pythia-1B’s 0.666.

   • Future Work: We plan to optimize the token selection strategy and explore different attention mechanisms to achieve more significant improvements.

   • Broader Impact: The stutter mechanism has shown consistent benefits across various tasks, indicating its broad applicability and potential for enhancing model performance.

   We are committed to further refining our approach to achieve more substantial and convincing improvements.

2. The experimental design is not quite fair: To clarify, the additional 1 billion tokens used for the stutter models are selected from the original Pythia training dataset (Biderman et al.), which is derived from the Pile dataset. Consequently, the stutter models do not encounter any new tokens compared to the base models. Instead, they are exposed to the same 1 billion tokens once more, as they have already been trained on these tokens previously. To address the concern about the fairness of the experimental design, we conducted an experiment where the base model was also trained with the same 1 billion tokens and plan to provide result next week. This ensures that any performance fluctuations can be attributed to the inclusion of the extra 10% token-retrospect layers rather than the additional training data.

3. The details of implementation are missing (sequence length problem): For training, our settings is to stutter at every token. Our method employs a two-pass process for training. For a given sample ( $X = {x_1, x_2, ..., x_n}$ ):

   • First Pass: We store the hidden states (${h^1_l*, h^2_l*, ..., h^n_l*}$) for each token in the sequence.

   • Second Pass (Stutter Phase): During this phase, we utilize the stored hidden states by passing both the stored hidden states and the current token to the model. For instance, to generate the final prediction for the next token of ( $x_3$ ), we input ( $x_2$ ) and ( $h^2_l$* ) to the stutter model. The original Feed-Forward (FF) and Attention (Attn) mechanisms perform the same as in the base model, where ( $x_2$ ) will attend on ( $x_1$ ). Our token retrospect mechanism performs linear attention using ( $x_2$ ) and ( $h^2_l*$ ) to extract information from the first pass to generate the next token.

   This approach ensures that the sequence length remains unchanged, and the model can effectively utilize the stutter mechanism without increasing the computational complexity associated with longer sequences but still able to utilize the information in the first pass.

Questions

(1) For the second pass, is the token retrospect allowed to attend over previously generated token hiddens in the second pass? Why or why not?

No, in our current implementation, the token retrospect is not allowed to attend over previously generated token hiddens in the second pass. We aimed to keep our mechanism simple and effective in this prototype. However, we acknowledge that allowing the token retrospect to attend over previously generated token hiddens could be a valuable enhancement. We plan to explore different kinds of attention mechanisms in future work to further improve the performance and flexibility of our approach.

(2) Can you provide more info on the continual pretraining? What hardware is used? For the 1B token training dataset, how many epochs are used? Do you see any benefit with more tokens or longer training?

Here are the details:

• Hardware: We utilized 4 NVIDIA A6000 GPUs for the continual pretraining process.

• Training Dataset: The 1B token training dataset was processed for 1 epoch.

• Hyperparameters:

  • Learning rate: 5e-5
  • Scheduler: Cosine scheduler with warmup
  • Warmup ratio: 0.01
  • Optimizer: Adam
  • Gradient accumulation steps: 8

• Token and Training Duration:We observed that increasing the number of tokens to 2B provided a slight improvement in performance. However, the performance gains were marginal and did not justify the additional computational cost. Therefore, we found that training with 1B tokens strikes a good balance between efficiency and performance.

Thank you again for your valuable feedback. We hope that the additional clarifications and improvements we have made to our paper address your concerns. We kindly request that you consider these enhancements when re-evaluating our paper and feel free to let us know if you have further questions.

Weakness

• Performance vs computation costs:

  • Inference: In the benchmark dataset we use for inference, the only token we use for stutter is the choice of the samples. In the Lambada-openai dataset, it only predicts the last token. In the multiple-choice dataset, the only tokens we stutter are the choice tokens. For example, in the Lambada dataset, given the context:
    • Given: "He heard Rihanna speak 'The Queen wants you in her carriage.' Tom spoke 'No, I’m not going in some asylum.' Ran was seen standing next to him spoke 'It’s just for a private talk with you that’s all.' Tom groaned and went inside the carriage to sit down next to the"
    • The goal is to predict "Queen"
    • So we actually just stutter at the last token.

   For another example from multiple choice dataset:

  • Question: "George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat?"
  • Choice: ["dry palms", "wet palms", "palms covered with oil", "palms covered with lotion"]

  The test was to concatenate the question and each answer, so there will be four inputs, and the model calculation is given by the sum of the lowest log probability of the choice part:

  1. "George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat? dry palms"
  2. "George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat? wet palms"
  3. "George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat? palms covered with oil"
  4. "George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat? palms covered with lotion"

  For inference, we will only stutter the choice part (which is "dry palms", "wet palms", "palms covered with oil", and "palms covered with lotion".) or the last token. Therefore, the time complexity for inference with the stutter mechanism is  O(n) for the first pass and  O(1) for each stuttered token in the second pass, the overall time complexity remains O(n), similar to the base model. This ensures that the stutter mechanism does not significantly increase the computational complexity during inference.

• Hyperparameter robustness:  We will provide more details on the selection and sensitivity of the few-shot examples and the 1B token training dataset from Pile. In our current setting. For the current 1B dataset used, it was randomly selected from the shuffled Pile dataset. The Pile is a diverse and comprehensive dataset that includes a wide range of domains such as web pages, academic papers, books, code repositories, legal documents, forums, patents, medical abstracts, literature, subtitles, encyclopedic entries, mathematical content, chat logs, and more. This variety ensures that the training data covers a broad spectrum of topics and styles, contributing to the robustness and generalizability of the model

• Inclusion of LMs from other families: We understand the importance of evaluating our approach against prominent models. We are currently experimenting with the stutter mechanism on LLAMA 1B and plan to provide an update with detailed results and comparisons next week.

Dear Reviewer r8Jr,

Thank you for your valuable feedback and for taking the time to review our paper. We appreciate your insightful comments and suggestions, which have helped us improve the quality of our work. Below, we address each of your concerns in detail.

Weakness

1. The exposition of the paper requires substantial improvements:

• Cross layer parameter sharing: Thank you for pointing out the need to discuss cross-layer parameter sharing. While the primary goal of the work by Lan et al. in "ALBERT: A LITE BERT FOR SELF-SUPERVISED LEARNING OF LANGUAGE REPRESENTATIONS" is to reduce memory usage, our focus is on improving performance through the stutter mechanism. We will include a citation and discussion of this relevant work in our updated paper to provide a comprehensive context for our approach.

• Citation format: We will update the citations in our paper to follow the ICLR recommendations, using \citep where appropriate.

• Proper citations for models/methods/datasets: We will provide proper citations for all models, methods, and datasets used in the paper, including those mentioned on line 066 and line 287. This will ensure that the experimental setup is clear and the comparisons are meaningful.

• Figure 1 update: We will update Figure 1 to reflect that the entire input token sequence is used for the second pass, providing a clearer representation of our method.

2. The experiment setup is problematic without enough convincing evidence:

• Inconsistent Improvements and lack of insights: We appreciate the reviewer's feedback regarding the varying improvements observed across different LMs and datasets. We acknowledge that while there are certain improvements, such as in the WSC benchmark, our paper does not delve deeply into the specific cases or statistical significance of these results. We are actively investigating these aspects to gain a better understanding of the improvements and their underlying causes and will provide detailed insights next week.

• Performance vs computation costs:

  • Training: In our current setting in training, we stutter at every token. Our method employs a two-pass process. For a given sample ( $X = {x_1, x_2, ..., x_n}$ ):

  1. First Pass: We store the hidden states (${h^1_l*, h^2_l*, ..., h^n_l*}$) for each token in the sequence.

  2. Second Pass (Stutter Phase): During this phase, we utilize the stored hidden states by passing both the stored hidden states and the current token to the model. For instance, to generate the final prediction for the next token of ( $x_3$ ), we input ( $x_2$ ) and ( $h^2_l*$ ) to the stutter model. The original Feed-Forward (FF) and Attention (Attn) mechanisms perform the same as in the base model, where ( $x_2$ ) will attend on ( $x_1$ ). Our token retrospect mechanism performs linear attention using ( $x_2$ ) and ( $h^2_l*$ ) to extract information from the first pass to generate the next token.

  If the base model has time complexity O(n), the time complexity of stutter model has time complexity 2 x O(n) since we run each sample twice, remaining the same time complexity as the base model.

(4) The authors did not include discussion/comparison to a very relevant method: pause tokens (Goyal et al.): Thank you for pointing out the relevance of the pause tokens method by Goyal et al. We appreciate the opportunity to discuss and compare our approach with theirs.

While pause tokens are indeed an efficient method that allows encoding the prefix in parallel, there are some key differences and advantages to our stutter mechanism:

1. Efficiency in Pretraining: The pause tokens method involves pretraining on the entire C4 dataset, which requires substantial computational resources and time. In contrast, our stutter mechanism is designed to be more efficient, requiring only 1 billion tokens for pretraining. This makes our approach more accessible and less resource-intensive.

2. Parameter Updates: The pause tokens method updates all parameters during training, which can be computationally expensive. Our stutter mechanism, on the other hand, updates only a small subset of parameters, making it more efficient in terms of both computation and memory usage.

3. Pause tokens: The pause tokens method introduces special tokens into the input sequence, which Goyal et al. themselves noted may not provide meaningful information to the model. In contrast, our stutter mechanism uses the hidden state of the previous token, allowing the model to "think again" and continue reasoning from the previous context. This approach provides a more coherent and continuous reasoning process, enhancing the model's ability to utilize prior information effectively.

We acknowledge the strengths of the pause tokens method and will include a discussion and citation of this relevant work in our updated paper. We believe that our approach offers a complementary and efficient alternative, particularly for scenarios with limited computational resources.

Thank you again for your valuable feedback. We hope that the additional clarifications and improvements we have made to our paper address your concerns. We kindly request that you consider these enhancements when re-evaluating our paper and feel free to let us know if you have further questions.

Dear Reviewer fJGV,

Thank you for your valuable feedback and for taking the time to review our paper. We appreciate your insightful comments and suggestions, which have helped us improve the quality of our work. Below, we address each of your concerns in detail.

(1) The selected tasks mostly don't require complex reasoning: The benchmark datasets we used are aligned with those reported in the Pythia paper, and they do include reasoning tasks. However, we acknowledge the importance of evaluating our method on more complex reasoning tasks. We are currently exploring more complex datasets like GSM8K and synthetic tasks that require multi-hop reasoning to further assess the effectiveness of our approach. We appreciate your suggestion and are committed to providing a more comprehensive evaluation in future work.

(2) The gain is not very consistent or significant across different tasks. In this case, the authors should also report variance to show the significance of the results: To address this concern, we plan to report the variance of our results to provide a clearer picture of their statistical significance. We will measure the variance by training multiple models with different random seeds and reporting the variance of the results. This approach will help us understand the stability and reliability of the performance improvements introduced by the stutter mechanism. We are currently conducting these additional experiments and will include the variance in our updated results next week to provide a more comprehensive evaluation of our method's effectiveness.

(3) Inefficiency in inference: In the benchmark dataset we use for inference, the only token we use for stutter is the choice of the samples. In the Lambada-openai dataset, it only predicts the last token. In the multiple-choice dataset, the only tokens we stutter are the choice tokens. For example, in the Lambada dataset, given the context:

• Given: "He heard Rihanna speak 'The Queen wants you in her carriage.' Tom spoke 'No, I’m not going in some asylum.' Ran was seen standing next to him spoke 'It’s just for a private talk with you that’s all.' Tom groaned and went inside the carriage to sit down next to the"

• The goal is to predict "Queen"

• So we actually just stutter at the last token.

    For another example from multiple choice dataset:

• Question: "George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat?"

• Choice: ["dry palms", "wet palms", "palms covered with oil", "palms covered with lotion"]

The test was to concatenate the question and each answer, so there will be four inputs, and the model calculation is given by the sum of the lowest log probability of the choice part:

1. "George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat? dry palms"

2. "George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat? wet palms"

3. "George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat? palms covered with oil"

4. "George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat? palms covered with lotion"

For inference, we will only stutter the choice part (which is "dry palms", "wet palms", "palms covered with oil", and "palms covered with lotion") or the last token. Therefore, the time complexity for inference with the stutter mechanism is  O(n) for the first pass and  O(1) for each stuttered token in the second pass. The overall time complexity remains O(n), similar to the base model. This ensures that the stutter mechanism does not significantly increase the computational complexity during inference.

Dear Reviewer fJGV,

We have updated the paper's figures and descriptions to better reflect the mechanism. Additionally, we have included a discussion and references of the pause token (Goyal et al.) in the updated version. We also include the time complexity of training and inference in the appendix. Thank you again for your valuable insights.