Current implementation: token-level surprisal-share weighting

Back to implementation losses

Methodology note: Token-level surprisal-share weighting

This note records how the current token-level weighting mechanism is implemented in code.

Core definition

The current implementation computes aligned z-free per-token base-model surprisal

$$b_t(x, y) = -\log p_{\mathrm{base}}(y_t \mid x, y_{<t}),$$

then converts it into normalized per-example token weights

$$\pi_t^{(\alpha)}(x, y) = \frac{m_t \, b_t(x, y)^\alpha}{\sum_s m_s \, b_s(x, y)^\alpha}.$$

So the current code supports tempered surprisal-share weighting, not only the linear (\alpha = 1) case.

Reference-model source

The weighting signal comes from a frozen reference model in:

• path:core/disentanglement/token_weighting.py

That helper:

• scores the z-free input sequence under torch.no_grad()
• computes masked per-position NLL
• aligns it to the conditioned sequence using alignment_offsets
• raises aligned base losses to alpha
• normalizes them over the aligned supervised region
• returns detached token weights plus summary statistics

The returned bundle currently includes:

• token_weights
• inverse_token_weights
• aligned_base_token_nll
• example_base_loss
• token_weighted_base_loss
• inverse_token_weighted_base_loss
• token_concentration_factor
• inverse_token_base_scale_factor

How the weighted loss enters the objective

The token-level weighted term is now one explicit reconstruction component:

• token_weighted_reconstruction_loss

It is still an additional scalar reconstruction term, not a replacement for the raw reconstruction term.

The current trainer therefore lets the user assign an independent coefficient to this term inside the six-term reconstruction stack.

Shared implementation structure

The current reconstruction helpers use a two-stage pattern:

1. compute masked per-token NLL once
2. aggregate that tensor with either uniform weights or token weights

In code:

• masked_token_nll_per_position(...) returns token NLL with shape [B, T-1]
• token_weighted_masked_nll_per_example(...) aggregates that tensor using token_weights

This same token-level weighting signal is also reused by the current inter-latent divergence loss in the discrete implementation when that term is enabled.

Configuration surface

The current configuration contract lives in:

• path:core/disentanglement/config.py

under:

• ReconstructionLossConfig.token_weighted_weight
• ReconstructionLossConfig.token_weight_alpha

The current CLI surface is:

• --token-weighted-reconstruction-weight
• --token-weighted-reconstruction-alpha

What is not implemented here

The current token-level implementation still does not support:

• predictive-entropy token weights
• dynamic token weights from the current training model
• interpolation between uniform and token-weighted reconstruction inside one unified reconstruction term

Related notes

• Current reconstruction loss stack
• Current implementation: normalized concentration-corrected token-weighted reconstruction