I have a rather large ML framework that takes multiple conditional probability terms that are computed via classifiers/neural networks. This arbitrary loss function is computed via a function:

loss_value = arbitrary_loss(probability1, probability2, ..., P(Y|Z))

I wish to have an end-to-end framework that computes everything and trains everything together. So I do not want to have an independently trained classifier. Say at some point I develop some intermediate values (embeddings) Z from the input samples X. I wish to model the conditional probability P(Y|Z) via an MLP softmax layer.

This term P(Y|Z) is then estimated and plugged into the final loss which is the sum and product of other probabilities.

P(Y|Z) = MLP(input_Z) #probability given input Z over labels

My issue is that if I simply take the value of the softmax layer to estimate this probability and plug it in, at no point are the true labels taken into account for a supervised machine learning problem.

How can I fix this without modifying the final loss function?

TLDR: I need a probability term P(Y|Z) modeled via an MLP softmax layer to be used in a complex arbitrary loss function. How do i ensure this term is accurate via the true label values, so that it can be used in the final loss?