Token Assorted: Mixing Latent and Text Tokens for Improved Language
  Model Reasoning
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.
67a448b69ca42c642a723a7d
DiJia Su
67a448b69ca42c642a723a7e
Hanlin Zhu
admin_assigned
2025-02-06T16:10:24.407Z
6467bc59b990713c50339d2d
/avatars/064aba45e37040f7b1de8f76169f5174.svg
Hanlin Zhu
user
hanlinzhu
67a448b69ca42c642a723a7f
Yingchen Xu
admin_assigned
2025-02-06T16:09:40.111Z
6481333f8c6a3b8f11fc4114
/avatars/15c194296506b32e3f218530382c9f78.svg
Yingchen Xu
user
xuyingchen
67a448b69ca42c642a723a80
Jiantao Jiao
admin_assigned
2025-02-06T16:09:29.786Z
653b306986b88947d5cacfa4
/avatars/21ebd4daf35ec67c7d5f9b0a53628b00.svg
Jiantao Jiao
user
nexus-jt-llm
67a448b69ca42c642a723a81
Yuandong Tian
admin_assigned
2025-02-06T16:09:19.670Z
6344cf73ee1504dbcd5bdfe7
/avatars/6dd2bf1f9c5679e5c8c85d62c9836aac.svg
Yuandong Tian
user
tydsh
67a448b69ca42c642a723a82
Qinqing Zheng
admin_assigned
2025-02-06T16:09:11.277Z
64d27579dafee18faf9308ac
/avatars/8914a47244017c3541d3d5ac5b2d0372.svg
Qinqing Zheng
user
goodsleep
Token Assorted: Mixing Latent and Text Tokens for Improved Language
  Model Reasoning
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.
A hybrid representation using latent tokens generated by VQ-VAE reduces the length of reasoning traces in large language models, improving performance on both reasoning and fine-tuning tasks.
Large Language Models (LLMs) chain-of-thought (CoT) VQ-VAE latent discrete tokens Keys-Finding Maze problem logical reasoning mathematical reasoning mixed training procedure