∇-Reasoner：在潜空间中通过测试时梯度下降实现大语言模型推理 nabla-Reasoner: LLM Reasoning via Test-Time Gradient Descent in Latent Space

完整流程如 Algorithm 1/2/3 所示，主要包含四步。第一步是初始 rollout：$y, z \sim \pi_{LLM}(\cdot|x)$，得到完整回答和它的 logits 矩阵。第二步是 DTO 优化（Algorithm 2）：把 token 用 straight-through Gumbel-Softmax 表示为 $y^{(t)}_i = \delta_{j^*} + \text{softmax}(z^{(t)}_i/\tau) - \text{StopGrad}(\text{softmax}(z^{(t)}_i/\tau))$，然后计算总损失 $\mathcal{L} = \mathcal{L}_{nll} + \lambda \mathcal{L}_{reward}$，其中 $\mathcal{L}_{nll} = -\sum_i \log \pi_{LLM}(y^{(t)}_i | y^{(t)}_{\leq i-1}, x)$ 提供分布正则，$\mathcal{L}_{reward} = -r(y^{(t)}|x)$ 提供方向信号；用梯度下降 $z^{(t+1)} = z^{(t)} - \eta \nabla_z \mathcal{L}$ 迭代 $T$ 次，得到 $\tilde{z}$。第三步是再采样与拒绝：从 $\tilde{z}_1/\tau$ 采出候选 $\tilde{y}_1$，若 $\tilde{y}_1 = y_1$ 则直接接受并进入下一步；否则让 LLM 用 $\tilde{y}_1$ 作前缀做一次新 rollout $\tilde{y}$，比较 $r(\tilde{y},\tilde{y}_1|x)$ 与 $r(y|x)$，只保留 reward 更高的那个。第四步是加速策略：(a) Gradient Caching，因为 argmax 的结果往往连续多步不变，把 $\partial \mathcal{L}/\partial y$ 缓存下来，下一轮用 surrogate loss $\sum y^\top_i g_i$ 直接复用，可以跳过 63.8% 的梯度反传；(b) Rollout Reusing，拒绝时直接把当前 token 之后的轨迹滑到下一个位置作为新 rollout，省掉一次额外生成；(c) Confidence/Gradient-Guided Token Selection，用熵阈值 $\epsilon_{ent}=0.25$ 和梯度范数阈值 $\epsilon_{grad}=8$ 跳过那些 logits 已经自信或梯度太小的 token，可以省掉 89.2% 的 DTO 调用；(d) 早停，当 rollout 总数达到 $N_{max}=8$ 时切换到普通自回归解码。

本文的技术新颖性主要体现在三点。第一，把一阶梯度信息显式地引入 LLM 测试时推理的解码循环中，并用 straight-through estimator 把离散 token 与连续 logits 的优化无缝衔接，避开了以往 EBM-based 方法端到端优化不稳定的问题。这一点在方法论上是范式转变——从'采样+打分'到'打分+梯度+采样'。第二，提出 token 级的 DTO 目标 $\mathcal{L} = \mathcal{L}_{nll} + \lambda \mathcal{L}_{reward}$，并推导了 Proposition C.1 把梯度分解为 $\delta_{prefix}$、$\delta_{postfix}$、$\delta_{reward}$ 三项，揭示了 attention 机制天然支持的双向信号传播，这是把 ToT 类搜索算法'编译'进一个并行梯度步的关键。第三，给出 Theorem 4.1 的理论结果：在分布空间做 Wasserstein gradient flow 最小化 PPO 目标，等价于在样本空间对每个样本做 DTO 的随机梯度流 $dx_t/dt = -\nabla \mathcal{L}(x_t) + \sqrt{2}\epsilon_t$。这把'训练时 RL 对齐'与'测试时解码优化'放在同一个数学框架下，统一了 parametric inference（训练）与 particle-based inference（测试时）的视角，是论文最具理论深度的贡献。

我认为主要有三处可以进一步加强。第一，KL 正则强度 $\lambda$ 固定为单一标量（实际靠 $\mathcal{L}_{nll}+\lambda \mathcal{L}_{reward}$），但不同 benchmark、不同难度下 $\lambda$ 的最优值可能差异很大；论文所有实验共用一组超参（$\epsilon_{ent}=0.25, \epsilon_{grad}=8, lr=0.01, T=20$），缺少对超参敏感度的报告。建议改成自适应 $\lambda$ 调度，例如随 reward 变化幅度动态调整、或在训练时用元学习得到任务相关的 $\lambda$。第二，DTO 的并行能力虽然靠 attention 实现，但每次解码仍要执行一次完整前向 + 后向，在 batch size 不够大时 wall-clock 不一定优于 BoN；附录 D.3 也承认 FLOPs 远高于 BoN（差 26 倍），只是恰好硬件能并行才打平。改进方向是把 DTO 与 speculative decoding 结合：用小模型快速预填草稿，对不确定 token 调用 DTO，可以进一步摊薄单 token 的成本。第三，依赖 outcome reward 意味着 reward 信号稀疏、且只有最终答案对的题才能拿到梯度；对于多步推理，中间步骤错了也无信号，作者虽提到可以扩展到 process reward，但实验里并未验证。建议引入 process reward（PRM）配合 DTO，并对每一个 thought boundary 做局部优化，从根本上解决长链推理里 reward vanishing 的问题。

作者在 Related Work 末尾和 Conclusion 里给出了几个明确方向：(1) 把 DTO 应用到 o1、DeepSeek-R1 这类已经过 RL 训练的 long-CoT 模型上，把它们的推理链再用一阶梯度细化；(2) 与 prompt 优化方法结合，因为 DTO 在输出空间、prompt 优化在输入空间，两者互补；(3) 与 continuous latent space reasoning（如 Geiping et al. 的 recurrent depth）协同，把隐空间推理与 token 空间推理做层次化。从结果延伸，我还能想到几条：(a) 用 DTO 训练出更小的 RM 蒸馏版，把 4B RM 压到 1B 同时保留梯度质量，进一步降低部署成本；(b) 把 Theorem 4.1 的 Wasserstein 视角推广到多模态（图像/视频生成），用相同框架做文/图生成的对齐；(c) 设计 reward-gradient 本身的正则化（如 Jacobian 谱约束）来缓解 reward hacking；(d) 在 agentic / tool-use 场景下，把 tool 调用结果当作额外 reward 信号纳入 DTO，让 agent 在测试时也能边走边优化。

论文整体可复现性较好。作者明确给出代码仓库 https://github.com/VITA-Group/Nabla-Reasoner，附录 D.1 列出全部超参（温度 0.5、top-p 0.95、最大长度 1024/3072、$\epsilon_{ent}=0.25$、$\epsilon_{grad}=8$、学习率 0.01、$T=20$、Adam-W + cosine scheduler），Algorithm 1-4 提供完整伪代码，Appendix C 给出 Theorem 4.1 的推导。所用模型 Qwen-2.5-Math-7B/7B-Instruct、Llama-3.1-8B-Instruct、Skywork-V2-Qwen3-4B/8B、Llama-3.1-8B 均为 HuggingFace 开源。数据集 MATH-500、AMC、AIME24/25、Open-thoughts、Numina-Math 也全部公开。算力方面，附录 D.3 提到在 8x H200 (140GB) 上做 AMC 上 83 prompts 的 wall-clock 评测；考虑到 DTO 单次前向+反向要过两个 7-8B 模型，且要做 20 步梯度下降 + 最多 8 次 rollout，完整跑一遍表 1 的全部实验估计需要 4-8 张 A100/H100 跑数天，对个人研究者门槛中等偏高。最大难点其实是工程实现：(i) 要把 Gumbel-Softmax 直通估计器与 HuggingFace logits 对齐；(ii) Gradient caching 需要在自定义 forward 里拦截 argmax 变化；(iii) Rollout reusing 必须维护正确的 KV cache，否则会引入 off-by-one 错误。论文给到了伪代码和 hyperparameter，但具体的工程 hook 仍需读源码。