Research map / Large language models / Reasoning and chain of thought
MUX: Continuous Reasoning via Multiplexed Tokens
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.
Proposes MUX, a training-time local distillation method that supervises each latent token with a position-weighted superposition of one-hot vocabulary encodings of a discrete reasoning span, and proves lossless recovery, non-collapse, and parallel BFS capability.
Problem
Continuous latent reasoning in language models is hard to train: trajectory-level losses yield shortcut or uninformative latent tokens, while existing local distillation methods need auxiliary decoders or KV-cache compression.
Approach
Target each latent token with a vocabulary-space multiplexed target: a normalized, position-weighted linear superposition of one-hot encodings of its aligned reasoning span (geometric, sinusoidal, or rotary weights), matched via a linear-softmax head with KL loss. Weightings satisfying a subset-sum separation condition make the target lossless, so no autoregressive decoder is needed.
Key results
- MUX achieves the best latent reasoning accuracy in all 32 math settings; e.g., 48.1% ID on GSM8K-AUG with GPT-2 vs 43.7% for CODI, and 58.0% with LLaMA 3.2 1B (Jacobi setting) vs 56.5% for KaVa.
- On LLaMA 3.1 8B GSM8K-AUG, MUX reaches 68.1% ID vs 64.1% for SIM-CoT.
- On search, MUX reaches 99.6% on MNNS and 88.7% on Game of 24, above CoT2 (98.9%, 85.0%).
Surprising findings
- MUX outperforms its own discrete SFT-CoT teacher in 15 of 32 settings despite using about 2.4-5.9x fewer reasoning tokens.
- Uniform (lossy) weighting still allows nontrivial demultiplexing accuracy, showing subword order is partly recoverable from occurrence counts.
Limitations
- Gains over SFT-CoT and baselines on OOD sets are modest in some settings, and the gap between lossless and uniform weighting is modest on GSM8K-AUG.
- Sequential latent token budget is small (six tokens); the parallel setting relies on Jacobi iterations.
- Losslessness is proven under exact arithmetic; finite precision results depend on the separation margin and span lengths (S<=11 used).
Experimental setup
Datasets: GSM8K-AUG GSM8K-AUG-NL GSM8K
Benchmarks: GSM8K SVAMP GSM-Hard MultiArith MNNS Game of 24
Metrics: accuracy Recall@5 Step Alignment MRR
Baselines: SFT-CoT No-CoT iCoT Coconut CODI SIM-CoT PCCoT KaVa CoT2
Builds on: Coconut CODI SIM-CoT KaVa iCoT
Closest papers by meaning
- Reasoning by Superposition: A Theoretical Perspective on Chain of Continuous Thought (NeurIPS 2025 · similarity 0.89)
- Geometric Latent Reasoning Induces Shorter Generations in LLMs (NeurIPS 2026 · similarity 0.88)
- Efficient Reasoning via Constrained Optimization in Latent Space (NeurIPS 2026 · similarity 0.88)
- Why Struggle with Continuous Latents? Interpretable Discrete Latent Reasoning via Rendered Compression (NeurIPS 2026 · similarity 0.88)
- Token Assorted: Mixing Latent and Text Tokens for Improved Language Model Reasoning (ICML 2025 · similarity 0.87)
- Emergence of Superposition: Unveiling the Training Dynamics of Chain of Continuous Thought (ICLR 2026 · similarity 0.87)
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.
