We present a latent variable model for classification that provides a novel probabilistic interpretation of neural network softmax classifiers. We derive a variational objective to train the model, analogous to the evidence lower bound (ELBO) used to train variational auto-encoders, that generalises the softmax cross-entropy loss. Treating inputs to the softmax layer as samples of a latent variable, our abstracted perspective reveals a potential inconsistency between their anticipated distribution, required for accurate label predictions, and their empirical distribution found in practice. We augment the variational objective to mitigate such inconsistency and induce a chosen latent distribution, instead of the implicit assumption found in a standard softmax layer. Overall, we provide new theoretical insight into the inner workings of widely-used softmax classifiers. Empirical evaluation on image and text classification datasets demonstrates that our proposed approach, variational classification, maintains classification accuracy while the reshaped latent space improves other desirable properties of a classifier, such as calibration, adversarial robustness, robustness to distribution shift and sample efficiency useful in low data settings.

翻译：我们提出了一种用于分类的潜变量模型，该模型为神经网络softmax分类器提供了新颖的概率解释。我们推导了一个变分目标来训练该模型，类似于用于训练变分自编码器的证据下界（ELBO），该目标推广了softmax交叉熵损失。将softmax层的输入视为潜变量的样本，我们的抽象视角揭示了其预期分布（用于准确标签预测所需）与实际中发现的经验分布之间可能存在的矛盾。我们增强了变分目标以缓解这种不一致性，并诱导出选定的潜分布，而非标准softmax层中隐含的假设。总体而言，我们为广泛使用的softmax分类器的内部工作机制提供了新的理论见解。在图像和文本分类数据集上的实证评估表明，我们提出的方法——变分分类——在保持分类准确性的同时，通过重塑潜空间改善了分类器的其他理想特性，如校准性、对抗鲁棒性、对分布偏移的鲁棒性以及在低数据场景中有用的样本效率。