SelfAdaptiveWeighting

Module for Self-Adaptive Weighting class.

class SelfAdaptiveWeighting(update_every_n_epochs=1)

Bases: BaseWeighting

The self-adaptive weighting strategy based on gradient norm balancing.

This scheme dynamically adjusts the weights assigned to each loss term by computing the norm of their gradients with respect to the model parameters. The resulting weights are chosen to counterbalance disparities in gradient magnitudes, promoting a more uniform contribution of all loss components during optimization.

In practice, loss terms with smaller gradient norms are assigned larger weights, while those with larger gradients are down-weighted. This helps mitigate training imbalance and prevents dominance of specific loss terms.

See also

Original reference: Wang, S., Sankaran, S., Stinis., P., Perdikaris, P. (2025). Simulating Three-dimensional Turbulence with Physics-informed Neural Networks. DOI: arXiv preprint arXiv:2507.08972.

Initialization of the SelfAdaptiveWeighting class.

Parameters:

update_every_n_epochs (int) – The number of training epochs between weight updates. If set to 1, the weights are updated at every epoch. This parameter is ignored by static weighting schemes. Default is 1.

update_weights(losses)

Update the weights based on the current losses.

This method defines how the weighting strategy adapts over time. It is responsible for computing and storing updated weights that will be used during aggregation.

Parameters:

losses (dict) – The mapping from loss names to loss tensors.

Returns:

The updated weights.

Return type:

dict