Story Boarding

Abstract

• Motivate problem: a reasoning language model defines a distribution $p_\theta(y|x)$ over response sequences $y$ given a prompt or task $x$. This distribution implicitly encodes an entangled mixture over multiple high-level abstract strategies for solving a given problem $x$. This is only implicitly represented. Uncovering these latent abstract strategies
• We formulate and explore the problem of learning a latent-variable router-generator factorization $(r_\phi(z|x), g_\phi(y|x,z))$ of a language model $p_\theta(y|x)$, where the router $r_\phi(z|x)$ selects a strategy $z$ given a task $x$ and the strategy-conditioned generator generates a solution $y$ guided by the strategy.
• This setting presents unique challenges that stem from the expressivity of the generator class and the optimization dynamics at-initialization where the generator already captures the $y$ given $x$ distribution without the latent strategy variable, leading to a form of posterior collapse.
• Our core methodology addresses this by formulating a variational training objective that converts absolute reconstruction improvement into fractional improvement relative to the base model, while also applying token-level pressure for $z$ to encode additional information about $y$ beyond $x$, leading to separation in strategy space.
• We empirically evaluate this on a suite of multi-strategy algorithmic tasks demonstrating the viability of the methodology.

Introduction

Motivate the problem

• Key setting of interest: language reasoning models (language models applied to reasoning or problem solving). Problem $x$. Model generates solution response $y ~ p_\theta(y|x)$. Solution $y$ has reward/cost/correctness measured by, say, $r(x,y)$ (e.g., 0)
• The question we ask: How does a model generate a solution $y$ across different implicit latent reasoning modes.
• A language model is trained on large scale data. For a given problem type, it sees many different approaches or strategies.
  • Examples
• The model implicitly represents an entangled mixture over
• Language models are trained on large scale

Figure 1: A cartoon depicting the core problem. (Left) Cartoon depiction of multi-modal p(y|x) distribution with multiple modes across reasoning strategies (proof by induction, by contradiction, etc). (Right) Depiction of strategy-aligned router-generator latent-variable factorization of a language model

Related Work

• Motivation: Steering Active, directed exploration in abstract strategy space
  • In cotnrast to: passive on-policy exploration; or exploration in low-level token-space
• (Mechanistic) Interpretability:
  • Sparse VAEs
  • …
• Latent variable models
  • Variational inference
  • Statistical literature on learning finite mixturez
  • VAEs
    • conditional VAEs
    • beta VAEs
    • VAEs for language models or sequence models

[Problem Statement] Latent-Variable Factorization of Abstract Strategies in Language Models

• Defining the problem. Introduce terminology, notation, etc.
• Disederate of latent variable $z$
  • Separate reasoning "modes"
  • Refer to [Figure 1]
• Emphasize how this problem differs from standard latent-variable generative models (e.g., standard VAEs).
  • Factorization: Take an existing conditional generative model $p_\theta(y|x)$ and factorize it into $\int_z r(z|x) g(y|x,z)$
  • not learn a latent-variable generative model of a data distribution $\mathcal{D}$ over data $x$ or $(x,y)$
  • This introduces unique challenges

[Methodology Pt1] How to parameterize and represent the router-generator $(r(z|x), g(y|x,z))$ factorization of a language model $p(y|x)$

• Special token <z> conditions/triggers model to reveal a latent variable $z ~ r(z|x) = p_\theta(\cdot | x_1, \ldots, x_{t_x}, \mathtt{\langle z \rangle})$
• Strategy-conditioned generator now conditions on this latent variable in its generation of $y$: $y_t \sim g_\theta(\cdot | x_1, \ldots, x_{t_x}, \mathtt{\langle z \rangle}, z, \mathtt{\langle/ z \rangle}, y_{<t})$

[Methodology Pt2] Latent-Variable Factorization Objective and Optimization Procedure

Motivate the core challenge

• A standard approach to learning latent variable models, in neural nets in particular, is variational inference and variational autoencoders, which use the evidence lower bound objective

  wrapfigure: Show that this fails…

  • ? What kind of figure shall we use?

• Interpret and explain why this challenge arises

  • Unique to this setting:
    • Base model already encodes the $y|x$ data distribution
    • The model is initialized at
    • Generator already captures the data distribution; in fact, it defines the data distribution
    • Related to posterior collapse problems caused by expressive generator class, but even more pathological due to initialized parameters defining the data distribution
    • The goal is not to learn a latent-variable generative model (i.e., not merely to capture the data distribution, which we already do), but to factorize and separate/disentangle generative modes.
    • There does not exist a unique factorization that captures
    • Naive variational inference objective does not prefer one factorization over another, and due to the expressiveness of the generator class and its initialization the nearest solution is the one that ignores the latent variable $z$

• Goal: We need a modified objective that explicitly places a preference over factorizations, preferring disentangled factorizations that separate latent modes and increasing pressure for z to encode information about y beyond what is already predictable from x.

Solution: Proposed methodology

1. Carefully control KL regularization strength
   • Related to $\beta$-VAE-like objective (except applied to conditional autoregressive sequence modeling)
2. Rather than standard reconstruction loss, Latent-Free Surprise Excess Reconstruction Loss $(y|x) \to (y|x,z)$
   • Story emphasis: move from absolute reconstruction improvement to fractional improvement relative to the base model
3. Misc tweaks and ablations

Experiments

Tasks & Data

• List tasks; List set of strategies from each task.
• Details of format

Setup: Entanglement

• We include
  1. small-scale models trained from scratch with random initialization
  2. Pre-trained language models (e.g., Qwen) adapted to task data
• This provides a controlled test bed with which to

Figure: Position-wise strategy decodability. A heatmap of task x (relative) position depicting strategy-decodability. Strategy information is present in the latent embeddings; just not explicitly separable or available in the generative procedure $p(y|x)$. This indicates the in-principle possibility of disentanglement (the information is there to be disentangled). The model implicitly "knows" about the existence of different strategies.

Strategy Disentanglement: $p(y|x) \to (r(z|x), g(y|x,z))$

Figure: Latent space disentanglement throughout training. Show frames of strategy color-coded $z \sim r(z|x), y \sim g(y|x,z)$. Show separation of strategies in latent space across training.

Figure: Normalized Reconstruction Loss versus Linear Strategy-Decodability of latent variable $z$