PaperlayerView exact revision on arXiv

Source-linked AI summary

SLPO: Scaling Latent Reasoning via a Surrogate Policy

Runyang You, Zhiyuan Liu, Yongqi Li, Wenjie Li

arXiv:2607.19691v2cs.CLcs.AIcs.LG

TL;DR

Latent reasoners lack a tractable per-step likelihood and adaptive stopping interface for outcome-reward reinforcement learning. SLPO supplies both through a surrogate latent-transition policy and learned stopping head, improving Pass@8 and Pass@16 across all 12 evaluated settings with gains up to 12.07 percentage points.

  • Problem

    Latent trajectories lack a tractable per-step likelihood and adaptive stopping interface, limiting outcome-reward reinforcement learning for latent test-time scaling.

  • Method

    SLPO uses a differentiable surrogate policy over latent transitions and a correctness-supervised stopping head refined by outcome-reward optimization.

  • Results

    SLPO improves Pass@8 and Pass@16 in all 12 evaluated backbone–dataset settings, with gains of up to 12.07 percentage points.

  • Takeaways & Limitations

    SLPO produces latent test-time scaling across continuous and soft-token inference, allocating longer trajectories to harder instances alongside improved deterministic accuracy.

Abstract

from arXiv · show

Reinforcement learning with verifiable rewards has become the predominant recipe for eliciting test-time scaling in explicit Chain-of-Thought reasoners. Yet this scaling path remains computationally costly, since every intermediate step must be decoded as a language token. Latent reasoning instead carries intermediate computation as continuous vectors and already matches or surpasses explicit CoT at far shorter horizons. Despite this promise, latent reasoners remain largely imitation-bound, while explicit CoT has already moved past imitation via outcome-reward RL. Latent trajectories lack a tractable per-step likelihood and an adaptive stopping interface under fixed thinking budgets, so outcome rewards cannot elicit latent test-time scaling. We introduce Surrogate Latent Policy Optimization (SLPO) to bring outcome-reward RL to autoregressive latent reasoners: an empirical surrogate policy density over latent transitions for trajectory-level credit assignment, and a correctness-supervised stopping head that outcome-reward optimization refines into a variable-horizon policy. Across continuous and soft thinking settings, SLPO improves Pass@$k$ under parallel sampling and allocates longer latent computation to harder instances with higher deterministic accuracy.

1 Introduction

SLPO brings outcome-reward reinforcement learning to latent reasoning by combining a surrogate policy density for latent transitions with a learned stopping policy. It improves latent test-time scaling and allocates longer computation to harder problems.

  • Explicit CoT scales test-time computation expensively because every intermediate reasoning step is decoded as a natural-language token.Many decoded tokens support linguistic coherence, exposition, or redundancy rather than the underlying problem-solving state.
  • Latent reasoners remain imitation-bound, while outcome-reward RL has enabled explicit CoT to scale beyond teacher imitation.Latent transitions lack the tractable per-step policy likelihood needed for trajectory-level credit assignment.
  • SLPO uses a surrogate hidden-space likelihood to assign rollout-advantage credit over continuous latent transitions and a correctness-supervised stopping head for learnable horizons.The stopping head is initialized through a correctness-supervised cold start before outcome-reward optimization refines the variable-horizon policy.
  • SLPO’s scaling effect persists across RLOO and GRPO, transfers to soft-token latent inference, and assigns longer latent trajectories to harder problems.The learned stopping policy converts a fixed thinking budget into difficulty-adaptive computation.
  • 12.07 percentage points: SLPO improves both Pass@8 and Pass@16 across all 12 backbone–dataset settings.The gains span two latent reasoners, two backbones, and three held-out benchmarks.

2 Related Work

Prior work contrasts explicit Chain-of-Thought’s discrete-token traces with latent reasoning’s compressed continuous thoughts. Outcome-reward RL scales explicit CoT through tractable token-policy likelihoods, whereas latent reasoning lacks exposed action log-probabilities and requires alternative policy interfaces.

  • Latent reasoning: Latent reasoning autoregresses in hidden space, replacing linguistically redundant discrete traces with compressed continuous thoughts.COCONUT, CODI, and CoLaR are representative latent-reasoning approaches cited in the related work.
  • Reinforcement learning: Outcome-reward RL has become central to scaling explicit CoT, using token policies with tractable sequence likelihoods.GRPO uses group-relative objectives, while RLOO uses leave-one-out baselines.
  • Reinforcement learning: Latent reasoning lacks exposed action log-probabilities because its continuous hidden transitions do not provide tractable token-policy likelihoods.Recent latent-policy methods address this gap through vocabulary-mediated or architecture-specific policy interfaces.

3 Preliminaries

Explicit CoT provides token-level action probabilities that support policy-gradient credit assignment, whereas latent CoT performs autoregressive computation in hidden space and lacks this token-policy likelihood. Dropout-induced stochastic latent trajectories provide an exploration substrate for reinforcement learning.

  • Explicit CoT and RL: Explicit CoT samples discrete reasoning tokens autoregressively and then generates the final answer conditioned on the sampled trace.The reasoning distribution is π_θ(· | x_i, y_<t).
  • Explicit CoT and RL: Its rollout likelihood is directly available from token probabilities, enabling policy-gradient methods to weight sampled rollouts by an advantage.The advantage is described as bA_i = R(â_i, a★_i) − b_i.
  • Explicit CoT and RL: Thus explicit CoT exposes the policy interface needed for RL because sampled reasoning trajectories already carry token-level action probabilities.This supports assigning credit to sampled reasoning rollouts.
  • Latent CoT: Latent CoT replaces discrete traces with autoregressive hidden-space trajectories, bypassing vocabulary sampling and removing the token-policy likelihood available in explicit CoT.A trajectory is represented as h_i,1:Tmax = (h_i,1, ..., h_i,Tmax), with h_i,t ∈ R^d.
  • Latent CoT: Dropout inference can make latent reasoning stochastic, supplying the exploration substrate for RL through sampled latent trajectories.The sampled trajectory is denoted h̃_i,1:Tmax.

4 Methodology

SLPO adds a correctness-supervised stopping head and a Gaussian surrogate for latent transitions, enabling reward-weighted optimization of latent computation, answer generation, and stopping decisions. At inference, the learned gate adaptively stops sequential latent reasoning without enumerating candidate prefixes.

  • Stopping head: SLPO first supervises a stopping head using final-answer correctness across candidate trajectory lengths, enabling learned continue-versus-stop decisions before outcome-reward optimization.The method evaluates each sampled trajectory at candidate stopping lengths and trains the gate to favor answer-valid stopping times.
  • Rollout policy objective: SLPO combines latent-transition, answer-token, and stopping-time likelihoods in a reward-weighted objective whose empirical update targets expected rollout reward.The behavior law comes from MC-dropout trajectories, answer decoding, and the stopping gate; RLOO or GRPO provide the stop-gradient advantage.
  • Surrogate transition likelihood: SLPO constructs a tractable isotropic Gaussian surrogate for stochastic latent transitions from repeated MC-dropout evaluations, providing per-step scores for advantage-weighted credit assignment.The surrogate uses empirical means and floored variances, while sampled rollout states serve as targets for the recomputed moments.
  • Inference: At inference, the stopping head emits probabilities after each latent state and halts at the first threshold crossing, so computation varies by instance without candidate-prefix enumeration.The final answer is decoded only after the stopping event.

5 Experiment Setup

Section 5 evaluates whether SLPO improves outcome performance across latent reasoners and optimization algorithms, transfers across latent interfaces, and enables difficulty-adaptive computation. Experiments use mathematical reasoning benchmarks, COCONUT and CODI checkpoints, and controlled training and evaluation protocols.

  • Experiment goals: The experiments test SLPO’s outcome-level gains, compatibility with policy-optimization algorithms, transfer across continuous and soft latent interfaces, and difficulty-adaptive computation.These objectives assess whether a fixed latent budget can support adaptive test-time computation.
  • Datasets: Experiments train on GSM8K-Aug and evaluate on GSM8K-Test, GSM-Hard, and MultiArith.The soft-latent transfer experiment additionally uses MATH500, AIME 2025, and AMC23.
  • Models and baselines: SLPO is applied to COCONUT and CODI checkpoints with GPT-2 (124M) and Llama-3.2-1B-Instruct backbones.Released baselines use T_max = 6 latent steps, whereas +SLPO raises the maximum to T_max = 12 and lets the stopping gate select the realized length.
  • Training and evaluation protocols: Unless otherwise specified, +SLPO trains the stopping head first, then optimizes with RLOO using (K, G) = (4, 8) on GSM8K-Aug.At test time, models use the learned stopping gate with its threshold selected by validation sweep.
  • Training and evaluation protocols: The evaluation reports Acc from one deterministic rollout and Pass@k for k∈{8, 16} from independent MC-dropout evaluation seeds.Acc disables dropout, while Pass@k averages over three seeds with p = 0.1.

6 Main Results

SLPO consistently elicits latent test-time scaling across backbones, reasoners, and held-out benchmarks, improving parallel-sampling performance and deterministic accuracy. It also outperforms latent and explicit baselines while reducing average reasoning length.

  • Latent test-time scaling: SLPO raises both Pass@8 and Pass@16 across all 12 backbone–dataset combinations.The gains span COCONUT and CODI, GPT-2 and Llama-3.2-1B, and three held-out benchmarks without changing the vector recurrence.
  • Latent test-time scaling: Deterministic accuracy improves on average for every backbone–reasoner pair.This accompanies the consistent Pass@8 and Pass@16 gains reported across the 12 combinations.
  • Latent test-time scaling: 12.07 percentage points, or 26.8% relative, is the Pass@8 improvement on MultiArith with Llama-3.2-1B COCONUT.This expands the number of successful parallel rollouts produced by the latent policy.
  • Broader model comparisons: CODI+SLPO achieves the strongest average accuracy among latent methods on Llama-3.2-1B with 5.79 latent steps on average.Compared with Latent-SFT, it raises average accuracy by 1.05 percentage points while using 50.3% fewer latent steps.
  • Broader model comparisons: On GSM8K, CODI+SLPO surpasses CoT-SFT by 1.17 percentage points while using 75.2% fewer reasoning steps.The comparison includes explicit CoT-SFT, iCoT, and latent baselines.

7 Analysis

SLPO’s performance depends more on sampling complete trajectories than on surrogate-estimation passes, while its surrogate likelihood transfers across outcome-reward optimizers and soft latent inference. SLPO also reshapes latent geometry and enables difficulty-adaptive computation through longer trajectories for harder problems.

  • Rollout scaling: Increasing group size G consistently improves Pass@2 on GSM-Hard and MultiArith, whereas performance is substantially less sensitive to K.K controls stochastic forward passes for surrogate estimation; G controls complete trajectories sampled per problem.
  • Optimizer transfer: RLOO and GRPO show closely aligned scaling curves across COCONUT and CODI, preserving SLPO’s gains over released backbone initialization.The comparison uses the same surrogate, rollout budget, and latent backbones across MultiArith, GSM-Hard, and GSM8K.
  • Soft latent inference: SLPO transfers to soft latent inference on Llama3.2–1B and Llama3.2–3B, compared with vocabulary-routed baselines.The soft-token 1B setup also tracks mean rollout length during training.
  • Latent geometry: Inter-step cosine distance increases for every backbone–dataset pair after SLPO, with the largest change on GSM-Hard.This indicates that successive latent states become more differentiated and produce clearer stage-wise progression.
  • Latent geometry: Effective rank decreases in every backbone–dataset–method combination, while increased inter-step separation and rank compression steer rollouts through task-relevant directions.The prefix consequently concentrates into a lower-dimensional subspace after SLPO.
  • Adaptive stopping: Harder problems receive longer latent trajectories, with Pearson r= 0.30 on validation and r= 0.26 on test; both p< 0.001.Easy bins terminate earlier, while the hardest bins receive the longest hidden computation.

8 Conclusion … A.3 SLPO Training

SLPO enables outcome-reward optimization for autoregressive latent reasoners through a surrogate latent-transition policy and correctness-supervised adaptive stopping. The appendix details its initialized checkpoints, stopping-gate cold start, rollout estimation, optimization, and hardware setup.

  • 8 Conclusion: SLPO introduces a differentiable surrogate policy over MC-dropout latent transitions for trajectory-level reward credit in vector-space reasoning.A correctness-supervised stopping-gate cold start provides a stopping-time prior for jointly optimizing adaptive computation and latent reasoning.
  • 8 Conclusion: Across continuous and soft thinking, SLPO improves Pass@𝑘 under parallel sampling and produces longer latent trajectories on harder instances.
  • A.1 Base Models and Initialization: The stopping-gate cold start and SLPO are applied directly to initialized weights without retraining the latent reasoners’ imitation stage.Explicit CoT baselines follow the training protocols reported by Hao et al. [10], Shen et al. [12], and Tan et al. [13].
  • A.2 Stopping-gate Cold Start: The stopping-gate cold start samples N=4 stochastic latent trajectories with rollout dropout p=0.1 up to Tmax=12, with candidate stopping lengths in [Tmin,Tmax] and Tmin=3.It uses the GSM8K-Aug split and raises the maximum budget from the ungated COCONUT/CODI horizon of six.
  • A.2 Stopping-gate Cold Start: The cold-start stage uses Adafactor at learning rate 10−4 for 15 epochs and restores the checkpoint with highest validation accuracy under patience 20.GPT-2 and Llama-3.2-1B configurations use different per-device batch and accumulation settings, while the stop gate uses dropout p=0.1.
  • A.3 SLPO Training: SLPO estimates advantages with RLOO over G=8 sampled latent trajectories per problem and uses K=4 independent MC-dropout forwards for each surrogate transition likelihood.Rollouts use MC-dropout p=0.1; realized latent states are stop-gradient targets, with gradients flowing through recomputed moments.
  • A.3 SLPO Training: SLPO uses Adafactor at learning rate 10−6 for up to 40,000 steps, with 2,000 warmup steps and effective prompt batch 64, while omitting PPO clipping and KL penalties.Training uses β=0 and runs over 4 processes on 4×RTX 5880 Ada GPUs.

A.4 Computational Cost Analysis

SLPO adds training-time computation for surrogate transition estimation, with cost controlled by group size G and stochastic forwards K. At inference, it uses the original latent recurrence and stopping head, so compute depends on adaptive trajectory length.

  • Training-time cost: Under (K, G) = (4, 8), an optimizer step takes 1.130 seconds, with surrogate construction consuming 0.274 seconds, or 24% of total time.Rollout generation and backward pass take 0.294 and 0.545 seconds, respectively.
  • Training-time cost: Reducing G= 8 to G= 4 at K= 4 lowers step time to 0.676 seconds (0.60×), while increasing Kfrom 4 to 8 at G= 8 raises it to 2.021 seconds (1.79×).These controls expose how group size and stochastic-forward count affect SLPO training cost.
  • Inference cost: At inference, SLPO uses the original latent recurrence and learned stopping head, with compute governed by realized adaptive trajectory length rather than K or G.The surrogate estimator and K stochastic forwards are used only during training.

A.5 Surrogate Optimization Guarantees … C Rollout Hyperparameter Sweeps

The appendix formalizes SLPO’s surrogate optimization guarantees, specifies deterministic stopping-gate inference and evaluation protocols, and extends the method to soft latent inference. It also documents baseline configurations, policy-optimization comparisons, and rollout hyperparameter sweeps.

  • A.5 Surrogate Optimization Guarantees: SLPO optimizes an explicit reward-weighted surrogate score on a detached rollout batch, establishing its empirical optimization target.The implemented gradient matches the detached empirical surrogate gradient at the rollout parameters.
  • A.5 Surrogate Optimization Guarantees: Under score alignment, the surrogate direction approximates the expected-reward gradient, and a sufficient local condition guarantees strictly increased expected reward.MC-dropout supplies exploratory trajectories, while moment-matched Gaussian scores provide the differentiable interface; the variance floor ensures a finite optimum.
  • A.6 Stopping-gate Inference: At inference, latent reasoning iterates until the stopping probability reaches a threshold or T_max, then deterministically decodes the final answer with dropout disabled.Unless otherwise noted, T_max = 12; candidate thresholds are {0.5, 0.6, 0.7, 0.8, 0.9}, with 0.5 as the default.
  • A.7 Inference and Evaluation: Accuracy uses one deterministic forward pass, whereas Pass@k samples k independent MC-dropout latent rollouts and counts a problem solved when any parsed answer is correct.Pass@k uses dropout probability p=0.1, with repeated evaluation because dropout seeds change sampled trajectories.
  • A.8 Baseline Configurations: Ungated COCONUT and CODI use six latent thoughts, while +SLPO uses the same backbones with T_max = 12 and reports realized stopping time as #L.CoLaR uses thinking speed 2× with a maximum latent budget of 64 steps; ReGuLaR follows its released evaluation protocol.
  • A.9 Soft Latent Inference Transfer: SLPO transfers to soft-token inference by scoring latent embedding-space transitions, retaining ordinary token likelihoods for answer generation.The transfer uses variance floor 𝜖=10^-6, latent blocks of length 32, an 8-step surrogate prefix by default, and K=2 MC-dropout samples per step.
  • B Policy-Optimization Algorithm Comparison: With the SLPO surrogate, rollout budget, and backbones fixed, the policy-optimization comparison swaps only RLOO versus GRPO and reports full Pass@k values.The comparison covers COCONUT and CODI.
  • C Rollout Hyperparameter Sweeps: Rollout sweeps vary one hyperparameter over {2, 4, 8} while fixing the other at K=4 or G=4, reporting full Pass@2 results.Figure 8 presents separate K and G sweeps, with the larger marker at 4 denoting the value used in each one-factor sweep.

D Latent Geometry Analysis … D.5 Figure Construction

The appendix defines the sampling, representation, and geometry metrics used to compare pre- and post-SLPO latent trajectories. Across full dataset–backbone evaluations, post-SLPO effective-rank curves lie lower while preserving their progression shape across latent steps.

  • D.1 Sampling Protocol: The analysis compares pre- and post-SLPO checkpoints on identical evaluation instances using stochastic latent rollouts with MC-dropout.Unless otherwise stated, rollouts use p=0.1, latent budget T=6, 16 dropout samples, and the first 32 problems per dataset.
  • D.2 Step Representations: Dropout-sampled latent states are aggregated at each step into a shared sequence of step-wise representations used by both geometry metrics.The resulting sequence is (c_i,1, . . . , c_i,T).
  • D.3 Inter-Step Separation: Inter-step separation measures the cosine distance between successive latent step representations, averaged across consecutive steps and then problems.The analysis reports relative pre- to post-SLPO change for COCONUT across GSM8K, GSM-Hard, MultiArith, and SVAMP.
  • D.4 Prefix Effective Rank: Prefix effective rank tracks how many independent hidden-space directions are spanned by the first t latent representations.The centered representation matrix’s stable rank equals 1 for collinear steps and increases as the prefix spans more independent directions.
  • D.4 Prefix Effective Rank: Effective-rank curves are averaged over problems at each prefix length, so the x-axis represents prefix length rather than a single latent state.Adding steps raises the maximum attainable rank, so slope along t should not be interpreted alone as geometric expansion.
  • D.5 Figure Construction: The main figure reports inter-step change for COCONUT and effective rank for representative CODI settings: GPT-2 on SVAMP and Llama on MultiArith.The representative rank panels correspond to equation 5, while inter-step separation uses equations 3–4.
  • D.5 Figure Construction: Full COCONUT and CODI extensions show that post-SLPO effective-rank curves lie lower at fixed t while preserving progression shape across latent steps.Figures 9 and 10 cover all dataset–backbone combinations; Figure 9 covers COCONUT and Figure 10 covers CODI.
Loading 2607.19691v2…