mtl_backward

torchjd.autojac.mtl_backward.mtl_backward(losses, features, aggregator, tasks_params=None, shared_params=None, retain_graph=False, parallel_chunk_size=None)

In the context of Multi-Task Learning (MTL), we often have a shared feature extractor followed by several task-specific heads. A loss can then be computed for each task.

This function computes the gradient of each task-specific loss with respect to its task-specific parameters and accumulates it in their .grad fields. Then, it computes the Jacobian of all losses with respect to the shared parameters, aggregates it and accumulates the result in their .grad fields.

Parameters:

  • losses (Sequence[Tensor]) – The task losses. The Jacobian matrix will have one row per loss.

  • features (Union[Sequence[Tensor], Tensor]) – The last shared representation used for all tasks, as given by the feature extractor. Should be non-empty.

  • aggregator (Aggregator) – Aggregator used to reduce the Jacobian into a vector.

  • tasks_params (Optional[Sequence[Iterable[Tensor]]]) – The parameters of each task-specific head. Their requires_grad flags must be set to True. If not provided, the parameters considered for each task will default to the leaf tensors that are in the computation graph of its loss, but that were not used to compute the features.

  • shared_params (Optional[Iterable[Tensor]]) – The parameters of the shared feature extractor. The Jacobian matrix will have one column for each value in these tensors. Their requires_grad flags must be set to True. If not provided, defaults to the leaf tensors that are in the computation graph of the features.

  • retain_graph (bool) – If False, the graph used to compute the grad will be freed. Defaults to False.

  • parallel_chunk_size (int | None) – The number of scalars to differentiate simultaneously in the backward pass. If set to None, all coordinates of tensors will be differentiated in parallel at once. If set to 1, all coordinates will be differentiated sequentially. A larger value results in faster differentiation, but also higher memory usage. Defaults to None. If parallel_chunk_size is not large enough to differentiate all tensors simultaneously, retain_graph has to be set to True.

Return type:

None

Example

A usage example of mtl_backward is provided in Multi-Task Learning (MTL).

Note

shared_params should contain no parameter in common with tasks_params. The different tasks may have some parameters in common. In this case, the sum of the gradients with respect to those parameters will be accumulated into their .grad fields.

Warning

mtl_backward relies on a usage of torch.vmap that is not compatible with compiled functions. The arguments of mtl_backward should thus not come from a compiled model. Check https://github.com/pytorch/pytorch/issues/138422 for the status of this issue.

Warning

Because of a limitation of torch.vmap, tensors in the computation graph of the features parameter should not have their retains_grad parameter set to True.