Current implementation: normalized inverse-surprisal token-weighted reconstruction ablation

Methodology note: Inverse-surprisal token-weighted reconstruction ablation

This note records how the normalized inverse-surprisal token-weighted reconstruction ablation is currently implemented in code.

Current status

This ablation is implemented as the sixth explicit reconstruction term:

normalized_inverse_surprisal_token_weighted_reconstruction_loss

Current definition

The implementation reuses the aligned frozen-reference token NLL tensor $b_t$ computed for token-level surprisal-share weighting. It builds inverse token weights

\rho_t^{(\alpha,\epsilon)} = \frac{m_t \max(b_t,\epsilon)^{-\alpha}}{\sum_s m_s \max(b_s,\epsilon)^{-\alpha}},

where $m_t$ is the aligned supervised-token mask, $\alpha$ is the shared token-weighting exponent, and $\epsilon$ is the inverse-token clipping floor.

The latent-conditioned token NLL tensor is then aggregated with $\rho_t$ to produce a per-example inverse-token reconstruction loss $L_{\mathrm{tokinv},b}$ . The implementation also computes the inverse-token z-free base loss $L^{\mathrm{base}}_{\mathrm{tokinv},b}$ and the inverse-token base scale factor

\eta_b = \frac{L^{\mathrm{base}}_{\mathrm{tokinv},b}}{\mu_b},

where $\mu_b$ is the uniformly averaged z-free supervised-token base loss for the same example.

The scalar component is

L_{\mathrm{tokinv}} = \frac{1}{c_{\mathrm{base}}} \cdot \frac{1}{B} \sum_b \frac{L_{\mathrm{tokinv},b}}{\eta_b}.

Where the supporting statistics come from

The shared frozen-reference pass lives in:

path:core/disentanglement/token_weighting.py

The returned statistics now include:

inverse_token_weights
inverse_token_weighted_base_loss
inverse_token_base_scale_factor

The high-surprisal token weights and inverse-surprisal token weights are computed from the same aligned z-free token NLL tensor, so the ablation does not add another reference-model forward pass.

Where the term is assembled

The scalar helper lives in:

path:core/disentanglement/losses.py

and is integrated in:

path:core/disentanglement/shared_discrete_cvae.py
path:core/disentanglement/continuous_cvae.py

In the discrete path, each latent-specific token NLL tensor is aggregated three ways: uniform, high-surprisal token-weighted, and inverse-surprisal token-weighted. In the continuous path, the sampled-z token NLL tensor is aggregated with the same three reconstruction aggregators.

Configuration surface

The core configuration fields are:

ReconstructionLossConfig.normalized_inverse_surprisal_token_weighted_weight
ReconstructionLossConfig.inverse_surprisal_token_weight_epsilon
ReconstructionLossConfig.token_weight_alpha
ReconstructionLossConfig.zfree_reference_scale

The trainer CLI exposes:

--normalized-inverse-surprisal-token-weighted-reconstruction-weight
--ablation-inv-surp-token-weight
--inverse-surprisal-token-weight-epsilon
--token-weighted-reconstruction-alpha

The alias --ablation-inv-surp-token-weight writes to the same trainer config field as the canonical weight flag. The exponent --token-weighted-reconstruction-alpha is intentionally shared with the high-surprisal token-weighted terms.

Diagnostics

The implementation logs:

inverse_token_weighted_base_loss
inverse_token_base_scale_factor
inverse_token_weighting/weight_max
inverse_token_weighting/weight_ratio_max_min
inverse_token_weighting/effective_token_count_mean

The readable reconstruction tag uses tokinv when the inverse-token term has positive weight. The run identity keeps token_weighted_reconstruction_alpha and inverse_surprisal_token_weight_epsilon only when an active token-weighted term makes those knobs behaviorally relevant.