I am building a neural network (using tensorflow/keras) that attempts to classify into one of 5 categories (0, 1, 2, 3, 4). I have been using sparse categorical cross-entropy as my loss function. This works as is expected, but my data is "somewhat" ordinal. When a label is 3 a prediction of 2 is much better than a prediction of 1 (which is better than 0, etc.). Basically, the loss function would use the predicted probabilities of the tensorflow model rather than the predicted classification ([.1, .2, .2, .4, .1] rather than 3).
Is there any loss function that takes advantage of this?
I could imagine something like:
where 𝑁 is the number of samples, 𝑦𝑖,𝑗=𝑖 is the probability of predicting the correct class, 𝑦𝑖,|𝑖−𝑗|=1 is the probability of predicting one class off, etc.
In short, there is a penalty for the probability of prediction one unit away, a larger penalty for two units away, and so on. However, I am unsure of how to implement this myself in tensorflow as a custom loss function because most loss functions I have seen do not utilize the predicted probabilities. Any help pointing me to a loss function that accomplishes this or help on a way to implement something like the above would be appreciated. As well as any dialogue on alternatives.