NeuralTangentKernelWeighting

Module for Neural Tangent Kernel Class

class NeuralTangentKernelWeighting(update_every_n_epochs=1, alpha=0.5)

Bases: BaseWeighting

The Neural Tangent Kernel (NTK) weighting strategy.

This weighting scheme dynamically adjusts the contribution of each loss term during training by leveraging gradient information with respect to the model parameters. For each loss component, the norm of its gradient is computed and used to derive relative importance weights. The resulting weights are smoothed over time using an exponential moving average controlled by the parameter alpha.

See also

Original reference: Wang, Sifan, Xinling Yu, and Paris Perdikaris. When and why PINNs fail to train: A neural tangent kernel perspective. Journal of Computational Physics 449 (2022): 110768. DOI: 10.1016.

Initialization of the NeuralTangentKernelWeighting 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.

• alpha (float) – The parameter controlling the exponential moving average of the weights. It must be in the range [0, 1], where a value of 0.0 means that only the current gradient norms are used to compute the weights, and a value of 1.0 means that only the last saved weights are used. Default is 0.5.

Raises:

• ValueError – If alpha is not a float.

• ValueError – If alpha is not between 0.0 and 1.0 (inclusive).

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