我有一个基本的LSTM模型，通过输出6个类别的概率分布来对输入数据进行分类。没什么特别复杂的。

现在，模型运作正常并给我一个输出，我会取最大值来对输入进行分类。然而，我认为我们可以做得更多。

对我来说，与其看到实际类别，不如看到LSTM输出的概率分布更加有用；比如说

[ 0.0528042 ,  0.11904617,  0.27744624,  0.37874526,  0.13942425,         0.03253399]

因为这些信息可以告诉我第二好的猜测，第三好的猜测等等，以及LSTM对其标记为某个类别的信心程度。

有趣的是，我的类别之间非常相关；比如说，如果类别1和2对应于“非常大值”和“大值”，我知道我的值很大，而类别3和4对应于“小值”和“非常小值”。

有没有什么方法可以利用类别之间的密切关系来获得更好的输出？当然，我不希望简单地只有两个输出（“大”和“小”），因为概率分布对我来说非常重要。

回答：

需要实现一个自定义损失函数来编码类间关系。

假设你的6个类别是按顺序排列的（比如，["极大", "很大", "大", "小", "很小", "极小"]），一个合适的损失函数可能是1D-Wasserstein距离（也称为地球移动距离）。

一维EMD有一个闭合形式的公式。例如，你可以尝试实现这篇论文中描述的内容。