언어 모델의 잠재적 추론 전략 발견

초록

추론 작업에 대해 학습된 언어 모델 \(p_\theta(y \mid x)\)는 여러 개의 뚜렷한 전략을 통해 문제를 해결하는 방법을 학습하지만, 이러한 전략은 모델의 응답 분포 내에서 암시적이고 얽혀 있다. 본 연구는 주어진 사전 학습된 언어 모델의 응답 분포를 구조화된 전략 조건부 표현으로 분해하는 문제를 다룬다. 구체적으로, 잠재 변수 분해 \(p_\theta(y \mid x) \leadsto (r_\phi(z \mid x), g_\phi(y \mid x, z))\)를 학습하는데, 여기서 라우터 \(r\)은 각 입력을 잠재 전략 \(z\)에 대한 분포로 매핑하고, 생성기 \(g\)는 해당 전략에 조건화된 응답을 생성한다. 핵심 난제는 기본 모델에서 초기화된 생성기가 이미 \(z\)를 사용하지 않고 \(p_\theta(y \mid x)\)를 표현한다는 점이다. 따라서 표준 변분 추론은 모델이 \(z\)를 통해 정보를 전달할 유인을 제공하지 않으며, 심각한 형태의 사후 분포 붕괴를 초래할 수 있다. 이를 해결하기 위해 우리는 기본 모델의 응답 손실에 대한 부분적 정보 이득을 측정하고, 기본 모델의 놀라움(surprisal)이 높은 토큰에 재구성 압력을 집중시키는 변분 목적 함수를 제안한다. 이는 \(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.