Token Assorted: Mixing Latent and Text Tokens for Improved Language Model Reasoning

We thank the reviewer for the comment and reply below.

#1

The experimental designs are overall sound. However, I have one question regarding Section 4.2.2, where you mention selecting the learning rate based on the lowest validation error. However, the MetaMathQA dataset only provides a training set and does not include a validation set. This raises some confusion regarding how the validation process was conducted. Could you clarify this?

Thanks for this question. Yes, for the MetMathQA dataset, we split it into 80% train and 20% validation sets. We tuned our hyper-parameter based on the validation set. Once the hyper-parameters were chosen, we retrained the model using the complete dataset (100% of available data) and reported test results from this final model. We will update the paper to clarify this.

#2

The authors of the MetaMath paper chose 3 epochs for fine-tuning, whereas you opted for 1 epoch. I am concerned about whether the baseline models are sufficiently trained in your experiments. Could you clarify this choice?

We actually regenerated the MetaMath dataset with the Llama3-405B-inst model (as described by line 251 on our paper), instead of the original MetaMath dataset . The original dataset was generated using GPT-3.5 Turbo, whereas we enhanced the quality by regenerating responses with the more powerful Llama-405B-inst model. In doing this, we distill the knowledge of Llama3-405B-inst into our smaller Llama3.1 ~ 3.2 series model. In this enhanced dataset, we observed that the smaller Llama models (1B/3B/8B) started to overfit after epoch 1 as is shown in the validation set, and thus, we train with 1 epoch for all of them.

#3.

Could you provide an efficiency analysis that accounts for both the reduced number of generated tokens and the additional parameters introduced by the VQ-VAE?

Thank you for raising this important point about efficiency analysis. Our method offers a favorable trade-off between token efficiency and additional parameters: The VQ-VAE introduces only 50M parameters (0.05B), which represents just 1.7% overhead for the Llama-3.2-3B model and an even smaller 0.6% for the Llama-3.1-8B model. This modest parameter increase is significantly outweighed by the efficiency gains of 20% reduction in token length for Llama3.2-1B and 3B models. And for Llama-3.1-8B model, it achieves 10% reduction in token length

Importantly, the VQ-VAE is only used during training. During inference, the LLM directly generates the latent tokens without requiring the VQ-VAE, resulting in pure computational savings with no additional inference overhead. These token reductions translate directly to proportional decreases in both inference time and computational cost, making our approach particularly advantageous for deployed reasoning systems.

For the ablation study on the compression-ratio (r), on the Llama3.2-3B:

With Avg Acc = the average math score across all math benchmarks as in Table4.2. The graphical result is here: https://imgur.com/a/iGB2TvU. Importantly, all settings of latent compression (r=2/16/32) outperform the CoT baseline. In general, a smaller r results in less abstract representations, leading to longer token sequences. Conversely, larger r values cause over-compression, which reduces sequence length but also degrades accuracy.

#4.

The use of latent tokens reduces explainability and readability. How do you address this concern?

Thank you for raising this concern. We have looked at the output of our latent-LLM. It seems that it strategically utilizes latent tokens at the beginning of decoding, serving as compact, high-level guides for the subsequent reasoning process. However, we emphasize that text-based reasoning still follows after these latent tokens.

Furthermore, to directly address readability and interpretability of these latent codes, we can explicitly transform the latent tokens back into the text tokens using the decoder from the VQ-VAE. In fact, we have explored this and observed that the decoded latent representations are indeed meaningful and interpretable, providing additional insight into the reasoning abstractions captured by the latent codes.

Please check the examples on our responses to Reviewer 2JVK (#4) (due to character limit)

#5.

Dart-math and MetaMath comparison

Thanks for pointing this out. Despite that both datasets have similar accuracy (with our latent approach being +0.4 point better), we still see an overall improvement of token efficiency of 16%. Using the latent approach still shows a good advantage.

#6.

Mean and variance

With 3 seeds, we compute the pass@1 for our math benchmark metrics. The improvement still holds, please see our results here: https://imgur.com/a/V3ijL3g

We thank the reviewer for the comment and reply below.

#1.

I could not find the final training and testing loss of the VQ-VAE. Also, some ablation on the codebook size or the chunk size would be useful.

Yes, the training loss and the testing loss of the VQ-VAE is 1.21 and 1.25, respectively. In general, the greater the compression ratio (r), the lower the accuracy. We have conducted an ablation study to examine the effect of codebook size (see #1 on our response to Reviewer yHAN). Regarding compression-ratio, we conducted an ablation study of this on Llama3.2-3B model, please see the result on our response to Reviewer rPRY (point # 3).

#2.

How many latent tokens does the model produce on average? Is there a tradeoff between number of latent tokens and accuracy, and if so, how could the tradeoff be managed?

On average, the model outputs 3.11 latent tokens on average with the overall distribution between 0 ~ 8 latent tokens. In general, we see that the longer the latent tokens, the higher the accuracy. One way to manage / trade-off this number is to manually trigger / cancel the start and end of the latent-tag.

#3.

It would be interesting to see which samples the model uses more latent tokens for and which it uses fewer. Are there any patterns to these samples, and has the model learned how to effectively leverage latent tokens?

Yes, we have analyzed the distribution of latent token lengths across different datasets. We found that in the gsm8k dataset (considered the "easiest" benchmark since it achieves the highest scores in our LLaMA model experiments), there is generally a larger number of latent tokens, averaging 3.21. In contrast, the more challenging benchmarks, such as the olympiad-math which is marked by the lowest scores, show an average of just 3.00 latent tokens. This pattern indicates that the model tends to leverage more latent tokens in relatively easier benchmarks where it can already form a strong representation in the text space, while it invokes fewer latent tokens in more difficult domains. For future directions, we could enhance model performance by incorporating more latent samples into the training data weighted by the difficulty of the sample.

#4

If the learned discrete latent tokens are decoded back to text, what is the produced text and is it interpretable?

To examine the interpretability of the latent-codes, we use the decoder from the trained VQ-VAE to convert the latent-code back to the text space. In general, it is interpretable. We bolded the decoded part of the text from the latent code, which is delimited by the <boLatent> ... <eoLatent> tags. In general, we find the latent tokens captures long/verbose sentences which may trigger useful follow-up thinking pattern.

Example#1. Here, the decoded latent token "To convert the number from in in a 8 ounces" captures the core concept of proportional conversion from the 8-ounce bottle to the 12-ounce bottle. Despite its grammar mistake, it still precisely states the mathematical operation needed to solve the problem.

Q: In a 8 fluid ounce bottle of Vitamin Water, there are 125 calories. How many calories would be contained in a 12 fluid ounce bottle? Express your answer in decimal form.

A: <boLatent> To convert the number from in in a 8 ounces <eoLatent>, we need to find the ratio of calories to fluid ounces in the 8 fluid ounce bottle and then multiply it by 12.\n\nFirst, let's find the ratio of calories to fluid ounces in the 8 fluid ounce bottle:\n\nCalories per fluid ounce = Total calories / Total fluid ounces\n= 125 calories / 8 fluid ounces\n= 15.625 calories per fluid ounce\n\nNow, let's multiply this ratio by 12 fluid ounces to find the total calories in the 12 fluid ounce bottle:\n\nTotal calories in 12 fluid ounces = Calories per fluid ounce * Total fluid ounces\n= 15.625 calories per fluid ounce * 12 fluid ounces\n= 187.5 calories \nThe answer is 187.5.

---

Example#2. The decoded latent token "We find the greatest of CD of::63 common find" abstracts the initiation of the Euclidean algorithm, leading directly into the process that determines the GCD.

Q: What is the greatest common divisor of 315 and 108?

A: <boLatent> We find the greatest of CD of::63 common find <eoLatent>, we can use the Euclidean algorithm. Here's how to do it:\n\n1. Divide 315 by 108: 315 = 108 × 2 + 99\n2. Divide 108 by 99: 108 = 99 × 1 + 9\n3. Divide 99 by 9: 99 = 9 × 11 + 0\n\nSince the remainder is 0, the GCD is the divisor in the last step, which is 9. \nThe answer is 9.

---

#5.

Figure 4.1 is so small that all text is unreadable. Increasing font size or decreasing the number of tokens shown would significantly improve the figure.

Thanks for pointing this out! Yes, we have fixed this figure and will update the paper to reflect the changes.

We thank the reviewer for the comment and reply below.

#1.

In the paper, the codebook size also need to be tuned, making the proposed method problem-dependent

We would like to clarify that the codebook size is not problem dependent and the model performance remains robust across different codebook sizes. To further verify this, we conducted additional experiments on the ProsQA and ProntoQA varying the codebook size from 64 to 128 and 256.

We see that the performance is:

As these results clearly show, performance remains remarkably stable across different codebook sizes, with minimal variation in both accuracy and token efficiency. This robustness demonstrates that our method is not dependent on fine-tuning this hyperparameter for each problem domain.

For study on compression-ratio, see response to Reviewer rPRY (#3).

#2.

Report only the number of generated tokens can be misleading as they do not fully reflect the entire computational complexity: the data are still be fed to the entire network

We would like to clarify that the VQVAE is only used during training. After we trained the VAVAE, we then convert the text-tokens into the discrete latent code, then we save the data offline. Note, the training data passes through the VQVAE only once and gets encoded into latent tokens. That’s it. Then, the next stage is to train the LLM with these saved latent tokens and text tokens, and the data will not pass through the VQVAE during both the LLM training and inference time. In summary, we don’t need VQVAE at all during the training of reasoning models.

During inference, the LLM directly generates the latent tokens without requiring the separate VQ-VAE, resulting in pure computational savings with no additional inference-time overhead. The one-time training cost of the VQ-VAE is negligible compared to the full LLM fine-tuning process.

The efficiency benefits are substantial:

• The VQ-VAE introduces a minimal parameter overhead of just 50M (0.05B) to LLMs, which adds
  • 5% to Llama-3.2-1B
  • 1.7% to Llama-3.2-3B
  • 0.6% to Llama-3.1-8B
• Despite the slight increase in parameters, it significantly enhances token efficiency, it reduces tokens by:
  • 20% for both Llama-3.2-1B and Llama-3.2-3B
  • 10% for Llama-3.1-8B

For real-world deployment scenarios processing millions of reasoning tasks, these efficiency gains translate to significant reductions in computation costs and inference time.

#3.

From my perspective, I cannot figure out where the improvement over COT comes from .. Instead, the proposed method is more likely to compress language-based thoughts, which, in my opinion, does not improve upon the COT baseline.

The advantage of our method is that it compresses these high-level abstractions into discrete latent tokens (as some form of information distillation). During inference, the LLM conditions on these high-level latent tokens generated during the beginning of decoding, which effectively guides the reasoning process by:

• Providing a more abstract representation that helps the model focus on relevant information while spending less efforts on verbose and high-level CoTs.
• Creating better initial conditions that influence the entire downstream reasoning process

We provide examples of these on our responses to Rev 2JVK (#4), please check them out.

Recently, [1] demonstrates that conditioning on the first token ( representing different reasoning paths) significantly enhances a model's reasoning capabilities. Similarly, our model generates the learned latent tokens early in the sequence, these latent tokens encapsulate high-level reasoning abstraction, and sets a good initial conditions for the LLM, guiding the entire reasoning process from a higher level of abstraction. In contrast to prior works like Coconut, which maintain multiple explicit reasoning trajectories simultaneously, our approach implicitly encodes multiple reasoning possibilities into a compressed latent representation. Although it does not explicitly explore multiple reasoning paths simultaneously at the token level, it effectively captures diverse reasoning strategies in a latent manner, resulting in a more streamlined and efficient decoding process. Because of this, this structured latent conditioning enables more effective reasoning compared to standard token-by-token generation.

Reference: [1] Xuezhi Wang & Denny Zhou. "Chain-of-thought reasoning without prompting." arXiv preprint arXiv:2402.10200.