Tasks

Select a task with the learner’s head leaf (learner.params.head.name). The head maps a representation to predictions and defines the loss and the compatible metrics. The representation comes directly from a frozen encoder or from a trained aggregator or decoder; see Modeling for the paths that produce it.

Task Zoo

Config name	Loss	Default metrics	When to use
binary_classification	Cross-entropy	auroc, balanced_accuracy, auprc, f1	Two-class labels; see Classification
multiclass_classification	Cross-entropy	auroc_macro, balanced_accuracy, f1_macro	Two or more classes; see Classification Use binary_classification when the problem is strictly binary
branch_aware_classification	Cross-entropy	auroc_macro, balanced_accuracy, f1_macro	The multiclass head over clam_mb’s per-class branches; soma uses it in place of multiclass_classification when the aggregator is clam_mb
ordinal_classification	MSE	qwk, balanced_accuracy	Ordered integer grades; see Classification
regression	MSE	mae, r2	Continuous targets; see Regression
survival	Discrete-time hazard NLL	c_index	Time-to-event with right censoring, over time bins; see Survival
cox_survival	Cox partial likelihood	c_index	Time-to-event with right censoring, one risk score; see Survival
segmentation	CE + soft-Dice	mean_dice, mean_iou	Dense per-pixel classification (the dense learner); see Segmentation
detection	Foreground-weighted MSE	mean_f1	Cell / nucleus point detection (the dense learner); see Detection

Head dropout

The classification, ordinal-classification, regression and survival heads take an optional dropout probability, applied to the head’s input — the aggregated bag representation, or the frozen embedding itself when no aggregator leaves the head as the only trainable component:

learner:
  params:
    head:
      name: binary_classification
      params:
        dropout: 0.2

The default is 0.0. Aggregator dropout, set in the aggregator leaf’s params, is independent.

Regression

The regression head maps a feature vector to continuous predictions and trains with MSE. With one output it reads the value target, predicts predicted_value, and reports mae and r2 by default.

num_outputs: G regresses G values at once, such as a gene panel. The head then reads the targets value_0 to value_<G-1>, predicts columns of the same names, and reports mean_pearson: the Pearson correlation of each output over the samples, averaged over the outputs. The HEST scores a closed-form Ridge+PCA probe with the same columns and metric.

class soma.models.heads.regression.RegressionHead(input_dim, num_outputs=1, metrics=None, dropout=0.0)

Bases: TaskHead

Linear regression head for one or many outputs.

Parameters:

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

• num_outputs (int) – Number of regressed values G. With one output the head reads the cohort’s value target and predicts predicted_value; with G > 1 it reads value_0 .. value_{G-1} and predicts the same-named columns (a gene panel, say), scored by mean_pearson.

• metrics (list[str] | None) – Metrics to compute. Empty uses the regression defaults (mae, r2) for one output and mean_pearson for several.

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

Task-head interface

soma.components("heads") lists the registered task heads with their params and defaults. To bring a custom loss, metric reduction or post-processing, subclass a built-in head and register it; see A custom component through the registry.

class soma.models.heads.base.TaskHead(*args, **kwargs)

Bases: ABC, Module

Abstract base class for task heads.

Maps an aggregated slide representation to task-specific predictions.

Subclasses must implement forward, compute_loss, postprocess, compute_metrics, and extract_targets, and should set target_dtypes. Values decided by the cohort (class counts, label maps, survival bins) come from the infer hook the head is registered with (ADR 0018). prediction_columns / to_columns declare what soma.predict() writes; the base implementation covers soma’s set families.

The head owns its target contract: target_dtypes declares the keys it consumes (and their tensor dtypes), extract_targets maps a SampleRecord to a dict of raw target values, and compute_loss / compute_metrics receive a dict[str, Tensor] keyed by those names.