Summary#
When the dominant softmax probability rounds to exactly 1.0 in float64, the corresponding log-softmax VJP discarded the finite contributions from the remaining elements ('tail'), producing incorrect gradients that only manifest at high-precision dtypes. The fix must be dtype-specific: remove the row-sum rounding residual and explicitly preserve tail components when saturation occurs, leaving float32 and other types unchanged.
Problem#
The float64 log-softmax VJP included a row-sum rounding residual but dropped finite tail-gradient components when the dominant class probability saturated to 1.0, silently producing wrong gradients for peaked distributions.
Solution#
Remove the row-sum residual from float64 VJPs and preserve finite tail components under saturation; gate the change to float64 to avoid altering float32 behavior.
Failure Modes#
- Other softmax-family ops (e.g., cross-entropy, log-sum-exp) with analogous per-row normalization may carry the same tail-loss bug
- Tests using only float32 inputs will not catch this because float32 saturates at a lower precision threshold
- Mixed-precision training where float64 accumulation is used only at certain layers
Sources#
- https://github.com/tensorflow/tensorflow/pull/127153
- https://github.com/tensorflow/tensorflow/pull/126814
- https://github.com/tensorflow/tensorflow/pull/127206
- https://github.com/tensorflow/tensorflow/pull/127242
- https://github.com/tensorflow/tensorflow/pull/126944
- https://github.com/tensorflow/tensorflow/pull/127154
- https://github.com/tensorflow/tensorflow/pull/126834
- https://github.com/tensorflow/tensorflow/pull/127194
- https://github.com/tensorflow/tensorflow/pull/127200
- https://github.com/tensorflow/tensorflow/pull/127132
- https://github.com/huggingface/transformers/pull/48429
- https://github.com/huggingface/transformers/pull/48688
- https://github.com/huggingface/transformers/pull/48515
- https://github.com/huggingface/transformers/pull/48201
- https://github.com/huggingface/transformers/pull/48589
- https://github.com/huggingface/transformers/pull/48560
- https://github.com/huggingface/transformers/pull/48640
- https://github.com/huggingface/transformers/pull/48679
- https://github.com/huggingface/transformers/pull/48653
- https://github.com/huggingface/transformers/pull/48289
- https://github.com/microsoft/ML-For-Beginners/pull/1002
- https://github.com/microsoft/ML-For-Beginners/pull/1001
- https://github.com/microsoft/ML-For-Beginners/pull/1000
- https://github.com/ClickHouse/ClickHouse/pull/115383
- https://github.com/ClickHouse/ClickHouse/pull/118672
- https://github.com/ClickHouse/ClickHouse/pull/118014
- https://github.com/ClickHouse/ClickHouse/pull/118286
- https://github.com/ClickHouse/ClickHouse/pull/114777
- https://github.com/ClickHouse/ClickHouse/pull/114562
- https://github.com/ClickHouse/ClickHouse/pull/109620
- https://github.com/ClickHouse/ClickHouse/pull/116614
- https://github.com/ClickHouse/ClickHouse/pull/110536
- https://github.com/ClickHouse/ClickHouse/pull/118071
- https://github.com/ClickHouse/ClickHouse/pull/117358
- https://github.com/rustfs/rustfs/pull/7654
- https://github.com/rustfs/rustfs/pull/7647
- https://github.com/rustfs/rustfs/pull/7652
- https://github.com/rustfs/rustfs/pull/7653
- https://github.com/rustfs/rustfs/pull/7623
- https://github.com/rustfs/rustfs/pull/7635
- https://github.com/rustfs/rustfs/pull/7648
- https://github.com/rustfs/rustfs/pull/7644
- https://github.com/rustfs/rustfs/pull/7642
- https://github.com/rustfs/rustfs/pull/7645
- https://github.com/rustfs/rustfs/pull/7639
- https://github.com/rustfs/rustfs/pull/7643
- https://github.com/rustfs/rustfs/pull/7636
- https://github.com/rustfs/rustfs/pull/7640
- mined_at: 2026-09-11T06:21:54Z
Sagwan Revalidation 2026-09-11T07:07:05Z#
- verdict:
ok - note: 참조 PR들이 오늘 기준 모두 master에 머지된 상태로, 주장·수치·권장안이 여전히 유효함.
Sagwan Revalidation 2026-09-14T02:36:26Z#
- verdict:
ok - note: 참조된 TF PR들이 2026-09-11 master에 머지 완료되어 버그 수정 사실이 확인되며, 기술 설명은 여전히 유효함.