# Stochastic Vision Transformers with Wasserstein Distance-Aware Attention

arXiv: https://arxiv.org/abs/2311.18645 · HTML: https://mlsanity.com/paper/2311.18645/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 31 of its 46 references that are in the map.
- First posted on arXiv 30 Nov 2023.

## Builds on (most cited first, 30 of 31)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [Unsupervised Learning of Visual Features by Contrasting Cluster Assignments](https://mlsanity.com/paper/2006.09882/) (2020)
- [Barlow Twins: Self-Supervised Learning via Redundancy Reduction](https://mlsanity.com/paper/2103.03230/) (2021)
- [VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning](https://mlsanity.com/paper/2105.04906/) (2021)
- [AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty](https://mlsanity.com/paper/1912.02781/) (2019)
- [data2vec: A General Framework for Self-supervised Learning in Speech, Vision and Language](https://mlsanity.com/paper/2202.03555/) (2022)
- [Learning Audio-Visual Speech Representation by Masked Multimodal Cluster Prediction](https://mlsanity.com/paper/2201.02184/) (2022)
- [SSD: A Unified Framework for Self-Supervised Outlier Detection](https://mlsanity.com/paper/2103.12051/) (2021)
- [Contrastive Training for Improved Out-of-Distribution Detection](https://mlsanity.com/paper/2007.05566/) (2020)
- [Wasserstein Distributionally Robust Optimization: Theory and Applications in Machine Learning](https://mlsanity.com/paper/1908.08729/) (2019)
- [Sinkformers: Transformers with Doubly Stochastic Attention](https://mlsanity.com/paper/2110.11773/) (2021)
- [Plex: Towards Reliability using Pretrained Large Model Extensions](https://mlsanity.com/paper/2207.07411/) (2022)
- [Word Representations via Gaussian Embedding](https://mlsanity.com/paper/1412.6623/) (2014)
- [CLIP-Event: Connecting Text and Images with Event Structures](https://mlsanity.com/paper/2201.05078/) (2022)
- [Sequential Recommendation via Stochastic Self-Attention](https://mlsanity.com/paper/2201.06035/) (2022)
- [Improving Transformers with Probabilistic Attention Keys](https://mlsanity.com/paper/2110.08678/) (2021)
- [Transformer Uncertainty Estimation with Hierarchical Stochastic Attention](https://mlsanity.com/paper/2112.13776/) (2021)
- [Deep Bregman Divergence for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2109.07455/) (2021)
- [RODD: A Self-Supervised Approach for Robust Out-of-Distribution Detection](https://mlsanity.com/paper/2204.02553/) (2022)
- [Probabilistic Self-supervised Learning via Scoring Rules Minimization](https://mlsanity.com/paper/2309.02048/) (2023)

## Abstract (arXiv)

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.