Technical Summary: MUX - Continuous Reasoning via Multiplexed Tokens

1. Problem Statement

Modern autoregressive language models solve complex problems (mathematics, coding, commonsense) by articulating intermediate reasoning steps in natural language (Chain-of-Thought). While effective, this discrete approach is computationally inefficient. Each reasoning step transmits only a single subword, creating a strict bandwidth bottleneck. Furthermore, many reasoning tokens are redundant, mirroring human communication patterns rather than optimized computation.

Continuous latent reasoning offers an alternative where models predict continuous vectors instead of subwords, theoretically allowing higher bandwidth by encoding multiple subwords simultaneously via superposition. However, existing latent reasoning methods face significant hurdles:

• Global methods (supervising only the final answer or trajectory endpoint) often lead to "shortcut behaviors" where latent tokens become uninformative placeholders.
• Local methods (supervising each latent token against a span of discrete reasoning) often require complex auxiliary architectures (e.g., autoregressive decoders) or fail to maintain diverse hypotheses necessary for search tasks.

The core challenge is defining a supervision target for continuous latent reasoning that is faithful to discrete reasoning, prevents latent collapse, and supports parallel exploration without excessive architectural complexity.

2. Methodology: MUX

The authors propose MUX, a local distillation method that trains language models to perform continuous latent reasoning using multiplexed tokens.

Core Concept

MUX treats the local distillation task as learning a fixed-dimensional continuous signal that encodes a variable-length categorical signal (a span of discrete reasoning steps). Drawing inspiration from code-division multiplexing in communication systems, MUX represents a span of discrete subwords as a weighted linear superposition of their one-hot encodings in the vocabulary space.

Technical Implementation

1. Multiplexed Targets: For a discrete reasoning span $r_i = (r^1_i, \dots, r^{S_i}_i)$, the learning target is constructed as:
   $$\text{mux}(r_i) := \sum_{j=1}^{S_i} \alpha^{(i)}_j \cdot \text{onehot}(r^j_i)$$
   where $\alpha^{(i)}_j$ are normalized positive positional weights.
2. Positional Weighting: To ensure the superposition is lossless (i.e., the original span can be uniquely recovered), the authors propose specific weighting schemes:
   • Geometric: $w_j = \rho^{j-1}$ (exponential decay).
   • Sinusoidal: Based on sine functions peaking near the end of the span.
   • Rotary: Based on averaged cosine components analogous to rotary position embeddings.
   • Theoretical Guarantee: These weightings satisfy a "subset-sum separation condition," ensuring that different sets of positions produce distinct total masses for any subword, preventing ambiguity even if a subword appears multiple times.
3. Training Objective: The model is trained to match these multiplexed targets using a linear-softmax head and a KL divergence loss ($L_{local}$). This is combined with a standard answer prediction loss ($L_{answer}$) and an optional trajectory-level loss ($L_{global}$) to align hidden features with discrete reasoning models.
4. Inference: During inference, the model autoregresses latent tokens in the hidden space. The vocabulary projection is used only for supervision during training, not for generating the final output tokens.

3. Key Contributions

(1) Latent Reasoning via Multiplexed Tokens

MUX introduces a simple local distillation framework where latent tokens represent spans of discrete reasoning via weighted linear superpositions. Unlike prior local methods, it does not require an auxiliary autoregressive decoder or compressed cache targets, relying instead on vocabulary-space projections.

(2) Lossless Multiplexing

The paper provides theoretical proofs that specific positional weightings (geometric, sinusoidal, rotary) guarantee lossless multiplexing. This means the discrete reasoning span is fully preserved in the continuous representation. The authors prove that this losslessness prevents latent collapse (where tokens become semantically homogenized), a common failure mode in global latent reasoning methods.

(3) Parallel Search via Superposition

Theoretically, the authors demonstrate that multiplexed tokens can naturally implement parallel breadth-first search (BFS). Because the tokens are superposed, a single latent token can represent and update an entire set of hypotheses (a search frontier) simultaneously. This allows the model to explore multiple problem-solving paths in parallel within a serial supervision framework.

(4) Empirical Superiority

Across 32 evaluation settings (spanning four language models, two training corpora, and four test sets), MUX outperforms strong baselines in both discrete and continuous reasoning. Notably, MUX often surpasses the discrete Chain-of-Thought (SFT-CoT) baseline it is distilled from, despite using significantly fewer tokens (approx. 2.4x to 5.9x fewer).

4. Experimental Results

• Mathematical Reasoning: On GSM8K-AUG, GSM8K-AUG-NL, SVAMP, GSM-Hard, and MultiArith, MUX achieved the best latent reasoning performance across GPT-2, LLaMA 3.2 1B, 3B, and 3.1 8B models. It frequently outperformed SFT-CoT, demonstrating that compact superposed reasoning can generalize better than the original discrete traces.
• Search Benchmarks: On the Minimum Nonnegative Sum (MNNS) and Game of 24 tasks, MUX significantly outperformed baselines (including Coconut and CoT2), validating the theoretical claim that multiplexed tokens enable parallel search capabilities.
• Ablation Studies:
  • Local Loss: Removing global loss ($\gamma=0$) and relying solely on the multiplexed local loss still yielded superior results compared to SIM-CoT, confirming the efficacy of the local target.
  • Weighting: Lossless weightings (Geometric, Sinusoidal, Rotary) consistently outperformed lossy uniform weighting, though the gap was modest on highly structured datasets.
  • Chunking: Randomized chunking of discrete spans into latent tokens performed best, acting as structured data augmentation.

5. Significance and Claims

The paper claims that lossless superposition as a local learning target constitutes a sufficient condition for achieving strong, efficient, and interpretable latent continuous reasoning.

• Efficiency: MUX achieves high performance with fewer reasoning steps than discrete CoT, addressing the computational bandwidth bottleneck.
• Simplicity: The method avoids the architectural complexity of prior local distillation approaches (e.g., auxiliary decoders) while maintaining or improving performance.
• Interpretability: Probing analyses confirm that the learned latent tokens encode faithful and interpretable reasoning content, countering the notion that latent reasoning is merely an uninterpretable placeholder.
• Parallelism: The results suggest that parallel search dynamics can naturally emerge from serial supervision when using multiplexed targets, offering a new pathway for improving planning and search in language models.

The authors conclude that simple, interpretable local targets can make latent reasoning stronger and easier to train, bridging the gap between the efficiency of continuous spaces and the fidelity of discrete reasoning.