Stochastic Vision Transformers with Wasserstein Distance-Aware Attention
Franciskus Xaverius Erick Mina Rezaei Johanna Paula Müller Bernhard Kainz
Self-supervised learning is one of the most promising approaches to acquiring
knowledge from limited labeled data. Despite the substantial advancements made
in recent years, self-supervised models have posed a challenge to
practitioners, as they do not readily provide insight into the model's
confidence and uncertainty. Tackling this issue is no simple feat, primarily
due to the complexity involved in implementing techniques that can make use of
the latent representations learned during pre-training without relying on
explicit labels. Motivated by this, we introduce a new stochastic vision
transformer that integrates uncertainty and distance awareness into
self-supervised learning (SSL) pipelines. Instead of the conventional
deterministic vector embedding, our novel stochastic vision transformer encodes
image patches into elliptical Gaussian distributional embeddings. Notably, the
attention matrices of these stochastic representational embeddings are computed
using Wasserstein distance-based attention, effectively capitalizing on the
distributional nature of these embeddings. Additionally, we propose a
regularization term based on Wasserstein distance for both pre-training and
fine-tuning processes, thereby incorporating distance awareness into latent
representations. We perform extensive experiments across different tasks such
as in-distribution generalization, out-of-distribution detection, dataset
corruption, semi-supervised settings, and transfer learning to other datasets
and tasks. Our proposed method achieves superior accuracy and calibration,
surpassing the self-supervised baseline in a wide range of experiments on a
variety of datasets.