Base-Directed Reconstruction Objective

This note records notation and reference language for the paper's proposed $p_\theta$-directed reconstruction objective. The core move is to replace the uniform reconstruction term in the negative conditional ELBO with a reconstruction term defined relative to the fixed base model $p_\theta(y|x)$.

Reference phrasing:

We replace the uniform reconstruction term in the negative conditional ELBO with a $p_\theta$-directed reconstruction term $\mathcal R_\theta$, which normalizes reconstruction by the z-free base-model reference scale and adaptively concentrates token-level reconstruction pressure where the base model has high residual surprisal.

Naming

Use $p_\theta$-directed reconstruction objective as the formal term for the objective modification.

Use base-directed reconstruction as the short prose term.

Use base-directed reconstruction pressure when emphasizing the optimization role: the term increases reconstruction pressure in units defined by the base model and its residual surprisal.

Use reference-normalized surprisal-weighted reconstruction for the mechanism when the sentence needs to name both parts: global normalization by a z-free base-model reference scale and token-level weighting by z-free base-model surprisal.

Avoid making implementation labels such as tokcc or distribution-normalized reconstruction the main paper term. Those labels are useful for configs, ablations, and analysis tables, but the paper method should be named by its role in the methodology.

Motivation

The ordinary reconstruction term does not directly measure how informative $z$ is beyond $x$ on a stable scale. The loss $\mathbb E_{z\sim q}[-\log g_\phi(y|x,z)]$ is a code-length objective for predicting $y$ from both $x$ and $z$; by itself, it does not distinguish prediction supplied by $x$ from prediction supplied by the latent variable.

Reference phrasing:

Standard reconstruction pressure does not ask whether $z$ explains $y$ beyond $x$. It only asks whether the generator predicts $y$ from the pair $(x,z)$.

This distinction is especially important in the adapted-from-base setting. The base model already predicts much of $y$ from $x$, so the factorized generator can maintain marginal fit while leaving $z$ inert. A useful reconstruction objective must create pressure for $z$ to explain residual variation in $y$ beyond what is already captured by the z-free base model.

Reference phrasing:

The relevant question is not absolute reconstruction quality, but whether conditioning on $z$ explains residual surprise beyond the input $x$.

Base-directed reconstruction addresses this by expressing reconstruction in base-model reference units and directing token-level pressure using z-free base-model surprisal. This makes the reconstruction pressure adaptive to both $p_\theta$ and the data distribution.

Parameter And Reference Convention

Use semicolon arguments for optimized parameters and subscripts for fixed reference objects.

The main objective should be written as

$$\mathcal J_\theta(x,y;\phi,\xi) = \mathcal R_\theta(x,y;\phi,\xi) + \beta\, \mathrm{KL} \left( q_\xi(\cdot|x,y) \Vert r_\phi(\cdot|x) \right).$$

Here $\phi$ denotes the optimized router-generator parameters, $\xi$ denotes the optimized training-time posterior parameters, and $\theta$ denotes the fixed base model used to define reference scores, token weights, and normalization constants.

The $\theta$ subscript marks dependence on this fixed base-model reference. Reference scores, token weights, and normalization constants such as $b_{\theta,t}$, $b_\theta$, $w_{\theta,t}^{(\gamma)}$, $a_{\theta,t}$, and $c_\theta$ are computed using the frozen reference copy of $p_\theta$ and treated as constants when optimizing $\phi$ and $\xi$; the objectives $\mathcal R_\theta(x,y;\phi,\xi)$ and $\mathcal J_\theta(x,y;\phi,\xi)$ remain functions of the optimized parameters indicated after the semicolon.

Use $\mathcal R_\theta(x,y;\phi,\xi)$ in the main text. Use $\mathcal R_{\theta;\alpha,\gamma}(x,y;\phi,\xi)$ only when comparing hyperparameters or discussing limiting cases.

Baseline Negative Conditional ELBO

Call the minimized baseline objective the negative conditional ELBO. Use explicit router-generator notation in this definition:

$$\mathcal J_{\mathrm{cELBO}}(x,y;\phi,\xi) = \mathbb E_{z\sim q_\xi(\cdot|x,y)} \left[ -\log g_\phi(y|x,z) \right] + \beta\, \mathrm{KL} \left( q_\xi(\cdot|x,y) \Vert r_\phi(\cdot|x) \right).$$

This is clearer than using compact $p_\phi$ notation at this point because it makes the router $r_\phi$ and generator $g_\phi$ explicit. Use $p_\phi$ when referring to the induced factorized-model distribution family, especially the latent marginal $p_\phi(y|x)$.

Reconstruction Unit

In the main paper, index response tokens by $t=1,\ldots,T_y$, where $T_y$ is the response length. Avoid $S(y)$ in the core objective unless the text specifically needs to discuss implementation masks.

Define the per-token generator loss as

$$\ell_t(x,y,z;\phi) = -\log g_\phi(y_t|x,z,y_{<t}).$$

This follows the semicolon convention: $\ell_t$ is a differentiable function of the current generator parameters $\phi$, not a fixed reference object.

The uniform response-token reconstruction term is

$$\mathcal R_{\mathrm{unif}}(x,y;\phi,\xi) = \mathbb E_{z\sim q_\xi(\cdot|x,y)} \left[ \frac{1}{T_y} \sum_{t=1}^{T_y} \ell_t(x,y,z;\phi) \right].$$

It is acceptable to write $-\log g_\phi(y|x,z)$ in the generic negative conditional ELBO. Before introducing the proposed objective, state that the experimental reconstruction unit is the response-token mean NLL.

Implementation note language:

In implementation, the response-token mean is computed with a supervision mask that removes prompt, latent-control, and padding positions. In the paper notation, $y$ denotes the response sequence, so the main equations use $t=1,\ldots,T_y$.

Base-Model Surprisal And Reference Scale

Define token-level z-free base-model surprisal as

$$b_{\theta,t}(x,y) = -\log p_\theta(y_t|x,y_{<t}).$$

Define the matching response-token mean base loss as

$$b_\theta(x,y) = \frac{1}{T_y} \sum_{t=1}^{T_y} b_{\theta,t}(x,y).$$

Define the z-free base-model reference scale as

$$c_\theta = \mathbb E_{x\sim \mathcal D_X,\; y\sim p_\theta(\cdot|x)} \left[ b_\theta(x,y) \right].$$

Use $\widehat c_\theta$ for the empirical estimate used in experiments.

Preferred prose: z-free base-model reference scale. The quantity can be described as an entropy-scale estimate or expected base-model response-token surprisal, but avoid using $\widehat H(Y|X)$ as the primary notation unless the surrounding text is explicitly making the entropy interpretation.

Adaptive Token Weights

Use $\gamma\ge 0$ for the concentration exponent controlling how strongly token weights focus on high-surprisal tokens.

Define uniform token weights

$$u_t=\frac{1}{T_y}.$$

Define the raw normalized surprisal weights as

$$\widetilde w_{\theta,t}^{(\gamma)}(x,y) = \frac{ b_{\theta,t}(x,y)^\gamma }{ \sum_{s=1}^{T_y} b_{\theta,s}(x,y)^\gamma }.$$

The raw weights sum to one. When $\gamma=0$, they reduce to uniform weights. When $\gamma>0$, they concentrate mass on tokens that have higher z-free base-model surprisal.

Define the concentration factor

$$\kappa_\theta^{(\gamma)}(x,y) = \frac{ \sum_{t=1}^{T_y} \widetilde w_{\theta,t}^{(\gamma)}(x,y) b_{\theta,t}(x,y) }{ b_\theta(x,y) }.$$

Define the final concentration-corrected weights as

$$w_{\theta,t}^{(\gamma)}(x,y) = \frac{ \widetilde w_{\theta,t}^{(\gamma)}(x,y) }{ \kappa_\theta^{(\gamma)}(x,y) }.$$

The final corrected weights generally do not sum to one. Instead, they satisfy the base-scale property

$$\sum_{t=1}^{T_y} w_{\theta,t}^{(\gamma)}(x,y) b_{\theta,t}(x,y) = b_\theta(x,y).$$

This is the point of the concentration correction: under the z-free base model, the corrected token-weighted loss has the same per-example scale as the uniform response-token mean loss.

Proposed Reconstruction Term

Use $\alpha\in[0,1]$ for the interpolation between uniform token pressure and corrected surprisal-directed token pressure.

Define the final reconstruction weights as

$$a_{\theta,t}(x,y) = \alpha u_t + (1-\alpha)w_{\theta,t}^{(\gamma)}(x,y).$$

Use $a_{\theta,t}^{(\alpha,\gamma)}(x,y)$ only when the hyperparameter dependence needs to be explicit.

Define the $p_\theta$-directed reconstruction term as

$$\mathcal R_\theta(x,y;\phi,\xi) = \frac{1}{c_\theta} \mathbb E_{z\sim q_\xi(\cdot|x,y)} \left[ \sum_{t=1}^{T_y} a_{\theta,t}(x,y) \ell_t(x,y,z;\phi) \right].$$

This folds the global reference normalization into $\mathcal R_\theta$ itself, so the objective can be described as replacing the standard uniform reconstruction term with a $p_\theta$-directed reconstruction pressure.

Under the z-free base model, the interpolation preserves the per-example base scale:

$$\sum_{t=1}^{T_y} a_{\theta,t}(x,y) b_{\theta,t}(x,y) = b_\theta(x,y).$$

Therefore, in expectation over $x\sim \mathcal D_X$ and $y\sim p_\theta(\cdot|x)$,

$$\mathbb E[\mathcal R_\theta^{\mathrm{base}}(x,y)] = 1.$$

This is the central normalization property: the reconstruction term is expressed in base-model reference units, so its baseline scale is $O(1)$ instead of potentially vanishing when the base model already predicts $y$ from $x$ well.

Interpretation

The objective is base-directed in two senses.

First, the global scale $c_\theta$ is defined by the z-free base model's expected response-token surprisal. This restores reconstruction pressure relative to the KL term in regimes where the adapted generator already has low z-free loss.

Second, the token weights are directed by the base model's residual surprisal. Tokens that are predictable from $x$ under $p_\theta$ receive less additional emphasis; tokens where the base model remains surprised receive more reconstruction pressure, subject to concentration correction.

The concentration correction should not be described as a harmless re-expression. It preserves the z-free base scale, but it also changes how examples with concentrated surprisal contribute relative to the raw token-weighted objective.

For paper-facing metric and figure language around $\mathcal R_\theta$, use Metrics and frontier language.

Usage Across The Paper

In the methodology section, introduce the negative conditional ELBO, explain why its uniform reconstruction term is weak in the adapted-from-base regime, then introduce $\mathcal R_\theta$ as the $p_\theta$-directed replacement.

In the experiments section, describe $\alpha$ and $\gamma$ as objective hyperparameters: $\alpha$ controls the mixture between uniform and surprisal-directed token pressure, and $\gamma$ controls the concentration of token weighting toward high-surprisal positions.

In captions and figure labels, prefer short names such as base-directed reconstruction, surprisal-weighted reconstruction, reference-normalized reconstruction, or base-directed objective depending on space.

In implementation-facing discussion, it is fine to mention distribution-normalized reconstruction and normalized concentration-corrected token-weighted reconstruction as the two components unified by the paper notation.

Terms That Need Care

Base-normalized reconstruction is concise but emphasizes only the global scale, not the token-level adaptive weighting.

Surprisal-weighted reconstruction captures the token-level mechanism but omits the reference normalization.

Residual-surprisal reconstruction is suggestive but can sound as if the objective optimizes an explicit residual difference rather than a weighted reconstruction loss.

Adaptive reconstruction pressure is useful explanatory prose, but too broad as the formal method name.

Entropy normalization should be used cautiously because $c_\theta$ is an empirical response-token reference scale in the same units as the implemented reconstruction loss.