Solution of multiple losses on mse_loss, bce_loss, KLD, and loss_contra.

[linjing-lab]

Here is a loss function composed of multiple loss combinations from one spatial transcriptomics task of spatial domain identification. Can losses calculated by different functions be used for multi-task learning?

import torch
import torch.nn.functional as F
from torch.optim import Adam

+ from torchjd import mtl_backward
+ from torchjd.aggregation import UPGrad

  optimizer = Adam(params, lr=0.1)
+ aggregator = UPGrad()

for ...
      mse_loss = F.mse_loss(decoded, x)
      bce_loss = F.binary_cross_entropy_with_logits(preds, labels)
      KLD = -0.5 / n_nodes * torch.mean(torch.sum(1 + 2*logvar -mu.pow(2)-logvar.exp().pow(2), 1))
      loss_contra = contrastive_loss(...)
      optimizer.zero_grad()
-     loss = mse_loss + bce_loss + KLD + loss_contra
-     loss.backward()
+     mtl_backward(losses=[mse_loss, bce_loss, KLD, loss_contra], features=features, aggregator=aggregator)
      optimizer.step()

[PierreQuinton]

Certainly! However, I am not sure about the usage of mtl_backward (rather than backward) unless you do have an explicit feature that is then fed into several parametrized models each giving you one of the losses (they should be independent too). If you can pinpoint this feature tensor, then use mtl_backward, otherwise simply use backward(losses=[mse_loss, bce_loss, KLD, loss_contra], aggregator=aggregator)

[linjing-lab]

Thank you. My understanding of Jacobian involves both mathematical theory and optimization research and development, so I am very interested in the torchjd project. Features of spatial transcriptomics are the result of feature compression, which is consistent with the mtl_backward example. Using loss=mse_loss+bce_loss+KLD+loss_conta in training with autograd can converge to good results. I would like to know the generalization of the prediction after training with mtl_backward.

[PierreQuinton]

Then it is definitely worth trying!