Reconstruction Loss Stack API Design

Reconstruction Loss Stack API Design

Back to reconstruction weighting redesign working map

This note records the current proposed API for the next reconstruction-loss redesign. It is an ideation note rather than an implementation note because the design is not yet adopted in code.

Design goals

The next implementation pass should:

  • keep a simple linear combination of explicit reconstruction terms,
  • support all currently considered reconstruction variants directly,
  • always compute and log all reconstruction variants regardless of which weights are nonzero,
  • keep the loss surface easy to read in compute_total_loss,
  • avoid introducing special dispatch logic such as a token-level kind switch unless that later proves clearly better.

Five reconstruction terms

The future reconstruction stack should expose the following five scalar reconstruction terms:

  1. reconstruction_loss
    • the ordinary unnormalized reconstruction loss
  2. distribution_normalized_reconstruction_loss
    • the distribution-level z-free-relative normalized reconstruction loss
  3. example_inverse_surprisal_weighted_reconstruction_loss
    • the example-level inverse-surprisal weighted reconstruction loss
  4. token_weighted_reconstruction_loss
    • the raw tempered token-level surprisal-share weighted reconstruction loss
  5. normalized_concentration_corrected_token_weighted_reconstruction_loss
    • the normalized concentration-corrected token-weighted reconstruction loss

The future total reconstruction contribution should be

Jrecon=wraw Rraw+wdist Rdist+wex Rex+wtok Rtok+wtok,cc Rtok,cc.\mathcal{J}_\mathrm{recon} = w_\mathrm{raw} \, R_\mathrm{raw} + w_\mathrm{dist} \, R_\mathrm{dist} + w_\mathrm{ex} \, R_\mathrm{ex} + w_\mathrm{tok} \, R_\mathrm{tok} + w_\mathrm{tok,cc} \, R_\mathrm{tok,cc}.

Even though some of these terms overlap strongly, the current design choice is to keep them explicit rather than collapsing them into a more abstract family.

Reconstruction config

The current proposal is to introduce a nested reconstruction config under the existing disentanglement loss config.

In core/disentanglement/config.py, the proposed dataclass shape is:

@dataclass(frozen=True)
class ReconstructionLossConfig:
    raw_weight: float = 1.0
    distribution_normalized_weight: float = 0.0
    example_inverse_surprisal_weight: float = 0.0
    token_weighted_weight: float = 0.0
    normalized_concentration_corrected_token_weighted_weight: float = 0.0
    token_weight_alpha: float = 1.0
    example_inverse_surprisal_epsilon: float = 1e-8
    zfree_reference_scale: float = 1.0

and then:

@dataclass(frozen=True)
class DisentanglementLossConfig:
    reconstruction: ReconstructionLossConfig = field(
        default_factory=ReconstructionLossConfig
    )
    beta: float = 1.0
    router_support_weight: float = 1.0
    posterior_entropy_weight: float = 0.0
    posterior_variance_weight: float = 0.0
    router_marginal_kl_to_prior_weight: float = 0.0
    router_supervision_weight: float = 0.0
    enable_router_support_loss: bool = True
    inter_latent_divergence_weight: float = 0.0

Meaning of the config fields

  • raw_weight
    • coefficient on reconstruction_loss
  • distribution_normalized_weight
    • coefficient on distribution_normalized_reconstruction_loss
  • example_inverse_surprisal_weight
    • coefficient on example_inverse_surprisal_weighted_reconstruction_loss
  • token_weighted_weight
    • coefficient on token_weighted_reconstruction_loss
  • normalized_concentration_corrected_token_weighted_weight
    • coefficient on normalized_concentration_corrected_token_weighted_reconstruction_loss
  • token_weight_alpha
    • tempering exponent used by both token-level terms
  • example_inverse_surprisal_epsilon
    • denominator stabilizer for the example-level inverse-surprisal term
  • zfree_reference_scale
    • the fixed distribution-level z-free reference scale (c_\mathrm{base})

CLI surface

The current proposed trainer flags are:

  • --raw-reconstruction-weight
  • --distribution-normalized-reconstruction-weight
  • --example-inverse-surprisal-reconstruction-weight
  • --token-weighted-reconstruction-weight
  • --normalized-concentration-corrected-token-weighted-reconstruction-weight
  • --token-weighted-reconstruction-alpha
  • --example-inverse-surprisal-epsilon

The old boolean --normalize-reconstruction-loss is not part of the proposed redesign. Its role is replaced by a direct scalar weight on distribution_normalized_reconstruction_loss.

Likewise, the current single token_weighted_reconstruction_weight becomes the raw token-level term specifically, not a generic token-level family selector.

Returned scalar losses and diagnostics

compute_total_loss should always return the five scalar reconstruction terms listed above.

In addition, the approved extra scalar diagnostics are:

  • token_weighted_base_loss
    • batch mean of the token-weighted z-free baseline value
  • token_concentration_factor
    • batch mean of the token concentration factor (\kappa_\alpha(x, y))

These are scalar batch means, not per-example [B] tensors. The per-example quantities are needed internally to construct the losses, but the standard returned bundle should remain scalar-oriented.

token_weight_alpha should be logged through config / run metadata rather than returned as a runtime loss-bundle scalar.

Expected internal structure

Internally, the reconstruction code will still need per-example tensors such as:

  • uniform per-example z-free base loss (b_u(x_i, y_i)),
  • raw token-weighted z-free base loss (R_{\mathrm{tok},\alpha}^{\mathrm{base}}(x_i, y_i)),
  • concentration factor (\kappa_\alpha(x_i, y_i)).

But those are implementation intermediates, not standard public outputs from compute_total_loss.

Migration stance

The current design preference is hard migration rather than preserving the old boolean / special-case API indefinitely.

That means:

  • prefer replacing normalize_reconstruction_loss with explicit weighted reconstruction terms,
  • prefer replacing the existing single reconstruction-special-case logic with the new nested reconstruction config,
  • prefer updating downstream analysis code to the new names rather than adding long-lived backward-compatibility shims.

Known redundancy

The proposed API intentionally keeps two overlapping pairs explicit:

  • reconstruction_loss vs distribution_normalized_reconstruction_loss
  • token_weighted_reconstruction_loss vs normalized_concentration_corrected_token_weighted_reconstruction_loss

This redundancy is acceptable at the current stage because:

  • it makes ablations trivial,
  • it keeps the implementation simple,
  • it mirrors the current design style,
  • it supports always-on tracking of all candidate reconstruction objectives.

In practice, most runs should likely set at most one member of each overlapping pair to a nonzero weight.

Relationship to other notes

This note is the current implementation-design target for the next coding pass.

Links to this page:

  1. Ideation
  2. Reconstruction Weighting Redesign Working Map
Built with LogoFlowershow