Unified Reconstruction Aggregation Ideas

Back to reconstruction weighting and normalization

This note records deferred ideas for replacing a simple linear combination of reconstruction variants with one unified reconstruction objective. These ideas are not yet adopted. The current redesign direction is to keep a linear combination of explicit reconstruction losses and revisit unified aggregation later.

Current holding decision

For the next loss-stack redesign, keep reconstruction aggregation in the simple form

\mathcal{J}_\mathrm{recon} = \lambda_\mathrm{std} R_\mathrm{std} + \lambda_\mathrm{dist} R_\mathrm{dist} + \lambda_\mathrm{ex} R_\mathrm{ex} + \lambda_\mathrm{tok} R_\mathrm{tok,\alpha} + \lambda_\mathrm{tok,cc} \widehat R_\mathrm{tok,\alpha},

where:

$R_\mathrm{std}$ is the unnormalized standard reconstruction loss,
$R_\mathrm{dist}$ is the distribution-level z-free-relative reconstruction loss,
$R_\mathrm{ex}$ is the example-level inverse-surprisal weighted reconstruction loss,
$R_\mathrm{tok,\alpha}$ is the token-level tempered surprisal-share weighted reconstruction loss,
$\widehat R_\mathrm{tok,\alpha}$ is the normalized concentration-corrected token-weighted reconstruction loss,
$\alpha$ is the token-level tempering parameter, with default $\alpha = 1$ recovering the current linear surprisal-share weighting.

During training, all five reconstruction losses should be tracked and logged regardless of whether their coefficients are zero.

This design is intentionally redundant at the level of expressive power, especially between $R_\mathrm{std}$ and $R_\mathrm{dist}$ and between the two token-level terms. That redundancy is acceptable for now because it keeps the redesign simple, symmetric, and transparent while the conceptual picture is still under discussion.

Why unified aggregation is deferred

The main reason to defer a unified reconstruction family is that there is not yet a single parameterization that is clearly preferable along all of the following axes:

conceptual interpretation,
stable global scaling relative to KL,
support for all currently considered reconstruction variants,
clean extensibility to future weighting signals.

The current discussion identified a recurring tension:

if the aggregation family is made fully policy matched, the normalizer and the weighting policy become self-referential;
if the aggregation family uses one fixed global reference scale, some variants keep a cleaner interpretation than others.

Because that tradeoff is not yet resolved, it is better to keep the active implementation design simple for now.

Candidate unified directions

Multi-level weighting family

One natural family is to separate:

token aggregation policy $\pi_t(x, y)$ ,
example weight $\omega(x, y)$ ,
optional global rescaling constant $C$ .

This gives a generic form

L_\mathrm{recon}[\pi,\omega,C] = \frac{1}{C} \mathbb{E}_{(x, y) \sim D} \left[ \omega(x, y) \sum_t \pi_t(x, y)\,\ell_t(x, y) \right].

This is attractive because it factorizes three different roles:

which token positions matter more,
which examples matter more,
how reconstruction is scaled relative to KL.

However, this family still leaves open how $C$ should be chosen.

Policy-matched baseline normalization

Another idea is to choose $C$ from the z-free reference under the same weighting policy:

C_{\pi,\omega} = \mathbb{E}\left[ \omega(x, y; b)\sum_t \pi_t(x, y; b)\,b_t(x, y) \right].

This has one appealing property:

when the current model matches the z-free reference, the resulting aggregated reconstruction loss equals $1$ for every configuration.

But this design also has a substantial drawback:

the normalizer depends on the same weighting policy whose effect we are trying to study, so the reference scale becomes self-referential and loses interpretive clarity.

This concern is the main reason it is not currently favored.

Fixed-reference scaling plus local weighting

A less self-referential direction is:

keep one fixed global z-free reference scale,
let token-level and example-level weighting only redistribute reconstruction emphasis locally.

This preserves a simpler notion of what the reconstruction loss is measured against, but it does not make every reconstruction config equal exactly $1$ at the z-free baseline.

Open design questions

Can one unified reconstruction family preserve the clean fractional interpretation of example-level inverse-surprisal weighting?
Can one unified reconstruction family keep token-level weighting on a stable (O(1)) scale without making the denominator policy dependent?
If a unified family is adopted later, should it replace the explicit linear combination entirely, or sit underneath it as one configurable reconstruction term?

Relationship to active notes

This note is intentionally forward-looking. It should be revisited only after the simpler explicit linear-combination redesign has been implemented and evaluated.