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, and the empirical distribution they follow in practice. We then devise a variational objective to mitigate such inconsistency and encourage a specified 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; and empirical evaluation on image and text classification datasets demonstrates that our proposed remedy, variational classification, maintains classification accuracy while the reshaped latent space improves other desirable classifier properties, such as calibration, adversarial robustness, robustness to distribution shift and sample efficiency useful in low data settings.

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