Optimized Uncertainty Estimation for Vision Transformers: Enhancing Adversarial Robustness and Performance Using Selective Classification

Deep Learning models often exhibit undue confidence when encountering out-of-distribution (OOD) inputs, misclassifying with high confidence. The ideal outcome, in these cases, would be an "I do not know" verdict. We enhance the trustworthiness of our models through selective classification, allowing the model to abstain from making predictions when facing uncertainty. Rather than a singular prediction, the model offers a prediction distribution, enabling users to gauge the model’s trustworthiness and determine the need for human intervention. We assess uncertainty in two baseline models: a Convolutional Neural Network (CNN) and a Vision Transformer (ViT). By leveraging these uncertainty values, we minimize errors by refraining from predictions during high uncertainty. Additionally, we evaluate these models across various distributed architectures, including new AI architectures, Cerebras CS-2, and SambaNova SN30.