Hi all,
I have a multiclass classification problem and my network structure is a bit complex than usual.  In a nutshell, I have 2 types of sets for labels. The ground-truth is always one label from one of the sets. (think like, labels from 0 to C are from one set and labels from C+1 to N are from another set)
My network calculates 2 diferent logits for each set with different architectures. What I do is:
I convert these logits to probability distributions via softmax and now I have 2 probability distributions one for each target set: p1 and p2.
I have a learnable scalar s(in range [0,1], which weights the learnt probability distributions.
I calculate p1 = s * p1 and p2 = (1-s) * p2
Basically, depending on the case, I generate the label from one of the sets. At the end I concetanete both distributions and it is still a valid probability distribution since s is in range [0,1]. Lastly, I want to compute the loss on that.
I cannot use cross entropy loss because it requires raw logits. Therefore I have to use NLLLoss, However, it requires log probabilities. Since I already applied softmax, I have to apply log manually but this leads to instabilities and numerical issues as mentioned in the documentation.
But I need to apply the softmax in the first place because if i apply log_softmax then I do not have valid probability distributions and I can not weight them in a reasonable way with parameter s.
Any suggestions ?