fix(decision): preserve gradients for confidently wrong predictions

[YingchaoX]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Fixes #10354

Decision cross-entropy currently clamps probabilities to 1e-9 before taking their logarithm. For a confidently wrong hard-label prediction, this caps the loss at approximately 20.7233 and removes the corrective gradient. For example, logits [30, -30] with the second class as the target produce zero gradients instead of [1, -1].

Compute masked log-probabilities directly with FP32 log_softmax. Handle empty option rows and clear padded log-probabilities before multiplying by targets to avoid NaNs. Update the paired decision documentation and add three focused regression tests (four parametrized cases): FP32/BF16 extreme wrong predictions through DistillRPSLoss, masked soft-target CE value/gradient equivalence, and all-padding rows. Existing decision head structural tests cover ordinary loss and gradient behavior.

Experiment results

Local CPU validation: Python 3.12.13, PyTorch 2.14.1, Transformers 5.19.0, editable repository installation on macOS arm64.

The four new cases against the original loss method give 2 failed, 2 passed; both failures assert the missing corrective gradient (FP32/BF16). With this fix, the same extreme input produces loss 60 and gradient [1, -1].

python -m pytest tests/train/test_decision_loss.py \
  tests/test_align/test_decision_jev_align.py::TestStageDStructural \
  tests/test_align/test_decision_clef_align.py::TestStageDStructural \
  tests/test_align/test_decision_omnijev_align.py::TestStageDStructural -q -rs
# 7 passed, no skips (4 new cases + 3 existing structural tests).

pre-commit run --all-files
# All hooks passed.

git diff --check
# Passed.

No full model training or GPU validation was rerun for this local patch. The change addresses the shared CE term; additional calibration terms can still contribute their own gradients.

Checklist

• I have linked the related issue above, if there is one.
• I have reviewed and understand every changed line, and I can explain how the change works and why it is correct.
• I have run the relevant tests or reproductions locally, and included the commands and results above (if applicable).
• I am responsible for the content of this PR, regardless of how it was produced.

[z0o0ey]

Thanks for the revisions, you were right, and I've merged them.