Normalized Concentration-Corrected Token-Weighted Reconstruction

Back to reconstruction weighting and normalization

This note records an under-consideration token-level reconstruction variant that combines:

• tempered token-level surprisal-share weighting to focus reconstruction inside each example,
• concentration correction to cancel the scale inflation induced when surprisal is concentrated on a short prefix,
• the same distribution-level z-free reference scale $c_\mathrm{base}$ used by Distribution-Level Z-Free-Relative Reconstruction Rescaling.

The intended role of this objective is:

focus reconstruction pressure on the most informative tokens within each example, while keeping the overall reconstruction term on the same scale as the ordinary uniform-token reconstruction loss.

Setup

Let the supervised y-region contain $|S(x, y)|$ active token positions. For one example $(x, y)$, define:

• z-conditioned per-token loss: $$\ell_t(x, y) = -\log p_\psi(y_t \mid x, z, y_{<t}),$$
• z-free base-model surprisal: $$b_t(x, y) = -\log p_\mathrm{base}(y_t \mid x, y_{<t}),$$
• uniform mean z-free loss: $$\mu(x, y) = b_u(x, y) = \frac{1}{|S(x, y)|}\sum_{t \in S(x, y)} b_t(x, y),$$
• distribution-level z-free reference scale: $$c_\mathrm{base} = \mathbb{E}_{(x, y)\sim D}[b_u(x, y)].$$

For compactness, define masked raw moments of the z-free token surprisal:

$$m_r(x, y) = \frac{1}{|S(x, y)|}\sum_{t \in S(x, y)} b_t(x, y)^r.$$

Then $m_1 = \mu$.

Raw token-weighted reconstruction

Let the tempered token-level surprisal-share weights be

$$\pi_t^{(\alpha)}(x, y) = \frac{b_t(x, y)^\alpha}{\sum_{s \in S(x, y)} b_s(x, y)^\alpha}, \qquad \alpha \ge 0,$$

with $\alpha = 0$ interpreted by continuity as uniform token weighting.

The corresponding token-weighted reconstruction is

$$R_{\mathrm{tok},\alpha}(x, y) = \sum_{t \in S(x, y)} \pi_t^{(\alpha)}(x, y)\,\ell_t(x, y).$$

At the z-free baseline, where $\ell_t(x, y)=b_t(x, y)$, this becomes

$$R_{\mathrm{tok},\alpha}^{\mathrm{base}}(x, y) = \sum_t \pi_t^{(\alpha)} b_t = \frac{m_{\alpha+1}(x, y)}{m_\alpha(x, y)}.$$

This is the base-model token loss after reweighting toward tokens with large $b_t^\alpha$.

Concentration factor

The corresponding token concentration factor is defined as

$$\kappa_\alpha(x, y) = \frac{R_{\mathrm{tok},\alpha}^{\mathrm{base}}(x, y)}{b_u(x, y)} = \frac{m_{\alpha+1}(x, y)}{m_\alpha(x, y)\,m_1(x, y)}.$$

This measures how much the $\alpha$-tempered token weighting inflates the z-free base loss relative to ordinary uniform token averaging.

Important properties:

• $\kappa_0(x, y)=1$,
• $\kappa_\alpha(x, y)\ge 1$ for all $\alpha \ge 0$,
• $\kappa_\alpha$ increases as surprisal becomes more concentrated,
• increasing $\alpha$ makes $\kappa_\alpha$ weakly larger, because more weight is placed on high-surprisal tokens.

For $\alpha=1$,

$$R_{\mathrm{tok},1}^{\mathrm{base}}(x, y) = \frac{m_2(x, y)}{m_1(x, y)} = \mu(x, y)\left(1+\mathrm{CV}(b)^2\right),$$

where $\mathrm{CV}(\cdot)$ is the coefficient of variation $\mathrm{CV} = \sigma / \mu$, so

$$\kappa_1(x, y)=1+\mathrm{CV}(b)^2.$$

Thus $\kappa_1$ is the factor by which linear surprisal-share weighting amplifies the z-free base loss relative to uniform token averaging.

Concentration-corrected reconstruction

Define the concentration-corrected token-weighted reconstruction

$$\widetilde R_{\mathrm{tok},\alpha}(x, y) = \frac{R_{\mathrm{tok},\alpha}(x, y)}{\kappa_\alpha(x, y)}.$$

At the z-free baseline,

$$\widetilde R_{\mathrm{tok},\alpha}^{\mathrm{base}}(x, y) = \frac{R_{\mathrm{tok},\alpha}^{\mathrm{base}}(x, y)}{\kappa_\alpha(x, y)} = b_u(x, y).$$

So the concentration correction restores the same per-example baseline scale as the uniform-token reconstruction loss.

Equivalently, the corrected loss can be written as

$$\widetilde R_{\mathrm{tok},\alpha}(x, y) = b_u(x, y)\, \frac{\sum_t b_t(x, y)^\alpha \,\ell_t(x, y)}{\sum_t b_t(x, y)^{\alpha+1}}.$$

This form makes the main caveat explicit:

• within each example, the relative token emphasis is still determined by $b_t^\alpha$,
• but across examples, the correction introduces an implicit example-level weight $1 / \kappa_\alpha(x, y)$.

So this is not a purely token-level intervention. It is a token-level weighting rule plus a concentration-dependent example-level correction.

Distribution-normalized form

Finally, define the normalized concentration-corrected token-weighted reconstruction

$$\widehat R_{\mathrm{tok},\alpha}(x, y) = \frac{\widetilde R_{\mathrm{tok},\alpha}(x, y)}{c_\mathrm{base}}.$$

At the z-free baseline,

$$\widehat R_{\mathrm{tok},\alpha}^{\mathrm{base}}(x, y) = \frac{b_u(x, y)}{c_\mathrm{base}}.$$

Taking expectation over the data distribution gives

$$\mathbb{E}_{(x, y)\sim D}\left[\widehat R_{\mathrm{tok},\alpha}^{\mathrm{base}}(x, y)\right] = \frac{\mathbb{E}[b_u(x, y)]}{c_\mathrm{base}} = 1.$$

So this objective has two useful normalization properties:

• per example, concentration correction restores the same z-free baseline scale as ordinary uniform reconstruction;
• in expectation over examples, the additional division by $c_\mathrm{base}$ makes the baseline value equal to $1$.

Interpretation

This objective can be read as follows:

• the raw token-weighted loss focuses attention on high-residual-surprisal positions;
• the concentration factor $\kappa_\alpha$ removes the scale inflation caused purely by surprisal being concentrated on a small subset of tokens;
• the outer division by $c_\mathrm{base}$ expresses the result in the same distribution-level z-free-relative units used by the standard normalized reconstruction loss.

That makes $\widehat R_{\mathrm{tok},\alpha}$ a natural candidate if the design goal is:

concentrate reconstruction pressure inside each example without changing the basic reconstruction scale just because surprisal mass is concentrated on a short prefix.

Useful structural properties

This objective family has several useful structural properties.

One of them is continuity at $\alpha = 0$:

$$\pi_t^{(0)} = \frac{1}{|S|}, \qquad \kappa_0 = 1, \qquad \widehat R_{\mathrm{tok},0}(x, y) = \frac{b_u(x, y)}{c_\mathrm{base}}.$$

So the family reduces exactly to distribution-level normalized uniform-token reconstruction when $\alpha = 0$.

This makes $\alpha$ a clean concentration knob:

• $\alpha=0$: ordinary distribution-normalized reconstruction,
• $\alpha=1$: concentration-corrected linear surprisal-share weighting,
• $\alpha>1$: stronger focus on the highest-surprisal tokens.

Two other important properties are:

• the raw token-weighted loss still concentrates reconstruction pressure within each example according to the token-level surprisal-share rule defined by $\alpha$;
• the concentration correction and outer division by $c_\mathrm{base}$ keep the resulting objective on the same baseline scale as ordinary distribution-normalized reconstruction: per example at the z-free baseline it reduces to $b_u(x, y)$, and in expectation it reduces to $1$ after dividing by $c_\mathrm{base}$.

Main caveat

The main caveat is that the concentration correction changes more than scale.

Because each example is divided by $\kappa_\alpha(x, y)$, examples with more concentrated surprisal are downweighted relative to the raw token-weighted loss. This may be good if the uncorrected objective overemphasizes short strategy-revealing prefixes simply because they induce a large raw scale inflation. But it may also be undesirable if those same concentrated examples are exactly the ones where the latent variable should matter most.

So the concentration-corrected objective should be understood as a distinct methodological proposal, not merely a harmless re-expression of raw token weighting.

Current status in this repository

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

The current implementation computes the aligned per-token z-free baseline losses needed for:

• $R_{\mathrm{tok},\alpha}$,
• $b_u$,
• $\kappa_\alpha$,
• and therefore $\widehat R_{\mathrm{tok},\alpha}$.

The current code touchpoints are:

• token-weight helper statistics in path:core/disentanglement/token_weighting.py,
• scalar loss helper in path:core/disentanglement/losses.py,
• discrete and continuous loss composition in path:core/disentanglement/shared_discrete_cvae.py and path:core/disentanglement/continuous_cvae.py,
• CLI / trainer wiring in path:experiments/synthetic_sequences/disentangled/config.py, path:experiments/synthetic_sequences/disentangled/train_step.py, and path:experiments/synthetic_sequences/disentangled/train_disentangled.py.

In the current reconstruction stack, this objective is the normalized high-surprisal token-weighted term alongside the raw token-weighted term rather than a replacement for it. See Current implementation: normalized concentration-corrected token-weighted reconstruction.

Neighboring variants

• Token-level surprisal-share weighting The currently implemented raw token-weighted loss.
• Distribution-Level Z-Free-Relative Reconstruction Rescaling The $\alpha=0$ endpoint of the normalized concentration-corrected family.
• Unified reconstruction aggregation ideas A deferred broader design space that would combine multiple reconstruction weightings inside one configurable family.