Two documents, two ways to decide how much they count
A team prepares domain documents to adapt an AI model. Some are short, others long. If every document is intended to matter equally, is mixing them in one batch and averaging loss enough? Not necessarily: averaging over all supervised positions gives more weight to documents with more positions. The denominator is part of the learning objective, not a cosmetic programming detail.
We ask what quantity is being averaged and how the gradient changes. Two synthetic groups contain 2 and 8 valid positions, and a one-parameter binary model lets us calculate cross-entropy completely. We then add padding positions used only to align lengths. We do not train a language model: a controllable model isolates a mechanism relevant to adapting domain-specific models.
What one loss measures
A token is a unit of text representation; a supervised position is a location where we request a prediction and have a target. This example needs neither real words nor a particular tokenizer. Every position has label y equal to 0 or 1. The model assigns the same probability p to class 1 at every position. It is deliberately simple: its purpose is to expose aggregation rules, not learn language.
Theta is the logit, a real-valued parameter the model can change. At theta=0, p=0.5. Natural logarithms give loss in nats. For y=1 we penalise small p; for y=0 we penalise small 1−p. Loss is not a percentage of correct responses. Changing observation weights can yield the same loss value but different gradients.
The global mean implicitly chooses document weights
Let n_i be the number of valid positions in document i and L_i its mean loss. A token mean sums all losses and divides by the total valid count. A document mean first computes each L_i and gives documents equal weight. In the first objective, L_i receives a coefficient proportional to n_i; in the second, equal coefficients, provided every document has at least one valid position.
D is the document count. With n_A=2 and n_B=8, the token objective is 0.2 L_A+0.8 L_B; the document objective is 0.5 L_A+0.5 L_B. Document B weighs four times A in the first objective but not the second. Token averaging is not automatically wrong: it consistently gives positions equal importance. It conflicts with the stated intention only when documents, rather than positions, should receive equal weight.
The gradient exposes the difference
Assign label 1 to both A positions and label 0 to all eight B positions. These are constructed targets, not corpus frequencies. To see how theta changes, differentiate loss with respect to it. The sigmoid derivative is p(1−p); substituting into the logarithmic derivatives simplifies everything to p−y. The gradient gives the local direction of increasing loss.
At theta=0 and p=0.5, the token gradient is 0.3 and the document gradient is 0. A descent step subtracts a positive step size times the gradient: the first objective pushes theta negative, while the second is already stationary. Both initial losses equal log 2, about 0.693147 nats. The loss value alone would not reveal that the objectives suggest different updates.
We can also locate the optima. The gradient is p minus the weighted fraction of label 1, so it vanishes at p=0.2 for tokens and p=0.5 for documents. The second derivative p(1−p) is positive strictly between 0 and 1, giving unique finite optima in this scalar model. This does not prove convexity of a large language model’s loss landscape; it concerns only the stated model.
Padding is not a new class-zero example
To obtain two rows of length 8, add six artificial positions to A. They have no task target; they organise tensor shapes. If we label them 0 and include them as real data, there are 16 positions with only 2 positives. The optimum becomes p=2/16=0.125. Document content did not change; six invented examples entered the objective.
A mask m equal to 1 at valid positions and 0 at padding must affect both numerator and denominator. Adding artificial masked positions then leaves the correct loss and gradient unchanged. The program checks this with 0,1,6 and 20 added positions and three theta values. This is a property of the constructed calculation, not a field-robustness measurement.
A second error removes padding losses from the numerator but still divides by 16. The result is 10/16 of the correct loss. At theta=0 it falls from 0.693147 to 0.433217 without improving any probability; the gradient falls from 0.3 to 0.1875. The scalar optimum remains p=0.2 because the whole function was multiplied by a positive constant. Update scale changes, however: a lower plotted loss does not prove a better prediction.
A third error treats label 0 as synonymous with padding and ignores every 0, including all eight legitimate B targets. Only two label 1 positions remain: the initial gradient becomes−0.5 and pushes p towards 1. The minimum is now a limit as theta increases without bound, not a finite optimum. A class value differs from a position’s artificial role; the mask must describe the role.
| Objective | Loss at theta=0(nats) | Gradient at theta=0 | Optimal p |
|---|---|---|---|
| Token mean | 0.693147 | 0.3000 | 0.200 |
| Document mean | 0.693147 | 0.0000 | 0.500 |
| Padding included | 0.693147 | 0.3750 | 0.125 |
| Wrong denominator | 0.433217 | 0.1875 | 0.200 |
| Class zero ignored | 0.693147 | -0.5000 | 1(limit) |
The rows describe different objectives, not five training runs evaluated on one test. Token versus document averaging is an explicit weighting choice. Other rows introduce errors relative to the agreed valid-data definition. A constant scale factor could be absorbed into the optimisation step size; if padding proportions vary by batch, the factor varies and changes relative update contributions too.

Splitting a batch should not silently change the objective
Suppose A and B are processed as two microbatches, successive portions used to reduce memory. At theta=0 their mean gradients are−0.5 and+0.5. A simple average gives 0, reconstructing the document objective. Reproducing the global token mean requires weighting each mean gradient by its valid-position count: (2 ×−0.5+8 × 0.5)/10=0.3, matching the unsplit batch.
B here counts microbatches. Equal valid counts make the formulas agree, as can special cancellations. Equal padded tensor lengths are insufficient. This checks accumulation algebra; it does not allege a defect in a particular Trainer or framework version. An implementation must be compared with its declared objective and actual reduction rules.
From formula to library without mixing the levels
Official PyTorch 2.14 CrossEntropyLoss documentation distinguishes class-index targets from probability-distribution targets. For class indices, ignore_index excludes targets from loss and gradient; without class weights, mean reduction uses the non-ignored target count. Our example uses discrete labels and no class weights. We do not extend that formula unchanged to probability targets or other weighting schemes.
In a causal language model, valid positions also depend on aligning inputs with next-token targets and deciding whether prompts, answers, or both are supervised. Full document length need not equal n_i. An attention mask and a loss mask answer different questions, even when an implementation connects them. Define which predictions actually contribute to the objective before counting; tensor occupancy alone is insufficient.
With no valid position, the mean denominator is zero and our formula is undefined. The program reports this explicitly rather than returning zero as if prediction were perfect. A document without targets also has undefined L_i; excluding it or treating it otherwise needs a declared policy. We do not attribute our code’s behaviour to every library, whose handling may differ.
Reproducing the argument without training an LLM
The program uses only Python’s mathematics library for loss and analytic gradients, with a numerically stable logarithmic formulation. It saves all five objectives, checks invariance under correctly masked padding, and compares the gradient with a central finite difference at theta=0.3. The observed check error is about 4.1 × 10⁻¹¹, below its chosen 1 × 10⁻⁸tolerance. This does not validate an automatic-differentiation system.
In the short snippet, lengths counts contributing positions and gradients stores each document’s mean contribution. The final two means reconstruct the objectives without GPUs, model weights, or external data. Summing positions has linear cost in the processed count, but this does not measure inference or training cost. This scalar operation does not simulate an actual model’s expensive computation.
Which unit of importance should be optimised?
Length changes training when it changes contributions’ objective weights. Equal document importance requires document averaging or equivalent weights; equal position importance requires token averaging. Neither is universally better. Every choice must distinguish padding from valid targets, and splitting data into microbatches should not introduce undeclared weights. Actual model quality ultimately requires independent evaluation, which this example does not replace.
Source and reproducibility materials
import math
p=0.5
lengths=[2,8]
labels=[1,0]
gradients=[p-y for y in labels]
token_gradient=sum(n*g for n,g in zip(lengths,gradients))/sum(lengths)
document_gradient=sum(gradients)/len(gradients)
assert abs(token_gradient-0.3)<1e-12
assert document_gradient==0
print(token_gradient,document_gradient)
Code, data, and instructions · JSON. Educational calculations executed with Python 3.14.0; figures with Matplotlib 3.11.2. AI-assisted analysis, without claiming peer review or human review. Original illustrative ImageGen cover: it does not document EL-AI people, premises, or installations. Sources accessed 21 October 2026.

