Current implementation: normalized concentration-corrected token-weighted reconstruction

Back to implementation losses

Methodology note: Normalized concentration-corrected token-weighted reconstruction

This note records how the normalized concentration-corrected token-weighted reconstruction term is currently implemented in code.

Current status

This idea is now implemented as one explicit reconstruction term:

• normalized_concentration_corrected_token_weighted_reconstruction_loss

Current definition

The implementation starts from the raw token-weighted per-example reconstruction loss

$$L_{\mathrm{tok},b}.$$

It then divides by the per-example token concentration factor

$$\kappa_b = \frac{L^{\mathrm{base}}_{\mathrm{tok},b}}{b_{\mathrm{base},b}},$$

and finally divides the batch mean by the shared global z-free reference scale

$$c_{\mathrm{base}}.$$

So the current scalar term is

$$L_{\mathrm{tok,cc}} = \frac{1}{c_{\mathrm{base}}} \cdot \frac{1}{B} \sum_b \frac{L_{\mathrm{tok},b}}{\kappa_b}.$$

Where the supporting statistics come from

The required per-example quantities come from:

• path:core/disentanglement/token_weighting.py

The relevant fields are:

• token_weighted_base_loss
• token_concentration_factor

These are computed from aligned frozen-reference z-free token losses using the same token weights as the raw token-weighted reconstruction term.

Where the term is assembled

The low-level 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

Configuration surface

The current configuration contract lives in:

• path:core/disentanglement/config.py

under:

• ReconstructionLossConfig.normalized_concentration_corrected_token_weighted_weight
• ReconstructionLossConfig.token_weight_alpha
• ReconstructionLossConfig.zfree_reference_scale

The current CLI surface is:

• --normalized-concentration-corrected-token-weighted-reconstruction-weight
• --token-weighted-reconstruction-alpha

Current relationship to raw token-weighted reconstruction

The implementation keeps both token-level terms explicit:

• token_weighted_reconstruction_loss
• normalized_concentration_corrected_token_weighted_reconstruction_loss

They are tracked separately and each has its own objective weight.

Related notes

• Current implementation: token-level surprisal-share weighting
• Current implementation: normalized inverse-surprisal token-weighted reconstruction ablation
• Current implementation: empirical z-free-relative reconstruction rescaling
• Current reconstruction loss stack