KLDivLoss: Why the Loss Can Turn Negative

黎 浩然/ 17 10 月, 2026/ 大语言模型/LARGELANGUAGEMODEL/LLM, 教程/TUTORIAL, 计算机/COMPUTER/ 0 comments

When KLDivLoss is substantially negative, first check whether input contains log-probabilities. By default, its pointwise expression is target*(log(target)−input). Putting ordinary softmax probabilities in the input position computes something else, without the nonnegativity guarantee of KL divergence.

Contents
  1. Write the two distributions before the API call
  2. API argument order differs from KL notation
  3. A negative counterexample exposes the wrong units
  4. Twelve checks reproduce direction and representation
  5. Mean and batchmean average different units
  6. Padding makes the averaging unit explicit
  7. Keep enough context to explain the scalar
  8. Sources and validation scope

For teacher q=(0.75,0.25) and student p=(0.5,0.5), the correct KL is about 0.130812. Supplying p directly as log-input instead gives about −1.062335. The counterexample leads into direction, log_target and reduction denominators, rather than merely fixing a sign. 中文版.

Write the two distributions before the API call

Let q denote the teacher or target distribution and p the student. This article compares D(q||p)=sum_i q_i(log q_i−log p_i), summing over classes. It weights terms by q and generally differs from D(p||q). Writing that expression explicitly is more reliable than assuming variable names establish the intended direction.

Initially, every probability is strictly positive and each distribution sums to one. Logs are natural. These conditions let the standard-library example call log directly; Mathematically, a term with q_i=0 is defined as zero by a limiting convention; q_i>0 with p_i=0 gives positive infinity. Ordinary multiplication after log(0) does not automatically implement those conventions. The teaching function is not a complete loss-library implementation.

Nonnegativity belongs to the complete KL sum for normalized distributions, not every classwise contribution. An individual q_i*log(q_i/p_i) can be negative while other terms compensate. A negative element from reduction=none therefore differs from a substantially negative total for valid distributions. Diagnose the quantity actually observed before deciding that the mathematical guarantee has failed.

API argument order differs from KL notation

The PyTorch 2.14 reference requires log-space input. The target is a probability by default, or log-space with log_target=True. To compute D(q||p), put log p in input and q in the default target, rather than following the left-to-right order of mathematical KL notation.

When the target contains log q, the pointwise expression becomes exp(log q)*(log q−log p). That changes representation without reversing direction. The switch does not normalize arbitrary values into a valid distribution; consistency still depends on the quantities supplied.

Student logits are commonly transformed with log_softmax along the class axis. Its official reference explains the combined, numerically stable formulation. There is no reason to introduce a potentially underflowed probability intermediate solely to follow a literal softmax-then-log sequence.

Student logits → log_softmax → log p
Teacher distribution → q
Per class: q(log q − log p)
Sum classes; choose averaging unit
Original input map: by default target is a probability and input a log-probability; API argument order differs from mathematical KL notation.

A negative counterexample exposes the wrong units

For q=(0.75,0.25) and p=(0.5,0.5), the correct expression is 0.75*log(1.5)+0.25*log(0.5), approximately 0.130812. Reversing the distributions gives a different value; reduction choices do not explain that directional difference.

If input incorrectly contains p, the sum becomes sum(q_i*log q_i)−sum(q_i*p_i). The first term is −H(q), and the second sum is 0.5 here. The result is therefore −H(q)−0.5≈−1.062335. This does not show that KL permits negative values. It shows that the computation no longer represents the stated KL.

Applying abs or clamping the result conceals the contract error and can change gradients. Check normalization, representation and direction first. Tiny negative values near zero need precision-aware tolerance analysis; they are a different diagnostic situation from the large negative value deliberately constructed here. A visual sign correction is not a substitute for identifying the expression being differentiated.

Twelve checks reproduce direction and representation

The code uses only Python’s standard library and the two strictly positive distributions above. No PyTorch execution, distillation training or GPU test was performed. zip receives known equal-length tuples here; a production interface must validate shapes rather than silently accepting a truncated pairing.

# Positive, normalized teaching distributions; no PyTorch/model run.
from math import log, exp, isclose

def point_sum(log_input, target):
    return sum(q*(log(q)-x) for x,q in zip(log_input,target))

q, p = (0.75,0.25), (0.5,0.5)
lq, lp = tuple(map(log,q)), tuple(map(log,p))
k = point_sum(lp,q)
wrong = point_sum(p,q)  # Probabilities in the log-input position.
h = -sum(x*log(x) for x in q)
assert isclose(sum(q),1.0)
assert isclose(sum(p),1.0)
assert isclose(k,0.13081203594113697)
assert isclose(point_sum(lq,q),0.0,abs_tol=1e-12)
assert not isclose(point_sum(lq,p),k)
assert wrong < 0
assert isclose(wrong,-h-0.5)
assert isclose(sum(exp(y)*(y-x) for x,y in zip(lp,lq)),k)
total, batch, classes = 2*k, 2, 2
assert isclose(total/batch,k)
assert isclose(total/(batch*classes),k/2)
# Same two distributions stored as [B=1,L=2,C=2].
assert isclose(total/1,2*k)  # batchmean over sequences
assert isclose(total/2,k)   # mean over two valid tokens
print("12 KL checks passed", k, wrong)

Twelve checks cover normalization, the expected value, zero self-divergence, directional asymmetry, the negative wrong-input result, equivalent log-target representation and denominators for two tensor layouts. The last two checks explicitly calculate denominators; they do not invoke a framework and should not be presented as framework tests.

Mean and batchmean average different units

For shape [B,C] with one distribution per row, summing classes gives each sample’s KL. The documented reductions sum all elements for sum, divide by the first dimension for batchmean, and divide by all elements for mean.

With two identical samples, two classes and per-sample KL K, sum=2K, batchmean=K and mean=K/2. Mean neither reverses direction nor repairs an input mismatch. It adds a class-count divisor, changing the relative weight when this term is combined with another loss.

Batchmean is not automatically mean token KL for every layout. For [B,L,C], it divides only by B and still sums length L. Moving the same two distribution pairs from [2,2] into [1,2,2] preserves their total but changes the first-dimension divisor from two to one. A reshape can change the statistical unit of the returned scalar even when the underlying pairs are identical.

Shape [2,2]: two samples, two classes
sum: 2K
batchmean: K
mean: K/2
Original reduction map: two identical distribution pairs, each with KL=K. Element averaging and sample averaging use different denominators.

Padding makes the averaging unit explicit

Language-model distillation often calls for summing the vocabulary dimension, aggregating valid tokens through a mask and dividing by the valid-token count. Equal sequence weighting is different: first average within each sequence, then across sequences. Those objectives disagree when sequence lengths differ. Neither is established merely by naming the result an average loss.

This connects to the earlier gradient accumulation normalization article. Microbatch KL means cannot always be combined with equal weight; their valid-item counts matter. Two fields called loss can represent different denominators, so a matching field name is insufficient evidence that the objective has been preserved.

Define how to skip or reject a batch with no valid tokens before dividing by zero. A mask is also not a general NaN sanitizer. Computing an invalid logarithm first and multiplying it by zero afterward need not produce a safe value. Choose a valid computational path rather than hoping masking will repair invalid arithmetic later.

Keep enough context to explain the scalar

On a small batch, inspect the class axis, shapes, q row sums and exp(log p) row sums. Then check that target representation agrees with log_target. Also verify classwise meaning: different vocabularies or class orders do not become aligned merely because shapes match. Successful normalization along one axis does not establish normalization along the class axis required by the task. A wrong axis can create plausible-looking tensors while changing the intended distributions.

Retain classwise sums per distribution, valid-item counts, the final denominator and mixture weights. A change in vocabulary, padding fraction or layout can make raw scalar losses incomparable. Whether the teacher is fixed and how temperature is configured are separate training decisions; this numeric KL example does not determine them.

Start with the two-class distributions to validate the input contract, then restore real logits and masks. This keeps a representation mistake from being confused with model behavior. If the counterexample helps locate a negative loss or unexpectedly small KL term, voluntary support through the section below is welcome.

Sources and validation scope

Checked October 17, 2026 against the official KLDivLoss and log_softmax references. This is an interface and mathematical explanation. Both diagrams and twelve standard-library checks are original teaching material. No model-quality experiment or device-specific numerical error bound was measured.

Support

If the explanations or examples helped you solve a problem or save time, you’re welcome to support this site with a voluntary tip.

WeChat support QR code; click to enlarge
WeChat
Alipay support QR code; click to enlarge
Alipay
Buy Me a Coffee; support this site
Buy Me a Coffee

Click a QR code to enlarge. More options: support page。

Share this Post

Leave a Comment

您的邮箱地址不会被公开。 必填项已用 * 标注

*
*