Reconstruction Weighting And Normalization

This note organizes the current design space for reconstruction weighting and normalization in the entangled-initialized CVAE setting. The shared problem is that the base model already predicts $y$ from $x$ well, so a plain reconstruction term may be too weak to make $z$ carry additional information about $y$ beyond $x$. This is related to the posterior collapse problem.

A useful high-level framing is that unnormalized reconstruction measures absolute predictive improvement, while normalization can turn that into fractional predictive improvement relative to the z-free base model. In that sense, normalization turns absolute information gain into fractional information gain.

Shared notation

Let $\bar{\ell}_\theta(x, y)$ denote the current model's per-example y-region reconstruction loss, averaged over supervised positions and over the posterior on $z$:

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

Let $b_t(x, y)$ denote the per-token z-free base-model surprisal at supervised position $t$:

Deferred unified aggregation

One possible future direction is to replace the current explicit linear combination of reconstruction variants with a single unified reconstruction family that supports weighting at multiple levels. That direction is currently deferred rather than adopted.

See Unified reconstruction aggregation ideas for the current sketch of that design space.

Absolute vs fractional information gain

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

is the code-length reduction obtained by conditioning on $z$ in addition to $x$. This is the cleanest information-theoretic quantity in the cluster: it measures absolute improvement over the z-free baseline.

By contrast, normalization forms ratios such as

These are naturally read as fractional quantities: how much of the baseline surprisal remains after conditioning on $z$, or equivalently what fraction of the residual surprise in $y$ given $x$ is explained by $z$.

Current and proposed variants

• Distribution-Level Z-Free-Relative Reconstruction Rescaling Rescales reconstruction by one global constant $1 / \mathbb{E}[b]$. This is a global rescaling of reconstruction relative to KL and auxiliary terms.
• Token-sum versus token-mean NLL aggregation Clarifies that the compact reconstruction notation $-\log p(y \mid x, z)$ denotes a formal sequence likelihood, while the current code implements a supervised-token mean NLL; the distinction changes weighting when target lengths vary.
• Example-level inverse-surprisal weighting Weights each example by the inverse of its own base-model surprisal. This is example reweighting and can be interpreted as an equalization force across examples, not just a global normalization.
• Token-level surprisal-share weighting Uses per-token base-model surprisal shares as a weighting signal inside each example. This is within-example token weighting toward positions where $z$ has room to help.
• Normalized concentration-corrected token-weighted reconstruction Adds a concentration correction to token-level weighting so that the z-free baseline scale matches the ordinary uniform-token reconstruction loss before applying the outer distribution-level normalizer.
• Inverse-surprisal token-weighted reconstruction ablation A planned directional negative control for token-level surprisal-share weighting. It reverses the within-example token pressure toward low-surprisal tokens while using a baseline-scale correction so the raw inverse weighting scale does not confound the ablation.
• Unified reconstruction aggregation ideas Records deferred ideas for replacing the explicit linear combination with one unified reconstruction family.

Key interpretation distinctions

Surprisal

The quantity $b(x, y)$ can be viewed as base-model surprisal on one observed target. Surprisal can act as a proxy for residual uncertainty, but it also mixes together intrinsic ambiguity, model mismatch, and outlier status.

Global rescaling vs reweighting

The dataset-level objective

is a ratio of means. It rescales reconstruction globally.

The example-level objective

is a mean of ratios. It changes how examples are weighted.

Difference vs ratio

When the baseline $b(x, y)$ is frozen, the absolute-difference quantity

is mainly an interpretation or diagnostic: it has the code-length-reduction meaning above, but it leaves the reconstruction gradient unchanged.

The ratio

does change optimization, because it rescales the reconstruction gradient by $1 / (b(x, y) + \epsilon)$. So the main optimization effect comes from reweighting, while the main information-theoretic interpretation comes from the underlying code-length difference.

Equalization vs high-surprisal emphasis

These variants use base-model surprisal to approximate where the z-free model still has residual unpredictability. The main design question is which direction to use that signal:

• larger surprisal can be used to emphasize examples or tokens where $z$ has room to help. This is the logic behind token-level weighting.
• inverse surprisal can be used to convert absolute improvement into fractional improvement. This is best read as an equalization force that normalizes away absolute difficulty, so high-surprisal examples do not dominate simply because they carry more raw loss mass.
• inverse token-level surprisal can be used as an ablation pressure: it deliberately emphasizes low-surprisal tokens to test whether the high-surprisal direction is important, rather than because low-surprisal tokens are expected to be strategy-informative.
• concentration correction can be used to remove the scale inflation caused when token-level weighting focuses on a short high-surprisal prefix. This preserves within-example token emphasis while introducing a second effect: examples with more concentrated surprisal receive a smaller overall weight.

Current implementation status

• Dataset-level normalization is implemented.
  • Core config surface: path:core/disentanglement/config.py
  • Trainer and normalizer loading: path:experiments/synthetic_sequences/disentangled/train_disentangled.py
  • CLI flag: --normalize-reconstruction-loss in path:experiments/synthetic_sequences/disentangled/config.py
• Token-level weighting is implemented with a frozen reference model.
  • Weight computation: path:core/disentanglement/token_weighting.py
  • Weighted losses: path:core/disentanglement/losses.py
  • CLI flag: --token-weighted-reconstruction-weight in path:experiments/synthetic_sequences/disentangled/config.py
• Example-level inverse-surprisal weighting is implemented as an explicit reconstruction term.
  • Core helper: path:core/disentanglement/losses.py
  • CLI flag: --example-inverse-surprisal-reconstruction-weight in path:experiments/synthetic_sequences/disentangled/config.py
• Inverse-surprisal token-level weighting is implemented as an explicit ablation term.
  • Weight computation: path:core/disentanglement/token_weighting.py
  • Core helper: path:core/disentanglement/losses.py
  • CLI flags: --normalized-inverse-surprisal-token-weighted-reconstruction-weight and --ablation-inv-surp-token-weight in path:experiments/synthetic_sequences/disentangled/config.py

Open design pressure

The central question across these variants is:

This hub note is descriptive. The concrete mathematical claims, interpretations, and code links live in the three atomic notes it indexes.

For the corresponding current code behavior, see Implementation losses index.