Skip to content

Metric Losses API¤

Loss functions for differentiable metric optimization.

DifferentiableAUROC¤

diffbio.losses.metric_losses.DifferentiableAUROC ¤

DifferentiableAUROC(temperature: float = 1.0)

Bases: Module

Differentiable approximation of the Area Under the ROC Curve.

This is a smooth training surrogate. For exact AUROC evaluation use :class:ExactAUROC, which delegates to calibrax's trapezoidal-rule implementation.

Approximates AUROC by replacing the hard indicator in the Wilcoxon-Mann-Whitney statistic with a sigmoid function, making it fully differentiable and JIT-compatible.

For every (positive, negative) pair the hard AUROC checks whether the positive score exceeds the negative score. This module replaces that indicator with sigmoid((pos - neg) / temperature), yielding a smooth surrogate whose gradient can drive optimisation.

Parameters:

Name Type Description Default
temperature float

Controls sharpness of the sigmoid approximation. Lower values approach the hard indicator; higher values give smoother gradients. Default 1.0.

 1.0
Example 
auroc_loss = DifferentiableAUROC(temperature=1.0)
predictions = jnp.array([0.9, 0.8, 0.1, 0.2])
labels = jnp.array([1.0, 1.0, 0.0, 0.0])
value = auroc_loss(predictions, labels)

Parameters:

Name Type Description Default
temperature float

Sigmoid temperature. Lower values produce a sharper (closer to hard) approximation.

 1.0

__call__ ¤

__call__(
    predictions: Float[Array, " n"],
    labels: Float[Array, " n"],
) -> Float[Array, ""]

Compute the differentiable AUROC approximation.

Parameters:

Name Type Description Default
predictions Float[Array, ' n']

Model output scores, shape (n,).

 required
labels Float[Array, ' n']

Binary ground-truth labels (0 or 1), shape (n,).

 required

Returns:

Type Description
Float[Array, '']

Scalar AUROC approximation in [0, 1].

ExactAUROC¤

diffbio.losses.metric_losses.ExactAUROC ¤

ExactAUROC()

Bases: Module

Exact AUROC metric using calibrax's trapezoidal-rule implementation.

Delegates to :func:calibrax.metrics.functional.classification.roc_auc to compute the exact Area Under the ROC Curve via threshold-sweep and the trapezoidal rule.

Use this for evaluation; use :class:DifferentiableAUROC for training (the sorting-based trapezoidal rule has zero gradients w.r.t. predictions because argsort is not differentiable).

Example 
exact = ExactAUROC()
predictions = jnp.array([0.9, 0.8, 0.1, 0.2])
labels = jnp.array([1.0, 1.0, 0.0, 0.0])
value = exact(predictions, labels)  # 1.0

__call__ ¤

__call__(
    predictions: Float[Array, " n"],
    labels: Float[Array, " n"],
) -> Float[Array, ""]

Compute the exact AUROC via calibrax.

Parameters:

Name Type Description Default
predictions Float[Array, ' n']

Model output scores, shape (n,).

 required
labels Float[Array, ' n']

Binary ground-truth labels (0 or 1), shape (n,).

 required

Returns:

Type Description
Float[Array, '']

Scalar AUROC in [0, 1].