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.

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) → torch.Tensor[source]

Forward pass with dynamic weight scheduling.

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:

Weighted total loss (scalar during training, tuple with log during eval).

Return type:

torch.Tensor or 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.