Token-Level Surprisal-Share Weighting

Back to reconstruction weighting and normalization

Implementation note: Current implementation: token-level surprisal-share weighting

This note identifies a teacher-forcing dilution failure mode in the standard reconstruction objective and derives a token-level weighting scheme that reallocates reconstruction pressure according to each token's share of the residual code length under the z-free base model.

The main derivation is written in the discrete CVAE setting, because that is where the exact latent expectation is easiest to analyze. The same weighting principle is also used in the current continuous implementation.

Motivation: Teacher-Forcing Dilution

Setup and notation

Consider a discrete CVAE with $K$ latent values trained under teacher forcing. For one example $(x, y)$ with $y = (y_1, \dots, y_{T})$ , define:

per-token z-conditioned NLL: $\ell_{t} (z) = -\log p_\phi (y_{t} | x, z, y_{< t}),$
per-token z-free base-model surprisal: $b_{t} = -\log p_\phi (y_{t} | x, y_{< t}).$

Both use the same backbone $\phi$ ; the baseline simply omits the $angle.l z angle.r z_{k} angle.l / z angle.r$ section from the input.

The current reconstruction loss averages uniformly over all supervised tokens $\mathcal{S}$ with $|\mathcal{S}| = T$ :

$R(z) = 1 / T \sum_{t \in \mathcal{S}} \ell_{t} (z).$

Two-regime structure under teacher forcing

This is simple interpretation is approximate and applies only to

TODO: restructure this note to make it more concise and more clear; this simple "strategy revealed by $t_0$ interpretation should be separated into a brief callout and it should be made clear that this is a simple model that applies only to simple settings."

Under teacher forcing, the solution sequence has two regimes:

Strategy-sensitive prefix ( $t \leq t_0$ ). The teacher-forced history $y_{< t}$ does not yet determine the strategy. The base model has non-trivial uncertainty: $b_{t} > 0$ . The z-conditioned model can reduce this uncertainty when $z$ matches the strategy used in $y$ , so $\ell_{t} (z)$ varies across latent values.
Strategy-determined suffix ( $t > t_0$ ). The partial solution $y_{< t}$ reveals the strategy. Continuation is nearly deterministic: $b_{t} \approx 0$ and $\ell_{t} (z) \approx 0$ for all $z$ , because the teacher-forced prefix renders $z$ redundant.

The boundary $t_0$ is not a sharp cutoff but a useful abstraction. For synthetic-sequence tasks, it corresponds roughly to the first few solution tokens where the procedural strategy is manifested.

The dilution

Under the two-regime model, the token-averaged reconstruction becomes:

R(z) = 1 / T \sum_{t=1}^{t_0} \ell_{t} (z) + \underbrace{1 / T \sum_{t=t_0 + 1}^{T} \ell_{t} (z)}_{\approx 0}.

The variation of $R(z)$ across latent values:

R(z) - R(z') = 1 / T \sum_{t=1}^{t_0} [\ell_{t} (z) - \ell_{t} (z')] = t_0 / T \cdot \delta_{z z'},

where $\delta_{z z'} = 1 / t_0 \sum_{t \leq t_0} [\ell_{t} (z) - \ell_{t} (z')]$ is the mean per-informative-token difference. The factor $t_0 / T$ is the dilution: the $T - t_0$ easy tokens contribute zero to the numerator but inflate the denominator.

Gradient Analysis

Posterior parameters $\xi$

The reconstruction gradient on the posterior parameters $\xi$ of $q_\xi (z | x, y)$ is:

(\partial \mathcal{L}_\mathrm{recon}) / (\partial \xi) = \sum_{z} (\partial q_\xi) / (\partial \xi) \cdot R(z).

Because $\sum_{z} \partial q_\xi / \partial \xi = 0$ (the probabilities sum to one), any component of $R(z)$ that is constant across $z$ cancels. The effective gradient magnitude is proportional to the variation:

\left\lVert \frac{\partial \mathcal{L}_\mathrm{recon}}{\partial \xi} \right\rVert \sim \max_{z \neq z'} |R(z) - R(z')| = \frac{t_0}{T} |\delta|.

The KL gradient on $\xi$ is $O(\beta)$ , independent of $T$ . The reconstruction signal dominates when:

$t_0 / T |\delta| > \beta \cdot C$

for a constant $C$ set by the KL curvature. When $t_0 / T << 1$ and $\beta = O(1)$ , the KL wins and the posterior collapses to the prior.

Token-Weighted Reconstruction Loss

Definition

Define per-token surprisal-share weights from the stop-gradiented z-free baseline NLL:

\pi_t = \frac{\operatorname{sg}(b_t)}{\sum_{s \in \mathcal{S}} \operatorname{sg}(b_s) + \epsilon}

for supervised positions $t \in \mathcal{S}$ .

Equivalently, replace the uniform token average with a weighted average:

R^\pi (z) = \sum_{t \in \mathcal{S}} \pi_t \, \ell_t(z).

This is algebraically equivalent to the implementation form

R^w (z) = \frac{\sum_{t \in \mathcal{S}} w_{t} \cdot \ell_{t} (z)}{\sum_{t \in \mathcal{S}} w_{t} + \epsilon},

with $w_t = \operatorname{sg}(b_t)$ .

The full weighted reconstruction loss under the posterior:

\mathcal{L}_\mathrm{recon}^\pi = \mathbb{E}_{z \sim q_\xi (z | x, y)}[R^\pi (z)].

The total objective:

\mathcal{J} = \mathcal{L}_\mathrm{recon}^\pi + \beta \cdot \mathrm{KL}(q_\xi (z | x, y), p_\phi (z | x)) + \mathrm{auxiliary\ terms}.

Intuition

The latent variable should prioritize encoding information that helps predict tokens that remain uncertain under the z-free model $p(y \mid x)$ ; i.e., branching tokens that are not already determined by $x$ and the teacher-forced prefix $y_{<t}$ . The current weighting scheme operationalizes this by allocating reconstruction pressure in proportion to each token's share of the residual base-model surprisal inside the example.

Gradient improvement

Since $w_{t} \approx 0$ for $t > t_0$ , the weighted reconstruction concentrates on the strategy-sensitive prefix:

R^w (z) \approx (\sum_{t=1}^{t_0} b_{t} \cdot \ell_{t} (z)) / (\sum_{t=1}^{t_0} b_{t}).

The variation across latent values:

R^w (z) - R^w (z') = (\sum_{t=1}^{t_0} b_{t} \cdot [\ell_{t} (z) - \ell_{t} (z')]) / (\sum_{t=1}^{t_0} b_{t}) = O(|\delta|),

with no $t_0 / T$ dilution factor. The posterior gradient is amplified by $T / t_0$ relative to the uniform-average loss. The z-parameter gradient is similarly amplified: the denominator is $\sum_{t \leq t_0} b_{t}$ instead of $T$ , and late-token noise gradients are suppressed.

The reconstruction gradient on $\xi$ now dominates the KL when:

$|\delta| > \beta \cdot C,$

which is $T / t_0$ times easier to satisfy than the uniform-average condition.

Conditions For Effectiveness

When it works

Nonzero initial z-sensitivity. At initialization, z-token embeddings are random, producing small random differences in $\ell_{t} (z)$ across $z$ at early positions. Under the standard loss, this signal is diluted by $t_0 / T$ and may be too weak to overcome the KL pressure. Under the weighted loss, the signal is preserved at $O(|\delta|)$ , potentially sufficient for the optimizer to bootstrap z-usage.
Non-trivial base-model uncertainty at early positions. The synthetic- sequence tasks have genuine strategy ambiguity: different strategies produce different early solution tokens for the same input $x$ . The base model (initialized from the entangled checkpoint) should have $b_{t} > 0$ at these positions.
Self-adapting behavior. Early in training, the base model may be uncertain at many positions, making $w_{t}$ roughly uniform. As the backbone improves and late tokens become easy, the weights concentrate on the strategy-sensitive prefix. The reweighting strengthens exactly when dilution becomes the binding constraint.

When it may not suffice

Complete z-insensitivity. If the decoder produces $\ell_{t} (z) = \ell_{t} (z')$ for all $z, z'$ at every position (including early tokens), the weighted variation is still zero. Reweighting amplifies existing z-variation but cannot create it from nothing. This limitation can be addressed by pairing the weighted loss with an architectural change that ensures nonzero $\partial \ell_{t} / \partial \psi$ (for example, z-conditioned layer normalization or a multi-token z-prefix).
Degenerate weights. If $b_{t} \approx 0$ for all $t$ (including early tokens), the denominator $sum w_{t} \approx 0$ and the loss degenerates. This would indicate that the task has no token-level strategy ambiguity, in which case disentanglement via reconstruction alone may be impossible.
Single-token z-conditioning. The z-token is a single token in the input sequence. Even with correct gradient scaling, its influence on early-position predictions is mediated by attention and may remain small. Reweighting is necessary for correct gradient balance but may not be sufficient to overcome an architectural bottleneck.

Implementation

Computing the baseline $b_{t}$

The z-free baseline requires one additional forward pass per batch through the z-free input template:

<bos> <x> {x} </x> <y> {y} </y> <eos>

This is the original entangled-model format. The per-token NLL $b_{t}$ is computed under teacher forcing at supervised positions and detached before use as weights.

Two implementation policies are possible:

Frozen policy. Compute $b_{t}$ from a frozen reference model initialized from the entangled checkpoint. This keeps the weighting signal stable while the disentangled model is trained.
Dynamic policy. Compute $b_{t}$ from the current training backbone at each step (still stop-gradiented before use), so weights evolve with the model.

Current implementation status in this repository: the frozen policy is implemented; the dynamic policy is intentionally deferred.

Current status in this repository

This token-level weighting mechanism is currently implemented.

Core weighting helper: path:core/disentanglement/token_weighting.py
Weighted loss helpers: path:core/disentanglement/losses.py
Discrete CVAE integration: path:core/disentanglement/shared_discrete_cvae.py
Continuous CVAE integration: path:core/disentanglement/continuous_cvae.py
Trainer/runtime wiring:
- path:experiments/synthetic_sequences/disentangled/train_step.py
- path:experiments/synthetic_sequences/disentangled/evaluation_runtime.py
CLI flag: --token-weighted-reconstruction-weight in path:experiments/synthetic_sequences/disentangled/config.py

Relation to neighboring variants

Distribution-Level Z-Free-Relative Reconstruction Rescaling Restores reconstruction scale globally, but does not change which token positions matter more.
Example-level inverse-surprisal weighting Reweights whole examples by inverse base-model surprisal, rather than emphasizing high-surprisal positions inside each example.
Inverse-surprisal token-weighted reconstruction ablation A low-surprisal token-level ablation that reverses the within-example weighting direction to test whether high-surprisal branching-token pressure is the important ingredient.

Among the current reconstruction variants, token weighting is the one most directly aligned with the intuition that $z$ should help where the z-free model still has room to improve.

Modifying `masked_token_nll_per_example`

The current implementation computes:

denom = active.sum(dim=1).clamp_min(1.0)
return (token_loss * active).sum(dim=1) / denom

The weighted variant replaces the uniform active mask with active * weights, where weights is the detached per-token baseline NLL:

weighted_active = active * weights.detach()
denom = weighted_active.sum(dim=1).clamp_min(eps)
return (token_loss * weighted_active).sum(dim=1) / denom

Cost

One additional forward pass per batch (z-free baseline). For the exact discrete CVAE with $K$ latent values, the training loop already performs $K + 1$ forward passes per batch ( $K$ latent-conditioned and $1$ posterior). The baseline adds one more, increasing cost by $1 / (K + 1)$ .