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 cross-entropy loss used to train classification models. 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 to be output, and the empirical distribution found in practice. We augment the variational objective to mitigate such inconsistency and encourage a chosen latent distribution, instead of the implicit assumption in off-the-shelf softmax classifiers. Overall, we provide new theoretical insight into the inner workings of widely-used softmax classification. 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分类的内部机制提供了新的理论见解。在图像和文本分类数据集上的实证评估表明，我们提出的变分分类方法在保持分类精度的同时，重塑后的潜空间改善了分类器的其他理想特性，如校准性能、对抗鲁棒性、分布偏移鲁棒性以及低数据情境下的样本效率。