揭示语言模型中的潜在推理策略

摘要

在推理任务上训练的语言模型 p_θ(y mid x) 能够通过多种不同的策略来解决问题，但这些策略隐含在模型的响应分布中，相互纠缠。我们研究如何将给定预训练语言模型的响应分布分解为结构化的、以策略为条件的表示。具体而言，我们学习一个潜变量分解 p_θ(y mid x) leadsto (r_ϕ(z mid x), g_ϕ(y mid x,z))，其中路由器 r 将每个输入映射到潜策略 z 的分布，生成器 g 则根据该策略生成响应。一个关键挑战在于：生成器初始化自基模型，本身已能表示 p_θ(y mid x) 而无需使用 z。因此，标准变分推断无法激励模型通过 z 传递信息，并可能导致严重的后验坍塌。为解决此问题，我们提出一种变分目标函数，该函数衡量相对于基模型响应损失的分数量增益，并将重构压力集中在基模型惊讶度高的词元上，从而鼓励 z 编码与策略相关的响应变异。我们引入一个多策略算法任务的基准测试，并证明该目标函数能够在保留基模型响应分布的同时，恢复出与不同参考策略对齐的潜编码。

English

A language model p_θ(y mid x) trained on reasoning tasks learns to solve problems via multiple distinct strategies, yet these strategies are implicit and entangled within the model's response distribution. We study the problem of decomposing the response distribution of a given pretrained language model into a structured, strategy-conditioned representation. Specifically, we learn a latent-variable factorization p_θ(y mid x) leadsto (r_ϕ(z mid x), g_ϕ(y mid x,z)), where a router r maps each input to a distribution over latent strategies z and a generator g produces the response conditioned on that strategy. A key challenge is that the generator, initialized from the base model, already represents p_θ(y mid x) without using z. Standard variational inference therefore gives the model no incentive to route information through z and can yield a severe form of posterior collapse. To address this, we propose a variational objective that measures fractional information gain relative to the base model's response loss and concentrates reconstruction pressure on tokens with high base model surprisal, encouraging z to encode strategy-relevant response variation. We introduce a benchmark of multi-strategy algorithmic tasks and show that this objective recovers latent codes aligned with distinct reference strategies while preserving the base model's response distribution.