学习细化隐藏状态以实现可靠的LLM推理

# 学习优化隐藏状态以实现可靠的 LLM 推理
来源: https://arxiv.org/html/2606.17524
###### 摘要

大型语言模型展现出强大的推理能力，但在复杂的多步设置中，其内部推理过程仍可能不稳定，早期的隐藏状态错误可能会传播到不正确的预测中。我们提出 ReLAR，一种强化引导的潜在细化框架，能够在解码前迭代地更新隐藏表示。ReLAR 维护一个紧凑的潜在推理状态，并使用学习到的深度控制器和动作控制器来自适应地确定细化步骤的数量和方向。这些控制器基于逐步似然改进的策略梯度目标进行训练，使得在不生成显式思维链的情况下，也能实现高效的输入依赖推理。在医学、数学、多跳推理以及开放生成基准上的实验表明，ReLAR 在推理开销显著低于显式推理基线的情况下，提高了准确性、生成质量和推理稳定性。代码可在 tongyu0924/Learning-to-Refine-Hidden-States (https://github.com/tongyu0924/Learning-to-Refine-Hidden-States-for-Reliable-LLM-Reasoning) 获取。

## 1 引言

大型语言模型（LLMs）在医学问答、临床总结和诊断推理方面展现了强大的能力（Singhal 等，2023 (https://arxiv.org/html/2606.17524#bib.bib1)；Thirunavukarasu 等，2023 (https://arxiv.org/html/2606.17524#bib.bib2)；Lucas 等，2024 (https://arxiv.org/html/2606.17524#bib.bib23)），凸显了它们作为未来临床决策支持系统核心组件的潜力。

然而，在复杂的多步设置中，可靠推理仍然更具挑战性。输入可能不完整、异构或内部冲突，即使是微小的逻辑不一致也可能在推理步骤间传播，并在高价值环境中导致错误结论（Chen 等，2025 (https://arxiv.org/html/2606.17524#bib.bib24)；He 等，2025 (https://arxiv.org/html/2606.17524#bib.bib25)）。

在这种情况下，问题往往不是缺乏知识，而是模型如何在多个推理步骤间内部整合证据的不稳定性。模型可能会过度锚定一个显著信号，低估其他相关信息，从而偏离到错误结论。因此，确保稳定且可控的多步推理对于可靠部署至关重要。

一种主要的 LLM 推理方法是显式推理，例如思维链（CoT）提示，它鼓励模型以自然语言生成中间推理步骤（Wei 等，2022 (https://arxiv.org/html/2606.17524#bib.bib5)；Kojima 等，2022b (https://arxiv.org/html/2606.17524#bib.bib26)；Wang 等，2023a (https://arxiv.org/html/2606.17524#bib.bib27)；Yao 等，2024 (https://arxiv.org/html/2606.17524#bib.bib20)；Shinn 等，2023 (https://arxiv.org/html/2606.17524#bib.bib28)）。这些方法通常能提高任务性能，并且在医学环境中被广泛采用，因为它们看起来具有可解释性。然而，它们在生成的文本层面运作，并不直接调节模型的内部推理过程。先前工作表明，即使最终答案看起来流畅或正确，推理痕迹也可能包含逻辑漏洞或幻觉内容（Lyu 等，2023 (https://arxiv.org/html/2606.17524#bib.bib29)；Lanham 等，2023 (https://arxiv.org/html/2606.17524#bib.bib30)）。此外，生成长推理痕迹增加了推理延迟和计算成本，这可能会限制在时间敏感的临床场景中的实用性。

参见图 1：ReLAR 与常规自回归推理的比较。ReLAR 在解码前迭代地优化隐藏状态，而自回归基线则从固定的 h0 解码，不进行修正。这种失败模式如图 1 所示。在此案例中，基线模型从固定的隐藏状态解码，导致其过早地确定一个推理结果。更可靠的推理过程需要在确定输出之前整合问题数量、算术关系和分数约束。关键区别不在于数学知识的可用性，而在于模型能否在积累证据的同时修正并稳定其内部表示。

因此，近期工作探索了潜在表示编辑和干预作为控制模型推理的机制（Wang 等，2025 (https://arxiv.org/html/2606.17524#bib.bib53)；Stolfo 等，2025 (https://arxiv.org/html/2606.17524#bib.bib54)；Helff 等，2026 (https://arxiv.org/html/2606.17524#bib.bib55)）。隐藏状态表示编码了结构化且语义上有意义的信息，对内部激活的干预可以比输出层监督更直接地影响模型行为（Meng 等，2022 (https://arxiv.org/html/2606.17524#bib.bib31)；Li 等，2023 (https://arxiv.org/html/2606.17524#bib.bib40)）。然而，现有的潜在方法对于复杂的多步推理仍然有限。大多数方法关注静态或单步干预，不支持迭代优化或对推理步骤间内部一致性的显式控制。因此，潜在状态中的早期错误可能持续存在并累积，导致不稳定的推理轨迹。因此，现有方法要么在输出层面运作而不控制内部推理，要么作用于潜在表示但不进行迭代和自适应的优化，使得推理不稳定性问题本质上未得到解决。

为了解决这一局限，我们提出了一种迭代的隐藏状态优化框架，能够在解码前通过强化学习控制的内部推理。我们的方法完全在隐藏状态空间内执行一系列优化步骤，使得内部表示能在生成任何输出之前逐步调整和稳定。一个学习到的控制器动态决定优化的方向和优化的迭代次数，实现了基于任务难度的推理深度自适应分配。通过在解码前直接干预隐藏表示，我们的框架提供了对内部推理轨迹的显式且细粒度的控制，这是仅靠输出级推理监督无法实现的。

我们的贡献可以总结如下：(1) 我们提出了一种迭代的隐藏状态优化框架，能够在解码前直接控制内部推理轨迹。(2) 我们引入了基于强化学习的控制器，能够动态调节优化方向和推理深度，实现内部推理的自适应分配。(3) 在临床推理基准上的实验证明了步级连贯性和整体可靠性的提升，同时推理时间开销低于显式推理基线。

## 2 相关工作

### 2.1 大型语言模型中的隐式推理

大型语言模型（LLMs）不仅可以通过显式的自然语言理由进行复杂推理，还可以通过其内部表示中的隐式计算来实现。虽然思维链提示以文本形式引出中间推理步骤（Wei 等，2022 (https://arxiv.org/html/2606.17524#bib.bib5)；Kojima 等，2022a (https://arxiv.org/html/2606.17524#bib.bib6)），但近期研究表明，即使这些推理没有明确表述出来，模型也可能在隐藏状态中编码与任务相关的推理信息（Schlag 等，2021 (https://arxiv.org/html/2606.17524#bib.bib46)；Geva 等，2021 (https://arxiv.org/html/2606.17524#bib.bib48)）。

隐式推理之所以有吸引力，是因为它避免了冗长文本理由的成本和潜在不忠实性，同时仍允许模型在产生答案之前整合证据。然而，标准 LLM 推理通常依赖于单次前向传播，使得隐式推理过程在很大程度上不受控制。因此，隐藏表示可能在解码前编码了不完整或不稳定的推理状态。

我们的工作基于这一观点，将推理视为一个可以在生成前进行优化的内部潜在过程。我们不要求模型在文本中暴露所有中间步骤，而是迭代地更新隐藏表示，使得隐式推理能够在解码前被稳定和控制。

### 2.2 潜在推理与表示级优化

语言模型中的推理通常通过思维链（CoT）提示（Wei 等，2022 (https://arxiv.org/html/2606.17524#bib.bib5)）及其扩展来引出，包括自一致性（Wang 等，2023b (https://arxiv.org/html/2606.17524#bib.bib7)）和树状探索（Yao 等，2024 (https://arxiv.org/html/2606.17524#bib.bib20)）。这些方法在生成的文本层面运作，需要显式地产生中间推理痕迹，这可能不稳定且计算成本高，尤其是在临床环境中。

近期工作探索了潜在推理，其中多步推理发生在隐藏状态空间内，而不是通过生成的 token 进行（Schlag 等，2021 (https://arxiv.org/html/2606.17524#bib.bib46)；Creswell 等，2023 (https://arxiv.org/html/2606.17524#bib.bib49)）。先前的方法研究了隐藏状态编辑或激活优化，但通常依赖于预定义或启发式干预，缺乏对内部推理动态的原理性控制（Elazar 等，2021 (https://arxiv.org/html/2606.17524#bib.bib47)；Geva 等，2021 (https://arxiv.org/html/2606.17524#bib.bib48)），这对于可靠的医学推理尤其成问题。相比之下，我们的方法直接干预隐藏表示，实现了对推理动态（包括优化深度和方向）的显式且细粒度的控制。

### 2.3 用于自适应推理控制的强化学习

强化学习（RL）已被广泛用于大规模语言系统中的策略优化、奖励塑造和自适应计算（Ouyang 等，2022 (https://arxiv.org/html/2606.17524#bib.bib8)；Bai 等，2022 (https://arxiv.org/html/2606.17524#bib.bib10)；Rafailov 等，2023 (https://arxiv.org/html/2606.17524#bib.bib11)）。深度自适应机制如自适应计算时间（ACT）（Graves, 2016 (https://arxiv.org/html/2606.17524#bib.bib22)）和动态执行策略证明了根据输入复杂度分配可变计算量的好处。RL 也被应用于控制模型编辑和模块化架构中的表示级干预。

然而，这些方法并非旨在稳定高风险领域（如医学）中的多步推理。相比之下，我们的工作利用强化学习直接控制潜在推理动态，训练专门的控制器自适应地选择优化深度和优化方向（Turner 等，2023 (https://arxiv.org/html/2606.17524#bib.bib41)；Meng 等，2022 (https://arxiv.org/html/2606.17524#bib.bib31)）。这种设计使得在保持稳定性和效率的同时实现输入依赖的内部推理，这对于可靠的医学决策支持至关重要。

参见图 2：模型流水线概览。
参见图 3：迭代潜在状态优化

## 3 方法

我们提出 ReLAR（强化引导的潜在优化），这是一个迭代的隐藏状态优化框架，能够在预训练语言模型的潜在空间内实现可控的多步推理。ReLAR 不是通过单次前向传播产生答案，而是在解码前执行一系列表示优化步骤，由两个学习到的控制器引导，自适应地确定隐藏状态应该被修正的程度和方向。图 2 给出了完整流水线的概览。

### 3.1 预备知识

设 x = (x_1, ..., x_n) 为输入 token 序列，p_θ 为基于 transformer 的语言模型，通过隐藏表示将 x 映射到输出分布。通过 p_θ 的单次前向传播得到最后一层隐藏状态 h_0 ∈ R^{L×D}，其中 L 是序列长度，D 是隐藏维度。

我们的框架为 p_θ 增加了一个低维的推理状态 s_t ∈ R^{d_s}，用于跟踪模型对输入的不断演变的内部潜在表示。与思维链理由不同，s_t 永远不会被解码成文本；它仅作为内部控制信号。从 (h_0, s_0) 开始，框架产生耦合的优化轨迹：

s_0 → s_1 → ... → s_T,   h_0 → h_1 → ... → h_T,

其中优化深度 T 由学习到的深度控制器自适应选择，如 3.4 节所述。

### 3.2 初始推理状态

给定基础隐藏表示 h_0，我们通过一个轻量级投影网络 f_extract 构建初始推理状态：

s_0 = f_extract(h_0).

向量 s_0 ∈ R^{d_s} 将来自最终 transformer 层的任务相关信息压缩成一个紧凑表示。它充当一个共享瓶颈：深度控制器 π_d 和动作控制器 π_a 都以 s_0 作为唯一输入，因此所有下游优化决策都由这个单一摘要控制。

### 3.3 动作引导的表示优化

在每个优化步骤 t ∈ {0, ..., T-1}，动作控制器预测一个优化方向和两个调制参数：

a_t = (γ_t, β_t, v_t) ~ π_a(a_t | s_t),

其中 v_t ∈ R^D 被归一化以满足 ||v_t|| = 1。标量参数 γ_t 和 β_t 用于计算有效的带符号步长：

α_t = f_α(γ_t, β_t),

它决定了更新的大小和符号。然后通过加法扰动来优化隐藏表示：

h_{t+1} = h_t + α_t v_t.

等价地，这个加法更新定义了优化函数：

f_refine(h_t, s_t, γ_t, β_t, v_t) = h_t + f_α(γ_t, β_t) v_t.

因此，v_t 决定了优化的方向，而 α_t 决定了表示沿该方向移动的距离。

然后刷新推理状态以反映更新后的隐藏状态：

s_{t+1} = g(s_t, h_{t+1}).

将这个过程迭代 T 步，使得模型能够在不产生任何中间 token 的情况下逐步稳定其内部表示。

在所有优化步骤之后，最终的推理状态 s_T 被实现为与原始编码 h_0 锚定的解码表示：

h_T = f_decode(s_T, h_0),   ŷ ~ p_θ(y | x, h_T).