Research map / Large language models / Reasoning and chain of thought
Token Assorted: Mixing Latent and Text Tokens for Improved Language Model Reasoning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Partially replaces leftmost chain-of-thought text tokens with discrete VQ-VAE latent codes, trained with randomized replacement counts, to shorten reasoning traces while improving accuracy on maze, logic, and math tasks.
Problem
Text chain-of-thought traces are long, with many tokens serving linguistic coherence rather than reasoning, making training and inference costly; prior latent-reasoning and CoT-compression methods need multi-stage curricula and often underperform full CoT on real LLM tasks.
Approach
Abstract the leftmost part of each reasoning trace into discrete VQ-VAE latent tokens (compression rate 16), mixing latent and text tokens, and randomize the replaced prefix length per sample so a single-stage fine-tune adapts to the extended vocabulary.
Key results
- On Llama-3.1-8B fine-tuned on MetaMathQA, latent approach averages 37.9% vs 33.4% for CoT across seven math benchmarks, with +13.3 points on Fresh-Gaokao-Math-2023.
- Generated responses are about 17% shorter on average than CoT (e.g., 513.7 vs 578.5 tokens for Llama-3.1-8B).
- On Keys-Finding Maze, latent approach reaches 62.8% 1-Feasible-10 vs 43% for CoT, and on ProsQA 96.2% vs 77.5%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Mix latent tokens (VQ-VAE abstractions of leading reasoning steps) with text CoT tokens | T5 / GPT-2 trained from scratch · Keys-Finding Maze, ProntoQA, ProsQA | accuracy / 1-Feasible-10: +19.8 maze, +1.2 ProntoQA, +18.7 ProsQA over best baseline | Helped |
| Mix latent and text tokens with fewer generated tokens | Llama-3.2/3.1 fine-tuned on MetaMathQA · Math reasoning avg over 7 datasets | average accuracy: +3.7 (1B), +2.9 (3B), +5.3 (8B) average points over CoT per text | Helped |
| Mix latent and text tokens | Llama-3.1-8B · Fresh-Gaokao-Math-2023 | accuracy: 16.7 to 30.0 | Helped |
| Mix latent and text tokens | Llama models on math reasoning · Average generated tokens across math benchmarks | average generated tokens: about 17% fewer tokens than CoT on average | Helped |
| Replace CoT with iCoT-style curriculum elimination | Llama models on math reasoning · MetaMathQA avg over 7 datasets | average accuracy: iCoT 7.7 vs CoT 14.1 on Llama-3.2-1B | Hurt |
| Replace all CoT tokens with latent tokens at once instead of partial left-to-right replacement (ablation) | Llama-3.2-3B · MetaMathQA avg over 7 datasets | average accuracy: 11.8 vs 28.1 for Latent-AR | Hurt |
| Curriculum replacement of CoT chunks over epochs instead of random-m partial replacement (ablation) | Llama-3.2-3B · MetaMathQA avg over 7 datasets | average accuracy: 10.9 vs 28.1 | Hurt |
| Random Poisson chunk replacement instead of left-to-right partial replacement (ablation) | Llama-3.2-3B · MetaMathQA avg over 7 datasets | average accuracy: 20.5 vs 28.1 (3B); 36.3 vs 37.9 (8B) | Hurt |
| Compression rate r=2 instead of r=16 (ablation) | Llama-3.2-3B latent model · MATH/GSM8K-style math avg accuracy | average accuracy: 27.08 vs 28.1 | Hurt |
| Compression rate r=32 instead of r=16 (ablation) | Llama-3.2-3B latent model · MATH/GSM8K-style math avg accuracy | average accuracy: 27.9 vs 28.1 (fewer tokens: 481.4 vs 513.6) | Hurt |
Surprising findings
- Fully replacing all text CoT with latent tokens (All-Replace) performs much worse than partial replacement, and iCoT-style curriculum degrades math accuracy substantially versus CoT.
- Compression rate 16 outperformed both rate 2 and rate 32 on accuracy, indicating a sweet spot rather than monotonic gains.
Limitations
- Gains on math are modest on some benchmarks (e.g., OlympiadBench and DeepMind-Math show near-parity or losses against baselines), and out-of-domain evaluation sets are small (e.g., Gaokao has 30 questions).
- Latent tokens are opaque without the VQ-VAE decoder; the VQ-VAE decoder interpretations are sometimes garbled.
- Method relies on a VQ-VAE trained per benchmark and only tested with supervised fine-tuning on relatively small base models.
Experimental setup
Datasets: Keys-Finding Maze (synthetic, 100k) ProntoQA ProsQA MetaMathQA Dart-MATH MATH GSM8K
Benchmarks: Keys-Finding Maze ProntoQA ProsQA MATH GSM8K Fresh-Gaokao-Math-2023 DeepMind-Mathematics College-Math OlympiadBench-Math TheoremQA
Metrics: 1-Feasible-10 accuracy exact match accuracy average accuracy number of generated tokens
Baselines: Sol-Only CoT iCoT Pause Token All-Replace Curriculum-Replace Poisson-Replace
Builds on: VQ-VAE iCoT Coconut MetaMath
Closest papers by meaning
- Exploring the Limits of Sub-Billion Language Model Reasoners with Open Training Recipes (ICLR 2026 · similarity 0.93)
- Reasoning via Test-Time Instance-Level Policy Gradient in Latent Space (NeurIPS 2026 · similarity 0.92)
- SwiReasoning: Switch-Thinking in Latent and Explicit for Pareto-Superior Reasoning LLMs (ICLR 2026 · similarity 0.91)
- Think Deep, Not Just Long: Measuring LLM Reasoning Effort via Deep-Thinking Tokens (ICML 2026 · similarity 0.91)
- LaDiR: Latent Diffusion Enhances LLMs for Text Reasoning (ICLR 2026 · similarity 0.91)
- Pushing on Multilingual Reasoning Models with Language-Mixed Chain-of-Thought (ICLR 2026 · similarity 0.91)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.