Survival

Two heads model time-to-event with right censoring; the head selects the loss:

• survival — discrete-time survival modeling. The continuous time axis is split into num_bins bins; the head emits one hazard logit per bin and trains with the sigmoid-hazard negative log-likelihood (the Gensheimer / HIPT NLLSurvLoss formulation).

• cox_survival — continuous-time CoxPH. The head emits a single risk scalar and trains with the Breslow partial-likelihood loss. The risk set must contain several samples, and the loader uses an event-balanced sampler so every batch/window contains at least one event (trainer.accumulate_grad_batches must be 1). Two modes, switched by the head’s cox_window param:

  • Padded mode (cox_window unset / 1): the risk set is the batch (batch_size >= 2). Works on single-embedding slide/patient features (no aggregator) or on MIL bags (any of abmil/transmil/mean_pool), where variable-length bags are padded and masking keeps the result exact.

  • Accumulation mode (cox_window >= 2): for large variable-size MIL bags. batch_size is pinned to 1 and an aggregator is required. The trainer forwards cox_window bags un-padded and computes one Cox loss and one optimizer step per window.

Both losses rank with Harrell’s C-index via scikit-survival.

Survival datasets carry these target columns; the head’s targets param maps a target to a differently named column, for example targets: {time: label}:

Column	Meaning
time	Continuous time-to-event (uncensored) or time-to-last-follow-up (censored).
event	1 if the event was observed, 0 if right-censored.
bin	Index of the discrete time bin containing time, for every sample including censored ones. Required for the survival head only.

Compute the bins yourself (e.g. qcut on the uncensored times). Set the head’s num_bins param to fix the head width: any subset of indices in [0, num_bins) is then valid, including empty intervals, and indices are never renumbered. Without an explicit width, observed bins must be contiguous from zero and num_bins is inferred as max(bin) + 1.

Survival samples are slides or patients (cohort.collapse: patient_id); a tile-level run is rejected. For patient runs, all slides of a patient must agree on the survival target. Survival MIL supports abmil, transmil, mean_pool, and hierarchical hipt features; other aggregators are rejected.

Task heads

class soma.models.heads.survival.SurvivalHead(input_dim, num_bins, alpha=0.15, metrics=None, dropout=0.0)

Bases: TaskHead

Discrete-time survival head (sigmoid-hazard NLL).

Parameters:

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

• num_bins (int) – Number of discrete time bins (model output width).

• alpha (float) – Up-weighting of the uncensored NLL term (HIPT default 0.15).

• metrics (list[str] | None) – Metrics to compute. Empty list uses the survival default (c_index).

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

class soma.models.heads.survival.CoxSurvivalHead(input_dim, ties='breslow', min_events_per_window=1, cox_window=1, metrics=None, dropout=0.0)

Bases: TaskHead

Continuous-time CoxPH survival head (Breslow partial likelihood).

Emits a single risk scalar per sample; the Cox partial likelihood is computed over a risk set of multiple samples. Unlike SurvivalHead, there is no binning and risk is the raw model output — higher means higher hazard / shorter survival, which is what concordance_index_censored expects. (Reusing the NLL head’s -sum(surv) derivation here would double-invert the C-index sign.)

The cox_survival head. Two training modes, switched by its cox_window param (carried here as accumulation_window):

• Padded mode (cox_window unset / 1): the risk set is the batch. Single-embedding slide/patient features (no aggregator, batch_size >= 2) or padded MIL bags (batch_size >= 2, masking handles padding).

• Accumulation mode (cox_window >= 2): for large variable-size MIL bags. batch_size is pinned to 1; the trainer forwards cox_window bags un-padded, keeps their risk scalars graph-connected, and computes one Cox loss over the window. accumulates_predictions signals this to the trainer.

In both modes the pipeline builds the training loader with an event-balanced sampler (needs_event_balanced_batches) and the tune loss is computed over the whole cohort (full_cohort_eval_loss). Config validation enforces the mode constraints.

Parameters:

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

• ties (str) – Tie-handling method. Only "breslow" is implemented.

• min_events_per_window (int) – Minimum events the event-balanced sampler guarantees per batch / window (>= 1).

• cox_window (int) – Prediction-accumulation window size. 1 (default) selects padded mode; >= 2 selects accumulation mode with this risk-set size.

• metrics (list[str] | None) – Metrics to compute. Empty list uses the survival default (c_index).

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