CoChat
Sign Up
PreprintCited by 2Topic ModelingMultimodal Machine Learning ApplicationsNatural Language Processing Techniques

The DOI matches its DataCite record.

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

Open access
Su, DiJia, Zhu, Hanlin, Xu, Yingchen, Jiao, Jiantao, Tian, Yuandong, Zheng, Qinqing·2025·arXiv
DOI 10.48550/arxiv.2502.03275arXiv 2502.03275
Open this paper in CoChat — free accountOpen PDF

Abstract

Large Language Models (LLMs) excel at reasoning and planning when trained on chainof-thought (CoT) data, where the step-by-step thought process is explicitly outlined by text tokens. However, this results in lengthy inputs where many words support textual coherence rather than core reasoning information, and processing these inputs consumes substantial computation resources. In this work, we propose a hybrid representation of the reasoning process, where we partially abstract away the initial reasoning steps using latent discrete tokens generated by VQ-VAE, significantly reducing the length of reasoning traces. We explore the use of latent trace abstractions in two scenarios: 1) training the model from scratch for the Keys-Finding Maze problem, 2) fine-tuning LLMs on this hybrid data with an extended vocabulary including unseen latent tokens, for both logical and mathematical reasoning problems. To facilitate effective learning, we introduce a simple training procedure that randomly mixes latent and text tokens, which enables fast adaptation to new latent tokens. Our approach consistently outperforms the baselines methods in various benchmarks.

Overview

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Integrity signals

Registry

DataCite · 10.48550/arxiv.2502.03275

Cites

Statistics

Verification

Claims the paper makes

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Claims other papers make about this one

Each is checked against this paper’s own text: does it say what the citing paper says it does?

Self-check

Percentages, sums and reported differences re-computed from the paper’s own numbers.

The self-check needs the paper’s full text.

Citations

Cites

No references could be resolved yet.

Part of this paper’s reference list could not be split into single entries.

Research works printed but not linked: 52

  1. Azerbayev, Z., Schoelkopf, H., Paster, K., Santos, M. D., McAleer, S., Jiang, A. Q., Deng, J., Biderman, S., and Welleck, S. Llemma: An open language model for math- ematics. arXiv preprint arXiv:2310.10631, 2023. · arXiv
  2. Barrault, L., Duquenne, P.-A., Elbayad, M., Kozhevnikov, A., Alastruey, B., Andrews, P., Coria, M., Couairon, G., Costa-jussà, M. R., Dale, D., et al. Large concept models: Language modeling in a sentence representation space. arXiv e-prints, pp. arXiv-2412, 2024.
  3. Chen, W., Ma, X., Wang, X., and Cohen, W. W. Program of thoughts prompting: Disentangling computation from reasoning for numerical reasoning tasks. arXiv preprint arXiv:2211.12588, 2022a. · arXiv
  4. Chen, W., Ma, X., Wang, X., and Cohen, W. W. Program of thoughts prompting: Disentangling computation from reasoning for numerical reasoning tasks. arXiv preprint arXiv:2211.12588, 2022b. · arXiv
  5. Chen, W., Yin, M., Ku, M., Lu, P., Wan, Y., Ma, X., Xu, J., Wang, X., and Xia, T. Theoremqa: A theorem-driven question answering dataset. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 7889-7901, 2023.
  6. Chung, H. W., Hou, L., Longpre, S., Zoph, B., Tay, Y., Fe- dus, W., Li, Y., Wang, X., Dehghani, M., Brahma, S., et al. Scaling instruction-finetuned language models. Journal of Machine Learning Research, 25(70):1-53, 2024.
  7. Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021a. · arXiv
  8. Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021b. · arXiv
  9. Deng, Y., Prasad, K., Fernandez, R., Smolensky, P., Chaud- hary, V., and Shieber, S. Implicit chain of thought reasoning via knowledge distillation. arXiv preprint arXiv:2311.01460, 2023. · arXiv
  10. Deng, Y., Choi, Y., and Shieber, S. From explicit cot to implicit cot: Learning to internalize cot step by step. arXiv preprint arXiv:2405.14838, 2024. · arXiv
  11. Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024. · arXiv
  12. Dziri, N., Lu, X., Sclar, M., Li, X. L., Jian, L., Lin, B. Y., West, P., Bhagavatula, C., Bras, R. L., Hwang, J. D., Sanyal, S., Welleck, S., Ren, X., Ettinger, A., Harchaoui, Z., and Choi, Y. Faith and fate: Limits of transformers on compositionality. Advances in Neural Information Processing Systems, 36, 2024.
  13. Feng, G., Zhang, B., Gu, Y., Ye, H., He, D., and Wang, L. Towards revealing the mystery behind chain of thought: a theoretical perspective. Advances in Neural Information Processing Systems, 36, 2024.
  14. Gandhi, K., Lee, D., Grand, G., Liu, M., Cheng, W., Sharma, A., and Goodman, N. D. Stream of search (sos): Learning to search in language. arXiv preprint arXiv:2404.03683, 2024. · arXiv
  15. Goyal, S., Ji, Z., Rawat, A. S., Menon, A. K., Kumar, S., and Nagarajan, V. Think before you speak: Train- ing language models with pause tokens. arXiv preprint arXiv:2310.02226, 2023. · arXiv
  16. Hao, S., Sukhbaatar, S., Su, D., Li, X., Hu, Z., Weston, J., and Tian, Y. Training large language models to reason in a continuous latent space. arXiv preprint arXiv:2412.06769, 2024. · arXiv
  17. Hart, P. E., Nilsson, N. J., and Raphael, B. A formal basis for the heuristic determination of minimum cost paths. IEEE transactions on Systems Science and Cybernetics, 4(2):100-107, 1968.
  18. He, C., Luo, R., Bai, Y., Hu, S., Thai, Z. L., Shen, J., Hu, J., Han, X., Huang, Y., Zhang, Y., et al. Olympiadbench: A challenging benchmark for promoting agi with olympiad- level bilingual multimodal scientific problems. arXiv preprint arXiv:2402.14008, 2024. · arXiv
  19. Hendrycks, D., Burns, C., Kadavath, S., Arora, A., Basart, S., Tang, E., Song, D., and Steinhardt, J. Measuring mathematical problem solving with the math dataset. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2).
  20. Hendrycks, D., Burns, C., Kadavath, S., Arora, A., Basart, S., Tang, E., Song, D., and Steinhardt, J. Measuring math- ematical problem solving with the math dataset. NeurIPS, 2021.
  21. Jiang, W., Shi, H., Yu, L., Liu, Z., Zhang, Y., Li, Z., and Kwok, J. Forward-backward reasoning in large language models for mathematical verification. In Findings of the Association for Computational Linguistics ACL 2024, pp. 6647-6661, 2024.
  22. Jiang, Z., Zhang, T., Janner, M., Li, Y., Rocktäschel, T., Grefenstette, E., and Tian, Y. Efficient planning in a com- pact latent action space. arXiv preprint arXiv:2208.10291, 2022. · arXiv
  23. Jiang, Z., Xu, Y., Wagener, N., Luo, Y., Janner, M., Grefen- stette, E., Rocktäschel, T., and Tian, Y. H-gap: Hu- manoid control with a generalist planner. arXiv preprint arXiv:2312.02682, 2023. · arXiv
  24. Kim, S., Joo, S. J., Kim, D., Jang, J., Ye, S., Shin, J., and Seo, M. The cot collection: Improving zero-shot and few- shot learning of language models via chain-of-thought fine-tuning. arXiv preprint arXiv:2305.14045, 2023. · arXiv
  25. Kojima, T., Gu, S. S., Reid, M., Matsuo, Y., and Iwasawa, Y. Large language models are zero-shot reasoners. Ad- vances in neural information processing systems, 35: 22199-22213, 2022.
  26. Lehnert, L., Sukhbaatar, S., Su, D., Zheng, Q., McVay, P., Rabbat, M., and Tian, Y. Beyond a*: Better planning with transformers via search dynamics bootstrapping. In First Conference on Language Modeling, 2024. URL https: //openreview.net/forum?id=SGoVIC0u0f.
  27. Li, Z., Liu, H., Zhou, D., and Ma, T. Chain of thought empowers transformers to solve inherently serial problems, 2024. URL https://arxiv.org/abs/ 2402.12875.
  28. Liu, L., Pfeiffer, J., Wu, J., Xie, J., and Szlam, A. Deliber- ation in latent space via differentiable cache augmenta- tion. 2024. URL https://arxiv.org/abs/2412. 17747.
  29. Nye, M., Andreassen, A. J., Gur-Ari, G., Michalewski, H., Austin, J., Bieber, D., Dohan, D., Lewkowycz, A., Bosma, M., Luan, D., et al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint arXiv:2112.00114, 2021a. · arXiv
  30. Nye, M., Andreassen, A. J., Gur-Ari, G., Michalewski, H., Austin, J., Bieber, D., Dohan, D., Lewkowycz, A., Bosma, M., Luan, D., et al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint arXiv:2112.00114, 2021b. · arXiv
  31. Pagnoni, A., Pasunuru, R., Rodriguez, P., Nguyen, J., Muller, B., Li, M., Zhou, C., Yu, L., Weston, J., Zettle- moyer, L., Ghosh, G., Lewis, M., Holtzman, A., and Iyer, S. Byte latent transformer: Patches scale better than tokens. 2024. URL https://arxiv.org/abs/ 2412.09871.
  32. Pfau, J., Merrill, W., and Bowman, S. R. Let's think dot by dot: Hidden computation in transformer language models. arXiv preprint arXiv:2404.15758, 2024. · arXiv
  33. Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019.
  34. Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Ex- ploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1-67, 2020. URL http://jmlr. org/papers/v21/20-074.html.
  35. Saparov, A. and He, H. Language models are greedy rea- soners: A systematic formal analysis of chain-of-thought. arXiv preprint arXiv:2210.01240, 2022. · arXiv
  36. Saxton, D., Grefenstette, E., Hill, F., and Kohli, P. Analysing mathematical reasoning abilities of neural models. arXiv preprint arXiv:1904.01557, 2019. · arXiv
  37. Su, D., Sukhbaatar, S., Rabbat, M., Tian, Y., and Zheng, Q. Dualformer: Controllable fast and slow thinking by learning with randomized reasoning traces. arXiv preprint arXiv:2410.09918, 2024. · arXiv
  38. Su, D., Gu, A., Xu, J., Tian, Y., and Zhao, J. Galore 2: Large- scale llm pre-training by gradient low-rank projection. arXiv preprint arXiv:2504.20437, 2025. · arXiv
  39. Tang, Z., Zhang, X., Wang, B., and Wei, F. Mathscale: Scaling instruction tuning for mathematical reasoning. arXiv preprint arXiv:2403.02884, 2024. · arXiv
  40. Tong, Y., Zhang, X., Wang, R., Wu, R., and He, J. Dart- math: Difficulty-aware rejection tuning for mathemati- cal problem-solving. arXiv preprint arXiv:2407.13690, 2024. · arXiv
  41. Van Den Oord, A., Vinyals, O., et al. Neural discrete rep- resentation learning. Advances in neural information processing systems, 30, 2017.
  42. Wang, X. and Zhou, D. Chain-of-thought reasoning without prompting. 2024. URL https://arxiv.org/abs/ 2402.10200.
  43. Wang, X., Wei, J., Schuurmans, D., Le, Q., Chi, E., Narang, S., Chowdhery, A., and Zhou, D. Self-consistency im- proves chain of thought reasoning in language models. arXiv preprint arXiv:2203.11171, 2022. · arXiv
  44. Wang, X., Caccia, L., Ostapenko, O., Yuan, X., Wang, W. Y., and Sordoni, A. Guiding language model reasoning with planning tokens. arXiv preprint arXiv:2310.05707, 2023. · arXiv
  45. Wei, J., Wang, X., Schuurmans, D., Bosma, M., Xia, F., Chi, E., Le, Q. V., Zhou, D., et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24837, 2022a.
  46. Wei, J., Wang, X., Schuurmans, D., Bosma, M., Xia, F., Chi, E., Le, Q. V., Zhou, D., et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24837, 2022b.
  47. Wen, K., Zhang, H., Lin, H., and Zhang, J. From sparse dependence to sparse attention: Unveiling how chain-of- thought enhances transformer sample efficiency. arXiv preprint arXiv:2410.05459, 2024. · arXiv
  48. Yang, A., Yang, B., Zhang, B., Hui, B., Zheng, B., Yu, B., Li, C., Liu, D., Huang, F., Wei, H., et al. Qwen2. 5 technical report. arXiv preprint arXiv:2412.15115, 2024. · arXiv
  49. Yao, S., Yu, D., Zhao, J., Shafran, I., Griffiths, T., Cao, Y., and Narasimhan, K. Tree of thoughts: Deliberate problem solving with large language models. Advances in Neural Information Processing Systems, 36, 2024.
  50. Yu, L., Jiang, W., Shi, H., Yu, J., Liu, Z., Zhang, Y., Kwok, J. T., Li, Z., Weller, A., and Liu, W. Metamath: Boot- strap your own mathematical questions for large language models. arXiv preprint arXiv:2309.12284, 2023. · arXiv
  51. Yu, P., Xu, J., Weston, J., and Kulikov, I. Distilling system 2 into system 1. arXiv preprint arXiv:2407.06023, 2024. · arXiv
  52. Yue, X., Qu, X., Zhang, G., Fu, Y., Huang, W., Sun, H., Su, Y., and Chen, W. Mammoth: Building math generalist models through hybrid instruction tuning. arXiv preprint arXiv:2309.05653, 2023. · arXiv

Cited by

Sorted by citations

  • Stop Overthinking: A Survey on Efficient Reasoning for Large Language Models

    Sui, Y., Chuang, Y.-N., Wang, G., Zhang, J., Zhang, T., Yuan, J., Liu, H., Wen, A., Zhong, S., Zou, N., et al. (2025) · arXiv · Cited by 517

    “…[62]; Learn to Skip [53]; Token-Budget [34]; ReasoningOutput-based Latent e.g. Coconut [35]; CODI [75]; CCoT [12]; Heima [73]; Token Assorted [79]; Loop [70]; SoftCoT [100]; DynamicReasoningParadigm e.g. Speculative Rejection [81]; Sampling-Efficient TTS [91]; DPTS [26]; Cer-taindex [30];…”

  • iCLP: Large Language Model Reasoning with Implicit Cognition Latent Planning

    Chen, S., & Niu, D. (2025) · arXiv · doi:10.48550/arxiv.2512.24014

    “Reasoning in a latent continuous space has been increasingly explored in recent research (Xu et al., 2024; Hao et al., 2024; Pagnoni et al., 2024; Su et al., 2025; team et al., 2024; Tack et al., 2025), demonstrating promising improvements in both accuracy and efficiency.”

Metadata

Authors
Su, D., Zhu, H., Xu, Y., Jiao, J., Tian, Y., & Zheng, Q.
ORCID
Venue
arXiv
Publisher
arXiv
Published
2025
ISSN
2331-8422 (ISSN-L)
DOI
10.48550/arxiv.2502.03275
arXiv
2502.03275
OpenAlex
W4407231268
Cited by
1
References
0
Citations by year
2026: 1
Type
posted content
Language
en
Open access
Licence
Locations
  • arXiv (Cornell University) (submittedVersion, cc-by)
  • arXiv (Cornell University) (cc-by)
Topics
  • Physical Sciences › Computer Science › Artificial Intelligence › Topic Modeling
  • Physical Sciences › Computer Science › Computer Vision and Pattern Recognition › Multimodal Machine Learning Applications
  • Physical Sciences › Computer Science › Artificial Intelligence › Natural Language Processing Techniques
Keywords
large language models, chain-of-thought prompting, vector quantized variational autoencoder, token mixing, reasoning traces, fine-tuning, mathematical reasoning, logical reasoning

About these checks

Registry, retraction, Cites and statistics results are computed, not judged, with one exception: a printed reference that no index lists is read once by a model before it counts as a citation error. Claims are extracted by a model and kept only when their quote appears verbatim in the text; each cited source is then checked, with a model’s help, for whether it says what the claim relies on. Those verdicts can be wrong — read the claim’s quote and the cited source before relying on one.

Statistics and arithmetic are recomputed, not judged.

Confidence bands are shown only for fields calibrated on our annotated test set. None is calibrated yet, so no band is shown.

What counts as a claim ›How we measure confidence ›How we check this ›

Are you an author, or do you see a problem with this page? Request a review