Loss functions (spec layer)

This file defines a small collection of common losses (and their gradients) in a way that is:

• shape-generic: a loss takes Tensor α s and reduces it to a scalar α,
• explicit about reduction: most losses here are "mean over all elements",
• easy to line up with PyTorch terminology when you read training code.

In PyTorch you'll often see two layers:

• a low-level, elementwise loss (for example, Huber loss),
• plus a reduction (mean or sum).

TorchLean's spec layer mirrors that idea: most definitions are written as an elementwise formula followed by a global mean over the shape.

The teacher-logit gradient uses logit differences, centering, and a residual correction as in Barron, Gnanapragasam and Tschannen, "A Fix for Softmax and Cross-Entropy" (https://arxiv.org/abs/2610.31642). Its near-field expression retains small gradients that would be lost by subtracting two separately rounded probability vectors.

def Spec.weightedDifference {α : Type} [Zero α] [Sub α] [Mul α] [BEq α] (weight left right : α) : α

Weight a difference, distributing the multiplication when subtraction overflows. Ordinary finite differences keep their evaluation order. For opposite-sign finite operands, weighting each operand first can keep the result representable without halving small tangents.

def Spec.axisMeanDenom (s : Shape) (axis : ℕ) [Shape.AxisInBounds axis s] : ℕ

Totalized denominator for a mean over slices orthogonal to axis.

Classification losses sum along the selected class dimension and average over every other dimension. Thus a class vector has one slice, a matrix of shape (batch, classes) with class dimension 1 has batch slices, and a tensor may instead place its class dimension anywhere in the shape. The shared meanDenominator supplies the empty-slice convention.

def Spec.meanOverAxisSlices {α : Type} [Div α] [NatCast α] {s : Shape} (axis : ℕ) [Shape.AxisInBounds axis s] (x : α) : α

Divide a classification loss by the number of slices orthogonal to axis.

def Spec.mseSpec {α : Type} [TorchLean.Storage α] [Add α] [Sub α] [Mul α] [Div α] [Zero α] [NatCast α] {s : Shape} (predicted target : TorchLean.Tensor α s) : α

Mean squared error: average of $(\mathtt{predicted}-\mathtt{target})^2$.

def Spec.mseDerivSpec {α : Type} [TorchLean.Storage α] [Add α] [Sub α] [Mul α] [Div α] [One α] [NatCast α] {s : Shape} (predicted target : TorchLean.Tensor α s) : TorchLean.Tensor α s

Derivative of mseSpec with respect to predicted.

def Spec.maeSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (predicted target : TorchLean.Tensor α s) : α

Mean absolute error: average of |predicted - target|.

def Spec.maeDerivSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (predicted target : TorchLean.Tensor α s) : TorchLean.Tensor α s

Derivative of maeSpec w.r.t. predicted (subgradient via sign).

def Spec.huberSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (predicted target : TorchLean.Tensor α s) (delta : α := 1) : α

Huber loss with transition parameter delta.

Elementwise, for residual $d=\mathtt{pred}-\mathtt{target}$:

• if $\lvert d\rvert<\delta$: $\tfrac12d^2$
• otherwise: $\delta(\lvert d\rvert-\tfrac12\delta)$

Then we take a mean over all elements.

This is PyTorch's HuberLoss convention. It differs from SmoothL1Loss by a factor of delta. The Huber interpretation requires delta > 0.

The quadratic branch halves before squaring. If a finite prediction and target have an overflowing difference, the linear branch instead weights both operands before subtraction. Nonfinite inputs and genuinely unrepresentable losses or reductions can still overflow.

def Spec.huberDerivSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (predicted target : TorchLean.Tensor α s) (delta : α := 1) : TorchLean.Tensor α s

Derivative of huberSpec w.r.t. predicted.

def Spec.crossEntropySpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (axis : ℕ) [Shape.AxisInBounds axis s] (predicted target : TorchLean.Tensor α s) (epsilon : α := Context.defaultEpsilon) : α

Cross-entropy between distributions (probabilities).

This is closest to PyTorch when you already have probabilities q (e.g. after a softmax) and a probability target p (e.g. one-hot or label-smoothed), and you want:

$$ \operatorname{CE}(p,q) =-\operatorname{mean}_r\sum_c p_{rc}\log q_{rc}, $$

where c ranges along the selected class dimension and r ranges over all remaining dimensions. A lone class vector is one distribution and is not divided by its number of classes.

Each predicted probability is clamped to [epsilon, 1 - epsilon] before the logarithm. A one-hot prediction that exactly matches a one-hot target therefore has loss -log(1 - epsilon), which is small but not 0.

PyTorch's F.cross_entropy typically takes logits and does log_softmax + NLLLoss; that is a different API surface than this "probabilities in, scalar out" spec.

def Spec.crossEntropyDerivSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (axis : ℕ) [Shape.AxisInBounds axis s] (predicted target : TorchLean.Tensor α s) (epsilon : α := Context.defaultEpsilon) : TorchLean.Tensor α s

Derivative of crossEntropySpec w.r.t. predicted.

def Spec.CrossEntropy.vectorTerms {α : Type} [TorchLean.Storage α] [Context α] {n : ℕ} (logits target : TorchLean.Tensor α [n]) : TorchLean.Tensor α [n]

Weighted log-probabilities for one class vector. Keep the original logits available to reconstruct a weighted term when the unweighted log-probability overflows.

def Spec.CrossEntropy.innermostTerms {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} : TorchLean.Tensor α s → TorchLean.Tensor α s → TorchLean.Tensor α s

Apply the weighted-log-probability calculation independently to each final-axis vector.

def Spec.CrossEntropy.terms {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (axis : ℕ) (logits target : TorchLean.Tensor α s) : TorchLean.Tensor α s

Weighted log-probabilities along axis, restored to the original tensor layout before the loss reduction. This preserves the existing accumulation order for non-trailing axes.

def Spec.crossEntropyLogitsSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (axis : ℕ) [Shape.AxisInBounds axis s] (logits target : TorchLean.Tensor α s) (mean : Bool := true) : α

Cross-entropy on logits, retaining the original logits when a weighted term needs rescaling.

This matches the common PyTorch decomposition:

$$ \operatorname{cross\_entropy}(\mathtt{logits},\mathtt{target}) =-\operatorname{mean}_r\sum_c \mathtt{target}_{rc}\operatorname{logsoftmax}(\mathtt{logits})_{rc}, $$

where the dimension selected by axis contains classes and r ranges over the remaining dimensions. This is PyTorch's reduction="mean" convention for one-hot or soft distribution targets.

Unlike crossEntropySpec, this takes logits and uses Activation.logSoftmaxSpec. Ordinary finite log-probabilities are multiplied as usual. When a log-probability is nonfinite, its weighted term is reconstructed from the shifted normalizer and a rescaled logit difference. Thus zero or small weights need not multiply an overflowing unweighted difference. No target mask discards target derivatives. Nonfinite inputs or overflowing reductions can still yield NaN.

Probability targets usually sum to one along axis, as in one-hot or label-smoothed targets. The same formula also accepts arbitrary target weights. Their sum scales the contribution of that slice; the mean still divides by the number of slices, not by the total target weight. Set mean := false to sum the contributions without averaging.

def Spec.crossEntropyLogitsDerivSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (axis : ℕ) [Shape.AxisInBounds axis s] (logits target : TorchLean.Tensor α s) (mean : Bool := true) : TorchLean.Tensor α s

Derivative of crossEntropyLogitsSpec with respect to the logits.

For each class slice, differentiating the log-normalizer contributes softmax(logits) * sum(target). The direct logit term contributes -target. Keeping that slice sum makes the derivative agree with the loss for weighted, unnormalized, and zero targets. For a probability target, its sum is one and the familiar softmax(logits) - target follows. Evaluation uses the shared log-softmax VJP rather than forming this difference directly; the pivot's complementary probability is retained when its exponential rounds to one.

The reduction drops the class axis and then restores that exact axis. Ordinary trailing-axis broadcasting would mix slices when classes occupy an outer or middle dimension. Empty class axes use the same zero-sum convention as the loss.

def Spec.CrossEntropy.targetGradient {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (axis : ℕ) [Shape.AxisInBounds axis s] (logits : TorchLean.Tensor α s) (seed : α := 1) (mean : Bool := true) : TorchLean.Tensor α s

Target-weight derivative of logits cross-entropy, multiplied by an upstream scalar seed. Reconstructing the weighted log-probability before reduction also handles zero or small seeds without first materializing a possibly overflowing unweighted target derivative.

def Spec.distillSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (axis : ℕ) [Shape.AxisInBounds axis s] (student teacher : TorchLean.Tensor α s) : α

Cross-entropy against a fixed teacher's logits, averaged over class slices.

Unlike probability targets, teacher logits retain the information needed to compute small student–teacher gradient differences. Zero teacher weights contribute zero even when a student log-probability overflows to negative infinity. The teacher is held fixed during differentiation.

def Spec.distillDerivSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (axis : ℕ) [Shape.AxisInBounds axis s] (student teacher : TorchLean.Tensor α s) : TorchLean.Tensor α s

Student gradient of distillSpec, evaluated from logit differences.

Center the difference under the teacher distribution, then compute its log-normalizer with expm1 and log1p. Near equal logits this avoids subtracting rounded probabilities. Coordinates far from the teacher use the ordinary probability difference instead, avoiding 0 * ∞ in a small teacher weight times a large exponential. The final correction removes reduction residuals. The mean uses the same class-slice denominator as crossEntropyLogitsSpec.

def Spec.hingeSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (predicted target : TorchLean.Tensor α s) : α

Hinge loss (binary margin loss), elementwise then mean-reduced:

$\operatorname{hinge}(x,y)=\operatorname{mean}_i\max(0,1-y_i x_i)$.

This matches the usual SVM-style hinge loss. (PyTorch exposes similar behavior via margin-style losses such as HingeEmbeddingLoss / MultiMarginLoss, but the exact signature differs.)

def Spec.hingeDerivSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (predicted target : TorchLean.Tensor α s) : TorchLean.Tensor α s

Derivative/subgradient of hingeSpec w.r.t. predicted.

def Spec.poissonSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (predicted target : TorchLean.Tensor α s) : α

Poisson negative log-likelihood (log-input form), elementwise then mean-reduced:

If predicted represents log(rate) and target is a nonnegative count, then (up to an additive constant that does not affect gradients):

$\operatorname{loss}_i=\exp(\mathtt{pred}_i)-\mathtt{target}_i\mathtt{pred}_i$.

This corresponds to PyTorch's PoissonNLLLoss(log_input=true, full=false) at the math level.

def Spec.poissonDerivSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (predicted target : TorchLean.Tensor α s) : TorchLean.Tensor α s

Derivative of poissonSpec w.r.t. predicted.

def Spec.cosineSimilaritySpec.normFactors {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (x : TorchLean.Tensor α s) (epsilon : α) : α × α

Scale and scaled, epsilon-clipped norm. For positive epsilon over the reals, both are positive and their product is the clipped L2 norm. Numerical callers divide by them separately.

def Spec.cosineSimilaritySpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (predicted target : TorchLean.Tensor α s) (epsilon : α := Context.defaultEpsilon) : α

Cosine similarity loss: 1 - cos(predicted, target) (reduced-to-scalar).

Normalize each tensor before the dot product, using a shared scale before squaring. This avoids forming raw dot products or products of large norms. The norm floor remains epsilon, not an additive regularizer. Pass a positive, finite epsilon representable in the scalar type. Nonfinite inputs and an overflowing scaled reduction can still produce nonfinite results.

def Spec.cosineSimilarityDerivSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (predicted target : TorchLean.Tensor α s) (epsilon : α := Context.defaultEpsilon) : TorchLean.Tensor α s

Derivative of cosineSimilaritySpec w.r.t. predicted.

If $\cos=(p\mathbin{\cdot}t)/(\lVert p\rVert\lVert t\rVert)$ and $\operatorname{loss}=1-\cos$, then (for nonzero norms):

$$ \frac{\partial\operatorname{loss}}{\partial p} =\frac{p\mathbin{\cdot}t}{\lVert p\rVert^3\lVert t\rVert}p -\frac{1}{\lVert p\rVert\lVert t\rVert}t. $$

The shared scaled norms avoid forming the fourth-degree denominator in this expression. When the predicted norm is at or below epsilon, its clipped denominator is held constant and the radial term is zero. The same positive, finite epsilon requirement applies as above.

def Spec.logCoshSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (predicted target : TorchLean.Tensor α s) : α

Log-cosh loss (reduced-to-scalar): log(cosh(predicted - target)).

def Spec.logCoshDerivSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (predicted target : TorchLean.Tensor α s) : TorchLean.Tensor α s

Derivative of logCoshSpec w.r.t. predicted.

def Spec.binaryCrossEntropySpec {α : Type} [Context α] (predicted target : α) (epsilon : α := Context.defaultEpsilon) : α

Binary cross-entropy on scalars (probabilities), with clipping to avoid log(0).

This matches the core formula behind PyTorch's BCELoss when predicted is already a probability (not a logit):

$$ \operatorname{BCE}(p,y) =-\left(y\log p+(1-y)\log(1-p)\right). $$

Assumption: target is in [0, 1]. We do not clip the target; we only clip predicted. The safeguard must remain positive in the chosen scalar type. At the upper clamp, retain it directly as the complementary probability: subtracting 1 - epsilon from one could otherwise produce zero when the upper clamp rounds to one. Inside the interval, log1p (-p) retains the negative-class loss when 1 - p would round to one.

def Spec.binaryCrossEntropyDerivSpec {α : Type} [Context α] (predicted target : α) (epsilon : α := Context.defaultEpsilon) : α

Selected derivative of binaryCrossEntropySpec w.r.t. predicted.

The clipped forward function is not differentiable at epsilon or 1 - epsilon; this definition chooses zero at those two kinks and on the clipped exterior branches.

def Spec.binaryCrossEntropyTensorSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (predicted target : TorchLean.Tensor α s) (epsilon : α := Context.defaultEpsilon) : α

Tensor BCE (probabilities), elementwise then mean-reduced.

def Spec.binaryCrossEntropyTensorDerivSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (predicted target : TorchLean.Tensor α s) (epsilon : α := Context.defaultEpsilon) : TorchLean.Tensor α s

Derivative of binaryCrossEntropyTensorSpec w.r.t. predicted.