TorchJD is a library extending autograd to enable Jacobian descent with PyTorch. It 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.
Jacobian descent is an extension of gradient descent supporting the optimization of vector-valued functions. This algorithm can be used to train neural networks with multiple loss functions. In this context, JD iteratively updates the parameters of the model using the Jacobian matrix of the vector of losses (the matrix stacking each individual loss' gradient). For more details, please refer to Section 2.1 of the paper.
Averaging the losses and computing the gradient of the mean is mathematically equivalent to computing the Jacobian and averaging its rows. However, this approach has limitations. If two gradients are conflicting (they have a negative inner product), simply averaging them can result in an update vector that is conflicting with one of the two gradients. Averaging the losses and making a step of gradient descent can thus lead to an increase of one of the losses.
This is illustrated in the following picture, in which the two objectives' gradients
With Jacobian descent,
TorchJD can be installed directly with pip:
pip install torchjdSome aggregators may have additional dependencies. Please refer to the installation documentation for them.
There are two main ways to use TorchJD. The first one is to replace the usual call to
loss.backward() by a call to
torchjd.autojac.backward or
torchjd.autojac.mtl_backward, depending
on the use-case. This will compute the Jacobian of the vector of losses with respect to the model
parameters, and aggregate it with the specified
Aggregator.
Whenever you want to optimize the vector of per-sample losses, you should rather use the
torchjd.autogram.Engine. Instead of
computing the full Jacobian at once, it computes the Gramian of this Jacobian, layer by layer, in a
memory-efficient way. A vector of weights (one per element of the batch) can then be extracted from
this Gramian, using a
Weighting,
and used to combine the losses of the batch. Assuming each element of the batch is
processed independently from the others, this approach is equivalent to
torchjd.autojac.backward while being
generally much faster due to the lower memory usage.
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.autojac 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)
+ aggregator = 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()
- loss = loss1 + loss2
- loss.backward()
+ mtl_backward(losses=[loss1, loss2], features=features, aggregator=aggregator)
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.
The following example shows how to use TorchJD to minimize the vector of per-instance losses with Jacobian descent using UPGrad.
import torch
from torch.nn import Linear, MSELoss, ReLU, Sequential
from torch.optim import SGD
+ from torchjd.autogram import Engine
+ from torchjd.aggregation import UPGradWeighting
model = Sequential(Linear(10, 5), ReLU(), Linear(5, 3), ReLU(), Linear(3, 1), ReLU())
- loss_fn = MSELoss()
+ loss_fn = MSELoss(reduction="none")
optimizer = SGD(model.parameters(), lr=0.1)
+ weighting = UPGradWeighting()
+ engine = Engine(model.modules())
inputs = torch.randn(8, 16, 10) # 8 batches of 16 random input vectors of length 10
targets = torch.randn(8, 16) # 8 batches of 16 targets for the first task
for input, target in zip(inputs, targets):
output = model(input).squeeze(dim=1) # shape [16]
- loss = loss_fn(output, target) # shape [1]
+ losses = loss_fn(output, target) # shape [16]
optimizer.zero_grad()
- loss.backward()
+ gramian = engine.compute_gramian(losses) # shape: [16, 16]
+ weights = weighting(gramian) # shape: [16]
+ losses.backward(weights)
optimizer.step()More usage examples can be found here.
TorchJD provides many existing aggregators from the literature, listed in the following table.
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}
}