TorchJD: Training with multiple losses in PyTorch [P]

Hi everyone! I wanted to share some recent progress on TorchJD that might be useful to the machine learning community.

When training models with multiple losses (multiple tasks, constraints, auxiliary losses, regularization terms, etc.), you typically have two options:

• Scalarization: Various ways to combine those losses into a single loss (e.g. average them or combine them with trainable weights); then you can do gradient descent on it.
• Jacobian descent: Compute the Jacobian of the vector of losses (i.e. one gradient per loss), and aggregate it into an update vector that will decrease each individual loss (rather than just the average loss). There are many ways to do this aggregation step.

Scalarization methods are generally cheaper in memory, but in some cases there is so much disagreement between your objectives that it's better to use a Jacobian descent method. In any case, thanks to our amazing new contributors, we've now finally implemented most existing methods of the literature from both categories into our library TorchJD, so that you can try anything in just a few line changes!

Recently, TorchJD has been accepted into the PyTorch ecosystem, and we're trying to make it become the go-to library for training with multiple losses. If you'd like to help build the future of the project, come join us on Discord (link can be found in the readme of the repo). New ideas, contributions, bug reports, experiments, and any form of feedback are all welcome. We have many ideas on how to make all this even more efficient, and we will need help for that.

If you want to support us, a star on GitHub also helps a lot!