揭示語言模型中的潛在推理策略

摘要

一個在推理任務上訓練的語言模型 p_θ(y mid x) 能透過多種不同的策略來解決問題，但這些策略是隱含且糾結在模型的反應分佈之中。我們研究如何將給定預訓練語言模型的反應分佈，分解為一個結構化的、以策略為條件的表徵。具體而言，我們學習一個潛在變數分解 p_θ(y mid x) 轉化為 (r_ϕ(z mid x), g_ϕ(y mid x,z))，其中路由器 r 將每個輸入映射到一個關於潛在策略 z 的分佈，而生成器 g 則根據該策略產生反應。一個關鍵挑戰在於，從基礎模型初始化的生成器，在沒有使用 z 的情況下就已經代表了 p_θ(y mid x)。因此，標準的變分推論無法激勵模型透過 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.