TorchJD is a library enabling Jacobian descent with PyTorch, to train neural networks with multiple objectives. It is based on the theory from Jacobian Descent For Multi-Objective Optimization and several other related publications. The main purpose is to jointly optimize multiple objectives without combining them into a single scalar loss. Using TorchJD to minimize a vector of objectives comes at the expense of some computational overhead compared to directly using PyTorch to minimize a scalarization of the objectives. However when the objectives are conflicting, this can be essential. To get started, check out our basic usage example.

Important

This library is currently in an early development stage. The API is subject to significant changes in future versions.

TorchJD is open-source, under MIT License. The source code is available on GitHub.