Planned implementation: reconstruction-loss token aggregation mode

Back to implementation losses

Methodology note: Token-sum versus token-mean NLL aggregation

This note records the planned focused ablation for reconstruction_loss_token_aggregation_mode, which will support mean and sum. The current implementation corresponds to mean.

Approved scope

• The mode applies only to reconstruction-family terms.
• KL, posterior entropy, posterior variance, router support, router marginal KL, inter-latent divergence, and other auxiliary losses are unchanged.
• mean remains the default.
• sum is allowed as an optimization setting, including for the raw reconstruction term.
• The z-free reference scale must be recomputed from validation records using the same aggregation mode as the training objective.
• The mode must be explicit in run config, slurm params, run naming, and aggregate tables.

This is an ablation, not a replacement of the default objective.

Aggregation contract

For example $i$, let $S_i$ be the supervised y-region, $n_i = |S_i|$, z-conditioned token NLL $\ell_{i,t}(z)$, and frozen z-free token NLL $b_{i,t}$.

The uniform per-example aggregation is:

$$A_{\mathrm{mean}}(v_i) = \frac{1}{n_i}\sum_{t \in S_i} v_{i,t}, \qquad A_{\mathrm{sum}}(v_i) = \sum_{t \in S_i} v_{i,t}.$$

Batch reduction remains an example mean:

$$\frac{1}{B}\sum_i R_i.$$

So sum means "per-example supervised-token sum, then batch mean." It does not mean global token mean over the whole batch.

Token-weighted aggregation contract

Token-weighting code currently normalizes token weights per example, so $\sum_t \pi_{i,t}=1$ over active supervised tokens. To make token-weighted losses match the selected uniform aggregation mode, define:

$$W_{\mathrm{mean}}(v_i;\pi_i)=\sum_{t \in S_i}\pi_{i,t}v_{i,t}, \qquad W_{\mathrm{sum}}(v_i;\pi_i)=n_i\sum_{t \in S_i}\pi_{i,t}v_{i,t}.$$

The factor $n_i$ is required for consistency. Without it, raw reconstruction would move to sequence-sum scale while token-weighted reconstruction would remain on token-mean scale. With it, $\alpha=0$ recovers the corresponding uniform aggregation mode exactly.

Loss-term adaptation

Raw reconstruction

The raw reconstruction component becomes:

$$L_{\mathrm{raw}}^a = \mathbb{E}_i\mathbb{E}_{z\sim q(z \mid x_i,y_i)} \left[A_a(\ell_i(z))\right], \qquad a \in \{\mathrm{mean}, \mathrm{sum}\}.$$

This is the component logged as reconstruction_loss.

Distribution-normalized reconstruction

The distribution-normalized component becomes:

$$L_{\mathrm{dist}}^a = \frac{L_{\mathrm{raw}}^a}{c_{\mathrm{base}}^a}, \qquad c_{\mathrm{base}}^a = \mathbb{E}_{i\in \mathrm{val}}\left[A_a(b_i)\right].$$

The reference scale must be mode-matched. A mean-scale validation loss cannot be reused as the denominator for a sum-scale reconstruction objective.

Example-level inverse-surprisal reconstruction

The example-level inverse-surprisal component becomes:

$$L_{\mathrm{ex}}^a = \mathbb{E}_i \left[ \frac{R_i^a}{B_i^a+\epsilon_a} \right], \qquad R_i^a=\mathbb{E}_{q(z\mid x_i,y_i)}[A_a(\ell_i(z))], \qquad B_i^a=A_a(b_i).$$

The epsilon is in the same units as $B_i^a$, so the implementation should treat it as mode-sensitive. The existing default is effectively a mean-scale epsilon; for sum, the safest first implementation is to reuse the configured value but document that it is in selected-mode units and log base-loss minima so pathological tiny denominators are visible.

This method still makes sense under sum, but if $\epsilon_a$ is negligible, the per-example ratio mostly cancels the length factor. So it is probably less informative as the first aggregation-mode ablation than distribution-normalized or token-weighted terms.

Raw token-weighted reconstruction

The high-surprisal token-weighted reconstruction component becomes:

$$L_{\mathrm{tok}}^a = \mathbb{E}_i\mathbb{E}_{q(z\mid x_i,y_i)} \left[ W_a(\ell_i(z);\pi_i) \right],$$

where $\pi_i$ is the detached high-surprisal token-weight vector derived from the frozen z-free token NLL.

For sum, this means multiplying the current weighted mean by $n_i$.

Normalized concentration-corrected token-weighted reconstruction

The concentration factor should be computed with the selected aggregation mode:

$$\kappa_i^a = \frac{W_a(b_i;\pi_i)}{A_a(b_i)}.$$

Because both numerator and denominator are multiplied by $n_i$ under sum, $\kappa_i^{\mathrm{sum}}=\kappa_i^{\mathrm{mean}}$ when the same normalized token weights are used. The normalized concentration-corrected component is:

$$L_{\mathrm{tok,cc}}^a = \frac{\mathbb{E}_i\left[W_a(\ell_i;\pi_i)/\kappa_i^a\right]}{c_{\mathrm{base}}^a}.$$

This method still makes sense under sum. At $\alpha=0$, it reduces to distribution-normalized uniform reconstruction in the selected aggregation mode.

Normalized inverse-surprisal token-weighted reconstruction

The low-surprisal inverse-token ablation follows the same matched-scale pattern:

$$\eta_i^a = \frac{W_a(b_i;\rho_i)}{A_a(b_i)}, \qquad L_{\mathrm{tokinv}}^a = \frac{\mathbb{E}_i\left[W_a(\ell_i;\rho_i)/\eta_i^a\right]}{c_{\mathrm{base}}^a}.$$

As with $\kappa_i$, the scale factor is unchanged between mean and sum if both numerator and denominator use the same aggregation mode. This ablation still makes sense under sum, but the strongest comparison remains against the high-surprisal token-weighted term with matched mode, alpha, and z-free reference scale.

Implementation requirements

• Add reconstruction_loss_token_aggregation_mode with allowed values mean and sum.
• Add the corresponding CLI flag, likely --reconstruction-loss-token-aggregation-mode {mean,sum}.
• Thread the mode through:
  • path:core/disentanglement/config.py
  • path:core/disentanglement/losses.py
  • path:core/disentanglement/token_weighting.py
  • path:core/disentanglement/shared_discrete_cvae.py
  • path:core/disentanglement/continuous_cvae.py
  • path:experiments/synthetic_sequences/disentangled/config.py
  • path:experiments/synthetic_sequences/disentangled/train_disentangled.py
  • path:experiments/synthetic_sequences/disentangled/train_step.py
• Recompute zfree_reference_scale from validation records with the selected aggregation mode. For entangled_run sources, do not silently load an old saved mean-scale validation loss when mode is sum.
• Record the mode in the saved run config and any run metadata artifacts.
• Include the mode in slurm family params and template recovery tests.
• Include the mode in run naming so mean and sum runs cannot collide.
• Include the mode in aggregate parsing and aggregate tables.
• Log supervised token-count diagnostics, at least mean/min/max over training and validation batches, so length effects can be interpreted.

Pilot collection

The first pilot should be narrow. The cleanest comparisons are mode-matched normalized variants:

• dist with mean versus dist with sum;
• tokcc with mean versus tokcc with sum;
• optionally the current best hybrid with mean versus the same hybrid with sum;
• raw sum is allowed, but should be interpreted as both a length-weighting intervention and a larger reconstruction-to-KL scale unless paired with explicit beta or normalizer controls.

The main outcome to watch is whether sum changes latent usage and strategy separation beyond what can be explained by a larger reconstruction scale or by upweighting longer examples.

Expected confounds

• Sequence length may correlate with task, strategy, parse success, or answer correctness.
• Raw sum changes the reconstruction-to-KL balance by approximately the active supervised token count.
• Distribution normalization removes the global average scale change but does not remove within-distribution length weighting.
• Example-level inverse-surprisal ratios may cancel much of the mean-vs-sum difference, making them less diagnostic for the first pilot.
• Token-weighted sum preserves within-example token emphasis but gives longer examples more total reconstruction weight.

Current status

This note records the planned ablation contract. The code has not yet been updated.