Token-Sum Versus Token-Mean NLL Aggregation

Back to reconstruction weighting and normalization

Implementation note: Current reconstruction loss stack

This note records a denominator choice that is easy to hide behind the shorthand $-\log p_\phi(y \mid x, z)$: sequence likelihood, token-mean supervised loss, example-mean loss, and token-mean batch loss are not interchangeable when sequence lengths vary.

Core distinction

For an autoregressive decoder over a supervised target span $S(y)$, define the per-token negative log-likelihood

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

The formal sequence negative log-likelihood is a token sum:

$$L_{\mathrm{seq}}(x, y, z) = -\log p_\phi(y \mid x, z) = \sum_{t \in S(y)} \ell_t(x, y, z).$$

The implementation-style supervised-token mean is instead

$$L_{\mathrm{tokmean}}(x, y, z) = \frac{1}{|S(y)|}\sum_{t \in S(y)} \ell_t(x, y, z).$$

These differ by a factor of $|S(y)|$ per example. That factor is harmless only when lengths are fixed or when the denominator is a global constant that does not change the relative weighting of examples or responses.

Batch reductions are separate choices

After choosing a per-example reduction, the batch reduction adds another weighting decision. Three common reductions are:

$$\frac{1}{B}\sum_{i=1}^{B}\sum_{t \in S_i}\ell_{i,t} \qquad \frac{1}{B}\sum_{i=1}^{B}\frac{1}{|S_i|}\sum_{t \in S_i}\ell_{i,t} \qquad \frac{\sum_i\sum_{t \in S_i}\ell_{i,t}}{\sum_i |S_i|}.$$

The first is an example mean of sequence NLLs, so longer examples carry larger loss magnitude inside each example. The second gives each example equal total weight after length-normalizing it. The third gives each active token equal weight globally, so longer examples receive more total batch weight because they contribute more tokens.

Sequence VAE convention

Formal sequence VAE work usually writes the reconstruction term as the likelihood of the whole observed sequence. For an autoregressive decoder this expands to a token sum.

Bowman et al. (2016) define the sentence VAE ELBO with $\mathbb{E}_{q(z \mid x)}[\log p(x \mid z)]$ and describe the decoder as an RNN language model conditioned on the latent code. The corresponding sequence model factorization is

$$p(x \mid z) = \prod_t p(x_t \mid x_{<t}, z),$$

so the log likelihood is a sum over token log probabilities. This source establishes the formal sequence-likelihood convention, even though the paper also discusses the reconstruction term operationally as cross entropy.

Chung et al. (2015), A Recurrent Latent Variable Model for Sequential Data, make the sequence factorization and timestep sum explicit. They write the RNN sequence probability as $\prod_t p(x_t \mid x_{<t})$, and the VRNN lower bound as a sum over timesteps of KL terms plus reconstruction log probabilities:

$$\sum_{t=1}^{T} \left[ -\mathrm{KL}(q(z_t \mid x_{\le t}, z_{<t}) \Vert p(z_t \mid z_{<t}, x_{<t})) + \mathbb{E}_{q(z_t \mid x_{\le t}, z_{<t})}\log p(x_t \mid z_{\le t}, x_{<t}) \right].$$

LM training convention

Standard language-model training and evaluation often report or optimize token-level cross entropy, usually averaged over active labels. PyTorch CrossEntropyLoss defaults to reduction="mean" and, with ignore_index, averages over non-ignored targets. Hugging Face's fixed-length perplexity documentation defines perplexity as the exponential of average token NLL, and its example multiplies model loss by the number of target tokens before accumulating total NLL for corpus-level perplexity.

This convention is useful for optimization stability, reporting comparability, and stable gradient scales across sequence lengths. It is not the same mathematical object as the whole-sequence log likelihood unless lengths are fixed or the denominator is held constant.

RLVR length-normalization parallels

The same denominator issue appears in GRPO-style RLVR objectives. The details differ from supervised NLL, but the lesson is the same: averaging over response tokens versus summing over response tokens changes the weighting of long and short responses.

DeepSeekMath introduces GRPO with a group average over outputs and a per-output average over generated tokens:

$$\frac{1}{G}\sum_{i=1}^{G}\frac{1}{|o_i|}\sum_{t=1}^{|o_i|}\mathrm{clipPG}_{i,t}.$$

This gives each sampled response roughly equal total weight regardless of response length.

Liu et al. (2025), Understanding R1-Zero-Like Training: A Critical Perspective, argue that the $1/|o_i|$ term creates a response-level length bias: positive-advantage short responses receive larger per-token updates, while negative-advantage long responses are penalized less per token. Their Dr. GRPO variant removes response-length normalization and reward-standard-deviation normalization; in implementation terms, they replace division by mask.sum(...) with division by a fixed constant such as MAX_TOKENS.

DAPO makes a related distinction. It describes original GRPO as a sample-level loss that averages losses within each response before aggregating across responses, then introduces a token-level policy-gradient loss that aggregates over all response tokens:

$$\frac{1}{\sum_i |o_i|}\sum_{i=1}^{G}\sum_{t=1}^{|o_i|}\mathrm{clipPG}_{i,t}.$$

Under this loss, tokens in longer responses are no longer diluted solely because they occur in longer samples.

Implication for this project

The clean formal notation $-\log p_\phi(y \mid x, z)$ is still appropriate when describing a baseline sequence likelihood. But the current implementation is more specific: it computes masked per-token NLL over the supervised y-span, averages over active supervised positions inside each example, then averages over examples. In code, this is masked_token_nll_per_example(...) in path:core/disentanglement/losses.py, followed by .mean() in the discrete and continuous CVAE loss paths.

So the implementation-matched reconstruction term is

$$\frac{1}{|S(y)|}\sum_{t \in S(y)}-\log p_\phi(y_t \mid x, z, y_{<t}),$$

not the formal sequence NLL

$$\sum_{t \in S(y)}-\log p_\phi(y_t \mid x, z, y_{<t}).$$

The denominator is not cosmetic once lengths vary. Moving from token mean to sequence sum would increase reconstruction magnitude for longer targets, alter the reconstruction-to-KL balance as a function of length, and change which examples dominate gradients. Moving from per-example token means to a global token mean would also change weighting: long examples would contribute proportionally more total batch weight.

Study questions

• Should the paper reserve $-\log p_\phi(y \mid x, z)$ for the formal sequence-likelihood objective, while naming the implemented objective "supervised-token mean NLL"?
• Should future experiments include a controlled aggregation ablation: per-example token mean, global token mean, example mean of sequence sums, and sequence-sum with a fixed global scale?
• If sequence-sum reconstruction is tested, should $\beta$ or the z-free normalization constants be rescaled so the comparison isolates length weighting rather than only changing the global reconstruction-to-KL scale?

Implementation follow-up

The approved focused ablation is recorded in Planned implementation: reconstruction-loss token aggregation mode. It introduces reconstruction_loss_token_aggregation_mode with default mean and optional sum, applies the mode only to reconstruction-family terms, recomputes z-free reference scales in the selected mode, and keeps raw sum available as an optimization setting.

References

• Bowman et al. (2016) - Generating Sentences from a Continuous Space; sentence VAE objective and latent-conditioned RNN decoder.
• Chung et al. (2015) - A Recurrent Latent Variable Model for Sequential Data; RNN factorization in Eq. 2 and timestep-wise VRNN lower bound in the learning section.
• PyTorch CrossEntropyLoss documentation; reduction="mean" and ignore_index behavior.
• Hugging Face Transformers - Perplexity of fixed-length models; perplexity as average token NLL and the token-count accumulation pattern.
• Shao et al. (2024) - DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models; GRPO objective in Eq. 3.
• Liu et al. (2025) - Understanding R1-Zero-Like Training: A Critical Perspective; response-level length bias in Sec. 3.1 and Dr. GRPO fix in Sec. 3.2 / Listing 1.
• DAPO (2025) - DAPO: An Open-Source LLM Reinforcement Learning System at Scale; sample-level versus token-level policy-gradient loss in Sec. 3.3 / Eq. 12.