BugZoo: numerically stable losses and domain-sensitive ops

TensorFuzz found rare-input failures that ordinary test sets missed, including broken loss functions that produced NaN and quantization/full-precision disagreements:

• Odena, Olsson, Andersen, and Goodfellow, “TensorFuzz: Debugging Neural Networks with Coverage-Guided Fuzzing”, ICML 2019. https://proceedings.mlr.press/v97/odena19a.html

The numerical-bugs study by Wang et al. gives a complementary source of real PyTorch/TensorFlow failures: invalid domains for log, sqrt, division, exp, and related math APIs:

• Wang et al., “An Empirical Study on Numerical Bugs in Deep Learning Programs”, ASE NIER 2022. https://doi.org/10.1145/3551349.3559561

TorchLean cannot repair an arbitrary hand-written unstable loss after the fact. The design instead gives stable primitives and domain-aware variants a named place in the spec. For example, crossEntropyLogitsSpec is the logits API users should reach for: it is defined through logSoftmaxSpec, rather than through a fragile softmax followed by log. Likewise, safedivSpec and safeDivOp make epsilon-protected division explicit in the graph.

Bug-shaped PyTorch sketches:

# Unstable: softmax can round to 0, then log(0) gives -inf and the loss can become NaN.
probs = torch.softmax(logits, dim=-1)
loss = -(target * torch.log(probs)).sum()

# Safer: PyTorch's cross_entropy/log_softmax path.
loss = -(target * torch.log_softmax(logits, dim=-1)).sum()

TorchLean equivalent:

Spec.crossEntropyLogitsSpec axis logits target

For division/domain bugs:

# Risky when denom can be zero or denormal-sized.
y = x / denom

TorchLean makes the protected variant visible:

TorchLean.Tensor.safedivSpec x denom

The checked definition and theorem below expose the shifted denominator. Adding epsilon does not prevent division by zero when denom = -epsilon, nor guarantee finite floating-point results. Callers must establish the domain conditions for their inputs.

theorem NN.Examples.BugZoo.StableLoss.crossEntropyLogits_uses_logSoftmax {s : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis s] {α : Type} [TorchLean.Storage α] [Context α] (logits target : TorchLean.Tensor α s) : Spec.crossEntropyLogitsSpec axis logits target = have logp := Activation.logSoftmaxSpec axis logits; have total := (target.mulSpec logp).sumSpec; Spec.meanOverAxisSlices axis (-total)

The logits cross-entropy spec is exactly the log-softmax form.

It weights log probabilities by targets, then averages over non-class dimensions.

theorem NN.Examples.BugZoo.StableLoss.crossEntropyLogitsDeriv_uses_target_mass {s : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis s] {α : Type} [TorchLean.Storage α] [Context α] (logits target : TorchLean.Tensor α s) : Spec.crossEntropyLogitsDerivSpec axis logits target = (((Activation.softmaxSpec axis logits).mulSpec (TorchLean.Tensor.broadcastAfterSum s axis (TorchLean.Tensor.reduceDim (fun {sliceShape : Spec.Shape} => TorchLean.Tensor.sumSpec) axis target))).subSpec target).scaleSpec (1 / ↑(Spec.axisMeanDenom s axis))

The logits-loss gradient keeps the total target weight in each class slice.

For weights t, the formula is softmax(logits) * sum(t) - t, averaged over the non-class dimensions. Probability targets have total weight one and give the familiar softmax - target. An all-zero target gives a zero loss and zero gradient. The same formula also handles weighted targets without normalizing them or changing the loss.

This theorem unfolds the supplied gradient spec. The reduction drops the selected class axis, and broadcastAfterSum restores that same axis before multiplying by the probabilities.

theorem NN.Examples.BugZoo.StableLoss.crossEntropyProbabilities_clips_before_log {s : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis s] {α : Type} [TorchLean.Storage α] [Context α] (predicted target : TorchLean.Tensor α s) (epsilon : α) : Spec.crossEntropySpec axis predicted target epsilon = have clamp01 := fun (x : α) => have x := if x > epsilon then x else epsilon; if x < 1 - epsilon then x else 1 - epsilon; have q := TorchLean.Tensor.mapSpec clamp01 predicted; have logq := q.logSpec; have total := (target.mulSpec logq).sumSpec; Spec.meanOverAxisSlices axis (-total)

Probability-space cross entropy clips the predicted probability before taking log.

Pass logits to crossEntropyLogitsSpec; use this form for probabilities.

theorem NN.Examples.BugZoo.StableLoss.safeDivSpec_unfold {s : Spec.Shape} {α : Type} [TorchLean.Storage α] [Context α] (x y : TorchLean.Tensor α s) : x.safedivSpec y = TorchLean.Tensor.map2Spec (fun (a b : α) => a / (b + Context.defaultEpsilon)) x y

Epsilon-protected division is a separate named tensor operation, not a hidden rewrite.