Current reconstruction loss stack

Back to implementation losses

This note documents the current reconstruction-related loss composition in the disentangled trainers.

Current shape

Both the discrete and continuous implementations now treat reconstruction as a linear combination of six explicit scalar terms:

$$\mathcal{J}_{\mathrm{recon}} = \lambda_{\mathrm{raw}} L_{\mathrm{raw}} + \lambda_{\mathrm{dist}} L_{\mathrm{dist}} + \lambda_{\mathrm{ex}} L_{\mathrm{ex}} + \lambda_{\mathrm{tok}} L_{\mathrm{tok}} + \lambda_{\mathrm{tok,cc}} L_{\mathrm{tok,cc}} + \lambda_{\mathrm{tokinv}} L_{\mathrm{tokinv}}.$$

These six reconstruction terms are:

• reconstruction_loss
• distribution_normalized_reconstruction_loss
• example_inverse_surprisal_weighted_reconstruction_loss
• token_weighted_reconstruction_loss
• normalized_concentration_corrected_token_weighted_reconstruction_loss
• normalized_inverse_surprisal_token_weighted_reconstruction_loss

The full training objective then adds KL and the other configured auxiliaries:

$$\mathcal{J} = \mathcal{J}_{\mathrm{recon}} + \beta L_{\mathrm{KL}} + \text{other auxiliary terms}.$$

Current design choices

• All six reconstruction terms are always computed and logged in both latent types.
• The trainer does not choose one reconstruction family implicitly.
• The objective uses a plain weighted sum of the six explicit reconstruction scalars.
• Token-level and example-level weighting both depend on one frozen z-free reference-model pass.
• zfree_reference_scale is a shared scalar used by:
  • distribution_normalized_reconstruction_loss
  • normalized_concentration_corrected_token_weighted_reconstruction_loss
  • normalized_inverse_surprisal_token_weighted_reconstruction_loss

Shared implementation structure

The current implementation factors reconstruction into two stages:

1. compute token-level masked NLL once
2. aggregate that same token-level tensor in different ways

The common low-level helper lives in:

• path:core/disentanglement/losses.py

The important current pattern is:

• masked_token_nll_per_position(...) computes masked token NLL with shape [B, T-1]
• masked_token_nll_per_example(...) aggregates that tensor uniformly over active positions
• token_weighted_masked_nll_per_example(...) aggregates that same tensor with token weights
• token_weighted_masked_nll_per_example(...) is also reused with inverse-surprisal token weights

So the current reconstruction variants differ primarily in aggregation, not in how token NLL is scored.

For the formal and implementation-level implications of this denominator choice, see Token-sum versus token-mean NLL aggregation.

Planned ablation note: reconstruction-loss token aggregation mode.

Current configuration surface

The core loss config now nests reconstruction-specific settings under:

• ReconstructionLossConfig in path:core/disentanglement/config.py

The current reconstruction config fields are:

• raw_weight
• distribution_normalized_weight
• example_inverse_surprisal_weight
• token_weighted_weight
• normalized_concentration_corrected_token_weighted_weight
• normalized_inverse_surprisal_token_weighted_weight
• token_weight_alpha
• example_inverse_surprisal_epsilon
• inverse_surprisal_token_weight_epsilon
• zfree_reference_scale

The trainer-side config wiring lives in:

• path:experiments/synthetic_sequences/disentangled/config.py
• path:experiments/synthetic_sequences/disentangled/train_step.py
• path:experiments/synthetic_sequences/disentangled/train_disentangled.py

Diagnostics

The current implementation always tracks token-weighting summary diagnostics:

• token_weighted_base_loss
• token_concentration_factor
• inverse_token_weighted_base_loss
• inverse_token_base_scale_factor
• token_weighting/weight_max
• token_weighting/weight_ratio_max_min
• token_weighting/effective_token_count_mean
• inverse_token_weighting/weight_max
• inverse_token_weighting/weight_ratio_max_min
• inverse_token_weighting/effective_token_count_mean

These come from aligned z-free reference statistics computed in:

• path:core/disentanglement/token_weighting.py

Related notes

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