揭示隐藏状态循环:基于on-policy强化学习的可切换潜在推理 Demystifying Hidden-State Recurrence: Switchable Latent Reasoning with On-Policy Reinforcement Learning
通过显式边界标记使隐藏状态循环推理可RL优化且可解释
前置知识
链式思维(Chain-of-Thought, CoT)
链式思维是一种让大语言模型通过逐步推理来解决问题的技术,模型在给出最终答案前会生成一系列中间推理步骤。传统的显式CoT将这些推理步骤作为可见文本输出,但会消耗大量token。潜在CoT(Latent CoT)试图将这些推理步骤压缩到模型的隐藏状态中,而不是输出为可见文本,从而减少token消耗。
本文提出的Switch框架是一种潜在CoT方法,理解CoT的基本概念和其token效率问题是读懂本文的基础。
隐藏状态循环(Hidden-State Recurrence)
隐藏状态循环是潜在CoT的一种实现方式,由Coconut方法引入。其核心思想是将前一个潜在步骤的最后一层隐藏状态作为下一个潜在步骤的输入嵌入,形成一个递归的计算过程。当位置是潜在token时,输入嵌入等于上一步的隐藏状态,这使得模型可以在不输出任何token的情况下进行多步推理计算。
这是本文方法的基础架构,理解其工作原理对于理解Switch如何在该架构上引入边界标记和GRPO优化至关重要。
GRPO(Group Relative Policy Optimization)
GRPO是一种on-policy强化学习算法,通过在输出文本的每个位置计算策略比率来优化模型。标准GRPO假设每个rollout位置都从分类token分布中采样并贡献策略密度到重要性比率。然而,隐藏状态循环的潜在位置不发出任何token,没有采样分布,因此GRPO在这些位置是未定义的。Switch-GRPO通过重新定义rollout执行和似然因式分解解决了这个问题。
本文的主要贡献之一就是让隐藏状态循环推理能够用标准on-policy RL(特别是GRPO)进行优化,理解GRPO的限制和Switch-GRPO的解决方案是理解本文技术贡献的关键。
策略比率(Policy Ratio)
策略比率是强化学习中用于计算重要性采样权重的关键概念,定义为当前策略与参考策略在给定状态下采取某个动作的概率之比。在语言模型中,这通常表示为新模型和旧模型在给定上下文下生成某个token的概率之比。GRPO等算法使用策略比率来估计梯度更新方向。对于隐藏状态循环的潜在位置,由于没有token输出,策略比率无法直接计算,这是现有RL方法难以直接应用于潜在推理的根本原因。
策略比率的可定义性是Switch能够使用on-policy RL的关键,边界标记作为离散token使得策略比率在每个文本位置都是定义良好的。
研究动机
现有隐藏状态循环的潜在推理方法面临两个核心挑战。首先,on-policy强化学习(如GRPO)难以直接应用于潜在推理设置。潜在位置不输出任何token,因此没有策略密度,导致GRPO的重要性比率在潜在块内部未定义。现有系统要么跳过RL训练,要么运行仅文本的rollout,但这与推理时的解码路径不同,造成训练-推理不一致。其次,潜在计算难以检查。潜在位置位于连续文本续写内部,没有分析师可以抓住的token,不清楚潜在步骤是执行任务相关的计算还是仅作为周围文本补偿的惰性占位符。这两个问题的共同根源是缺乏标记潜在计算开始和结束的显式边界。
本文的目标是本文的具体目标是解决隐藏状态循环潜在推理的两个核心问题:使其能够用标准on-policy强化学习(特别是GRPO)进行优化,并使其能够进行直接的机制分析。目标是设计一个框架,既能在推理任务上达到或超越现有方法的性能,又能提供对潜在推理过程的可解释性洞察,验证潜在推理步骤是否真正在做有意义的计算。
与已有工作不同的是,本文的独特切入角度是引入一对显式的边界标记来标记潜在推理块,这与现有方法形成鲜明对比。Coconut及其后续工作(如CODI、CoLaR)使用固定的或启发式的潜在推理触发机制,但都没有显式的边界标记。本文观察到,缺乏显式边界是RL可优化性和可解释性问题的共同根源,因此提出用边界token作为边界,一举解决两个问题:边界使得GRPO比率在每个文本位置都定义良好(潜在位置仅贡献零策略梯度项),同时边界作为分析锚点,让研究者可以直接探测切换概率、从内部激活探测切换状态、并对特定的潜在隐藏状态进行因果干预。
核心方法
Switch方法的核心思想是在模型词汇表中添加三个特殊token:进入潜在模式、退出潜在模式和潜在占位符。在推理时,模型正常解码直到发出进入token,然后运行至少最少潜在步数个潜在步骤(每个潜在步骤将前一步的最后一层隐藏状态作为当前步骤的输入嵌入),最后发出退出token恢复文本解码。训练分为三个阶段:Phase 1(SFT)教模型何时发出边界token对;Phase 2(课程)逐渐用潜在位置替换边界token块内的文本,同时保持边界信号;Phase 3(Switch-GRPO)使用相同的边界使GRPO比率在每个文本位置定义良好,允许通过包含潜在步骤的轨迹进行on-policy RL。Phase 1和Phase 2合称为Switch-SFT,Phase 3称为Switch-GRPO。
核心创新点是引入显式的边界标记来标记潜在推理块。这使潜在推理成为一个学习的、按步决策,模型选择是否以及何时调用它。边界是普通离散token,因此GRPO比率在每个位置、退出token和可见答案token位置都定义良好,潜在位置仅贡献零策略梯度项。同一边界也作为分析锚点,让研究者可以读取进入概率、从内部激活探测切换状态、对特定潜在隐藏状态进行因果干预。这与现有方法的关键区别在于:Coconut及其变体没有显式边界,因此难以用RL优化且难以解释;Switch通过边界标记一举解决这两个问题。
方法步骤详情
Switch的训练包含三个阶段。Phase 1(定位切换位置):训练模型在数学CoT语料库中用边界标记标记高熵片段。高熵位置定义为基础模型的下一个token分布具有高香农熵的位置,连续的高熵位置用边界标记标记。然后对标注的语料库进行标准的下一个token交叉熵监督微调。Phase 2(潜在课程):逐渐用潜在位置替换边界标记块内的文本,同时保持边界在损失中。采用并行课程:每个span同时转换,每个span的潜在计数按照参数增长,其中增长系数为2,最大潜在步数为8。潜在标签被屏蔽,损失仅应用于非潜在位置。Phase 3(Switch-GRPO):使用GRPO改进正确性和标签格式良好性。Switch-GRPO重新定义rollout:rollout执行与部署解码相同的多次前向传播,因此优化器在训练时看到的轨迹正是推理时产生的轨迹。似然因式分解:隐藏状态注入给定前文是确定性的,因此rollout似然仅因式分解于文本位置。奖励函数包含四项:来自math-verify的正负1正确性奖励、正负1标签格式奖励、0或1潜在使用奖励、可选的正确性门控简洁性项。
技术新颖性
Switch的技术新颖性体现在三个方面。首先,它首次使隐藏状态循环的潜在推理能够用标准on-policy强化学习(特别是GRPO)进行优化。Switch-GRPO通过重新定义rollout执行和似然因式分解,解决了潜在位置没有策略密度导致GRPO未定义的问题,使训练时和推理时的轨迹完全一致。其次,它首次使隐藏状态循环的潜在推理能够进行直接的机制分析。边界标记作为分析锚点,让研究者可以读取进入概率、从内部激活探测切换状态、对特定潜在隐藏状态进行因果干预,从而验证潜在推理步骤是否真正在做有意义的计算。第三,它首次展示了隐藏状态循环潜在推理的机制可解释性。通过边界标记的机制分析揭示了三个发现:进入token是一个局部化的、习得的切换策略而非风格人工痕迹;潜在步骤执行问题特定的、因果上重要的计算而非惰性占位符;该计算集中在进入时的单个隐藏状态转换上。
实验结果
Switch在两个数学推理基准上取得了显著优于基线的性能。在MATH-500上,Switch达到79.3%,比同等规模的最强Coconut风格基线(CoLaR的53.6%)高25.7个百分点。在GSM8K上,Switch达到89.2%,同样优于所有基线。Switch-GRPO相比仅课程的SFTcheckpoint带来了显著提升:潜在条件准确率跳跃了+12.6个百分点(从66.7%到79.3%),同时切换率从81%降至58%,表明模型学会了选择潜在推理真正有帮助的问题。在Switch-GRPO训练过程中,每个问题的潜在调用从约1.5降至约1,可见token使用从约2900降至约1900。机制分析通过边界标记揭示了三个发现:进入token是一个局部化的、习得的切换策略(在标注的进入位置,模型将进入token基本放在词汇表顶部,排名小于等于1.7,在随机非边界位置抑制进入token约3-4个数量级);潜在步骤执行问题特定的、因果上重要的计算(零化潜在状态将诊断子集的准确率从100%降至33.3%,替换为同范数随机向量仅损失9.5个百分点,跳过潜在步骤损失19.0个百分点);该计算集中在进入时的单个隐藏状态转换上(logit lens显示退出token概率在每个潜在步骤都约等于1,没有最小步数约束潜在块会坍缩到单次隐藏前向传播)。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| MATH-500数学推理 | 准确率(%) | 79.3 (Switch-GRPO) / 66.7 (Switch-SFT) | 53.6 (CoLaR) / 48.3 (CODI) / 46.6 (Coconut) | +25.7 vs CoLaR, +31.0 vs CODI, +32.7 vs Coconut |
| GSM8K数学推理 | 准确率(%) | 89.2 (Switch-GRPO) / 80.2 (Switch-SFT) | 78.5 (CoLaR) / 76.4 (CODI) / 76.1 (Coconut) | +10.7 vs CoLaR, +12.8 vs CODI, +13.1 vs Coconut |
| MATH-500潜在条件准确率 | 准确率(%) | 79.3 (After Switch-GRPO) | 66.7 (After SFT) | +12.6个百分点 |
| 潜在调用率 | 百分比(%) | 58 (After Switch-GRPO) | 81 (After SFT) | 降低23个百分点(更精准调用) |
| 可见token使用 | 平均token数 | 1777 (After Switch-GRPO) | 1433 (After SFT) | 更高效(用更少token达到更高准确率) |
| 潜在状态干预(零化) | 准确率变化 | 从100%降至33.3% (-66.7个百分点) | 正常推理 | 证明潜在计算的因果重要性 |
| 切换状态线性探测 | 探测准确率 | 91.9% (After SFT) / 88.4% (After Switch-GRPO) | 随机猜测50% | 表明切换决策在模型内部可解码 |
局限与改进
作者承认的局限性包括:实验仅限于8B参数的Qwen3模型和数学推理基准(MATH-500和GSM8K),尚未在多领域推理或更大模型规模上评估Switch,在这些场景中学习切换和潜在深度之间的平衡可能不同。Switch-GRPO的梯度仅通过每个rollout的文本段流动;潜在位置通过冻结的KV cache贡献,因此潜在表示主要由Phase 2课程而非直接由RL目标塑造。机制的机理解释侧重于模型编码的内容(切换边界、潜在因果效应)而非表征失败模式;论文中的logit-lens解码是定性的,不应被视为潜在推理轨迹的忠实重构。自身的观察包括:Switch需要手动设置最小潜在步数来防止潜在块坍缩到单步,这是一个超参数;训练过程分为三个阶段,相对复杂;奖励函数包含多个项,需要平衡权重;机制分析虽然揭示了一些有趣发现,但对潜在推理内部过程的解释仍然有限。
独立分析的弱点
Switch的一个独立分析弱点是依赖固定的最小潜在步数来防止模型过早退出潜在模式。机制分析显示退出token概率在每个潜在步骤都约等于1,如果没有最小步数约束,潜在块会坍缩到单次隐藏前向传播。这表明模型可能没有真正学会多步潜在推理,而是依赖约束来保持一定的深度。改进方向可以是设计自适应的深度机制,让模型根据问题复杂度自主决定潜在步数,或者引入内在奖励来鼓励有意义的深度推理。另一个弱点是Switch-GRPO的梯度仅通过文本段流动,潜在位置通过冻结的KV cache贡献,因此潜在表示主要由课程而非RL直接塑造。这限制了RL对潜在推理质量的直接优化。改进方向可以是探索让潜在位置也能参与梯度传播的方法,例如通过可学习的潜在表示或引入潜在空间的内在奖励。此外,Switch的评估仅限于数学推理任务,在需要多步推理的其他领域(如逻辑推理、代码生成、常识推理)上的性能未知。改进方向是在更多领域评估Switch,并根据不同领域的特点调整课程设计和奖励函数。
未来方向
作者提出的未来方向包括:将潜在token本身也设计为可采样的,桥接隐藏状态循环和词汇混合潜在,这是自然的下一步;在匹配规模和数据下进行头对头比较。基于论文成果可以延伸的方向包括:将Switch应用于更多推理领域(如代码生成、逻辑推理、科学推理)并评估其泛化能力;探索自适应的潜在深度机制,让模型根据问题复杂度和不确定性动态决定潜在步数;研究如何让RL更直接地优化潜在表示,例如通过内在奖励或对比学习;将Switch与其他推理增强技术(如检索增强、工具使用)结合,构建更强大的推理系统;进一步研究潜在推理的机制可解释性,例如开发更精细的探测工具来理解潜在空间中的计算过程;在更大规模的模型上评估Switch,观察学习切换和潜在深度之间的平衡如何随模型规模变化。
复现评估
Switch的复现性评估如下:代码方面,论文提供了GitHub页面链接,表明代码将开源。模型权重方面,论文提供了HuggingFace链接,表明训练好的模型权重将公开。数据方面,论文使用了OpenR1-Math的标注子集,按照SwiReasoning的标注流程将高熵CoT子片段用边界标记包裹。硬件方面,论文在附录中提供了详细的训练和硬件细节。难度方面,复现Switch需要理解三阶段训练流程、课程学习设计、Switch-GRPO的修改实现、以及机制分析的工具。相对复杂的训练流程和多个超参数可能增加复现难度。总体而言,Switch的开源程度较高,代码、模型权重和数据都将公开,但需要一定的技术门槛来完整复现所有实验和机制分析。
论文图表