Example-Level Inverse-Surprisal Weighting

Back to reconstruction weighting and normalization

Implementation note: Current implementation: example-level inverse-surprisal weighting

This note records an under-consideration variant of baseline-normalized reconstruction that weights each example by the inverse of its own base-model surprisal rather than using one dataset-level constant.

Setup

Let $\bar{\ell}_\theta(x, y)$ denote the current model's per-example y-region reconstruction loss:

$$\bar{\ell}_\theta(x, y) = \mathbb{E}_{z \sim q_\xi(\cdot \mid x, y)}\left[-\log p_\psi(y \mid x, z)\right].$$

Let $b(x, y)$ denote the z-free base-model surprisal on the same example:

$$b(x, y) = -\log p_\mathrm{base}(y \mid x).$$

Objective

The proposed example-level normalized reconstruction term is

$$R_\mathrm{ex}(x, y) = \frac{\bar{\ell}_\theta(x, y)}{b(x, y) + \epsilon},$$

with $\epsilon > 0$ included explicitly to prevent low-surprisal examples from producing extreme ratios.

The empirical objective becomes

$$\mathbb{E}_{(x, y) \sim D}[R_\mathrm{ex}(x, y)] = \mathbb{E}\left[\frac{\bar{\ell}_\theta(x, y)}{b(x, y) + \epsilon}\right].$$

Fractional Interpretation

When $\epsilon$ is small relative to $b(x, y)$, this admits the decomposition

$$\mathbb{E}\left[\frac{\bar{\ell}_\theta(x, y)}{b(x, y) + \epsilon}\right] \approx 1 + \mathbb{E}\left[\frac{\bar{\ell}_\theta(x, y) - b(x, y)}{b(x, y) + \epsilon}\right].$$

This is the mean relative excess reconstruction loss over the base model.

If negative log-likelihood is read as code length, then

$$b(x, y) - \bar{\ell}_\theta(x, y)$$

is the code-length reduction from conditioning on $z$ in addition to $x$. That is the natural absolute-improvement quantity.

The ratio form changes the interpretation from absolute improvement to fractional improvement:

• it asks what fraction of the residual surprise in $y$ given $x$ is explained by $z$;
• equivalently, minimizing $R(x, y) / b(x, y)$ means minimizing the leftover fraction of baseline surprisal after conditioning on $z$.

This is the strongest user-facing interpretation of the objective:

• the numerator $\bar{\ell}_\theta - b$ measures excess loss over the base model;
• the denominator expresses that excess relative to the base-model surprisal scale of that same example.

In this sense, normalization turns absolute information gain into fractional information gain at the example level.

Difference vs optimization effect

The code-length difference

$$\bar{\ell}_\theta(x, y) - b(x, y)$$

has the cleanest information-theoretic interpretation, but if $b(x, y)$ is frozen it does not change the reconstruction gradient.

The ratio

$$\frac{\bar{\ell}_\theta(x, y)}{b(x, y) + \epsilon}$$

does change optimization:

$$\nabla \left(\frac{\bar{\ell}_\theta(x, y)}{b(x, y) + \epsilon}\right) = \frac{1}{b(x, y) + \epsilon}\nabla \bar{\ell}_\theta(x, y).$$

So the optimization effect comes from inverse-surprisal reweighting, while the clean interpretation comes from the underlying code-length reduction.

Relation to dataset-level normalization

It is important to distinguish this mean of ratios from the current dataset-level ratio of means:

$$\frac{\mathbb{E}[\bar{\ell}_\theta]}{\mathbb{E}[b]} \neq \mathbb{E}\left[\frac{\bar{\ell}_\theta}{b + \epsilon}\right].$$

The dataset-level variant is a global rescaling of reconstruction. This example-level variant is example reweighting.

See Distribution-Level Z-Free-Relative Reconstruction Rescaling for the distribution-level form.

Example-dependent reconstruction-vs-KL balance

The per-example objective takes the form

$$\frac{1}{b(x, y) + \epsilon}\bar{\ell}_\theta(x, y) + \beta K(x, y) + \mathcal{L}_\mathrm{aux}(x, y).$$

Equivalently, after factoring out the reconstruction coefficient, the KL-to-reconstruction balance becomes example dependent:

$$\bar{\ell}_\theta(x, y) + \beta_\mathrm{eff}(x, y) K(x, y) + \cdots, \qquad \beta_\mathrm{eff}(x, y) \propto \beta (b(x, y) + \epsilon).$$

So this objective does not merely normalize scale. It multiplies the effective KL pressure by an example-specific factor derived from the base-model surprisal.

This makes the objective naturally interpretable as a heterogeneous $\beta$-VAE or a sample-dependent rate-distortion tradeoff: different examples face different effective prices on latent rate, or equivalently different effective latent budgets.

Interpretation and caution

The most important caveat is how to explain the weighting signal.

At the mechanism level, the objective does assign larger reconstruction gradients to examples with smaller $b(x, y)$. But the intended interpretation is not "give more latent budget to easy examples." The better interpretation is that the ratio acts as an equalization force: it normalizes away absolute difficulty so that high-surprisal examples do not dominate simply because they carry more raw reconstruction loss.

Equivalently, the objective prefers large fractional improvement rather than large absolute improvement.

Two cautions follow:

• this objective is not literally the conditional mutual information $I(z; y \mid x)$; the direct information-theoretic quantity is the code-length difference $b - \bar{\ell}$, not the ratio;
• very small denominators can make the weighting unstable or overly aggressive, which is why the $\epsilon$ term is essential.

Current status in this repository

This variant is implemented as example_inverse_surprisal_weighted_reconstruction_loss, one explicit term in the current reconstruction stack.

The current code touchpoints are:

• the per-example reconstruction helper in path:core/disentanglement/losses.py
• the frozen-reference-model machinery in path:core/disentanglement/token_weighting.py
• the loss composition paths in:
  • path:core/disentanglement/shared_discrete_cvae.py
  • path:core/disentanglement/continuous_cvae.py
• CLI / trainer wiring in path:experiments/synthetic_sequences/disentangled/config.py and path:experiments/synthetic_sequences/disentangled/train_disentangled.py

The implementation note is Current implementation: example-level inverse-surprisal weighting.

Neighboring variants

• Distribution-Level Z-Free-Relative Reconstruction Rescaling Global reconstruction rescaling by one distribution-level constant.
• Token-level surprisal-share weighting Within-example token weighting toward positions where the z-free model has high residual surprisal shares and $z$ has the most local room to help.