Repository navigation
Solution of multiple losses on mse_loss, bce_loss, KLD, and loss_contra. #421
Description
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()
Activity
PierreQuinton commentedon Sep 14, 2025
Certainly! However, I am not sure about the usage of
mtl_backward(rather thanbackward) unless you do have an explicitfeaturethat is then fed into several parametrized models each giving you one of the losses (they should be independent too). If you can pinpoint thisfeaturetensor, then usemtl_backward, otherwise simply usebackward(losses=[mse_loss, bce_loss, KLD, loss_contra], aggregator=aggregator)linjing-lab commentedon Sep 14, 2025
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_contain training with autograd can converge to good results. I would like to know the generalization of the prediction after training with mtl_backward.PierreQuinton commentedon Sep 14, 2025
Then it is definitely worth trying!