autogram

The autogram package provides an engine to efficiently compute the Gramian of the Jacobian of a tensor of outputs (generally losses) with respect to some modules’ parameters. This Gramian contains all the inner products between pairs of gradients, and is thus a sufficient statistic for most weighting methods. The algorithm is formally defined in Section 6 of Jacobian Descent For Multi-Objective Optimization).

Due to computing the Gramian iteratively over the layers, without ever having to store the full Jacobian in memory, this method is much more memory-efficient than autojac, which makes it often much faster.

The list of Weightings compatible with autogram is:

• UPGradWeighting

• AlignedMTLWeighting

• CAGradWeighting

• ConstantWeighting

• DualProjWeighting

• IMTLGWeighting

• KrumWeighting

• MeanWeighting

• MGDAWeighting

• PCGradWeighting

• RandomWeighting

• SumWeighting