Token Assorted: Mixing Latent and Text Tokens for Improved Language Model Reasoning
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866915473720868864 |
|---|---|
| author | Su, DiJia Zhu, Hanlin Xu, Yingchen Jiao, Jiantao Tian, Yuandong Zheng, Qinqing |
| author_facet | Su, DiJia Zhu, Hanlin Xu, Yingchen Jiao, Jiantao Tian, Yuandong Zheng, Qinqing |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2502_03275 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Token Assorted: Mixing Latent and Text Tokens for Improved Language Model Reasoning Su, DiJia Zhu, Hanlin Xu, Yingchen Jiao, Jiantao Tian, Yuandong Zheng, Qinqing Computation and Language Artificial Intelligence Machine Learning Logic in Computer Science 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. |
| title | Token Assorted: Mixing Latent and Text Tokens for Improved Language Model Reasoning |
| topic | Computation and Language Artificial Intelligence Machine Learning Logic in Computer Science |
| url | https://arxiv.org/abs/2502.03275 |