Challenges of Learning a Latent Variable Factorization of Pre-trained Language Models

This note records a writing-level framing of why the project is not just a standard conditional VAE setup. It sits between the paper-writing notes and the methodology/theory notes.

Related notes:

• Contributions
• Strategy-structure recovery problem
• Base-directed reconstruction objective
• CVAE objective and loss terms
• Formalizing strategy identification

Relation to latent variable generative modeling

Our approach is related to latent variable generative modeling approaches such as VAEs. In a VAE, we observe data $x \sim D$ from some data distribution and learn a generator $p(x \mid z)$ that maps a latent variable $z$ to an instance. Conditional VAEs apply this to learning a conditional distribution $p(y \mid x)$ from a dataset of pairs $(x, y) \sim \mathcal{D}$.

Core challenges

1. Learning a latent variable generative model over discrete sequence space with an autoregressive model class.
2. Working in the adaptation setting rather than the ordinary dataset-learning setting.

The second point is the more fundamental difference. Rather than learning from data with some dataset

$$\mathcal{D} = \{(x_i, y_i)\},$$

we start from a model that already captures the observable distribution $p(y \mid x)$ and ask for a useful latent factorization of that law.

Challenges introduced by this setting

• The generative model already captures $p(y \mid x)$ and therefore does not need a latent variable in order to fit the observable law.
• The distribution may already be a mixed and entangled superposition of multiple strategies, with the latent "modes" entangled in the hidden states.
• Posterior collapse remains a major risk, especially because the decoder is strong enough to model much of the conditional law directly.
• Adaptation data comes from the model rather than from an external dataset, which makes the setup overlap with distillation-like training even if "distillation" is not quite the right term.

The adaptation setting makes posterior collapse unusually severe. If the factorized model is initialized at or near the base model, then $z$ does not need to encode information about $y$ beyond $x$ in order to preserve the observable conditional law. Distributional fidelity therefore under-specifies the intended goal of latent strategy recovery.

Reference phrasing:

The goal is not to obtain another generative model of $p_\theta(y|x)$; we already have one. The goal is to decompose that model into a latent-variable factorization that recovers implicit strategy structure while preserving observable behavior.

Another way to state the optimization difficulty is that the plain reconstruction term is an absolute code-length objective. High-surprisal examples contribute more raw loss mass, so they dominate the reconstruction term even when the scientific question is not "where can we get the largest absolute improvement?" but rather "where is $z$ carrying information beyond $x$?" Base-relative normalization changes that emphasis by turning reconstruction into a fractional objective.

The proposed methodology responds by using the fixed base model $p_\theta(y|x)$ to define adaptive reconstruction pressure: global scaling restores reconstruction pressure relative to the KL term, while token-level weighting emphasizes residual base-model surprisal.

Why this note matters

This framing explains why the project has to combine:

• methodology for preventing trivial or collapsed latent usage,
• theory for understanding non-identifiability and selection among many factorizations,
• experiments that test whether recovered latent variables align with intended strategies rather than arbitrary predictive partitions.