GradCuit:透過梯度直連實現魯棒可解釋的測試時潛推理

本文提出GradCuit,一種創新的測試時推理優化方法,旨在解決現有基於優化的潛推理中信用分配間接且黑盒化的問題。傳統方法通過解碼後的詞元間接連接潛狀態與推理軌跡,導致難以追蹤潛狀態更新如何影響後續推理。GradCuit通過在Transformer層的隱藏表示與生成延續之間插入可優化潛狀態,利用因果自注意力機制建立從最終獎勵到潛狀態的完全可微路徑,實現直接梯度分配。在五個指令微調骨干模型、三個推理基準及兩種答案格式上的實驗表明,GradCuit平均準確率達64.5%,超越思維鏈提示6.6個百分點,領先最強競品2.4個百分點。此外,該方法展現出卓越的魯棒性,標準差從1.53降至0.82,且通過詞元級梯度歸因揭示了潛影響集中在推理連接詞元上,為測試時擴展提供了全新視角。

大型语言模型在复杂推理任务中的表现往往受限于其静态参数,而测试时优化(Test-Time Optimization)作为一种新兴范式,试图在不更新模型参数的情况下,通过优化实例特定的连续状态来提升输出质量。然而,现有的基于优化的潜推理方法存在一个根本性的缺陷:它们通常将潜状态与推理轨迹通过解码后的词元(decoded tokens)进行连接。这种间接的连接方式导致序列级的信用分配(credit assignment)变得模糊且低效,研究人员难以清晰地理解潜状态的更新究竟如何塑造了后续的推理过程,使得这一过程如同黑盒一般。针对这一痛点,本研究提出了GradCuit(Gradient through Circuit)方法。其核心贡献在于打破了解码词元的中介作用,直接在Transformer层的隐藏表示与生成的延续部分之间插入可优化的潜状态。这一设计使得因果自注意力机制能够确保每一个后续生成词元的对数概率(log-probability)都拥有一条通往之前所有潜状态的可微路径。通过这种架构,来自整个生成序列的奖励加权梯度可以直接、无阻碍地分配给潜状态,从而实现了端到端的直接优化,不仅提升了推理的准确性,更极大地增强了过程的可解释性。在技术实现层面,GradCuit巧妙地利用了Transformer内部的结构特性。

具体而言,该方法选定Transformer中的特定层,在该层的隐藏表示(即提示词处理后的内部状态)与生成的延续文本之间注入可优化的潜变量。由于Transformer采用因果自注意力机制,后续生成的每一个词元在计算注意力时,都能关注到之前的所有位置,包括这些被注入的潜状态。这意味着,从最终的输出结果反向传播梯度时,梯度可以通过剩余的Transformer块,直接回流到这些潜状态上,而无需经过解码器或中间词元的非线性映射干扰。训练策略上,模型参数保持冻结,仅优化这些插入的潜状态向量,利用强化学习中的奖励信号或损失函数的梯度来更新潜状态。这种设计使得模型能够根据最终的推理结果,动态调整其内部的"思考路径",即调整潜状态以引导后续词元的生成方向。与传统的通过调整解码概率或采样策略不同,GradCuit直接干预了模型内部的表示空间,使得优化过程更加精细和可控,同时也为分析模型内部的推理机制提供了清晰的梯度流向。为了验证GradCuit的有效性,研究者在五个不同的指令微调骨干模型上进行了广泛的实验,涵盖了三个主流的推理基准数据集,并测试了两种不同的答案格式。实验结果显示,GradCuit取得了平均64.5%的准确率,这一成绩不仅比传统的思维链(Chain-of-Thought)提示方法高出6.6个百分点,还超越了当前最强的竞争方法2.4个百分点,显示出其在复杂推理任务中的显著优势。

在鲁棒性方面,GradCuit表现尤为出色。研究者在七种不同的学习率设置下进行了测试,发现GradCuit的性能波动极小,其准确性的标准差从对比方法LatentSeek的1.53降低到了0.82,表明该方法对超参数的选择不敏感,具有更强的稳定性。值得注意的是,即使是GradCuit的随机游走变体,其性能依然与LatentSeek相当,进一步证明了其优化机制的稳健性。消融实验和层分析还揭示,早期到中间的Transformer层是优化潜状态最有效的空间,而词元级的梯度归因分析则显示,潜状态的影响力主要集中在那些起到"推理连接"作用的词元上,这为理解模型如何构建逻辑链条提供了微观视角。GradCuit的提出对开源社区和工业落地具有深远的意义。首先,它提供了一种无需重新训练模型即可提升推理能力的低成本方案,这对于资源受限的场景极具吸引力。其次,其增强的可解释性使得研究人员能够更清晰地洞察大模型内部的推理过程,通过观察梯度流向和潜状态变化,可以诊断模型在哪些环节出现了逻辑断裂或错误,从而指导后续的模型改进或提示工程优化。在工业应用中,这种测试时缩放(test-time scaling)技术可以作为一种即插即用的模块,嵌入到现有的推理系统中,显著提升复杂任务(如数学证明、代码生成、逻辑推理)的可靠性和准确性。此外,GradCuit开辟了大模型适应"如何思考"而非仅仅"生成什么"的新方向,鼓励后续研究探索更多基于内部状态直接优化的方法,推动大模型从被动响应向主动推理的范式转变,为构建更加透明、可控和智能的AI系统奠定了重要基础。

Sources