Reduction, Masking, Weighting, and Multi-Output Objectives

On 16 October 2024 Hugging Face published a fix for a loss-normalization bug in the Transformers Trainer. Nothing was wrong with the model, the data or the optimizer. The division was in the wrong place.

With gradient accumulation the code averaged each micro-batch's cross-entropy separately, then averaged those means. Every one of those steps was a mean over valid tokens. The result was still not the quantity anyone intended. Hugging Face said so in the announcement: “the correct loss should be computed by the total loss across all batches in a gradient accumulation step divided by the total number of all non padding tokens in those batches. This is not the same as the average of the per-batch loss values.”

Both computations are means over valid tokens. Only one of them is the loss.

Masking excludes positions where the target is undefined, unavailable or irrelevant. Padding tokens are the standard example. The claim it makes is that no supervised error exists there, so the denominator counts valid positions only, and the risk it carries is a hidden missingness pattern. Weighting changes the influence of a target that is perfectly valid — a rare-class coefficient, say. Its claim is that this error matters differently, its denominator is often the total weight, and its risk is distorted calibration. Sampling does neither. It changes which observations appear in the batch at all, as when rare events are oversampled. It alters the distribution of updates, its denominator is batch-dependent, and its risk is duplicated noisy cases. Confusing the three can hide missing supervision or manufacture biased gradients.

How much of a batch carries supervised error at all is a dial, not a measurement. BERT set it once in 2018 and the field inherited the setting: “In all of our experiments, we mask 15% of all WordPiece tokens in each sequence at random.” The authors were explicit that this bounds the supervision available per step, noting that “the masked LM only make predictions on 15% of tokens in each batch”. Wettig and colleagues at Princeton turned the dial in 2023 and found 15% is not optimal. Masking 40% beat it for BERT-large-size models: SQuAD 88.0 → 89.8, QNLI 90.9 → 91.6, MNLI 84.2 → 84.5. Pushed to 80%, validation perplexity went from 17.7 to 1141.4, and more than 95% of fine-tuning performance survived anyway. The number a mask moves first is not the number anyone deploys.

The weighting side has a matching trap, and it is a spelling trap. PyTorch's torch.nn.CrossEntropyLoss ships the signature (weight=None, size_average=None, ignore_index=-100, reduce=None, reduction='mean', label_smoothing=0.0). Under reduction='mean' with class weights it divides the weighted sum by Σ w_{y_n}·1{y_n ≠ ignore_index} — the sum of the weights over non-ignored targets. Keras 3 uses the same word for a different denominator, and documents the choice on the reduction argument of its base Loss class: “"sum" sums the loss, "sum_over_batch_size" and "mean" sum the loss and divide by the sample size, and "mean_with_sample_weight" sums the loss and divides by the sum of the sample weights.” Its default is sum_over_batch_size. Dividing by the weights is the separate opt-in. Two frameworks a team may already have in one repository disagree by default about what the word mean names once weights exist.

The order of these operations matters, because each one changes what the next one is dividing.

First, compute elemental losses: an error for each token, pixel, sample, pair or task output. Second, apply validity masks, removing positions without meaningful targets or observable outcomes. Third, apply weights, changing influence for classes, examples, tasks or confidence levels. Fourth, choose denominators — samples, valid elements, total weight, or another quantity you have declared out loud. Fifth, combine objectives into the single scalar that is backpropagated.

The Transformers gradient-accumulation bug lived entirely in step four. PyTorch and Keras 3 disagree there by default. It is the step most often inherited rather than chosen.

A mean per image, a mean per positive label and a mean per annotated pixel produce different gradient units. Batch composition can therefore change effective step size even when the optimizer learning rate is fixed.

DAPO isolated that effect. It changed only the loss reduction of GRPO — from averaging tokens within a sample and then averaging samples, to a single average over all tokens in the batch, 1/Σ|o_i|. The complaint is about a denominator, not a gradient: “Since all samples are assigned the same weight in the loss calculation, tokens within longer responses (which contain more tokens) may have a disproportionately lower contribution to the overall loss, which can lead to two adverse effects.” In the ablation, AIME 2024 avg@32 rose from 41 to 42 with the token-level loss, and the authors note the gain is in stability rather than score. The full recipe reached 50 points on Qwen2.5-32B, above DeepSeek-R1-Zero-Qwen-32B's 47, using 50% of the training steps. Hugging Face's TRL has since made this reduction its default, shipping loss_type="dapo" and marking the original sequence-length normalization "grpo" as “Not recommended due to length bias”.

The denominator also fixes what a learning rate means. Minibatch SGD divides the summed gradient by the batch size n, so a rule can be stated directly on that division. Goyal and colleagues at Facebook AI Research stated it in 2017: “Linear Scaling Rule: When the minibatch size is multiplied by k, multiply the learning rate by k.” With it they trained ResNet-50 on ImageNet at minibatch 8192 across 256 GPUs in one hour, with no loss of accuracy and roughly 90% scaling efficiency from 8 to 256 GPUs.

It is a rule about an arithmetic mean, and it holds only as far as that mean does. Shallue and colleagues, in the Journal of Machine Learning Research in 2019, measured 71,638,836 loss values over 168,160 models across 35 workloads. They found the rule is not general — “their learning rate heuristic broke down for even larger batch sizes” — and captioned their Figure 8 “Optimal effective learning rates do not always follow linear or square root scaling heuristics.”

The denominator should match whatever the project intends to keep stable. Log both numerator and denominator so scale changes are visible.

This procedure catches scale problems before a weight sweep. Log every component — unweighted loss, valid count, weighted numerator and final contribution. Inspect gradient norms, measuring how each component affects shared and task-specific parameters. Perturb batch composition with short, long, sparse, dense, positive-free and heavily masked batches. Verify invariants: check whether duplicating an example or changing padding alters the intended unit. Then tune with decision metrics, using validation behavior rather than comparable-looking loss numbers.

That last step is the one most often replaced by machinery, and two groups reported at NeurIPS 2022 that the machinery does not pay. Xin and colleagues at Google ran large-scale language and vision experiments on the specialised multi-task optimizers, and put the result in one sentence of their abstract: “We show that, despite the added design and computational complexity of these algorithms, MTO methods do not yield any performance improvements beyond what is achievable via traditional optimization approaches.”

Kurin and colleagues reached the same place from the other direction in 2022. They found that “unitary scalarization” — simply minimising the sum of the task losses — “matches or improves upon the performance of complex multi-task optimizers in popular supervised and reinforcement learning settings”, once standard single-task regularization and stabilization are applied. They argue that many specialised multi-task optimizers “can be partly interpreted as forms of regularization”.

A plain weighted sum, audited and tuned against decision metrics, is the baseline the specialised methods have to beat. On this evidence they do not.