torchjd 0.2.0

Library for Jacobian Descent with PyTorch.

TorchJD

TorchJD is a library enabling Jacobian descent with PyTorch, to train neural networks with multiple objectives. In particular, it can be used for multi-task learning, with a wide variety of algorithms from the literature. It also enables the instance-wise risk minimization paradigm, as proposed in Jacobian Descent For Multi-Objective Optimization. The full documentation is available at torchjd.org, with several usage examples.

Installation

TorchJD can be installed directly with pip:

pip install torchjd

Compatibility

TorchJD requires python 3.10, 3.11 or 3.12. It is only compatible with recent versions of PyTorch (>= 2.0). For more information, read the dependencies in pyproject.toml.

Contribution

Please read the Contribution page.

Citation

If you use TorchJD for your research, please cite:

@article{jacobian_descent,
  title={Jacobian Descent For Multi-Objective Optimization},
  author={Quinton, Pierre and Rey, Valérian},
  journal={arXiv preprint arXiv:2406.16232},
  year={2024}
}