neuralop.losses.meta_losses.SoftAdapt

class neuralop.losses.meta_losses.SoftAdapt(params, num_losses, eps=1e-08, weights=None)[source]

SoftAdapt algorithm for adaptive loss weighting and aggregation.

SoftAdapt automatically adjusts the weights of multiple loss components based on their relative magnitudes and rates of change. The algorithm uses exponential weighting to give higher importance to losses that are not decreasing or are relatively larger, based on their ratio to previous losses, helping to balance the training of different objectives.

Reference: “Heydari, A.A., Thompson, C.A. and Mehmood, A., 2019. Softadapt: Techniques for adaptive loss weighting of neural networks with multi-part loss functions. arXiv preprint arXiv: 1912.12355.”

Parameters:

paramsList[torch.Tensor]

List of model parameters used to determine the device for tensor operations.

num_lossesint

Number of loss components that will be aggregated.

epsfloat, optional

Small constant added to denominators to prevent division by zero. Should be positive and small (e.g., 1e-8), by default 1e-8

weightsOptional[Dict[str, float]], optional

Optional static weights for each loss component, by default None

Methods

forward(losses, step)

Weights and aggregates the losses using the original variant of the SoftAdapt algorithm.

Warning

The algorithm assumes that all loss components should be minimized. If some losses should be maximized, they should be negated before passing to SoftAdapt.

Notes

At step 0, all losses are simply summed with equal weights. For subsequent steps, the algorithm computes adaptive weights based on the ratio of current to previous loss values, with higher weights given to losses that are increasing or have higher relative magnitudes.

forward(losses: Dict[str, Tensor], step: int) → Tensor[source]

Weights and aggregates the losses using the original variant of the SoftAdapt algorithm.

Parameters:

lossesDict[str, torch.Tensor]

Dictionary of loss components to be aggregated. For instance, the keys could be ‘data_loss’ and ‘physics_loss’.

stepint

Current optimization step. Used to determine whether to initialize the algorithm (step 0) or apply adaptive weighting (step > 0).

Returns:

losstorch.Tensor

The aggregated loss value.

lmbdatorch.Tensor

The computed adaptive weights for each loss component.

Notes

At step 0, all losses are summed with equal weights. For subsequent steps, the algorithm computes adaptive weights based on the ratio of current to previous loss values.