neuralop.losses.meta_losses.Relobralo

class neuralop.losses.meta_losses.Relobralo(params, num_losses, alpha=0.95, beta=0.99, tau=1.0, eps=1e-08, weights=None)[source]

Relative Loss Balancing with Random Lookback (ReLoBRaLo) algorithm for adaptive loss weighting.

ReLoBRaLo combines relative loss balancing with random lookback to provide more stable and effective multi-objective optimization. The algorithm uses both initial loss values and previous loss values to compute adaptive weights, with a random lookback mechanism. This random lookback mechanism introduces stochasticity in the weighting by sometimes reverting to initial loss values, encouraging robustness across loss scales over time.

Reference: “Bischof, R. and Kraus, M., 2021. Multi-Objective Loss Balancing for Physics-Informed Deep Learning. arXiv preprint arXiv:2110.09813.”

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.

alphafloat, optional

Controls the smoothing factor for the exponential moving average of the adaptive weights. High alpha gives more inertia to past weights (more stability). Low alpha allows faster adaptation. By default 0.95

betafloat, optional

Probability for the random lookback mechanism, i.e. probability of sampling from initial loss values instead of previous loss values during adaptive weight computation. By default 0.99

taufloat, optional

Temperature parameter that controls the sharpness of the exponential weighting. Higher values make the weights more uniform, by default 1.0

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 ReLoBRaLo 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 ReLoBRaLo.

The random component means that results may vary between runs, due to the Bernoulli sampling in the weight computation.

Notes

At step 0, all losses are simply summed with equal weights, and both initial and previous loss buffers are initialized with the current loss values.

The random lookback mechanism (controlled by beta) helps the algorithm escape local minima by occasionally using initial loss values instead of recent ones. This is particularly useful in physics-informed neural networks where different loss components may have different convergence characteristics.

The temperature parameter tau controls how aggressive the weighting is. Lower values make the algorithm more selective in which losses to prioritize, while higher values make the weights more uniform.

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

Weights and aggregates the losses using the ReLoBRaLo 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.

lmbda_ematorch.Tensor

The exponential moving average of adaptive weights for each loss component.

Notes

At step 0, all losses are summed with equal weights, and both initial and previous loss buffers are initialized with the current loss values.

For subsequent steps, the algorithm computes weights using both initial and previous loss values, with a random component that helps escape local minima. The final weights are computed as an exponential moving average.