aggregation¶
When doing Jacobian descent, the Jacobian matrix has to be aggregated into a vector to store in the
.grad fields of the model parameters. The
Aggregator is responsible for these aggregations.
When using the autogram engine, we rather need to extract a vector
of weights from the Gramian of the Jacobian. The
Weighting is responsible for this.
Note
Most aggregators rely on computing the Gramian of the Jacobian, extracting a vector of weights
from this Gramian using a Weighting
[PSDMatrix], and then combining the rows of the Jacobian using these
weights. For all of them, we provide both the
Aggregator interface (to be used in autojac) and the
Weighting interface (to be used in autogram).
For the rest, we only provide the Aggregator
interface – they are not compatible with autogram.
Aggregators and
Weightings are callables that take a Jacobian matrix or a
Gramian matrix as inputs, respectively. The following example shows how to use UPGrad to either
aggregate a Jacobian (of shape [m, n], where m is the number of objectives and n is the
number of parameters), or obtain the weights from the Gramian of the Jacobian (of shape [m, m]).
>>> from torch import tensor
>>> from torchjd.aggregation import UPGrad, UPGradWeighting
>>>
>>> aggregator = UPGrad()
>>> jacobian = tensor([[-4.0, 1.0, 1.0], [6.0, 1.0, 1.0]])
>>> aggregation = aggregator(jacobian)
>>> aggregation
tensor([0.2929, 1.9004, 1.9004])
>>> weighting = UPGradWeighting()
>>> gramian = jacobian @ jacobian.T
>>> weights = weighting(gramian)
>>> weights
tensor([1.1109, 0.7894])
Abstract base classes¶
- class torchjd.aggregation.Aggregator[source]¶
Abstract base class for all aggregators. It has the role of aggregating matrices of dimension \(m \times n\) into row vectors of dimension \(n\).
- class torchjd.aggregation.WeightedAggregator(weighting)[source]¶
Aggregator that combines the rows of the input Jacobian matrix with weights given by applying a
Weighting[Matrix] to it.
- class torchjd.aggregation.GramianWeightedAggregator(gramian_weighting)[source]¶
WeightedAggregatorthat computes the gramian of the input Jacobian matrix before applying aWeighting[PSDMatrix] to it.
- class torchjd.aggregation.Weighting[source]¶
Abstract base class for all weighting methods. It has the role of extracting a vector of weights of dimension \(m\) from some statistic of a matrix of dimension \(m \times n\), generally its Gramian, of dimension \(m \times m\).