语义步骤预测：通过步骤采样实现LLM推理轨迹中的多步潜在预测

# 基于步骤采样的多步潜在预测在LLM推理轨迹中的应用  
**来源**：https://arxiv.org/html/2604.18464  

###### 摘要  
\[本arXiv论文完全来自NUS CS5260课程的研究工作。\]  
语义管道预测（STP）利用表示几何在微调期间将LLM隐藏状态轨迹正则化为局部线性测地线，从而大幅提高数据效率。原始STP方法随机采样token子跨度，这与基础大语言模型（LLM）的训练架构兼容。受STP启发，我们感兴趣的是研究采样位置是否能够进一步增强多步推理的语义结构，从而影响其几何效果。我们在连续的语义推理步骤边界上应用STP，在ProcessBench（3,400个样本）上实现了比冻结基线高168倍的多步潜在预测精度，而随机token STP仅提高4倍。通过训练好的非线性预测器探测潜在流形表明，STP塑造的轨迹是平滑曲线而非直线：一个3层MLP在步骤边界模型上将预测误差进一步降低33–12倍（相对于线性外推）。移除语言建模损失后，得到的轨迹的MLP可预测性比组合损失高2倍，揭示了生成质量与几何纯度之间的权衡。我们的结果将采样位置确定为几何正则化的关键变量，并将多步潜在预测MSE确立为此类方法的新评估指标。  

## 1 引言  
大语言模型（LLM）通过生成逐步解决方案进行推理，每个推理步骤产生许多token。每个token都需要一次完整的自回归前向传播，使得多步推理计算成本高昂。越来越多的研究探索潜在推理，即在模型的连续隐藏状态空间中进行计算，而非在每个步骤解码离散token（Hao et al., 2024 (https://arxiv.org/html/2604.18464#bib.bib4)；PLaT, 2026 (https://arxiv.org/html/2604.18464#bib.bib15)）。要使潜在推理可行，模型的隐藏状态轨迹必须可预测——即给定当前推理状态，未来状态能够以合理的精度被预测。语义管道预测（STP）（Huang et al., 2026 (https://arxiv.org/html/2604.18464#bib.bib8)）为此提供了理论和实践框架。基于测地线假设（该假设坚持无错误的token序列在平滑的语义流形上沿局部线性测地线运行），STP在微调期间添加一个额外的余弦位移损失，强制连续的隐藏状态位移向量平行。由此产生的“语义管道”将轨迹限制在测地线的管状邻域内，提高了信噪比，并实现了16倍的数据效率提升。STP在其损失中随机采样token子跨度，其中三个索引`<|step|>`作为训练数据中推理步骤之间的分隔符（图1 (https://arxiv.org/html/2604.18464#S3.F1)），生成如下形式的序列：  
\[问题\]<|step|>\[步骤1\]<|step|>\[步骤2\]<|step|>⋯<|step|>\[步骤K\]<|step|>  
(2)  
在`<|step|>`位置的隐藏状态形成轨迹z = (z_0, z_1, ..., z_K)。我们在连续三元组上计算STP损失：  
L_STP^step = 1/(K-1) * Σ_{k=1}^{K-1} (1 - ((z_k - z_{k-1})·(z_{k+1} - z_k)) / (||z_k - z_{k-1}||·||z_{k+1} - z_k|| + ε))  
(3)  
这在数学上与STP的损失相同，但区别在于索引的采样位置。我们在语义上有意义的步骤边界处采样，而非随机token位置。组合目标函数为：  
L = L_NTP + β·L_STP^step，其中β=1。  
微调使用LoRA (Hu et al., 2022 (https://arxiv.org/html/2604.18464#bib.bib6))，秩为16，应用于q, k, v, o投影（在1.5B基座模型上约4.4M可训练参数）。  

### 3.2 实验网格  
**表1：六模型实验网格。每个模型在损失组成和/或STP采样策略上不同。**  
该网格（表1 (https://arxiv.org/html/2604.18464#S3.T1)）隔离了三个因素：（i）token级别与步骤级别采样（C vs. A2），（ii）步骤级别随机与连续采样（A2 vs. A），以及（iii）L_NTP的必要性（A vs. A1）。  

### 3.3 多步预测MSE  
对于每个轨迹(z_0, z_1, ..., z_K)和每个有效位置k≥1，我们通过线性外推预测m步后的状态：  
ẑ_{k+m} = z_k + m·(z_k - z_{k-1})  
(4)  
并计算归一化预测误差：  
MSE_m = 1/N * Σ_{(k, 样本)} (||ẑ_{k+m} - z_{k+m}||^2 / ||z_{k+m}||^2)  
(5)  
MSE_m ≈ 0表示轨迹完全线性；MSE_m ≈ 1表示预测误差等于信号幅度。  

### 3.4 轨迹平滑度分数  
我们报告余弦分数（与训练目标匹配）：  
cos_score_k = 1 - ((z_k - z_{k-1})·(z_{k+1} - z_k)) / (||z_k - z_{k-1}||·||z_{k+1} - z_k||) ∈ [0,2]  
(6)  
以及垂直分数（几何解释）：  
perp_score_k = ||d_k - (d_k·ŝ_k)ŝ_k|| / ||d_k|| = sin(θ_k) ∈ [0,1]  
(7)  
其中d_k = z_k - z_{k-1}，s_k = z_{k+1} - z_{k-1}。这两个分数单调相关，并产生相同的模型排序（Spearman ρ=1.0）。  

## 4 实验  
### 4.1 设置  
**模型：** Qwen2.5-Math-1.5B（数学专用）、Qwen2.5-1.5B（通用）、Llama-3.2-1B（不同架构，2048维隐藏状态 vs. 1536维）。所有模型均使用秩为16的LoRA进行微调。  
**训练数据：** 6,132个MATH竞赛问题及解答，在段落边界处用`<|step|>`分隔符拆分。3个epoch（每个模型约1150个优化器步骤）。  
**评估数据：** ProcessBench (Zheng et al., 2024 (https://arxiv.org/html/2604.18464#bib.bib12))：3,400个逐步数学解答，涵盖四个难度级别，并带有手工标注的错误位置。  

## 5 结果  
### 5.1 预测MSE  
**表2：基于Qwen2.5-Math-1.5B在ProcessBench上的六模型消融实验。平均余弦分数（↓，与训练目标一致）、平均垂直分数（↓）和多步预测MSE（↓）。改进倍数相对于B1在m=1时计算。**  
**参考图表说明 图2：多步预测MSE（实线，左轴，对数刻度）和预测准确率（虚线，右轴）随跳跃距离m的变化。步骤边界模型（A、A1、A2）比冻结基线（B1）和随机token STP（C）提升数个数量级。模型A在m=8时仍保持可用（MSE<1）。**  
在连续步骤边界上采样（模型A）在m=1时比冻结基线实现168倍更精确的预测，而原始随机token方法（模型C）仅实现4倍（表2 (https://arxiv.org/html/2604.18464#S5.T2)，图2 (https://arxiv.org/html/2604.18464#S5.F2)）。模型A在m=3时的预测（MSE 0.093）比B1在m=1时的预测（MSE 0.955）更精确。L_NTP的作用似乎是一种权衡而非冗余。模型A1（仅L_STP^step）测试是否可以丢弃L_NTP。在线性外推下，A1与A的差异在12%以内。在使用训练好的MLP预测器时，A1在所有跳跃距离上比A提高47–57%（§6.1 (https://arxiv.org/html/2604.18464#S6.SS1)）。选择是在生成质量（L_NTP：GSM8K +3.7%）和几何可预测性（仅L_STP：MLP预测提升2倍）之间的权衡。  

### 5.2 跨模型与跨数据集验证  
**表3：B1 vs. A在三个模型家族上的ProcessBench结果。准确率=1 - sqrt(MSE)。**  
**参考图表说明 图3：几何改进的鲁棒性。(a) 跨模型：B1 vs. A在m=1时的预测MSE（对数刻度）。改进幅度从39倍（Qwen2.5-1.5B）到335倍（Llama-3.2-1B）。(b) 跨数据集：B1 vs. A在四个不同难度的ProcessBench子集上。改进幅度一致（145倍–183倍），在较难的数据集上略强。**  
为了研究几何表示微调是否可以迁移到非数学模型，我们同样微调了两个模型：Qwen2.5-1.5B（与Qwen2.5-Math-1.5B直接对比）以及Llama-3.2-1B（代表不同模型架构和不同维度）。从结果中我们观察到，这种效果跨架构、隐藏维度和预训练范式迁移（表3 (https://arxiv.org/html/2604.18464#S5.T3)，图3 (https://arxiv.org/html/2604.18464#S5.F3)a）。Llama-3.2-1B取得了最强的改进（335倍），尽管它没有数学专用预训练。这种改进在不同难度级别上也成立（图3 (https://arxiv.org/html/2604.18464#S5.F3)b），在四个ProcessBench子集（GSM8K 165倍、MATH 145倍、OlympiadBench 183倍、Omni-MATH 180倍）上均一致。几何塑造效果一致，且在较难问题上略强，从而排除了STP过拟合简单问题结构的担忧。  

### 5.3 任务准确率：无退化  
**表4：GSM8K准确率（贪心解码）。**  
总体而言，所有模型在准确率上表现相当，确认了步骤边界STP不会降低模型的推理能力。模型A的准确率与普通微调（B2）相同，同时实现了168倍更好的潜在预测MSE（表4 (https://arxiv.org/html/2604.18464#S5.T4)）。随机token STP是唯一相对于冻结基线准确率下降的变体。这表明推理的语义结构防止了对语言模型目标的干扰。准确率和MSE测试实际上揭示了两个损失相互正交贡献：L_NTP似乎影响准确率提升，而L_STP则带来几何对齐，塑造模型通过潜在空间思考的方式。  

### 5.4 解码验证  
**表5：从预测嵌入解码（200个ProcessBench样本）。**  
解码验证实验（表5 (https://arxiv.org/html/2604.18464#S5.T5)）表明，两个隐藏状态向量可能在欧氏距离上很近，但在LM头用于下一个token预测的特定子空间上不同。与解码相关的信息占据了全1536维隐藏状态空间内的一个较低维流形。为了测试我们的预测嵌入是否功能等价而不仅仅是几何接近，我们将预测的ẑ_{k+m}和实际的z_{k+m}都通过模型的最终RMSNorm和LM头，产生两个词汇表上的概率分布。在m=1时，两个分布在前置词（top-1）上的一致性为93.4%（冻结基线B1为19.1%），KL散度仅为0.019 nats（B1为4.581）。步骤检索度量提供了互补视角：预测嵌入在所有步骤边界嵌入中（排除源点z_k和z_{k-1}）的最近邻在92.2%的情况下是正确的目标步骤，确认了ẑ_{k+m}落在正确的推理状态。在m=1时，模型A和B1之间出现了有启发性的对比：B1实现了相对较高的步骤检索（70.2%）但较低的token一致性（19.1%），意味着冻结模型的预测嵌入在几何上接近正确步骤，但编码了不同的下一个token信息——几何位置和功能内容未对齐。STP训练使两者对齐：模型A的92.2%检索率和93.4% token一致性表明，在几何塑造后几何邻近性与功能等价性一致。在更长视距上，模型A的token一致性优雅地下降——m=2时为71.1%，m=3时为34.1%——而B1崩溃至接近随机（10.1%和8.3%）。值得注意的是，模型A在m=2时（71.1%）仍超过B1在m=1时（19.1%），与MSE结果相呼应，确认了预测优势远不止一步预测。  

## 6 分析  
### 6.1 流形结构  
**表6：MLP vs. 线性外推在ProcessBench上的结果（所有六个模型）。比率 = MLP MSE / 线性MSE；随m递减的比率表明平滑曲率。**  
我们训练一个残差多层感知器（MLP）后验预测器探针：ẑ_{k+m} = z_k + m(z_k - z_{k-1}) + g_φ(z_k, z_{k-1})，其中g_φ是一个3层MLP，最后一层零初始化（架构细节见附录A.3 (https://arxiv.org/html/2604.18464#A1.SS3)）。零初始化确保MLP从线性预测器的精确副本开始，并且只能改进它。MSE的任何减少都代表线性外推遗漏的真实非线性结构。MLP探针在80%的ProcessBench步骤边界对上训练，并在剩下的20%上评估。该探针纯粹作为测量微调模型潜在几何的仪器，不改变模型本身。为了表征模型的流形结构，我们将MLP与一个平凡的线性探针在相同任务上进行比较：给定步骤k和步骤k-1的隐藏状态，预测步骤k+m的隐藏状态。MLP MSE与线性MSE的比率揭示了超出直线近似之外可学习结构的存在。关键的是，随着m递减的比率表明平滑、复合的曲率，这是连贯弯曲流形的特征，而非随机噪声（表6 (https://arxiv.org/html/2604.18464#S6.T6)，图4 (https://arxiv.org/html/2604.18464#S6.F4)）。模型B1、B2和C显示出平坦或弱递减的比率，而经过步骤边界微调的模型（如A、A1和A2）显示出强烈递减的比率。模型A从0.285降至0.085（m=1到m=3下降了70%），A1从0.137降至0.029（下降了79%）。随m递减的比率是平滑弯曲路径的标志。曲率随视距二次方累积，因此更长的外推暴露了更多非线性结构供MLP学习。A1（纯L_STP，无语言建模损失）具有最陡峭的下降。