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.

inductive Spec.LossType : Type

Enumeration of supported loss families used by configuration records.

• mse : LossType
• mae : LossType
• huber : LossType
• crossEntropy : LossType
• hinge : LossType
• poisson : LossType
• cosineSimilarity : LossType
• logCosh : LossType

structure Spec.Loss : Type

Loss configuration record that names the selected loss family.

• lossType : LossType

  Selected loss family for this configuration.

def Spec.Loss.mse : Loss

Configuration selecting mean-squared-error loss.

def Spec.Loss.mae : Loss

Configuration selecting mean-absolute-error loss.

def Spec.Loss.huber : Loss

Configuration selecting Huber loss.

def Spec.Loss.crossEntropy : Loss

Cross-entropy loss configuration.

def Spec.Loss.hinge : Loss

Configuration selecting hinge loss.

def Spec.Loss.poisson : Loss

Poisson loss configuration.

def Spec.Loss.cosineSimilarity : Loss

Cosine similarity loss configuration.

def Spec.Loss.logCosh : Loss

Log-cosh loss configuration.

def Spec.toScalarSpec {α : Type} [TorchLean.Storage α] [Add α] [Zero α] {s : Shape} : TorchLean.Tensor α s → α

Sum all tensor elements into a single scalar.

def Spec.meanOver {α : Type} [Div α] [NatCast α] {s : Shape} (x : α) : α

Mean of a scalar that conceptually came from a tensor with shape s.

The denominator is TorchLean.Tensor.meanDenominator, the same totalized element count the tensor reductions use, so the loss layer and mean agree on what an empty tensor averages to.

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

Number of slices orthogonal to axis in a tensor shape.

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.

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

Totalized denominator for a mean over slices orthogonal to axis.

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.

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.

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.crossEntropyLogitsSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (axis : ℕ) [Shape.AxisInBounds axis s] (logits target : TorchLean.Tensor α s) : α

Cross-entropy on logits (stable log-softmax form).

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 for numerical stability.

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 reduction still divides by the number of slices, not by the total target weight.

def Spec.crossEntropyLogitsDerivSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (axis : ℕ) [Shape.AxisInBounds axis s] (logits target : TorchLean.Tensor α s) : 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.

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.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 {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (predicted target : TorchLean.Tensor α s) (epsilon : α := Context.defaultEpsilon) : α

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

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. $$

We use epsilon to avoid division by zero (similar to common "eps" handling in PyTorch code).

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.

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.