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

Yao Tang; Li Dong; Yaru Hao; Qingxiu Dong; Furu Wei; Jiatao Gu

Source-linked AI summary

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

Yao Tang, Li Dong, Yaru Hao, Qingxiu Dong, Furu Wei, Jiatao Gu

arXiv:2601.08808v1Jan 13, 2026cs.CLcs.AIcs.LG

TL;DR

Discrete CoT is effective but expensive, while deterministic continuous reasoning limits stochastic exploration for reinforcement learning. Multiplex Thinking samples K candidate tokens and merges their embeddings into one continuous token with a tractable rollout distribution, enabling direct RL. Across challenging math benchmarks, it outperforms strong discrete CoT and RL baselines from Pass@1 through Pass@1024 while producing shorter sequences.

Problem
Long discrete CoT rollouts are costly, while deterministic continuous token methods collapse policy stochasticity and limit exploration needed for RL.
Method
Multiplex Thinking independently samples K tokens, aggregates their vocabulary embeddings into one continuous token, and factorizes rollout probabilities for direct RL optimization.
Results
Multiplex Thinking consistently outperforms strong discrete CoT and RL baselines from Pass@1 through Pass@1024 across challenging math reasoning benchmarks.
Takeaways & Limitations
Multiplex Thinking improves accuracy while maintaining shorter reasoning sequences by compressing richer information into individual multiplex tokens.
Takeaways & Limitations
Soft Thinking’s deterministic mapping from logits to continuous tokens limits diverse reasoning-path exploration for trial-and-error RL.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Multiplex Thinking addresses the cost and limited exploration of discrete CoT and deterministic continuous reasoning by sampling multiple tokens and merging their embeddings into one stochastic continuous token. It enables direct RL over multiplex rollouts while improving accuracy and shortening trajectories.

Motivation: Discrete CoT and RL rollouts are costly because each alternative requires a long, explicit reasoning sequence.Their exploration resembles depth-first search, whereas continuous tokens can compactly encode multiple candidate paths.
Motivation: Deterministic continuous token methods collapse the token-level policy distribution, producing identical rollouts and limiting on-policy exploration.This conflicts with RL’s reliance on stochastic trial-and-error rollouts.
Method: Multiplex Thinking independently samples K tokens from the model distribution and aggregates their vocabulary embeddings into one continuous multiplex token.Low-entropy distributions yield nearly coincident samples, while high-entropy distributions preserve richer multi-token information in one step.
Method: Multiplex Thinking preserves discrete sampling behavior and defines a probability distribution over complete reasoning trajectories, enabling direct RL optimization without the full token cost of discrete CoT.The probability of each multiplex token factorizes over its independently sampled tokens.
Results: Multiplex Thinking consistently improves Pass@1–Pass@1024 over strong discrete CoT and RL baselines while producing shorter reasoning sequences.The method compresses high-entropy reasoning steps so one multiplex token can encode richer information than a standard discrete token.

2 BACKGROUND

This section introduces discrete CoT, deterministic Soft Thinking, and RLVR as foundations for reasoning and optimization. It emphasizes that deterministic continuous reasoning limits diverse path exploration needed for RL.

Chain-of-Thought Reasoning: Standard CoT generates intermediate thinking tokens before producing final answers, using the language model’s vocabulary embeddings.The reasoning sequence is sampled token by token before answer generation.
Soft Thinking: Soft Thinking replaces discrete thinking tokens with continuous concept tokens formed by weighting vocabulary embeddings with the next-token distribution.The resulting representation compresses reasoning information into continuous vectors without fine-tuning.
Soft Thinking: Soft Thinking is deterministic for a given context, so its fixed mapping from logits to concept tokens prevents diverse reasoning-path exploration.This limits optimization through reinforcement-learning objectives based on trial and error.
Reinforcement Learning with Verifiable Rewards: RLVR trains language models on verifiable question-answer data by maximizing rewards assigned according to the sampled answer and ground truth.The objective optimizes expected reward under the model’s answer-sampling policy.

3 MULTIPLEX THINKING

Multiplex Thinking combines stochastic discrete sampling with continuous token representations, allowing each reasoning token to encode multiple candidate paths while retaining a tractable trajectory distribution for on-policy RL. Its entropy and sampling structure adapt from near-discrete behavior under confidence to broader exploration under uncertainty.

3 MULTIPLEX THINKING: Multiplex Thinking independently samples K tokens at each reasoning step and aggregates them into one continuous multiplex token.The construction combines discrete sampling with continuous representations.
3 MULTIPLEX THINKING: When K = 1, the multiplex representation collapses to a standard discrete token, while larger K retains mixtures of sampled candidates.As K approaches infinity, the empirical sample distribution converges to the model’s language-model-head distribution.
3 MULTIPLEX THINKING: Vocabulary-space weighting supports uniform averaging or LM-head reweighting, with reweighting adopted by default because it reflects confidence over sampled candidates.The two weighting choices have comparable empirical performance in the reported experiments.
3 MULTIPLEX THINKING: Independent sampling factorizes multiplex-token probabilities, making the trace log-probability a sum over constituent samples and enabling direct on-policy RL optimization.The RL objective operates over the joint generation of the multiplex trace and final answer.
3 MULTIPLEX THINKING: Multiplex-token entropy scales linearly with K, corresponding to an effective exploration volume that expands from |V| to |V|^K.The continuous token can encode a superposition of K paths while deferring discrete decisions.

4 EXPERIMENTS

Experiments evaluate Multiplex Thinking across six mathematical benchmarks, two model sizes, multiple baselines, and Pass@1–Pass@1024 test-time scaling. The method generally achieves stronger accuracy and exploration, with especially pronounced gains on harder tasks and after expanding multiplex width beyond one.

4 EXPERIMENTS: On the 7B backbone, Multiplex Thinking outperforms Stochastic Soft Thinking across all six benchmarks, compared with four of six on the 1.5B backbone.The authors associate the larger-model advantage with resolving interference between superposed reasoning paths.
4 EXPERIMENTS: Multiplex Thinking achieves the best Pass@1 performance in 11 of 12 settings across two model sizes and six mathematical benchmarks.It surpasses Discrete RL using the identical GRPO training setup across all tasks.
4 EXPERIMENTS: At k = 1024 on AIME 2025 with the 7B model, Multiplex Thinking reaches approximately 55%, while Discrete RL plateaus around 40%.The widening gap on challenging setups indicates stronger scaling of exploration with additional sampled trajectories.
4 EXPERIMENTS: Multiplex Thinking’s advantage widens on sparse-solution tasks such as AIME 2025 and OlympiadBench, whereas methods saturate quickly on simpler MATH-500.The reported pattern is difficulty-dependent and is attributed to maintaining superposed reasoning states.
4 EXPERIMENTS: Multiplex Thinking requires significantly fewer samples than discrete baselines to reach a target accuracy, reducing test-time compute.The evaluation measures Pass@k for k from 1 through 1024 as a proxy for exploration potential.

5 ANALYSIS

The analysis examines multiplex width, sequence length, exploration, aggregation, and trajectory behavior. Multiplex Thinking generally improves accuracy while shortening trajectories and preserving exploration.

Multiplex Thinking-I provides substantial inference-only gains over standard Discrete CoT, isolating benefits from the representation itself.
The impact of token width K: K ≥2 substantially improves performance over K = 1 across benchmarks; on AMC 2023, precision rises from 44.7% to 49.6% (+4.9%).
The impact of token width K: Performance continues increasing from K = 2 to K = 3 and 6, but the marginal gain between K = 3 and K = 6 becomes considerably smaller.
Test-time compute: multiplex width vs. sequence length: Multiplex Thinking achieves higher accuracy than Discrete CoT with substantially shorter trajectories as response length increases.
Test-time compute: multiplex width vs. sequence length: Multiplex Thinking-I-4k consistently outperforms or matches Discrete CoT-5k despite using 20% shorter sequences.
Training dynamics: Multiplex tokens encode multiple potential paths, producing shorter trajectories while maintaining superior accuracy.
Training dynamics: Multiplex training shows lower entropy reduction than discrete RL, indicating less entropy collapse and more sustained exploration during training.
Trajectory visualization: The trajectory visualization shows consensus when samples agree and multiplex aggregation when divergent high-entropy candidates compete.

6 RELATED WORKS

Related work spans discrete CoT and RL, continuous reasoning tokens, and parallel reasoning methods. Multiplex Thinking complements parallel sampling by changing the per-step token distribution.

Discrete reasoning: Discrete CoT and RL methods improve reasoning but incur high computation costs from generating long explicit reasoning sequences.
Continuous reasoning: Continuous reasoning methods use hidden states or continuous representations as latent reasoning steps and can outperform discrete CoT on logical reasoning.
Parallel reasoning: Parallel reasoning methods explore multiple paths and aggregate their outcomes through approaches such as self-consistency and Best-of-N selection.
Parallel reasoning: Multiplex Thinking complements parallel reasoning by changing the per-step token distribution rather than the outer-loop sampling budget.

7 CONCLUSION

The paper presents Multiplex Thinking as a bridge between discrete CoT and continuous reasoning, combining richer token representations with stochastic sampling. Across challenging mathematics benchmarks, it improves accuracy and token efficiency.

Multiplex Thinking aggregates multiple independent discrete tokens into continuous representations while preserving probabilistic sampling for reinforcement learning.
Across challenging mathematical benchmarks, Multiplex Thinking consistently outperforms strong discrete CoT and RL baselines from Pass@1 to Pass@1024.
The method achieves these gains with improved token efficiency by compressing complex reasoning steps into shorter trajectories.

A.1 IMPLEMENTATION DETAILS

The implementation uses verl and SGLang, with experiments conducted on eight NVIDIA DGX B200 GPUs. The stopping rule uses the highest-probability </think> token to end thinking.

The framework is implemented with verl and SGLang, using SGLang version 0.4.9.post6.
All experiments are conducted on 8× NVIDIA DGX B200 GPUs.
Thinking ends when the highest-probability discrete token is the special token </think>, rather than a training-free low-entropy heuristic.

A.1.1 HYPER-PARAMETERS

This section provides the hyper-parameters used for GRPO training of discrete RL and Multiplex Thinking.

Table 6 summarizes the hyper-parameters used in GRPO training for discrete RL and Multiplex Thinking.

A.2 ADDITIONAL RESULTS

This section presents experimental results omitted from the main body because of space constraints.

The appendix reports complete experimental results omitted from the main body due to space constraints.

A.2.1 FULL PASS@1–PASS@1024 RESULTS

The appendix expands the evaluation across Pass@k trajectories, multiplex widths, training dynamics, and qualitative reasoning examples. These results show a substantial advantage for multiplex variants over discrete RL, with diminishing returns as width increases.

FULL PASS@1–PASS@1024 RESULTS: Figure 7 details Pass@k performance from k = 1 to 1024 across six benchmarks and both 1.5B and 7B model scales.
FULL PASS@1–PASS@1024 RESULTS: Multiplex variants with K ≥2 show a significant performance gap over the Discrete RL baseline with K = 1.
FULL PASS@1–PASS@1024 RESULTS: Increasing multiplex width beyond K ≥2 yields diminishing returns, with the K = 2, 3, and 6 curves remaining closely clustered.
FULL PASS@1–PASS@1024 RESULTS: Figure 9 reports training and validation score dynamics, validating the trained model every 25 steps with Pass@4 on MATH-500.
FULL PASS@1–PASS@1024 RESULTS: The qualitative trajectory applies inclusion-exclusion to count assignments in which all three lanes are occupied, obtaining 540 ways.