TorchJD is a library extending autograd to enable Jacobian descent with PyTorch. In can be used to train neural networks with multiple objectives. In particular, it supports multi-task learning, with a wide variety of aggregators from the literature. It also enables the instance-wise risk minimization paradigm. The full documentation is available at torchjd.org, with several usage examples.
TorchJD can be installed directly with pip:
pip install torchjdNote
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.
The main way to use TorchJD is to replace the usual call to loss.backward() by a call to
torchjd.backward or torchjd.mtl_backward, depending on the use-case.
The following example shows how to use TorchJD to train a multi-task model with Jacobian Descent, using UPGrad.
import torch
from torch.nn import Linear, MSELoss, ReLU, Sequential
from torch.optim import SGD
from torchjd import mtl_backward
from torchjd.aggregation import UPGrad
shared_module = Sequential(Linear(10, 5), ReLU(), Linear(5, 3), ReLU())
task1_module = Linear(3, 1)
task2_module = Linear(3, 1)
params = [
*shared_module.parameters(),
*task1_module.parameters(),
*task2_module.parameters(),
]
loss_fn = MSELoss()
optimizer = SGD(params, lr=0.1)
A = UPGrad()
inputs = torch.randn(8, 16, 10) # 8 batches of 16 random input vectors of length 10
task1_targets = torch.randn(8, 16, 1) # 8 batches of 16 targets for the first task
task2_targets = torch.randn(8, 16, 1) # 8 batches of 16 targets for the second task
for input, target1, target2 in zip(inputs, task1_targets, task2_targets):
features = shared_module(input)
output1 = task1_module(features)
output2 = task2_module(features)
loss1 = loss_fn(output1, target1)
loss2 = loss_fn(output2, target2)
optimizer.zero_grad()
mtl_backward(
losses=[loss1, loss2],
features=features,
tasks_params=[task1_module.parameters(), task2_module.parameters()],
shared_params=shared_module.parameters(),
A=A,
)
optimizer.step()Note
In this example, the Jacobian is only with respect to the shared parameters. The task-specific parameters are simply updated via the gradient of their task’s loss with respect to them.
More usage examples can be found here.
TorchJD provides many existing aggregators from the literature, listed in the following table.
The following example shows how to instantiate
UPGrad and aggregate a simple matrix J with it.
from torch import tensor
from torchjd.aggregation import UPGrad
A = UPGrad()
J = tensor([[-4., 1., 1.], [6., 1., 1.]])
A(J)
# Output: tensor([0.2929, 1.9004, 1.9004])Tip
When using TorchJD, you generally don't have to use aggregators directly. You simply instantiate
one and pass it to the backward function (torchjd.backward or torchjd.mtl_backward), which
will in turn apply it to the Jacobian matrix that it will compute.
Please read the Contribution page.
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}
}
