Applicability of UPGrad in Parallel Multi-Model Setup with Shared Auxiliary Loss

[ayitime]

Hi, thank you for your great work!

I have a slightly unconventional network setup and would like to ask whether UPGrad would be applicable in this case.

Specifically, I have two parallel models, model1 and model2, each processing its own dataset independently and producing predictions p1 and p2. The respective losses L1 and L2 are computed separately. Additionally, I compute a third loss L3 based on both p1 and p2.

Currently, I update model1 using the combined loss L1 + L3, and model2 using L2 + L3. However, I am unsure how to appropriately balance the weights between L1 and L3, and between L2 and L3.

Would it be feasible to apply UPGrad in this setup to adaptively balance these losses?

Looking forward to your insights — thank you in advance!

[PierreQuinton]

Yes that should be doable with two calls to torchjd.backward by setting appropriate inputs. Something like

import torch

from torchjd import backward
from torchjd.aggregation import UPGrad

# Define model1 and model2
# Compute L1, L2 and L3

backward([L1, L3], UPGrad(), inputs=model1.parameters())
backward([L2, L3], UPGrad(), inputs=model2.parameters())

It is possible that you have to use retain_graph=True for the first one in case some parameters are shared between the models.

[ayitime]

Thank you very much for your valuable advice! I will try UPGrad right away and hope that it will work in this situation.

[ValerianRey]

Hi! That's an interesting setup!

I would go for two different calls to torchjd.backward: one with L1 and L3 and one with L2 and L3.
The tricky thing is that L3 depends on both models' parameters. So you will need to specify the inputs to ensure that the first call only affects the first model's parameters, and that the second call only affects the second model parameters. Lastly, I think you will need to use retain_graph=True in the first call so that the computation graph for L3 is preserved (because it will be needed for the second call).

It would look something like:

# Do the usual stuff (forward pass, compute the losses, zero_grad, etc)
torchjd.backward([loss1, loss3], aggregator=UPGrad(), inputs=model1.parameters(), retain_graph=True)
optimizer1.step()
torchjd.backward([loss2, loss3], aggregator=UPGrad(), inputs=model2.parameters(), retain_graph=False)
optimizer2.step()
# ...

Here, I assumed you have one optimizer for model1 and one for model2, but this can be easily adapted if you just have one optimizer for both:

# Do the usual stuff (forward pass, compute the losses, zero_grad, etc)
torchjd.backward([loss1, loss3], aggregator=UPGrad(), inputs=model1.parameters(), retain_graph=True)
torchjd.backward([loss2, loss3], aggregator=UPGrad(), inputs=model2.parameters(), retain_graph=False)
optimizer.step()
# ...

I hope this helps!

[ayitime]

Thank you very much for your valuable advice! I will try UPGrad right away and hope that it will work in this situation.

[ValerianRey]

@PierreQuinton you got it first haha

[PierreQuinton]

Well, your was better!