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

arXiv:2606.13106v1cs.LGcs.CL

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–2
  • Method

    SWITCH uses learned <swi>/</swi> boundary tokens, a visible-to-latent curriculum, and Switch-GRPO to train and analyze recurrent latent reasoning.

    Source: §1 · p. 2, §3 · p. 3
  • 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. 6
  • Takeaways & 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–12

    Experiments 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 · show

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

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
Loading 2606.13106v1…