Switch Explained: How a Pair of <swi> Boundary Tokens Solves Both Latent-Reasoning RL Training and Hidden-State Interpretability
> Paper: *Demystifying Hidden-State Recurrence: Switchable Latent Reasoning with On-Policy Reinforcement Learning* > Authors: Jiayu Yang, Chao Chen, Shengen Wu, Yinhong Liu, Yuxuan Fan, Lujundong Li, Songning Lai, Chengwei Qin, Zhijiang Guo (HKUST(GZ), Cambridge, NTU, JoinQuant, HKUST) > Paper: https://arxiv.org/abs/2606.13106 > Code: https://github.com/LARK-AI-Lab/SWITCH
1. Background: Latent Chain-of-Thought and Its Two Hard Problems
Standard LLM reasoning uses visible Chain-of-Thought text, which is token-expensive, slow, and may leak internal reasoning. Latent CoT instead performs reasoning with continuous hidden states: Coconut (Hao et al., 2025) feeds the previous step's last-layer hidden state directly as the next step's input embedding, reasoning entirely in representation space without emitting tokens. This is elegant, but leaves two unsolved problems:
1. On-policy RL is not trainable. GRPO (used by DeepSeek-R1) requires a samplable token distribution and policy density at every position. Latent positions emit no tokens, so the GRPO ratio is undefined inside latent blocks. Prior work either skips RL or trains on text-only rollouts, creating a train-test mismatch between training trajectories and latent inference paths. 2. Hidden states are uninterpretable. Analysts cannot tell whether latent steps perform useful reasoning or act merely as inert placeholders whose work is done by surrounding text.
2. Core Insight: One Pair of Boundary Tokens Solves Both
The paper's insight is that both problems share the same root cause: latent blocks have no explicit boundary. With a boundary, RL gains an anchor for policy density, and analysis gains an intervention handle.
Switch introduces <swi> (switch in) and </swi> (switch out) — a pair of ordinary discrete tokens bracketing K <latent> positions that use Coconut-style hidden-state recurrence.
- Why this fixes RL:
<swi>/</swi>have standard policy distributions and densities, so the GRPO importance ratio is well-defined at boundary positions. Latent positions produce no policy-gradient terms, but the *boundary decisions* (when to enter/exit latent mode) are RL-optimizable. Rollout execution uses the identical multi-forward-pass decoder in training and inference, and since hidden-state injection is deterministic given preceding text, rollout likelihood factors only over text positions — eliminating train-test mismatch. - Why this fixes interpretability: the boundaries give analysts three handles: reading p(
<swi>), probing switch state from internal activations, and performing causal interventions on specific latent hidden states. - Switch is +16.8 points over the strongest Coconut-style baseline (CoLaR) and uses ~38% fewer visible tokens than text CoT.
- What RL adds: latent-conditional accuracy gains +12.6 points from Switch-GRPO, while the switch rate *drops* from 81% to 58% — RL teaches the model to invoke latent mode more selectively, not more often. During training, per-problem latent calls fall from ~1.5 to ~1 and visible tokens shrink from ~2900 to ~1900.
- Per-subject: Algebra 88.7%, Prealgebra 80.5%, Number Theory 79.0%; accuracy declines smoothly from 93.0% (Level 1) to 53.7% (Level 5).
- Accuracy-efficiency operating curve: tuning the brevity bonus yields a ~76% accuracy / ~1300-token point (33% shorter) with 0% max-length truncation — a distribution-level shift, not simple truncation.
- First work to make on-policy GRPO truly applicable to hidden-state recurrence without modifying the RL algorithm itself.
- First to make latent reasoning verifiable via a reproducible experimental protocol (probing + causal intervention) for the placeholder question.
- Switchable design philosophy: rather than forcing latent mode everywhere, the model learns when latent reasoning is worth its cost — analogous to human intuition vs. deliberate thinking.
- Yang, J., et al. (2026). Demystifying Hidden-State Recurrence: Switchable Latent Reasoning with On-Policy Reinforcement Learning. *arXiv preprint arXiv:2606.13106*.
- Hao, S., et al. (2025). Training Large Language Models to Reason in a Continuous Latent Space. *arXiv*.
- DeepSeek-AI. (2025). DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. *arXiv*.
3. Three-Stage Training
1. Switch-SFT Part 1: Measure base-model Shannon entropy per position on math CoT data; wrap consecutive high-entropy intervals with <swi>/</swi>; train with standard next-token cross-entropy to teach when to enter/exit latent mode.
2. Latent Curriculum (Part 2): Gradually replace text inside the spans with <latent> positions. A key design choice: parallel curriculum (replacing all spans simultaneously, incrementally increasing latent steps) vastly outperforms sequential (one span at a time), because sequential lets the model keep most computation in the standard text distribution and avoid learning real latent computation.
3. Switch-GRPO: GRPO with four weighted rewards — correctness (±1, via math-verify), tag-format (±1), latent-usage (0/1, rewarding <swi> usage on correctly answered problems), and an optional correctness-gated brevity bonus ([0,1]).
4. Experiments: 79.3% on MATH-500, +25.7% Over Baseline
All methods share a Qwen3-8B base with matched data and decoding:
| Method | MATH-500 | GSM8K | Visible Tokens | |--------|----------|-------|----------------| | No-CoT | 34.7% | 72.8% | 45 | | Text-CoT SFT | 49.2% | 82.3% | 3100 | | iCoT (Pause Tokens) | 37.2% | 75.1% | 120 | | Coconut | 53.6% | 83.7% | 3800 | | CODI | 56.8% | 84.5% | 3400 | | CoLaR | 62.5% | 85.2% | 2900 | | Switch | 79.3% | 89.2% | 1900 |
5. Mechanistic Analysis: Making Latent Reasoning Dissectable
Using the boundary tokens as anchors, the paper answers three questions:
1. Is <swi> a learned policy or a stylistic artifact? A sharply localized, learned switching policy: p(<swi>) concentrates at consistent positions (after the problem statement, before computation), and logit lens plus linear probes distinguish imminent-switch activations across problem difficulty and type.
2. Do latent steps do real computation? Yes — causal interventions (swapping latent states across problems) sharply reduce answer accuracy, and replacing the latent block with equal-length text underperforms latent mode. Latent states carry task-relevant information.
3. Where does computation concentrate? At the first hidden-state transition at the latent block entrance; the influence of subsequent latent steps decays, consistent with Coconut's observations.
6. Why It Matters
| Method | RL support | Interpretable | Switchable | |--------|-----------|---------------|------------| | Coconut | ❌ | ❌ | ❌ | | CODI/CoLaR | ⚠️ text-only | ❌ | ❌ | | Vocabulary Mixtures | ✅ | ⚠️ | ❌ | | Pause Tokens | ❌ | ⚠️ | ❌ | | Switch | ✅ | ✅ | ✅ |
7. Limitations
1. A minimum latent-step constraint (Kmin = 4) is required, or the model exits immediately — latent "inertia" is still weak. 2. Experiments are limited to Qwen3-8B; larger models remain unverified. 3. Only mathematical reasoning is evaluated; generalization to code, logic, or common-sense tasks is open. 4. Kmax = 8 may be insufficient for problems requiring very long reasoning chains.
8. Conclusion
Switch uses one minimal primitive — a pair of <swi>/</swi> boundary tokens — to simultaneously make on-policy GRPO well-defined for latent reasoning and provide analysis anchors for probing and causal intervention. With 79.3% on MATH-500 (+25.7 over baseline) at only ~1,900 visible tokens, plus mechanistic evidence that latent steps perform real, problem-relevant computation, the key takeaway for reasoning-model designers is that explicit boundaries are not a restriction but a liberation — the path from latent CoT as a toy to latent CoT as a practical tool.
References: