Inverse-Surprisal Token-Weighted Reconstruction Ablation

Back to reconstruction weighting and normalization

Neighboring note: Token-level surprisal-share weighting

This note records an ablation for token-level surprisal-share weighting. The main token-weighted reconstruction term puts more reconstruction pressure on high-surprisal tokens, because those tokens are hypothesized to be the strategy-branching positions where $z$ can help. This ablation reverses that token-level pressure by upweighting low-surprisal tokens and downweighting high-surprisal tokens.

Role of the ablation

The ablation is a directional negative control. If high-surprisal token weighting helps because it targets strategy-branching positions, then an objective that emphasizes low-surprisal tokens should help less, fail, or suppress the same mechanism. If the inverse objective performs similarly, the evidence for the branching-token explanation is weaker: the improvement may instead come from nonuniform weighting, reconstruction-scale changes, or some other change to the objective.

Setup

Let $S(x, y)$ be the active supervised y-region. For one example, 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) = \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}[\mu(x, y)].$$

Use a clipped surprisal

$$\tilde b_t(x, y) = \max(b_t(x, y), \epsilon)$$

with $\epsilon > 0$ so inverse weights remain finite.

Raw inverse-token weighting

Define inverse-surprisal token weights

$$\rho_t^{(\alpha,\epsilon)}(x, y) = \frac{\tilde b_t(x, y)^{-\alpha}}{\sum_{s \in S(x, y)} \tilde b_s(x, y)^{-\alpha}}, \qquad \alpha \ge 0.$$

The raw inverse-surprisal token-weighted reconstruction is

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

Here $\alpha=0$ recovers uniform token weighting, while larger $\alpha$ puts more mass on lower-surprisal tokens.

Scale confound in the raw objective

At the z-free baseline, where $\ell_t(x, y)=b_t(x, y)$, the raw inverse objective has baseline value

$$R_{\mathrm{inv},\alpha,\epsilon}^{\mathrm{base}}(x, y) = \sum_{t \in S(x, y)} \rho_t^{(\alpha,\epsilon)}(x, y)\,b_t(x, y).$$

For $\alpha=1$ with negligible clipping, this is the harmonic mean of the supervised token surprisals. When surprisal is uneven across tokens, this is below the ordinary arithmetic mean $\mu(x, y)$. So the raw inverse objective does not only reverse token emphasis; it also changes the reconstruction scale. That scale change would confound the ablation, because a worse result could come from weaker effective reconstruction pressure rather than from emphasizing the wrong tokens.

Selected normalized ablation

To isolate token-emphasis direction more cleanly, define the inverse base-scale factor

$$\eta_{\alpha,\epsilon}(x, y) = \frac{R_{\mathrm{inv},\alpha,\epsilon}^{\mathrm{base}}(x, y)}{\mu(x, y)}.$$

Then define the scale-corrected inverse-token reconstruction

$$\widetilde R_{\mathrm{inv},\alpha,\epsilon}(x, y) = \frac{R_{\mathrm{inv},\alpha,\epsilon}(x, y)}{\eta_{\alpha,\epsilon}(x, y)}.$$

Finally, put the term in the same distribution-level z-free-relative units used by Distribution-Level Z-Free-Relative Reconstruction Rescaling:

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

At the z-free baseline,

$$\widehat R_{\mathrm{inv},\alpha,\epsilon}^{\mathrm{base}}(x, y) = \frac{\mu(x, y)}{c_\mathrm{base}},$$

so the expected z-free baseline value is $1$ over the data distribution. This matches the baseline-scale behavior of normalized concentration-corrected token-weighted reconstruction while reversing the within-example token emphasis.

Interpretation

The selected ablation asks whether the high-surprisal direction matters. It should be compared against the high-surprisal normalized concentration-corrected token term with the same $\alpha$, the same frozen reference model, and the same distribution-level reference scale $c_\mathrm{base}$.

The strongest expected evidence for the branching-token hypothesis would be:

• high-surprisal token weighting improves latent usage or strategy separation;
• inverse-surprisal token weighting does not produce the same improvement under matched scale;
• diagnostics show that the inverse variant puts its mass on low-base-surprisal continuation tokens rather than high-surprisal branching tokens.

Caveats

This ablation is not a proof that high-surprisal tokens are strategy tokens. Base-model surprisal can also reflect model mismatch, rare formatting, noisy labels, or non-strategy ambiguity. Conversely, some low-surprisal tokens may still carry strategy information if the strategy was already revealed earlier in the teacher-forced prefix.

The scale correction also changes more than pure token emphasis: examples with very small inverse-weighted base loss receive a larger correction factor. This is preferable to the raw inverse objective for a controlled ablation, but it should still be interpreted as a token-weighting rule plus a per-example baseline-scale correction.

Alternatives considered

• Raw inverse weighting: closest to the simple rule $w_t \propto 1 / b_t$, but confounds token direction with a smaller base loss scale.
• Reverse-rank weighting: assigns the largest high-surprisal weights to the lowest-surprisal positions, preserving the weight concentration distribution more exactly, but is less direct and more complex.
• Shuffled token weights: tests whether alignment between weights and surprisal matters, but does not specifically test the opposite low-surprisal direction.
• Bottom-quantile weighting: uniformly weights the lowest-surprisal tokens, but introduces a quantile hyperparameter and a discontinuous selection rule.

Current status in this repository

This ablation is now implemented as normalized_inverse_surprisal_token_weighted_reconstruction_loss, the sixth explicit reconstruction term. It reuses the aligned frozen-reference per-token base losses computed in path:core/disentanglement/token_weighting.py, exposes the canonical CLI flag --normalized-inverse-surprisal-token-weighted-reconstruction-weight, and also exposes the alias --ablation-inv-surp-token-weight. See Current implementation: normalized inverse-surprisal token-weighted reconstruction ablation for the code-level details.