Classification

Choose binary classification for two classes, multiclass classification for an unordered set of classes, or ordinal classification for ordered integer grades.

Binary

class soma.models.heads.classification.BinaryClassificationHead(input_dim, num_classes, metrics=None, label_map=None, dropout=0.0)

Bases: TaskHead

Linear classification head for binary (two-class) tasks.

Parameters:

  • input_dim (int) – Dimension of the input representation.

  • num_classes (int) – Must be 2.

  • metrics (list[str] | None) – Metrics to compute. Empty list uses the default set for binary_classification: auroc, balanced_accuracy, auprc, f1.

  • dropout (float) – Probability of zeroing an element of the head’s input, applied before the linear layer. 0.0 (default) builds no dropout module.

Multiclass

class soma.models.heads.classification.MulticlassClassificationHead(input_dim, num_classes, metrics=None, label_map=None, dropout=0.0)

Bases: TaskHead

Linear classification head for multi-class tasks.

Parameters:

  • input_dim (int) – Dimension of the input representation.

  • num_classes (int) – Number of output classes (>= 2).

  • metrics (list[str] | None) – Metrics to compute. Empty list uses the default set for multiclass_classification: auroc_macro, balanced_accuracy, f1_macro.

  • dropout (float) – Probability of zeroing an element of the head’s input, applied before the linear layer. 0.0 (default) builds no dropout module.

Ordinal

class soma.models.heads.ordinal_classification.OrdinalClassificationHead(input_dim, num_classes, metrics=None, label_map=None, dropout=0.0)

Bases: TaskHead

Linear head for ordinal classification using MSE loss.

Treats ordered integer labels as continuous values during training (MSE loss) while producing integer predictions at inference by rounding the continuous output to the nearest class. Both the rounded prediction and the raw continuous score are reported.

Parameters:

  • input_dim (int) – Dimension of the input representation.

  • num_classes (int) – Number of ordinal classes (e.g. 6 for labels 0–5).

  • metrics (list[str] | None) – Metrics to compute. Empty list uses the default set for ordinal_classification: qwk, balanced_accuracy.

  • dropout (float) – Probability of zeroing an element of the head’s input, applied before the linear layer. 0.0 (default) builds no dropout module.

Metric compatibility

multiclass_classification accepts qwk as an opt-in metric when the class labels have an ordinal interpretation; the loss stays cross-entropy.

Benchmarks

  • EVA — frozen-tile-probe runs of these heads on the EVA patch-level datasets (bach, breakhis, crc, mhist, patch_camelyon).