Multiplex Thinking: Reasoning via Token-wise Branch-and-Merge

TL;DR

The paper tackles the inefficiency of long chain-of-thought reasoning by introducing Multiplex Thinking, a stochastic soft reasoning framework that, at each step, samples K discrete tokens and aggregates them into a continuous multiplex token. This preserves the vocabulary embedding prior and discrete sampling dynamics while enabling a tractable, on-policy RL optimization over multiplex rollouts. Empirically, Multiplex Thinking improves Pass@1 through Pass@1024 on challenging math benchmarks and achieves shorter responses, with inference-time gains even without RL. The work analyzes the role of multiplex width, compute trade-offs, and entropy dynamics, and demonstrates robust gains across model scales, establishing a scalable path to more efficient and capable reasoning in LLMs.

Abstract

Large language models often solve complex reasoning tasks more effectively with Chain-of-Thought (CoT), but at the cost of long, low-bandwidth token sequences. Humans, by contrast, often reason softly by maintaining a distribution over plausible next steps. Motivated by this, we propose Multiplex Thinking, a stochastic soft reasoning mechanism that, at each thinking step, samples K candidate tokens and aggregates their embeddings into a single continuous multiplex token. This preserves the vocabulary embedding prior and the sampling dynamics of standard discrete generation, while inducing a tractable probability distribution over multiplex rollouts. Consequently, multiplex trajectories can be directly optimized with on-policy reinforcement learning (RL). Importantly, Multiplex Thinking is self-adaptive: when the model is confident, the multiplex token is nearly discrete and behaves like standard CoT; when it is uncertain, it compactly represents multiple plausible next steps without increasing sequence length. Across challenging math reasoning benchmarks, Multiplex Thinking consistently outperforms strong discrete CoT and RL baselines from Pass@1 through Pass@1024, while producing shorter sequences. The code and checkpoints are available at https://github.com/GMLR-Penn/Multiplex-Thinking.

Figures (10)

• Figure 1: Illustration of Multiplex Thinking. The model first generates an initial probability distribution conditioned on the prompt and the begin of thinking token[bot]. Instead of committing to a single discrete token or a deterministic soft average, we conduct token-wise branching and merging by independently sampling $K$ discrete tokens and aggregate them into a continuous multiplex token. When one of the sampled discrete tokens is the end of thinking token [eot], LLM continues conduct discrete decoding to give the answer. The design of sampling-based continuous thinking bridges the gap between continuous representation and stochastic discrete sampling, allowing for effective on-policy exploration and further RL training.
• Figure 2: Pass@1--Pass@1024 performance on four representative datasets. Top row: 1.5B models; bottom row: 7B models. Full results on all six datasets are reported in Appendix \ref{['app:additional_results']}.
• Figure 3: Performance comparison under different multiplex widths $K$. The x-axis is the multiplex width $K$ and the y-axis is the Pass@1 performance on each dataset. The grey columns ($K=1$) indicate the discrete RL performance and the blue columns represent Multiplex Thinking. The best results are bolded in each dataset.
• Figure 3: Pass@1 accuracy averaged on six math reasoning datasets over 64 runs. We compare Multiplex Thinking-I-4k could match the performance of Discrete CoT-5k which has 25% more sequence token length budget.
• Figure 4: Length scaling comparison. The y-axis is the accuracy averaged on six challenging reasoning datasets.