Distribution-Level Z-Free-Relative Reconstruction Rescaling

Back to reconstruction weighting and normalization

Implementation note: Current implementation: empirical z-free-relative reconstruction rescaling

This note describes a distribution-level reconstruction rescaling that uses the expected loss of the latent-free conditional distribution $p(y \mid x)$ as its reference scale. Its role is to restore reconstruction scale in the entangled-initialized regime, where the z-conditioned model already predicts $y$ from $x$ accurately enough that plain reconstruction can be too weak relative to KL.

Standard ELBO Objective

For one supervised pair $(x, y)$ , the conditional ELBO is

\mathcal{L}_\mathrm{cELBO}(x, y) = \mathbb{E}_{z \sim q_\xi(\cdot \mid x, y)}[\log p_\psi(y \mid x, z)] - \mathrm{KL}(q_\xi(z \mid x, y), p_\phi(z \mid x)).

With a beta-weighted KL:

\mathcal{L}_{\beta}(x, y) = \mathbb{E}_{z \sim q_\xi(\cdot \mid x, y)}[\log p_\psi(y \mid x, z)] - \beta \mathrm{KL}(q_\xi(z \mid x, y), p_\phi(z \mid x)).

Equivalently, the minimized loss is

\mathcal{J}(x, y) = R(x, y) + \beta K(x, y) + \mathcal{L}_\mathrm{aux}(x, y),

where the first term is the reconstruction loss

R(x, y) = -\mathbb{E}_{z \sim q_\xi(\cdot \mid x, y)}[\log p_\psi(y \mid x, z)],

and the second term is the KL divergence tethering the posterior to the prior/router:

K(x, y) = \mathrm{KL}(q_\xi(z \mid x, y), p_\phi(z \mid x)).

In a standard VAE trained from scratch, $R(x, y)$ is initially large, so reconstruction gradients dominate early. In the current project this is not the regime: $p_\psi (y | x, z)$ is initialized from an already strong entangled model $p_\theta (y | x)$ , so the initial y-region NLL is often already very small. The KL term therefore starts on a comparable or larger scale and can push $q_\xi$ toward $p_\phi$ before $z$ becomes informative.

Core issue

The failure mode is that the reconstruction loss is already small (at initialization, model can predict $y$ from $x$ without $z$ ), so the encoder-side gradient that would make $z$ useful is weak from step 0.

Normalized Reconstruction Objective

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

\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 frozen z-free base-model y-region surprisal for that same example:

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

Here "z-free" means the latent-free conditional distribution $p(y \mid x)$ , i.e. the model's prediction of $y$ from $x$ without conditioning on a latent variable $z$ .

Implementation detail

In the current implementation, the compact notation

-\log p_\psi(y \mid x, z)

is implemented as the masked mean token NLL

\frac{1}{|S(x, y)|}\sum_{t \in S(x, y)} -\log p_\psi(y_t \mid x, z, y_{<t}),

where $S(x, y)$ is the supervised y-region. Likewise,

-\log p_\mathrm{base}(y \mid x)

is implemented as

\frac{1}{|S(x, y)|}\sum_{t \in S(x, y)} -\log p_\mathrm{base}(y_t \mid x, y_{<t}).

So the compact notation suppresses two implementation details: the loss is masked to the supervised y-span, and it is averaged across those supervised tokens rather than summed.

Let

c_\mathrm{base} = \mathbb{E}_{(x, y) \sim D}[b(x, y)] > 0

denote the distribution-level mean z-free loss used as the global normalization constant. In methodology terms this is a data-distribution expectation. In the current implementation it is estimated empirically from validation data. In the entangled-run setting this scalar is loaded from run artifacts; in the pretrained-source setting it is recomputed by scoring the current validation rows with the source model.

Define the normalized reconstruction term

R_\mathrm{norm}(x, y) = \frac{\bar{\ell}_\theta(x, y)}{c_\mathrm{base}}.

The optimized objective becomes

\mathcal{J}_\mathrm{norm}(x, y) = R_\mathrm{norm}(x, y) + \beta K(x, y) + \mathcal{L}_\mathrm{aux}(x, y) = \frac{1}{c_\mathrm{base}} \bar{\ell}_\theta(x, y) + \beta K(x, y) + \mathcal{L}_\mathrm{aux}(x, y).

So rescaling does not alter the shape of the reconstruction objective; it changes only its global weight. When the entangled baseline loss is tiny, the effective reconstruction coefficient $1 / c_\mathrm{base}$ becomes large, restoring a better reconstruction-to-KL balance.

Code-Length Interpretation

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

b(x, y) - R(x, y)

is the code-length reduction obtained by conditioning on $z$ in addition to $x$ . This is the natural absolute-improvement quantity in the problem: it measures how much shorter the description of $y$ becomes once the latent is provided.

The current distribution-level rescaling does not optimize this difference directly. Instead, it rescales reconstruction by the aggregate z-free loss. That converts the absolute code-length picture into a fractional one at the distribution level.

Relative-Excess Interpretation

The dataset-level normalized reconstruction objective can be rewritten as

\mathbb{E}_{(x, y) \sim D}[R_\mathrm{norm}(x, y)] = \frac{\mathbb{E}[\bar{\ell}_\theta(x, y)]}{c_\mathrm{base}}.

Since $c_\mathrm{base} = \mathbb{E}[b(x, y)]$ , this equals

\frac{\mathbb{E}[\bar{\ell}_\theta(x, y)]}{\mathbb{E}[b(x, y)]} = 1 + \frac{\mathbb{E}[\bar{\ell}_\theta(x, y)] - \mathbb{E}[b(x, y)]}{\mathbb{E}[b(x, y)]} = 1 + \frac{\mathbb{E}[\bar{\ell}_\theta(x, y) - b(x, y)]}{\mathbb{E}[b(x, y)]}.

This is a relative-excess reconstruction objective at the level of distribution means:

the numerator $\mathbb{E}[\bar{\ell}_\theta - b]$ is the excess reconstruction loss over the base model;
the denominator $\mathbb{E}[b]$ expresses that excess relative to the mean base-model surprisal.

This interpretation is close to, but importantly different from, an example-level relative-excess objective. The current global rescaling forms

1 + \frac{\mathbb{E}[\bar{\ell}_\theta - b]}{\mathbb{E}[b]},

not

1 + \mathbb{E}\left[\frac{\bar{\ell}_\theta - b}{b}\right].

So the current method rescales reconstruction globally; it does not reweight examples relative to one another.

Interpretation

The corresponding validation metric

\rho_\mathrm{recon} = R_\mathrm{val} / c_\mathrm{base}

has a direct interpretation:

$\rho_\mathrm{recon}$ is the leftover fraction of aggregate baseline surprisal after conditioning on $z$ .
$1 - \rho_\mathrm{recon}$ is the fraction of aggregate residual surprise in $y$ given $x$ explained by $z$ .
$\rho_\mathrm{recon} = 1$ : the disentangled model matches the entangled baseline.
$\rho_\mathrm{recon} < 1$ : the latent-conditioned model improves on the baseline, so conditioning on $z$ explains a positive fraction of the aggregate residual surprise.
$\rho_\mathrm{recon} > 1$ : the factorized model is even worse than the baseline at predicting y from x, z.

This is useful both as an optimization rescaling and as a comparison metric, because it reports reconstruction quality relative to the actual entangled initializer rather than on an absolute NLL scale that varies by task.

The critical point is that this is a global rescaling, not an example reweighting. It does not change which examples matter more. It only changes the overall reconstruction-to-KL balance.

Another way to say this is that rescaling turns absolute information gain into fractional information gain at the distribution level. The metric does not ask which examples receive more emphasis; it asks what fraction of the aggregate z-free loss remains.

Rate-Distortion View

The normalized objective can be written as

\frac{1}{c_\mathrm{base}} R(x, y) + \beta K(x, y) = \frac{1}{c_\mathrm{base}}\left(R(x, y) + \beta c_\mathrm{base} K(x, y)\right).

So in rate-distortion language, baseline normalization changes the global slope of the reconstruction-vs-KL tradeoff. Equivalently, it changes the global latent budget available to reconstruction improvement, but it does not redistribute that budget across examples. This is why it is best understood as scale restoration rather than reweighting.

Current status in this repository

This distribution-level objective family is currently implemented through an empirical validation-based estimator of the z-free reference loss.

Core configuration contract: normalize_reconstruction_loss and reconstruction_loss_normalizer in path:core/disentanglement/config.py
Reconstruction weight application:
- path:core/disentanglement/shared_discrete_cvae.py
- path:core/disentanglement/continuous_cvae.py
Normalizer loading / computation:
- path:experiments/synthetic_sequences/disentangled/train_disentangled.py
- path:experiments/synthetic_sequences/disentangled/setup.py
CLI flag: --normalize-reconstruction-loss in path:experiments/synthetic_sequences/disentangled/config.py
Validation metric emission: loss/val/baseline_normalized_reconstruction_loss in path:experiments/synthetic_sequences/disentangled/evaluation_runtime.py

Neighboring variants

Example-level inverse-surprisal weighting Uses inverse per-example base-model surprisal. Unlike distribution-level rescaling, that variant reweights examples and induces an example-dependent effective $\beta$ .
Token-level surprisal-share weighting Uses per-token base-model surprisal shares to emphasize strategy-sensitive positions inside each example.

Relationship to Other Interventions

Token-Weighted Reconstruction

Token-weighted reconstruction changes the distribution of the reconstruction gradient across token positions by upweighting positions where the z-free model has high residual surprisal. Baseline normalization changes only the global magnitude of the standard reconstruction term. These are complementary:

normalization says "make reconstruction matter again at all,"
token weighting says "within reconstruction, focus on strategy-sensitive tokens."

Both can be used simultaneously.

Inter-Latent JSD

Inter-latent JSD injects an exclusivity signal: different latent values should induce different token distributions. It does not by itself solve the weak-reconstruction problem. Normalization and JSD therefore address different bottlenecks:

normalization encourages $z$ to become useful,
JSD encourages the information in $z$ to be strategy-separating.

InfoVAE-Style Beta Reduction

Lowering $\beta$ weakens the KL pressure directly. Normalization strengthens reconstruction directly. Both rebalance the same competition, but in opposite ways:

beta reduction rescales $K(x, y)$ downward,
normalization rescales $R(x, y)$ upward.

Normalization is especially natural in the entangled-initialized regime because its scale is tied to the observed baseline loss rather than chosen manually.

Diagnostics to Watch

The main diagnostics are:

raw reconstruction_loss: did the conditional generator remain strong?
baseline_normalized_reconstruction_loss: is the factorized model beating the entangled baseline?
kl_prior_loss: is the posterior carrying information instead of collapsing?
loss_contrib/fraction_abs/reconstruction_loss: did the optimization balance actually shift toward reconstruction?
downstream structure metrics such as alignment or probe accuracy: does the stronger reconstruction signal translate into more informative latents?

Expected signature If normalization is working, the reconstruction contribution should grow, the normalized reconstruction ratio should fall below

1

, and the latent variables should become more informative without a large drop in task correctness.