PaperlayerView exact revision on arXiv

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.

  • Method

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

  • Results

    Across MATH-500 and GSM8K, SWITCH outperforms all Coconut-style baselines under matched-base-model evaluation, reaching 79.3% on MATH-500.

  • 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.

  • Takeaways & Limitations

    Experiments are restricted to 8B-parameter Qwen3 models and mathematical reasoning benchmarks, without evaluation of multi-domain reasoning or larger scales.

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.
  • 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.
  • 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.
  • 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.
  • 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.

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.
  • 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.
  • 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.
  • SWITCH overview: 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.

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.
  • 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.
  • 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.
  • Switch-SFT: 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>.
  • 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.
  • 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.

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.
  • 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.
  • 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.
  • 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.
  • 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.
  • 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.

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.
  • Q1: Switching policy: 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.
  • 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.
  • 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.

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.
  • 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.
  • 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.
  • Mechanistic verification: 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.
  • 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.

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.
  • 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.

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.
  • A. Implementation Details: 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.
  • 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.
  • 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.
  • 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.
  • 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.
  • 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.
  • 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.

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.
  • F. Mechanistic Analysis: Additional Details: 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.
  • F. Mechanistic Analysis: Additional Details: 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.
  • F. Mechanistic Analysis: Additional Details: 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.
  • F. Mechanistic Analysis: Additional Details: 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.
  • 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 θ.
  • G. Per-Subject and Per-Difficulty Visualisation: 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.

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.
  • Failure profile of wrong trajectories: 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.

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}.
  • 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.
  • 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.
  • J.1. Latent Chain-of-Thought Reasoning: 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.
  • J.1. Latent Chain-of-Thought Reasoning: 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.
  • J.1. Latent Chain-of-Thought Reasoning: 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.
  • J.1. Latent Chain-of-Thought Reasoning: 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.

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.
  • J.2. Switchable / Hybrid Reasoning: 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.
  • J.2. Switchable / Hybrid Reasoning: 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.
  • 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.
  • 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.
  • 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.
  • K. Extended Discussion: 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.
  • K. Extended Discussion: 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.
Loading 2606.13106v1…