FairGrad

class torchjd.aggregation.FairGrad(alpha, max_iters=None)[source]

GramianWeightedAggregator using the step decision of Algorithm 1 of Fair Resource Allocation in Multi-Task Learning.

Parameters:

  • alpha (float) – The parameter controlling the type of fairness in the alpha-fairness formulation.

  • max_iters (int | None) – The maximum number of iterations of the optimization loop. If set to None, the default value of scipy.optimize.least_squares (100 * m) will be used.

Note

This aggregator requires optional dependencies. When they are not installed, instantiating it raises an ImportError with installation instructions. To install them, use pip install "torchjd[fairgrad]".

__call__(matrix, /)[source]

Computes the aggregation from the input matrix and applies all registered hooks.

Parameters:

matrix (Tensor) – The Jacobian to aggregate.

Return type:

Tensor

class torchjd.aggregation.FairGradWeighting(alpha, max_iters=None)[source]

Weighting [PSDMatrix] giving the weights of FairGrad, as defined in Equation 4 of Fair Resource Allocation in Multi-Task Learning.

Parameters:

  • alpha (float) – The parameter controlling the type of fairness in the alpha-fairness formulation.

  • max_iters (int | None) – The maximum number of iterations of the optimization loop. If set to None, the default value of scipy.optimize.least_squares (100 * m) will be used.

Note

This implementation was adapted from the official implementation.

Note

This aggregator requires optional dependencies. When they are not installed, instantiating it raises an ImportError with installation instructions. To install them, use pip install "torchjd[fairgrad]".

__call__(gramian, /)[source]

Computes the vector of weights from the input Gramian and applies all registered hooks.

Parameters:

gramian (Tensor) – The Gramian from which the weights must be extracted.

Return type:

Tensor