torchcell.models.equivariant_cell_graph_transformer.MaskedMultitaskLoss

class MaskedMultitaskLoss(head_weights: dict[str, float] | None = None, loss_fn: str = 'mse', dist_heads: dict[str, Module] | None = None)[source]

Bases: Module

Masked multitask loss over config-selected heads plus graph regularization.

Each head’s loss is masked to the genotypes in the batch that actually carry that phenotype (sparse supervision): absent modalities contribute zero. The existing graph-regularization attention loss term is added UNCHANGED.

A head may carry a DistHead (dist_heads[name]), in which case its loss is the head’s proper scoring rule (Gaussian CRPS / pinball) over the predicted DISTRIBUTION parameters. Heads with no entry keep the plain elementwise point loss (mse/l1) – this is what gene_interaction and the synthetic dry-run use, so the pre-distributional behavior is preserved exactly when no dist_heads are passed.

forward(head_outputs: dict[str, Tensor], targets: dict[str, Tensor], masks: dict[str, Tensor] | None = None, graph_reg_loss: Tensor | None = None, feature_masks: dict[str, Tensor] | None = None) → tuple[Tensor, dict[str, Tensor]][source]

Compute the masked multitask loss.

Parameters:

  • head_outputs – {head_name: prediction [B, …]}.

  • targets – {head_name: target [B, …]} for supervised heads.

  • masks – {head_name: bool [B]} selecting supervised genotypes per head. A missing mask means all rows are supervised for that head.

  • graph_reg_loss – Optional scalar graph-regularization loss, added as-is.

  • feature_masks – {head_name: bool [B, F]} selecting which FEATURES to score. Used by the masked-label objective so a step is scored only on the genes still HIDDEN at that step; a revealed gene is model input, and scoring it would reward copying. A missing entry scores every feature, which is the pre-existing behaviour.

Returns:

scalar summed loss. per_head_loss: {head_name: detached scalar loss} for logging.

Return type:

total_loss