matcha.nn.losses

Loss functions for regression, classification, and multitask learning.

Attributes

LossRegistry

Classes

`BCEFocalLoss`	Focal loss implementation using binary cross entropy with logits.
`Poly1BCELoss`	Polynomial expansion of the binary cross entropy loss, which can lead to better
`MultitaskLoss`	Implementation of a customizable multitask loss. The same loss is
`MultiLoss`	Implementation of a customizable multi-loss function that supports
`BoundedLoss`	Loss function for handling bounded regression
`MSELoss`	Mean squared error loss (wraps `torch.nn.MSELoss`).
`L1Loss`	Mean absolute error loss (wraps `torch.nn.L1Loss`).
`HuberLoss`	Huber loss (wraps `torch.nn.HuberLoss`).
`SmoothL1Loss`	Smooth L1 loss (wraps `torch.nn.SmoothL1Loss`).
`BoundedMSELoss`	`BoundedLoss` with MSE as the inner loss.
`BoundedMAELoss`	`BoundedLoss` with MAE as the inner loss.
`BoundedHuberLoss`	`BoundedLoss` with Huber as the inner loss.
`BoundedSmoothL1Loss`	`BoundedLoss` with Smooth L1 as the inner loss.
`BCELoss`	Binary cross-entropy with logits (wraps `torch.nn.BCEWithLogitsLoss`).
`CrossEntropyLoss`	Cross-entropy loss (wraps `torch.nn.CrossEntropyLoss`).
`WeightedBCELoss`	Weighted binary cross entropy loss with logits.
`GradNormLoss`	Implementation of GradNorm for adaptive loss balancing in multitask learning.
`DropoutLoss`	Wrap a per-element loss and randomly mask a fraction of entries each step.
`DropoutMSELoss`	`DropoutLoss` with MSE as the inner loss.
`DropoutMAELoss`	`DropoutLoss` with MAE as the inner loss.
`DropoutHuberLoss`	`DropoutLoss` with Huber as the inner loss.
`DropoutSmoothL1Loss`	`DropoutLoss` with Smooth L1 as the inner loss.
`DropoutBCELoss`	`DropoutLoss` with BCE-with-logits as the inner loss.
`DropoutFocalBCELoss`	`DropoutLoss` with focal BCE as the inner loss.
`DropoutPoly1BCELoss`	`DropoutLoss` with Poly1 BCE as the inner loss.
`DropoutWeightedBCELoss`	`DropoutLoss` with weighted BCE as the inner loss.

Module Contents

matcha.nn.losses.LossRegistry

class matcha.nn.losses.BCEFocalLoss(gamma=2, alpha=None, reduction='mean')[source]

Bases: torch.nn.Module

Focal loss implementation using binary cross entropy with logits. Suitable for binary classification with class imbalance. Reference: https://arxiv.org/abs/1708.02002

gamma = 2

alpha = None

reduction = 'mean'

eps = 1e-07

forward(inputs, targets) → torch.Tensor[source]

Parameters:

inputs (torch.Tensor) – Raw logits.
targets (torch.Tensor) – Binary target labels.

Returns:

Focal loss value.

Return type:

torch.Tensor

class matcha.nn.losses.Poly1BCELoss(epsilon: float = 1.0, reduction: str = 'mean')[source]

Bases: torch.nn.Module

Polynomial expansion of the binary cross entropy loss, which can lead to better classification performance if epsilon is tuned. Suitable for binary classification. Reference: https://arxiv.org/abs/2204.12511

epsilon = 1.0

reduction = 'mean'

eps = 1e-07

forward(inputs, targets) → torch.Tensor[source]

Parameters:

inputs (torch.Tensor) – Raw logits.
targets (torch.Tensor) – Binary target labels.

Returns:

Poly1 BCE loss value.

Return type:

torch.Tensor

class matcha.nn.losses.MultitaskLoss(loss_fn: str = 'mse', loss_args: dict = {})[source]

Bases: torch.nn.Module

Implementation of a customizable multitask loss. The same loss is used for all tasks.

loss_fn is the name of the loss to broadcast, loss_args are its arguments.

loss

forward(outputs: torch.Tensor, targets: torch.Tensor) → torch.Tensor[source]

Parameters:

outputs (torch.Tensor) – Predictions of shape (batch, num_tasks).
targets (torch.Tensor) – Targets of shape (batch, num_tasks); NaN marks missing.

Returns:

Scalar loss averaged across valid entries and tasks.

Return type:

torch.Tensor

class matcha.nn.losses.MultiLoss(loss_configs: list)[source]

Bases: torch.nn.Module

Implementation of a customizable multi-loss function that supports dynamic weight scheduling during training.

Each loss configuration is a dictionary with: - loss_fn: name of the loss function - loss_args: arguments for the loss function - task_map: tuple/list indicating which columns this loss applies to (start, end) - init_w: initial weight at T=0 - final_w: final weight at T=end - T: total epochs to transition from init_w to final_w - warmup: epochs to keep init_w fixed before starting transition

loss_configs

losses

forward(outputs: torch.Tensor, targets: torch.Tensor, T_current: int = 0) → tuple[torch.Tensor, dict][source]

Forward pass with dynamic weight scheduling.

Always returns (total_loss, loss_log) so training and validation share one contract. See issue #41: the earlier training-mode / eval-mode split was a hotfix with no in-repo reproducer and left BaseClassicModel / MLPModel callers unpacking a 0-d tensor. Callers that do not need the per-task log can ignore the second element.

Parameters:

outputs (torch.Tensor) – Model predictions.
targets (torch.Tensor) – Ground truth targets; NaN marks missing entries.
T_current (int) – Current training epoch for weight scheduling.

Returns:

Tuple of (total_loss, loss_log) where loss_log maps each task name to a {"loss": float, "weight": float} record.

Return type:

tuple[torch.Tensor, dict]

class matcha.nn.losses.BoundedLoss(loss_fn: str = 'mse', **kwargs)[source]

Bases: torch.nn.Module

Loss function for handling bounded regression

Allows the use of bound information on the readout (e.g. IC50 < x), so that the model is not penalized when it predicts e.g. y_pred < x. Implementation is based on: https://chemprop.readthedocs.io/en/latest/_modules/chemprop/nn/metrics.html#BoundedMixin

loss

forward(outputs: torch.Tensor, targets: torch.Tensor) → torch.Tensor[source]

Parameters:

outputs (torch.Tensor) – Predictions.
targets (torch.Tensor) – Targets with bound info in the last dimension. Shape (batch, [num_tasks,] 2) where [..., 0] is the value and [..., 1] encodes the bound type (-1 = less-than, 1 = greater-than).

Returns:

Loss computed only on non-masked predictions.

Return type:

torch.Tensor

class matcha.nn.losses.MSELoss(size_average=None, reduce=None, reduction: str = 'mean')[source]

Bases: torch.nn.MSELoss

Mean squared error loss (wraps torch.nn.MSELoss).

class matcha.nn.losses.L1Loss(size_average=None, reduce=None, reduction: str = 'mean')[source]

Bases: torch.nn.L1Loss

Mean absolute error loss (wraps torch.nn.L1Loss).

class matcha.nn.losses.HuberLoss(reduction: str = 'mean', delta: float = 1.0)[source]

Bases: torch.nn.HuberLoss

Huber loss (wraps torch.nn.HuberLoss).

class matcha.nn.losses.SmoothL1Loss(size_average=None, reduce=None, reduction: str = 'mean', beta: float = 1.0)[source]

Bases: torch.nn.SmoothL1Loss

Smooth L1 loss (wraps torch.nn.SmoothL1Loss).

class matcha.nn.losses.BoundedMSELoss(**kwargs)[source]

Bases: BoundedLoss

BoundedLoss with MSE as the inner loss.

class matcha.nn.losses.BoundedMAELoss(**kwargs)[source]

Bases: BoundedLoss

BoundedLoss with MAE as the inner loss.

class matcha.nn.losses.BoundedHuberLoss(**kwargs)[source]

Bases: BoundedLoss

BoundedLoss with Huber as the inner loss.

class matcha.nn.losses.BoundedSmoothL1Loss(**kwargs)[source]

Bases: BoundedLoss

BoundedLoss with Smooth L1 as the inner loss.

class matcha.nn.losses.BCELoss(weight: torch.Tensor | None = None, size_average=None, reduce=None, reduction: str = 'mean', pos_weight: torch.Tensor | None = None)[source]

Bases: torch.nn.BCEWithLogitsLoss

Binary cross-entropy with logits (wraps torch.nn.BCEWithLogitsLoss).

class matcha.nn.losses.CrossEntropyLoss(weight: torch.Tensor | None = None, size_average=None, ignore_index: int = -100, reduce=None, reduction: str = 'mean', label_smoothing: float = 0.0)[source]

Bases: torch.nn.CrossEntropyLoss

Cross-entropy loss (wraps torch.nn.CrossEntropyLoss).

class matcha.nn.losses.WeightedBCELoss(w1: float = 0.5, reduction: str = 'mean')[source]

Bases: torch.nn.Module

Weighted binary cross entropy loss with logits.

Applies per-class weights to handle class imbalance in binary classification. The user specifies the weight for the positive (minority) class; the weight for the negative class is computed so that w0 + w1 = 1.

Parameters:

w1 (float) – Weight for class 1 (positive / minority class). Must be in (0, 1).
reduction (str) – Reduction mode: 'mean', 'sum', or 'none'.

w1 = 0.5

w0 = 0.5

reduction = 'mean'

forward(inputs: torch.Tensor, targets: torch.Tensor) → torch.Tensor[source]

Parameters:

inputs (torch.Tensor) – Raw logits.
targets (torch.Tensor) – Binary target labels.

Returns:

Weighted BCE loss value.

Return type:

torch.Tensor

class matcha.nn.losses.GradNormLoss(loss_fn: str = 'mse', loss_args: dict = {}, num_endpoints: int = 1, weight_lr: float = 0.025)[source]

Bases: torch.nn.Module

Implementation of GradNorm for adaptive loss balancing in multitask learning.

GradNorm automatically balances training by dynamically tuning gradient magnitudes. It adjusts task weights to ensure that all tasks train at similar rates. The weight updates are handled internally — no separate optimizer needed.

Reference: https://arxiv.org/abs/1711.02257

Parameters:

loss_fn (str) – Name of the loss function to use for all tasks (e.g., "mse").
loss_args (dict) – Arguments to pass to the loss function constructor.
num_endpoints (int) – Number of tasks.
weight_lr (float) – Learning rate for updating task weights internally.

Example:

loss_fn = GradNormLoss(loss_fn="mse", num_endpoints=3)

# Training loop (no changes needed):
loss = loss_fn(outputs, targets, shared_layer=model.backbone[-1])
optimizer.zero_grad()
loss.backward()
optimizer.step()

ALPHA = 1.5

loss

num_endpoints = 1

weight_lr = 0.025

initial_losses = None

forward(outputs: torch.Tensor, targets: torch.Tensor, shared_layer: torch.nn.Module = None) → torch.Tensor[source]

Forward pass computing weighted multitask loss.

Parameters:

outputs (torch.Tensor) – Predictions of shape (batch_size, num_endpoints).
targets (torch.Tensor) – Targets of shape (batch_size, num_endpoints).
shared_layer (torch.nn.Module or None) – The last shared layer of the network. Required during training for GradNorm weight updates.

Returns:

Weighted sum of task losses.

Return type:

torch.Tensor

reset_initial_losses() → None[source]: Reset the initial losses for a fresh start of GradNorm tracking.

class matcha.nn.losses.DropoutLoss(loss_fn: str = 'mse', dropout: float = 0.0, seed: int | None = None, reduction: str = 'mean', **kwargs)[source]

Bases: torch.nn.Module

Wrap a per-element loss and randomly mask a fraction of entries each step.

Intended as a regularizer for multi-endpoint pretraining (e.g. predicting many molecular descriptors at once): randomly dropping a fraction of labels from the loss on every forward pass discourages overfitting to any single endpoint.

Reference: https://github.com/JacksonBurns/how-to-train-your-chemeleon/blob/main/pretraining/random_dropout_mse.py

The inner loss is instantiated with reduction="none" so masking happens before reduction. The dropout mask is resampled every forward and composes with the existing NaN mask (NaN targets are always excluded, as in MultitaskLoss). In eval() mode the wrapper reduces to a plain NaN-masked mean of the inner loss, regardless of dropout.

Parameters:

loss_fn (str) – Alias of the inner per-element loss (resolved via LossRegistry). Defaults to "mse".
dropout (float) – Fraction of non-NaN entries to drop from the loss on each training forward pass. Must satisfy 0.0 <= dropout < 1.0.
seed (int or None) – Optional integer seed for a private torch.Generator. When set, the mask trajectory is reproducible across runs and does not perturb the ambient torch RNG. When None (default), masks are drawn from the ambient RNG (like torch.nn.Dropout).
reduction (str) – Reduction applied after NaN and dropout masking: "mean" (default) returns a scalar averaged over kept entries — the standalone contract; "sum" returns the scalar sum of the masked per-element loss; "none" returns the per-element loss tensor with dropped and NaN entries zeroed, so MultiLoss and MultitaskLoss (which construct inner losses with reduction="none") receive the shape they expect. The inner loss is always instantiated with reduction="none"; this parameter controls only the outer reduction after masking.
kwargs – Extra keyword arguments forwarded to the inner loss constructor.

dropout

reduction = 'mean'

loss

forward(outputs: torch.Tensor, targets: torch.Tensor) → torch.Tensor[source]

Parameters:

outputs (torch.Tensor) – Predictions of shape (batch, num_tasks).
targets (torch.Tensor) – Targets of shape (batch, num_tasks); NaN marks missing entries and is always excluded from the loss.

Returns:

Masked per-element loss when reduction="none" (same shape as outputs, with dropped and NaN entries zeroed). Scalar sum for "sum". Scalar mean over kept entries for "mean" (the default standalone behavior).

Return type: