Latent-GRPO: Group Relative Policy Optimization for Latent Reasoning

Paper: arXiv:2604.27998 Code: DJC-GO-SOLO/Latent-GRPO Code reference: main @ c0994fb7 (2026-05-12)

1. Motivation

Latent reasoning 把中间推理压缩进连续表示，理论上能显著缩短 reasoning chain，但现有工作主要停留在 SFT；一旦直接把 GRPO/Soft-GRPO 迁到 latent space，训练就会不稳定，甚至 collapse。 论文把问题拆成三个具体瓶颈：缺少内在 latent manifold、trajectory-level reward 和 token-level update 的方向不一致、多个正确 latent path 在共享前缀处会发生 harmful averaging。解决这三个点，才能把 latent reasoning 从“能做”推进到“能稳定 RL”。 这个问题值得做，因为一旦稳定，latent reasoning 可以同时保留正确率和更短的推理长度；在难题上还能带来更强的 pass@k 搜索能力。

2. Idea

核心不是“再做一个 latent RL 版本”，而是把 latent RL 重新约束到一个可优化的几何空间里：先用 Latent-SFT 提供有效 manifold，再用三道护栏分别处理无效轨迹、错误的噪声方向、以及共享起点上的 mode averaging。

和 Soft-GRPO 的根本差异是，Soft-GRPO 主要解决“怎么估计 latent token 的 density”，而 Latent-GRPO 进一步改变了 exploration 和 update 的结构，让优势信号、采样扰动、以及多路径竞争都对齐到 latent 机制本身。

3. Method

Figure 1 解读：图 (a) 说明 explicit reasoning 和 latent reasoning 在 probability density 与 sampling mechanism 上都不同；图 (b) 展示 naive Gumbel objective 会把正优势轨迹的某些成分往下推、把负优势轨迹的某些成分往上推；图 (c) 说明多个正确 latent path 在共享前缀处会平均成一个不一定合法的连续状态。

3.1 Overall framework

Latent-GRPO 先假设模型已经经过 Latent-SFT 初始化，拥有可用的 latent manifold。训练时，每个 prompt 采样一组 group trajectories；每条轨迹先在 top-$K$ vocabulary 上形成 latent token，再根据轨迹级 reward 做 GRPO-style 更新。

方法上有三层改动：

• Invalid Sample Advantage Masking：把超长/不终止的轨迹从 group baseline 里剔除。
• One-sided Noise Sampling：把 Gumbel 扰动改成单边正扰动，并在优化阶段保持方向一致。
• Optimal Correct Path First Token Selection：在多个正确轨迹里只保留一个最优轨迹的第一个 latent token，避免共享前缀上的 harmful averaging。

3.2 Invalid Sample Advantage Masking

设一组采样轨迹为 $mathcalG$，第 $j$ 条轨迹长度为 $L_j$，最大响应长度为 $$\begin{gathered} L_{ \\ max} \end{gathered}$$。论文把在预算内正常终止的轨迹定义为 $$\begin{gathered} mathcal{G}_{ \\ mathrmvalid}= \\ j \\ mid{L}_j<L_{ \\ max} \\ . \end{gathered}$$ 然后只在 valid subset 上计算 group baseline： $$\begin{gathered} m{u}_R= \\ frac1| \\ mathcal{G}_{ \\ mathrmvalid}| \\ su{m}_{j \\ in \\ mathcal{G}_{ \\ mathrmvalid}}{R}_j, \\ qquad \\ sigm{a}_R= \\ sqrt \\ frac1| \\ mathcal{G}_{ \\ mathrmvalid}| \\ su{m}_{j \\ in \\ mathcal{G}_{ \\ mathrmvalid}}(R_j- \\ m{u}_R{)}^2. \end{gathered}$$ 优势定义为： \\hat{A}_j^{\\mathrm{valid}}= \\begin{cases} \\dfrac{R_j-\\mu_R}{\\sigma_R}, & j\\in\\mathcal{G}_{\\mathrm{valid}}, \\\\ 0, & j\\notin\\mathcal{G}_{\\mathrm{valid}}. \\end{cases} 直觉是：无效轨迹不是“坏样本的正常一员”，而是已经偏离 latent manifold 的污染源。如果把它们留在 group 统计里，会扭曲 valid samples 的归一化优势。

def mask_invalid_advantage(rewards, lengths, max_len):
    valid = lengths < max_len
    adv = torch.zeros_like(rewards)
    if valid.any():
        mu = rewards[valid].mean()
        sigma = rewards[valid].std(unbiased=False).clamp_min(1e-6)
        adv[valid] = (rewards[valid] - mu) / sigma
    return adv

论文公式与 released code 实现差异：训练脚本里 algorithm.exclude_overlong_samples_from_advantage 不是全局单一配置，GSM8K 脚本为 False，Math500 脚本为 True；代码里也分成了 compute_latent_grpo_outcome_advantage_firstmask_best_exclude_advantage 和 ...include_advantage 两条路径。也就是说，paper 的概念是一致的，但 released code 采用了 task-specific 的开关。

3.3 One-sided Noise Sampling

Soft-GRPO 的问题在于，标准 Gumbel perturbation 是双边的：$$\begin{gathered} g_i^{\prime }= \\ log{p}_i+ \\ x{i}_i \end{gathered}$$，其中 $x{i}_i$ 可以正也可以负。这样一来，正优势轨迹里的某些 component 反而可能被往下调，负优势轨迹里的某些 component 反而可能被往上推。

Latent-GRPO 把扰动改成单边形式： $$\begin{gathered} x{i}_i^{+}= \\ mathrmclip( \\ x{i}_i,-a,b)+a+ \\ delta, \\ qquad \\ x{i}_i \\ sim \\ mathrmGumbel(0,1), \end{gathered}$$ 并定义 $$\begin{gathered} g_i^{\ast }= \\ log{p}_i( \\ thet{a}_{ \\ mathrmold})+ \\ x{i}_i^{+}. \end{gathered}$$ 这样在优化初期，target margin 是正的；如果多轮 PPO 更新后 margin 被跨过去，论文再用 conditional STE 把 backward direction 翻回去，维持单边约束。

def one_sided_target(logp_old, gumbel_noise, a=1.5, b=3.0, delta=1e-6):
    xi_pos = gumbel_noise.clamp(-a, b) + a + delta
    return logp_old + xi_pos

释放代码里，这个设计主要落在 sglang_latent_reasoning_pkg/python/sglang/srt/layers/sampler.py 和 rollout 管线里：采样阶段保存 top-k gumbels、原始概率和 token ids，再交给 trainer 计算 surrogate density。训练脚本默认 actor_rollout_ref.rollout.add_noise_gumbel_softmax=True，而评测脚本默认会关掉噪声做 deterministic decoding；pass@k 是单独打开 Gumbel sampling 后测的。

3.4 Optimal Correct Path First Token Selection

这里解决的是 latent mixture non-closure。多个正确路径一起优化时，explicit CoT 还能靠离散采样选中一条具体路径，但 latent token 是连续向量，会把这些正确路径平均成一个可能不合法的中间态，尤其在第一个 reasoning step 最明显。

做法是：先找出验证正确的轨迹集合 $$\begin{gathered} mathcal{G}_{ \\ textcorrect} \end{gathered}$$，再用轨迹级平均 surrogate log-probability 选出最优路径 $$\begin{gathered} bar{ \\ ell}_j= \\ frac1L_j \\ left[ \\ su{m}_{t=1}^{T_{ \\ mathrmlat}^{(j)}} \\ log \\ p{i}_{ \\ thet{a}_{ \\ mathrmold}}^{+}(s_{j,t} \\ mid \\ boldsymbolx,s_{j,<t})+ \\ su{m}_{t=T_{ \\ mathrmlat}^{(j)}+1}^{L_j} \\ log \\ p{i}_{ \\ thet{a}_{ \\ mathrmold}}(y_{j,t} \\ mid \\ boldsymbolx,s_{j,1:T_{ \\ mathrmlat}^{(j)}},y_{j,<t}) \\ right], \end{gathered}$$ 然后取 $$\begin{gathered} j^{\ast }= \\ arg \\ ma{x}_{j \\ in \\ mathcal{G}_{ \\ textcorrect}} \\ bar{ \\ ell}_j. \end{gathered}$$ 只屏蔽其他正确路径在 $t=1$ 的更新，后续 token 仍然保留。这一点很关键：它不是把训练收缩成单一路径，而是只修掉共享起点上的 harmful averaging。

def select_first_token_mask(correct_scores, correct_flags):
    mask = torch.ones_like(correct_flags, dtype=torch.float32)
    if correct_flags.sum() > 1:
        best = torch.argmax(correct_scores[correct_flags])
        correct_idx = torch.nonzero(correct_flags, as_tuple=False).flatten()
        for i, idx in enumerate(correct_idx):
            if i != best:
                mask[idx, 0] = 0.0
    return mask

3.5 Final objective

最终目标就是把上面三层 mask/target 合到 PPO surrogate 里： $$\begin{gathered} mathcal{J}_{ \\ mathrmLatent \\ text-GRPO}( \\ theta)= \\ mathbb{E}_{ \\ boldsymbolx} \\ left[ \\ frac1G \\ su{m}_{j=1}^G \\ frac1L_j \\ su{m}_{t=1}^{L_j} \\ left( \\ min \\ left(r_{j,t}^{\ast }( \\ theta) \\ tilde{A}_{j,t}, \\ mathrmclip(r_{j,t}^{\ast }( \\ theta),1- \\ epsilo{n}_{ \\ mathrmclip},1+ \\ epsilo{n}_{ \\ mathrmclip}) \\ tilde{A}_{j,t} \\ right)- \\ beta \\ , \\ mathbb{D}_{ \\ mathrmKL}^{(j,t)} \\ right) \\ right]. \end{gathered}$$ 其中 $tilde{A}_{j,t}$ 同时包含 invalid-sample masking 和 first-token selection。

def latent_grpo_loss(ratio, ratio_old, advantage, kl, eps_clip, beta):
    clipped = ratio.clamp(1 - eps_clip, 1 + eps_clip)
    ppo_term = torch.min(ratio * advantage, clipped * advantage)
    return -(ppo_term - beta * kl).mean()

3.6 Code-to-paper mapping

Code reference: main @ c0994fb7 (2026-05-12) — pseudocode and mapping based on this commit

Paper concept	Source file	Key class/function	Role
Latent rollout / top-k latent token construction	sglang_latent_reasoning_pkg/python/sglang/srt/layers/sampler.py	top-k sampling helpers	构造 latent token 的 top-k mixture，并保存采样所需的 scores / probs / gumbels
Rollout 输出回传训练端	verl-0.4.x/verl/workers/rollout/sglang_rollout/sglang_rollout.py	SGLangRollout, _post_process_outputs	把 token ids、logprobs、top-k gumbels、original probs 整理回 trainer
Invalid sample masking + first-token selection	verl-0.4.x/verl/trainer/ppo/core_algos.py	compute_latent_grpo_outcome_advantage_firstmask_best_exclude_advantage, compute_latent_grpo_outcome_advantage_firstmask_best_include_advantage	实现 group advantage、无效轨迹过滤、最佳正确路径首 token mask
GRPO / PPO loss	verl-0.4.x/verl/trainer/ppo/core_algos.py	compute_policy_loss, compute_rewards, kl_penalty	构造最终 policy objective 和 KL 正则

3.7 机制补充：三道护栏如何共同稳定 latent RL

Latent-GRPO 的三项改动分别对应三种失败模式。Invalid Sample Advantage Masking 处理的是 group 统计被污染的问题：过长或没有正常终止的样本往往不是“低 reward 的普通负例”，而是采样过程已经偏离 latent manifold 的异常轨迹。若把它们放进均值和方差，valid 样本的优势会被错误归一化，模型可能为了远离异常样本而破坏原本正确的推理方向。把无效样本优势置零并从 baseline 中剔除，本质上是在告诉优化器：先不要从不可信轨迹里学习方向，只用正常完成的轨迹比较好坏。One-sided Noise Sampling 处理的是 latent token 采样和更新方向不一致的问题。

Soft-GRPO 使用 Gumbel 扰动估计 latent density，但双边噪声会让某些正优势轨迹的 logit 分量被目标压低，也会让负优势轨迹的分量被意外抬高。Latent-GRPO 把扰动改成单边正 margin，再用 conditional STE 在 margin 被跨过后修正反向传播方向，使“采样时选中的 latent token”和“训练时增加其相对概率”保持同向。这个设计尤其适合连续 latent，因为 latent token 不是一个可以独立解释的离散词；一旦某些维度被反向推动，得到的向量可能落到语言模型未见过的中间区域。

Optimal Correct Path First Token Selection 处理的是多条正确路径之间的 non-closure。显式 CoT 中，多条正确推理链共享前缀通常不是大问题，因为模型最终仍然输出某个离散 token；但 latent reasoning 的第一个 latent state 是连续向量，多个正确路径在同一起点被同时拉近时，平均结果未必对应任何一条合法推理模式。论文只屏蔽其他正确路径的 first latent token，而不是屏蔽整条轨迹，说明作者希望保留后续推理多样性，只修复最容易发生 mode averaging 的分岔入口。

这三道护栏的组合比单独任一项更重要。Latent-SFT 先提供可用的表示流形，advantage masking 保证 reward baseline 来自可信轨迹，one-sided noise 保证 latent density surrogate 的梯度方向不反常，first-token selection 再避免多个正确模式在入口处互相平均。最终目标仍然是 GRPO/PPO 风格的 clipped surrogate 加 KL 约束，但被优化的对象已经不再是普通离散 token 序列，而是由 top-$K$ 词分布、Gumbel 扰动、latent embedding 和文本解码共同定义的混合轨迹。

论文没有说明所有数据混合比例和每个 benchmark 的完整超参，因此复现时应优先检查 released config 中 mask 开关、Gumbel sampling 开关和 deterministic evaluation 设置是否与实验一致。从实验设计看，low-difficulty、high-difficulty 和 pass@k 三组结果分别验证不同问题：低难度任务检查 latent RL 是否会破坏已有能力，高难度任务检查优化后是否真的提升复杂推理，pass@k 检查 latent sampling 是否带来可搜索的多样性。

Ablation 的意义也应按三种失败模式理解：去掉 advantage masking 会让异常轨迹进入统计，去掉 one-sided noise 会重新出现方向错配，去掉 first-token selection 则会在多正确路径问题上更容易平均。也就是说，Latent-GRPO 的贡献不是某个单独 trick，而是把 latent reasoning 的采样、奖励归一化和轨迹选择重新对齐到同一个几何假设上。

3.8 补充：实现层面的训练闭环

把 Latent-GRPO 落到实现时，需要同时维护三类状态：rollout 端保存每个 latent step 的 top-$K$ token、Gumbel 扰动和旧策略概率；trainer 端根据最终答案 reward 计算 group advantage；loss 端再把 advantage、one-sided surrogate ratio、KL 和 first-token mask 合成逐 token 更新。这个闭环说明 latent RL 的困难不只在公式，而在采样记录必须足够完整。如果 rollout 只返回最终文本，trainer 无法重建 latent token 的 surrogate density；

如果只记录 token id 而不记录扰动，one-sided target 就无法和当时的采样事件对齐。与普通 GRPO 相比，Latent-GRPO 的 credit assignment 更依赖“哪些 latent state 仍然在有效流形内”。文本 token 即使低概率，通常仍然是词表中的合法符号；latent token 则可能因为多个 embedding 混合、双边噪声或错误路径平均而落到难以解码的区域。Invalid sample mask 先过滤明显离开流形的轨迹，one-sided noise 让被选中的方向更像一个可重复的正 margin，first-token selection 再减少正确路径入口处的混合。

这样 trainer 更新的不是抽象的“推理能力”，而是保持 latent 表示可解码、可搜索、可被 reward 区分的局部结构。这一点也影响如何解读局限。论文没有把所有 latent manifold 诊断指标、不同 top-$K$ 选择或不同噪声截断范围的完整曲线都列出，因此无法断言某个单独超参是通用最优。更稳妥的结论是：在 latent reasoning 已经由 SFT 初始化的前提下，RL post-training 必须同时约束样本有效性、采样扰动方向和多正确路径竞争；否则即使 reward 定义正确，优化也可能把连续 latent 空间推向不可解释的平均态或无效态。

4. Experimental Setup

Datasets

低难度设置：

• 训练：GSM8K-Aug train split；README 里对应默认文件是 GSM8k-Aug-oss-dup-all.parquet。
• 验证/评测：GSM8K-Aug test、GSM-Hard、SVAMP、MultiArith。
• 规模：论文没有逐个给出精确样本数；训练脚本只明确了数据文件名。

高难度设置：

• 训练：DAPO-Math-17k-en-train.parquet，文件名直接表明是 17k 规模。
• 验证/评测：Math-500-test.parquet、AIME24、AIME25、GPQA。
• SFT 初始化：Open-R1 的 4K-length reasoning chains 子集。

Baselines

• SFT
• GRPO
• Latent-SFT
• Soft-GRPO

Metrics

• Pass@1：单次采样下的正确率。
• L：平均 reasoning length，越小表示推理更短。
• Pass@k：采样 $k$ 次后至少一次答对的概率；论文在 Gumbel sampling 下报告这个指标。

Training config

共同点：

• MODEL_PATH 必须是 Latent-SFT 初始化后的 checkpoint。
• trainer.n_gpus_per_node=8
• algorithm.adv_estimator=grpo
• actor_rollout_ref.actor.optim.lr=1e-6
• actor_rollout_ref.rollout.enable_latent=True
• actor_rollout_ref.rollout.gumbel_softmax_temperature=1.0
• actor_rollout_ref.rollout.max_topk=10

低难度脚本 Latent-GRPO-gsm8k-llama3.sh：

• data.train_batch_size=64
• data.val_batch_size=128
• data.max_prompt_length=192
• data.max_response_length=128
• actor_rollout_ref.actor.ppo_mini_batch_size=16
• actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2
• actor_rollout_ref.actor.use_kl_loss=False
• actor_rollout_ref.actor.ppo_max_token_len_per_gpu=2048
• actor_rollout_ref.rollout.max_model_len=1024
• actor_rollout_ref.rollout.max_num_batched_tokens=2048
• actor_rollout_ref.rollout.temperature=0.6
• actor_rollout_ref.rollout.latent_end_token_id=524
• actor_rollout_ref.rollout.add_noise_gumbel_softmax=True
• actor_rollout_ref.rollout.use_one_sided_gumbel_noise=True
• actor_rollout_ref.rollout.noise_scale=1.0
• actor_rollout_ref.rollout.gpu_memory_utilization=0.6
• algorithm.exclude_overlong_samples_from_advantage=False
• trainer.total_epochs=10

高难度脚本 Latent-GRPO-math500-qwen.sh：

• data.train_batch_size=32
• data.val_batch_size=500
• data.max_prompt_length=1024
• data.max_response_length=4096
• actor_rollout_ref.actor.ppo_mini_batch_size=32
• actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1
• actor_rollout_ref.actor.use_kl_loss=True
• actor_rollout_ref.actor.ppo_max_token_len_per_gpu=20480
• actor_rollout_ref.model.enable_gradient_checkpointing=True
• actor_rollout_ref.actor.fsdp_config.param_offload=True
• actor_rollout_ref.actor.fsdp_config.optimizer_offload=True
• actor_rollout_ref.rollout.max_model_len=12000
• actor_rollout_ref.rollout.max_num_batched_tokens=12000
• actor_rollout_ref.rollout.temperature=0.6
• actor_rollout_ref.rollout.latent_end_token_id=522
• actor_rollout_ref.rollout.add_noise_gumbel_softmax=True
• actor_rollout_ref.rollout.use_one_sided_gumbel_noise=True
• actor_rollout_ref.rollout.noise_scale=1.0
• actor_rollout_ref.rollout.gpu_memory_utilization=0.8
• algorithm.exclude_overlong_samples_from_advantage=True
• trainer.total_epochs=5

README 还明确提醒：不要直接从非 latent 初始化模型开始做 Latent-GRPO；non-reasoning base model 需要在 chat template 里带 <think>，否则不会进入 reasoning mode。

5. Experimental Results

Low-difficulty suite

Figure 2 解读：LLaMA-3.2-1B-Instruct 在 GSM8K-Aug 上的训练动态。Latent-GRPO 的验证分数上升更稳，同时长度几乎不增长；GRPO 早期更快，但后期更容易 collapse。

Method	GSM8K-Aug	GSM-Hard	SVAMP	MultiArith	Avg
SFT	49.51 / 77.46	11.67 / 85.78	57.38 / 44.10	92.23 / 45.49	52.70 / 63.20
GRPO	62.26 / 111.1	14.70 / 130.8	60.78 / 69.40	94.16 / 65.50	57.98 / 94.20
Latent-SFT	47.55 / 22.00	10.25 / 25.50	53.22 / 15.30	90.80 / 16.30	50.46 / 19.78
Soft-GRPO	48.34 / 21.60	10.30 / 25.60	53.19 / 15.30	91.07 / 16.40	50.73 / 19.73
Latent-GRPO	66.29 / 23.80	15.83 / 28.00	56.50 / 15.80	94.65 / 17.20	58.32 / 21.20

关键结论：

• 平均上，Latent-GRPO 比 Latent-SFT 高 7.86 Pass@1，比 GRPO 短 4.44x。
• 它在 GSM-Hard 和 MultiArith 上都优于 GRPO，同时长度远短于 GRPO。
• 论文文字特别指出：one-sided noise 是最关键的稳定性来源，first-token selection 是进一步提升。

High-difficulty suite

Figure 3 解读：Qwen2.5-Math-7B 在 DAPO-Math-En 上的训练动态。Latent-GRPO 一边提升 Math500 验证分数，一边保持更短的 response length。

Method	Math500	AIME24	AIME25	GPQA	Avg
SFT	67.40 / 2662	13.39 / 6423	6.77 / 6550	32.07 / 3555	29.91 / 4798
GRPO	75.15 / 2982	16.61 / 7500	19.79 / 6815	38.25 / 4568	37.45 / 5466
Latent-SFT	64.85 / 1070	6.67 / 1651	6.61 / 1299	29.67 / 1214	26.95 / 1309
Soft-GRPO	60.60 / 1133	6.20 / 1894	3.59 / 1511	27.12 / 1426	24.38 / 1491
Latent-GRPO	80.40 / 929	26.56 / 2548	23.23 / 1735	36.68 / 1384	41.72 / 1649

关键结论：

• 平均上，Latent-GRPO 比 Latent-SFT 高 14.77 Pass@1，比 GRPO 高 4.27 Pass@1，同时长短比是 3.31x。
• 在 Math500、AIME24、AIME25 上都是最优 Pass@1。
• Soft-GRPO 在高难度上甚至低于 Latent-SFT，说明只做 density matching 不够。

Pass@k

Figure 4 解读：开启 Gumbel sampling 后，Latent-GRPO 的 pass@k 曲线整体优于 explicit GRPO，尤其在更难的 AIME 任务上提升明显。

论文给出的关键数值：

• Math500：Latent-GRPO sampled reach 87.8，GRPO 为 84.2。
• AIME24：GRPO 在 pass@64 为 23.3；Latent-GRPO 在 noise=0.5 / 1.0 时分别到 46.7 / 50.0。
• AIME25：GRPO 在 pass@64 为 30.0；Latent-GRPO 在 noise=0.5 / 1.0 时分别到 46.7 / 53.3。

Ablation and limitations

消融结论是定性的，但方向很明确：

• 去掉 One-sided Noise Sampling，验证性能会明显掉，训练甚至会 collapse。
• 去掉 First Token Selection，也会掉，但幅度比去掉 one-sided noise 小。

限制也很明确：

• No Sampling 的 latent reasoning 是 deterministic 的，不适合直接看 pass@k。
• Gumbel sampling 能提高 pass@k，但会牺牲一部分 pass@1，尤其在最难的 AIME 上。
• one-sided STE 本质上是 surrogate backward rule，不是严格的概率密度。