v7 Cross-Entropy Parity: Why Small Math Choices Matter

This page explains the v7 cross-entropy (CE) kernel in detail: the p - one_hot gradient, PyTorch reduction semantics, long-horizon drift behavior, and the concrete changes that moved CE from fragile to production-grade for the v7 training harness.

The Core Gradient: Why It Is p - one_hot

For one token with logits z and target class y:

\[ p_i = \frac{e^{z_i}}{\sum_j e^{z_j}}, \quad \mathcal{L} = -\log(p_y) \]

\[ \frac{\partial \mathcal{L}}{\partial z_i} = p_i - \mathbf{1}[i = y] \]

So each row gradient is:

\[ \nabla_{z}\mathcal{L} = p - \text{one\_hot}(y) \]

For mean reduction over N_valid targets:

\[ \nabla_{z}\mathcal{L}_{\text{mean}} = \frac{p - \text{one\_hot}(y)}{N_{\text{valid}}} \]

What Can Go Wrong in Practice

PyTorch CE Semantics We Matched

• F.cross_entropy(...) routes to torch._C._nn.cross_entropy_loss.
• Index-target path is effectively log_softmax + nll_loss.
• Default ignore_index = -100 for index targets.
• Mean reduction divides by non-ignored target count (or summed target weights if class weights are used): \(\frac{\sum \ell_t}{N_{\text{valid}}}\).
• All targets ignored with mean reduction returns NaN loss; gradients are zero.

Primary source files in local PyTorch checkout:

• pytorch/torch/nn/functional.py (cross_entropy entrypoint)
• pytorch/aten/src/ATen/native/native_functions.yaml (cross_entropy_loss dispatch)
• pytorch/aten/src/ATen/native/LossNLL.cpp (CPU reduction/backward semantics)

CK v7 CE Variants (Current)

Implementation: src/kernels/loss_kernels.c

Change Log: What We Improved

1. Switched CE loss calculation to stable log-sum-exp form to avoid probability-clamp artifacts.
2. Added explicit PyTorch-style index-target mean reduction semantics with ignore_index=-100.
3. Added a strict reference CE entrypoint (softmax_cross_entropy_loss_ptref) for parity A/B.
4. Added guarded default routing to preserve all-valid long-horizon stability while handling ignore/invalid targets correctly.
5. Expanded CE unit tests to include mixed-ignore and all-ignored behavior checks.

Production Evidence (v7 Harness)

Additional training gates run clean in this validation pass:

• version/v7/.cache/reports/train_parity_drift_smoke_latest.json (PASS)
• version/v7/.cache/reports/train_parity_realistic_long_horizon_latest.json (PASS)
• version/v7/.cache/reports/optimizer_parity_latest.json (PASS)
• version/v7/.cache/reports/fd_gradients_latest.json (PASS)
• version/v7/.cache/reports/replay_determinism_latest.json (PASS)

What We Still Do Not Match (Yet)

How To Reproduce and Validate

# CE unit tests (includes ignore_index and all-ignored semantics)
source .venv/bin/activate
python unittest/test_cross_entropy.py

# Long-horizon parity stress (850 steps)
source .venv/bin/activate
CK_NUM_THREADS=8 python version/v7/scripts/train_parity_epochs_v7.py \
  --epochs 8 --seq-len 8 --total-tokens 131072 --grad-accum 8 \
  --optimizer adamw --lr 1e-3 --seed 42 --max-steps 850 \
  --ck-rmsnorm-backend c --ck-swiglu-backend c --ck-loss-backend c \
  --json-out /tmp/v7_ce_semantics_c_850_postfix.json

# A/B against strict CE reference path
source .venv/bin/activate
CK_NUM_THREADS=8 python version/v7/scripts/train_parity_epochs_v7.py \
  --epochs 8 --seq-len 8 --total-tokens 131072 --grad-accum 8 \
  --optimizer adamw --lr 1e-3 --seed 42 --max-steps 850 \
  --ck-rmsnorm-backend c --ck-swiglu-backend c --ck-loss-backend c_ptref \
  --json-out /tmp/v7_ce_semantics_c_ptref_850_postfix.json