言語モデルにおける潜在的な推論戦略の解明

要旨

推論タスクで学習された言語モデル \(p_\theta(y \mid x)\) は、複数の異なる方略を通じて問題を解くことを学習するが、これらの方略は暗黙的であり、モデルの応答分布の中に絡み合っている。本研究では、与えられた事前学習済み言語モデルの応答分布を、構造化された方略条件付き表現に分解する問題を扱う。具体的には、潜在変数分解 \(p_\theta(y \mid x) \rightarrow (r_\phi(z \mid x), g_\phi(y \mid x, z))\) を学習する。ここで、ルータ \(r\) は各入力を潜在方略 \(z\) 上の分布に写像し、ジェネレータ \(g\) はその方略に条件付けられた応答を生成する。重要な課題は、ジェネレータがベースモデルから初期化されているため、\(z\) を用いずに既に \(p_\theta(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.