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} [Context α] {s : Shape} : Tensor α s → α

Sum all tensor elements into a single scalar.

def Spec.meanDenom (s : Shape) : ℕ

Denominator for totalized mean reductions over a shape.

For nonempty shapes this is the real element count. For empty shapes the mathematical mean is undefined; TorchLean's scalar-polymorphic spec layer is total, so it uses denominator 1 and the empty sum contributes 0.

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

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

def Spec.lastAxisSliceCount : Shape → ℕ

Number of independent last-axis slices in a tensor shape.

Classification losses sum over the innermost class axis and apply mean over the remaining axes. Thus a class vector has one slice, a matrix of shape (batch, classes) has batch slices, and higher-rank inputs use the product of every axis except the last. Scalars are treated as one single-coordinate slice.

def Spec.lastAxisMeanDenom (s : Shape) : ℕ

Totalized denominator for a mean over last-axis slices.

def Spec.meanOverLastAxisSlices {α : Type} [Context α] {s : Shape} (x : α) : α

Divide a last-axis-summed classification loss by its number of independent slices.

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

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

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

Derivative of mse_spec w.r.t. predicted.

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

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

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

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

def Spec.huberSpec {α : Type} [Context α] {s : Shape} (predicted target : 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} [Context α] {s : Shape} (predicted target : Tensor α s) (delta : α := 1) : Tensor α s

Derivative of huber_spec w.r.t. predicted.

def Spec.crossEntropySpec {α : Type} [Context α] {s : Shape} (predicted target : Tensor α s) (epsilon : α := Numbers.epsilon) : α

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 is the last (class) axis and r ranges over all remaining axes. 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} [Context α] {s : Shape} (predicted target : Tensor α s) (epsilon : α := Numbers.epsilon) : Tensor α s

Derivative of cross_entropy_spec w.r.t. predicted.

def Spec.crossEntropyLogitsSpec {α : Type} [Context α] {s : Shape} (logits target : 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 last axis c contains classes and r ranges over the remaining axes. 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.

Note: this spec assumes each last-axis target slice is a probability distribution (sums to 1), as in one-hot or label-smoothed targets.

def Spec.crossEntropyLogitsDerivSpec {α : Type} [Context α] {s : Shape} (logits target : Tensor α s) : Tensor α s

Derivative of cross_entropy_logits_spec w.r.t. logits.

def Spec.hingeSpec {α : Type} [Context α] {s : Shape} (predicted target : 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} [Context α] {s : Shape} (predicted target : Tensor α s) : Tensor α s

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

def Spec.poissonSpec {α : Type} [Context α] {s : Shape} (predicted target : 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} [Context α] {s : Shape} (predicted target : Tensor α s) : Tensor α s

Derivative of poisson_spec w.r.t. predicted.

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

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

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

Derivative of cosine_similarity_spec 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^2\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} [Context α] {s : Shape} (predicted target : Tensor α s) : α

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

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

Derivative of log_cosh_spec w.r.t. predicted.

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

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 : α := Numbers.epsilon) : α

Selected derivative of binary_cross_entropy_spec 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} [Context α] {s : Shape} (predicted target : Tensor α s) (epsilon : α := Numbers.epsilon) : α

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

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

Derivative of binary_cross_entropy_tensor_spec w.r.t. predicted.