I am not very familiar with the paper but computing once the shared part and then using the same features for each classification head will work. You then sum all the losses you want and call backward on it !
Is there some complications I’m missing here?