Source-linked AI summary
Demystifying Hidden-State Recurrence: Switchable Latent Reasoning with On-Policy Reinforcement Learning
Jiayu Yang, Chao Chen, Shengen Wu, Yinhong Liu, Yuxuan Fan, Lujundong Li, Songning Lai, Chengwei Qin, Zhijiang Guo
TL;DR
Latent hidden-state reasoning is difficult to optimize with on-policy RL and difficult to inspect causally because latent positions lack explicit discrete boundaries. SWITCH adds learned entry and exit tokens with Switch-GRPO, outperforming prior approaches while exposing localized, causally necessary latent computation.
Problem
Existing latent reasoning lacks discrete boundaries, leaving standard on-policy RL undefined inside latent blocks and making their computation difficult to inspect causally.
Source: §1 · pp. 1–2Method
Results
Across MATH-500 and GSM8K, SWITCH outperforms all Coconut-style baselines under matched-base-model evaluation, reaching 79.3% on MATH-500.
Source: §4.2 · p. 6Takeaways & Limitations
The results show recurrent latent reasoning can be optimized with RL and directly interpreted through localized switching decisions and causally necessary latent computation.
Source: §6 · pp. 11–12Experiments are restricted to 8B-parameter Qwen3 models and mathematical reasoning benchmarks, without evaluation of multi-domain reasoning or larger scales.
Source: Limitations · p. 12
Abstract
from arXiv · showhide
Latent chain-of-thought compresses reasoning by replacing visible reasoning traces with continuous hidden-state recurrence, but existing formulations are difficult to optimize with standard on-policy reinforcement learning (RL) and hard to interpret causally. Our key insight is that a single pair of explicit boundary tokens can address both issues at once: discrete entry and exit anchors make the latent block compatible with standard on-policy RL, and the same anchors offer a natural foothold for mechanistic analysis. Motivated by this, we propose SWITCH, a switchable latent reasoning framework. The model emits <swi> to enter latent mode and </swi> to exit. Because the boundaries are ordinary discrete tokens, the GRPO policy ratio is well-defined at every decision point. The same anchors also expose the latent steps to direct probing and causal intervention. We train the model with a visible-to-latent curriculum and a Switch-GRPO objective that propagates gradients through recurrent latent computation. SWITCH consistently outperforms prior hidden-state-recurrence latent reasoning approaches at similar scale. Mechanistic analysis through the boundary tokens further reveals three findings: (i) <swi> is a sharply localised, learned switching policy rather than a stylistic artefact; (ii) the latent step it opens performs problem-specific, causally important computation rather than acting as an inert placeholder; and (iii) that computation is concentrated at a single hidden-state transition on entry. Together, these results show that hidden-state-recurrence latent reasoning is both RL-trainable and open to direct mechanistic analysis, including of how on-policy RL itself improves the model from the inside.
1. Introduction
SWITCH uses explicit <swi>/</swi> boundary tokens to make hidden-state-recurrence latent reasoning compatible with on-policy RL and directly inspectable. Its curriculum and Switch-GRPO training address prior optimization and interpretability challenges while improving performance and latent-use efficiency.
Motivation
- Hidden-state recurrence compresses visible reasoning by feeding each latent step’s last-layer hidden state back as the next input embedding.The computation remains in the LLM’s existing representation space and reuses its forward pass without additional architectural components.Source: §1 · p. 1
Challenges
- Latent positions lack token policy densities, leaving standard on-policy methods such as GRPO undefined inside latent blocks.Existing systems therefore either skip RL or use text-only rollouts that diverge from inference-time decoding.Source: §1 · pp. 1–2
Core idea
- SWITCH introduces <swi> to enter latent mode and </swi> to exit, making latent reasoning a learned per-step decision with discrete RL anchors.The same boundaries expose latent computation to direct probing and causal intervention, addressing concerns that latent steps may be inert placeholders.Source: §1 · p. 2
Training
- SWITCH trains through SFT, gradual replacement of wrapped text with <latent> positions, and Switch-GRPO gradient propagation through recurrent latent computation.This three-phase procedure is described as a visible-to-latent curriculum followed by reinforcement-learning optimization.Source: §1 · p. 2
Results
- 79.3% on MATH-500, +25.7 points above the strongest same-scale Coconut-style baseline; Switch-GRPO also halves latent invocation while raising invoked-problem accuracy by +12.6 points.These results are reported for Switch versus the strongest Coconut-style baseline and for Switch-GRPO versus the SFT-only checkpoint, respectively.Source: §1 · p. 2
2. Related Work
Prior latent chain-of-thought methods differ in how latent tokens are represented: hidden-state recurrence feeds back hidden states, whereas vocabulary mixtures sample convex combinations of vocabulary embeddings. Because vocabulary mixtures are samplable and support direct policy gradients, recent reinforcement-learning work has moved away from hidden-state recurrence.
Latent CoT representations
- Latent CoT methods split into hidden-state recurrence, which feeds back the previous step’s last-layer hidden state, and vocabulary mixtures, which sample top-k convex combinations of vocabulary embeddings.Vocabulary mixtures use Gumbel-Softmax for sampling.Source: §2 · p. 2
Reinforcement learning
- Vocabulary mixtures are samplable and admit direct policy gradients, motivating recent RL work to abandon hidden-state recurrence.The passage frames this shift as a consequence of optimization compatibility.Source: §2 · p. 2
SWITCH overview
- SWITCH training combines SFT, a curriculum replacing text with latent positions, and Switch-GRPO for on-policy RL using answer reward.The curriculum jointly constitutes Switch-SFT before Switch-GRPO.Source: Figure 1 · p. 3
- At inference, the model emits <swi> to enter latent mode, executes <latent> steps, and emits </swi> to resume text decoding.The figure also describes hidden-state recurrence within the latent block.Source: Figure 1 · p. 3
3. Method
SWITCH makes hidden-state recurrence trainable by using discrete boundary tokens to mark latent blocks across supervised and reinforcement-learning stages. Its method combines recurrent latent execution, a parallel latent curriculum, and switch-aware GRPO rewards.
Overall training recipe
- Boundary tokens <swi> and </swi> provide discrete handles for training latent blocks across SFT, curriculum replacement, and policy optimization.The three-phase procedure teaches when to switch, how to compute inside the block, and how to optimize correctness and tag well-formedness.Source: §3 · p. 3, §3.2 · p. 4, §3.3 · p. 5
Latent execution
- Latent computation recursively feeds each step’s previous last-layer hidden state as the next input embedding, requiring a separate forward pass per latent step.Unlike text positions, latent positions emit a Dirac mass rather than sampled tokens, so they lack a direct policy density.Source: §3.1 · pp. 3–4
Switch-SFT
- Phase 1 tags contiguous high-entropy visible-CoT runs with <swi>/</swi>, teaching the model where to enter and exit latent reasoning.High-entropy positions are defined by uncertainty in the base model’s next-token distribution on a mathematical CoT corpus.Source: §3.2 · p. 4
- Phase 2 progressively replaces text inside switch spans with <latent> positions while retaining boundary-token loss, avoiding the collapse associated with one-shot replacement.At inference, the model enters latent mode at <swi>, performs at least K_min latent steps, and may exit with </swi>.Source: §3.1 · p. 3, §3.2 · p. 4
Latent curriculum
- The parallel curriculum schedule is substantially better than the sequential schedule because replacing every span simultaneously forces surrounding text to condition on latent hidden states.The sequential schedule can satisfy next-token-prediction loss while avoiding computation in latent space.Source: §3.2 · p. 4, Figure 2 · p. 4
Switch-GRPO
- Switch-GRPO executes deployed-decoder rollouts and combines correctness, tag-format, latent-usage, and optional brevity rewards to optimize latent reasoning.The latent-usage reward encourages correct answers to invoke <swi> rather than fall back to plain text.Source: §3.3 · p. 5
4. Experiments
Experiments show that Switch outperforms matched Coconut-style baselines and that Switch-GRPO adds substantial latent reasoning ability beyond curriculum SFT. Reinforcement learning also calibrates when latent computation is invoked and exposes an accuracy–length tradeoff.
Experimental setup
- Experiments use Qwen3-8B with <swi>, </swi>, and <latent> on OpenR1-Math, evaluated on MATH-500 and GSM8K.Baselines use the same base model, matched data, and matched decoding settings.Source: §4.1 · p. 5, §4.2 · p. 6
Baselines
- The comparison re-implements direct-answer, text-CoT, iCoT, Pause Tokens, Coconut, CODI, and CoLaR under matched conditions.This provides apples-to-apples comparisons across standard and latent reasoning approaches.Source: §4.1 · p. 5, §4.2 · p. 6
Headline performance
- Switch reaches 79.3% on MATH-500 and 89.2% on GSM8K, exceeding all Coconut-style baselines under the matched-base-model protocol.The reported analyses use a representative run with complete training, decoding, and intervention logs.Source: §4.2 · p. 6
Switch-GRPO contribution
- Latent-conditional accuracy improves by +12.6 points after Switch-GRPO, while switch rate falls from 81% to 58% with identical weights, vocabulary, and decoding.The comparison isolates reinforcement learning from curriculum SFT alone using the same MATH-500 set and K_min =4 greedy decoding.Source: §4.2 · p. 6
Switch calibration
- During training, latent invocations per problem decline from ∼1.5 to ∼1 and visible-token usage contracts from ∼2900 to ∼1900.The pattern indicates selective use of latent reasoning rather than indiscriminate invocation.Source: §4.2 · p. 7
Accuracy–efficiency tradeoff
- A brevity-bonus operating point trades about three points of accuracy for ∼33% shorter outputs and 0% max-length truncation.Varying the Switch-GRPO reward lets users select an explicit accuracy–length operating point.Source: §4.2 · p. 7
5. How Does Latent Work in Reasoning?
The boundary token <swi> functions as a learned, sharply localized switching policy rather than a stylistic tag. The following latent step performs causally important reasoning, concentrated in a single hidden-state transition on entry and preserved by the K_min constraint.
Q1: Switching policy
- <swi> behaves as a learned switching policy, sharply localized at the boundary rather than spanning a stylistic window.At annotated positions, <swi> has rank ≤1.7 on both checkpoints, versus rank ∼103 at random non-boundary positions; its probability collapses several orders of magnitude one token later.Source: §5 · pp. 8, 10
- Switch-GRPO softens the boundary probability from p(<swi>) = 0.85 after SFT to 0.48, while preserving the localized spike.The softened peak has margin near zero, consistent with less aggressive switching under uncertainty and the halved switch rate.Source: §5 · pp. 8–9
Q2: Causal contribution
- On the diagnostic subset, zeroing the latent state reduces accuracy from 100% to 33.3%, while same-norm random replacement costs 9.5 points and skipping costs 19.0 points.These interventions show that the latent step carries specific computation rather than generic perturbation or redundant text.Source: §5 · p. 10
Q3: Location of computation
- The latent computation is concentrated at a single hidden-state transition on entry, while the K_min constraint prevents collapse to one hidden forward pass.The first latent step becomes more diffuse and problem-conditional under the logit lens, whereas p(</swi>) ≈1 at every latent step regardless of correctness.Source: §5 · p. 11
6. Conclusion
SWITCH combines learned switching, curriculum training, and Switch-GRPO to make recurrent latent reasoning both reinforcement-learning trainable and directly interpretable. Experiments show competitive performance, an adaptable accuracy–efficiency trade-off, and causally necessary latent computation verified through explicit boundaries.
Contributions
- SWITCH integrates a learned switch token, a three-phase curriculum, and Switch-GRPO into hidden-state-injection models.These components constitute the proposed switchable latent reasoning framework.Source: §6 · p. 11
Empirical results
- SWITCH outperforms competitive baselines while providing an adaptable accuracy–efficiency trade-off.The conclusion reports this result across extensive experiments without specifying a single aggregate metric.Source: §6 · p. 11
Mechanistic verification
- The switch decision is highly localized and linearly decodable, reaching 91.9% probe accuracy.The explicit boundary design enables this direct verification of switching behavior.Source: §6 · p. 11
- Causal analysis confirms that latent reasoning steps are functionally necessary, with zero intervention collapsing diagnostic accuracy from 100% to 33% (−66.7 points).Random-norm and Skip interventions are far less destructive, and the relevant computation is the specific hidden state of Eq. 1 rather than any non-zero perturbation.Source: §6 · p. 11, Figure 8 · p. 12
Conclusion
- Overall, recurrent latent CoT can be successfully optimized via RL and directly interpreted.The conclusion links trainability and interpretability as the central outcome of the framework.Source: §6 · p. 12
Limitations
The evaluation is limited to 8B-parameter Qwen3 models and mathematical reasoning benchmarks, while Switch-GRPO updates latent representations only indirectly through a frozen KV cache. Larger-scale, multi-domain settings and the effects of this training design remain unevaluated.
Evaluation scope
- Experiments cover only 8B-parameter Qwen3 models and the MATH-500 and GSM8K mathematical reasoning benchmarks.The study does not evaluate multi-domain reasoning or larger model scales, where the balance between learned switching and latent depth may differ.Source: Limitations · p. 12
Optimization scope
- Switch-GRPO’s gradients flow through rollout text segments, while latent positions contribute through a frozen KV cache.Consequently, the latent representation is shaped primarily by the Phase 2 training component described in the passage.Source: Limitations · p. 12
A. Implementation Details … E. Visible-Token CDF
The paper specifies SWITCH’s implementation, Switch-GRPO objective, training safeguards, and checkpoint selection, then characterizes visible-token efficiency and the segmented recurrent forward algorithm. Key results include a controllable compression–accuracy trade-off and early stopping before late-training reward hacking.
A. Implementation Details
- SWITCH registers <swi>, </swi>, and <latent> as special tokens, resizing Qwen3-8B embeddings and initializing them from a content-neutral seed to avoid rank-collapse.The embeddings grow from 151 936 to 151 672 entries, with IDs 151 669–151 671.Source: A. Implementation Details · p. 16
- Phase-2 training uses c=2, Kmax=8, a 48-token per-sample latent cap, p_unif=0.1, and three-epoch warm-started stages over k∈{0,...,8}.Parallel replacement is the default and supplies all Phase-3 initializations.Source: A. Implementation Details · p. 16
B. Switch-GRPO Loss, in Full
- Switch-GRPO uses G=5 rollouts, εc=0.2, β=10^-3, learning rate 10^-6, and three inner epochs, with πθold also serving as the KL anchor.This removes the separate reference model and saves roughly 18 GB/GPU.Source: A. Implementation Details · p. 16, B. Switch-GRPO Loss, in Full · p. 18
D. Algorithm Boxes
- Segmented backward splits rollouts at <swi>/</swi> boundaries, streams the key-value cache left-to-right, and backpropagates only through text segments while latent segments run under torch.no_grad().Text segments contribute to the clipped surrogate loss, while latent segments store no autograd state.Source: A. Implementation Details · p. 16, B. Switch-GRPO Loss, in Full · p. 18
A. Implementation Details
- 69.0% MATH-500 accuracy at 1 276 average visible tokens and 0% max-length truncation contrasts with the default 72.6% at 1 919 tokens and 18.4% truncation.The brevity bonus applies only to correct responses using at least one <swi> block, yielding a controllable Pareto trade-off.Source: A. Implementation Details · p. 17, B. Switch-GRPO Loss, in Full · p. 17
C. Per-Checkpoint Trajectory of Switch
- The representative 1,964-step trajectory enters reward hacking after roughly step 1,200, with switch rate reaching 100% and latent invocations rising from ∼1 to ∼13 per problem.Average reward declines as extra latent computation fails to improve correctness, so training early-stops at step 800.Source: C. Per-Checkpoint Trajectory of Switch · p. 18, Figure 9 · p. 19
E. Visible-Token CDF
- Figure 10 shows the brevity-bonus variant dominating the SFT baseline through the median visible-token count while losing few problems to the high-token tail.The comparison uses empirical CDFs on MATH-500 for the SFT baseline, Switch-GRPO endpoint, and brevity-bonus operating point.Source: E. Visible-Token CDF · p. 18
D. Algorithm Boxes
- The recurrent forward algorithm partitions positions into maximal constant-L segments, injects embeddings for visible tokens and prior hidden states for latent tokens, and carries a streaming KV cache.Each segment is processed sequentially before returning all hidden states and logits.Source: E. Visible-Token CDF · p. 19
F. Mechanistic Analysis: Additional Details · G. Per-Subject and Per-Difficulty Visualisation
The supplementary mechanistic details specify probe, teacher-forced switch, intervention, and training-trajectory protocols, while per-subject and per-difficulty analyses report headline MATH-500 performance and latent-use splits. Together, these details clarify evaluation conditions, robustness, and the Switch-GRPO training procedure.
F. Mechanistic Analysis: Additional Details
- Probe accuracies use balanced binary datasets, ℓ2-regularised logistic classifiers with C=1.0, and a single 80:20 train/test split.The probes sample equal numbers of non-boundary positions per swi-start position, and reported values are test-set accuracies.Source: F. Mechanistic Analysis: Additional… · pp. 19–20
- Teacher-forced switch metrics compare annotated <swi> prefixes with random non-boundary positions on paired MATH-500 problems.The reported readouts are entropy, p(<swi>), <swi> rank, and log-margin to the top token.Source: F. Mechanistic Analysis: Additional… · p. 19
- The switch-window collapse one token after the boundary remains robust across offsets −8, . . . , +8.At that position, p(<swi>) is ∼2×10−6 and the rank is ∼5 000.Source: F. Mechanistic Analysis: Additional… · p. 19
- The representative MATH-500 trajectory distinguishes curriculum-only After SFT checkpoints from After Switch-GRPO reinforcement-learning checkpoints and a + brevity bonus variant.The post-RL endpoint is bolded, while the strongest end-to-end Switch-GRPO run reaches 79.3% MATH-500 / 89.2% GSM8K.Source: Table 8 · p. 20, F. Mechanistic Analysis: Additional… · p. 20
- Interventions are evaluated only on normally correct problems that used at least one latent block, with greedy decoding, Kmin =4, and the headline max_new_tokens setting.Correctness is determined by math-verify before applying each intervention mode.Source: F. Mechanistic Analysis: Additional… · p. 20
G. Per-Subject and Per-Difficulty Visualisation
- One Switch-GRPO step generates G rollouts with real hidden-state injection, computes rewards and normalized advantages, then performs segmented backward passes through text segments while latent segments remain no_grad.The update applies clipped policy-ratio and KL terms, updates cached quantities, steps the optimizer, and replaces θold with θ.Source: G. Per-Subject and Per-Difficulty… · p. 21
- The algorithm accumulates segment losses and token counts across rollout groups before applying the optimizer step and returning the normalized loss and rewards.Latent segments update the retained hidden-state set without gradient computation, whereas text segments contribute gradients.Source: G. Per-Subject and Per-Difficulty… · p. 21
H. Generation Trace Analysis
Wrong post-RL trajectories on MATH-500 are substantially longer than correct trajectories, regardless of latent usage, and exhibit slightly less confident switch decisions.
Failure profile of wrong trajectories
- Wrong trajectories are substantially longer than correct ones, both with and without latent usage.The analysis stratifies Switch’s post-RL trajectories by correctness and latent usage on MATH-500.Source: H. Generation Trace Analysis · p. 21
- Switch decisions are less confident on wrong trajectories, with entropy 0.717 vs. 0.608 and p(<swi>) =0.669 vs. 0.763.The reported comparisons are for wrong versus correct trajectories, respectively.Source: H. Generation Trace Analysis · p. 21
I. Ablations · J. Full Related Work · J.1. Latent Chain-of-Thought Reasoning
The ablation identifies K_min=4 as the default because eliminating latent dwell reduces MATH-500 accuracy to 53.0%. Related work spans hidden-state recurrence, vocabulary-embedding mixtures, non-decoding thinking tokens, and prior smaller-scale systems, while Switch preserves interpretable text reasoning outside latent blocks.
I. Ablations
- 53.0% accuracy results when K_min=0 collapses latent blocks to one hidden forward pass, whereas K_min=4 recovers the training distribution and is the default.The inference sweep tested K_min ∈ {0, 2, 4, 8, 16}.Source: §I · p. 21
J. Full Related Work
- The full related-work discussion situates Switch among several approaches to latent chain-of-thought reasoning.The section expands the treatment of related work promised in §2.Source: J. Full Related Work · p. 21
J.1. Latent Chain-of-Thought Reasoning
- Coconut and CODI feed previous-step last-layer hidden states back as next input embeddings, placing entire reasoning steps in latent space between text tokens.Coconut uses a multi-stage curriculum replacing explicit CoT tokens with k·c latent positions; CODI retains the injection mechanism but changes the curriculum.Source: §J.1 · p. 22
- Vocabulary-embedding mixture methods represent latent tokens as convex combinations of vocabulary input embeddings rather than directly reusing hidden states.Soft-Thinking uses next-token softmax probabilities as mixture weights, while Latent-SFT uses top-k mixtures with stochastic Gumbel-Softmax targets.Source: §J.1 · pp. 22–23
- Latent-SFT reports 2.7×–5.5× shorter chains than explicit SFT on six math benchmarks, and Latent-GRPO proposes vocabulary superposition with one-sided Gumbel margins for RL.Latent-GRPO contrasts its method with Coconut’s direct hidden-state latent tokens; the supplied passage truncates the SofT-GRPO discussion.Source: §J.1 · p. 23
- Prior hidden-state-recurrence studies mainly use GPT-2 or 1–2B-scale models, while Switch targets a regime with learned switching, tuned curriculum, and RL.Coconut’s LLaMA-3-8B probe improved GSM8K by 1.4 points over SFT, but lacked a tuned curriculum, learned switching token, or RL; CoLaR is also hidden-state-recurrence-based.Source: §J.1 · p. 23
- Pause tokens, filler tokens, and implicit-CoT internalisation insert non-decoding thinking tokens without continuous-state feedback, whereas Switch preserves visible text CoT outside <swi> blocks.This design keeps the visible trajectory interpretable; the supplied passage also mentions multimodal IVT-LR but truncates its description.Source: §J.1 · p. 23
J.2. Switchable / Hybrid Reasoning … K. Extended Discussion
SWITCH contrasts with training-free adaptive switching by jointly learning when to enter latent reasoning and the latent dynamics inside it. The discussion argues that explicit boundary tokens make hidden-state recurrence compatible with on-policy RL and enable mechanistic evidence that latent computation is task-relevant.
J.2. Switchable / Hybrid Reasoning
- SwiReasoning reports 1.8–3.1 accuracy points and 57–79% token-efficiency gains across math, STEM, coding, and general benchmarks.It uses entropy-based switching with a hard switch budget while keeping the reasoning model frozen.Source: §J.2 · p. 23
- Unlike training-free entropy switching, SWITCH jointly optimizes the discrete <swi> entry point and latent dwell dynamics end-to-end, including with RL.The comparison motivates training because SwiReasoning applies latent steps to a model not trained for them and fixes switch locations externally.Source: §J.2 · pp. 23–24
- SWITCH combines adaptive-compute decisions and latent representations in one trained model rather than emitting extra thinking as text.This distinguishes it from adaptive test-time compute methods that always express additional computation through visible tokens.Source: §J.2 · p. 24
J.3. Reinforcement Learning for Reasoning and Latents
- Standard GRPO is a discrete-token policy optimizer, while vocabulary-mixture latent methods restore this interface using Gumbel-Softmax or Gumbel reparameterization.The cited examples include Latent-GRPO and SofT-GRPO, which use different Gumbel-based constructions for latent sampling and optimization.Source: §J.3 · p. 24
J.4. Interpretability of Internal Reasoning States
- Logit-lens analysis and linear probing provide complementary tools for testing what intermediate hidden states represent.Logit lens qualitatively reads hidden states through the language-model head, whereas linear probing tests whether a target property is encoded in frozen activations.Source: §J.4 · p. 24
K. Extended Discussion
- Switch-GRPO shows hidden-state recurrence is RL-compatible because the policy ratio requires tractable density only at discrete decision points emitting <swi> and </swi>.This provides a constructive counterexample to claims that recurrent latent states cannot be optimized with on-policy RL.Source: K. Extended Discussion · p. 24
- After Switch-GRPO, p(<swi>) falls from 0.85 to 0.48, entropy rises to ∼0.5, neighbour contrast remains ∼102, and latent-conditional accuracy nearly doubles.The passage also reports that the switch rate halves, indicating RL changes switching behavior without erasing the SFT-learned policy contrast.Source: K. Extended Discussion · p. 25
- Zeroing latent hidden states reduces accuracy by 66.7 points, whereas replacing them with equal-norm random vectors costs 9.5 points.The intervention shows that answers depend on the specific hidden state produced by the latent computation, not merely on arbitrary nonzero activity.Source: K. Extended Discussion · p. 25