# Probabilistic Attention for Interactive Segmentation

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

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 55 of its 73 references that are in the map.
- First posted on arXiv 23 Jun 2021.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [End-to-End Object Detection with Transformers](https://mlsanity.com/paper/2005.12872/) (2020)
- [Mask R-CNN](https://mlsanity.com/paper/1703.06870/) (2017)
- [Fully Convolutional Networks for Semantic Segmentation](https://mlsanity.com/paper/1411.4038/) (2014)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Deformable DETR: Deformable Transformers for End-to-End Object Detection](https://mlsanity.com/paper/2010.04159/) (2020)
- [DeepLab: Semantic Image Segmentation with Deep Convolutional Nets, Atrous Convolution, and Fully Connected CRFs](https://mlsanity.com/paper/1606.00915/) (2016)
- [Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation](https://mlsanity.com/paper/1802.02611/) (2018)
- [XLNet: Generalized Autoregressive Pretraining for Language Understanding](https://mlsanity.com/paper/1906.08237/) (2019)
- [Pyramid Scene Parsing Network](https://mlsanity.com/paper/1612.01105/) (2016)
- [Non-local Neural Networks](https://mlsanity.com/paper/1711.07971/) (2017)
- [Rethinking Atrous Convolution for Semantic Image Segmentation](https://mlsanity.com/paper/1706.05587/) (2017)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [Show, Attend and Tell: Neural Image Caption Generation with Visual Attention](https://mlsanity.com/paper/1502.03044/) (2015)
- [Deep High-Resolution Representation Learning for Visual Recognition](https://mlsanity.com/paper/1908.07919/) (2019)
- [Deep High-Resolution Representation Learning for Human Pose Estimation](https://mlsanity.com/paper/1902.09212/) (2019)
- [LVIS: A Dataset for Large Vocabulary Instance Segmentation](https://mlsanity.com/paper/1908.03195/) (2019)
- [Self-Attention with Relative Position Representations](https://mlsanity.com/paper/1803.02155/) (2018)
- [On the Variance of the Adaptive Learning Rate and Beyond](https://mlsanity.com/paper/1908.03265/) (2019)
- [CCNet: Criss-Cross Attention for Semantic Segmentation](https://mlsanity.com/paper/1811.11721/) (2018)
- [Universal Transformers](https://mlsanity.com/paper/1807.03819/) (2018)
- [Segmentation Transformer: Object-Contextual Representations for Semantic Segmentation](https://mlsanity.com/paper/1909.11065/) (2019)
- [Stand-Alone Self-Attention in Vision Models](https://mlsanity.com/paper/1906.05909/) (2019)
- [Hopfield Networks is All You Need](https://mlsanity.com/paper/2008.02217/) (2020)
- [Residual Attention Network for Image Classification](https://mlsanity.com/paper/1704.06904/) (2017)
- [Bottleneck Transformers for Visual Recognition](https://mlsanity.com/paper/2101.11605/) (2021)
- [Attention Augmented Convolutional Networks](https://mlsanity.com/paper/1904.09925/) (2019)

## Built on by (most cited first, 12 of 12)

- [Designing Robust Transformers using Robust Kernel Density Estimation](https://mlsanity.com/paper/2210.05794/) (2022)
- [Interactive Class-Agnostic Object Counting](https://mlsanity.com/paper/2309.05277/) (2023)
- [Unveiling the Hidden Structure of Self-Attention via Kernel Principal Component Analysis](https://mlsanity.com/paper/2406.13762/) (2024)
- [Elliptical Attention](https://mlsanity.com/paper/2406.13770/) (2024)
- [Attention: Marginal Probability is All You Need?](https://mlsanity.com/paper/2304.04556/) (2023)
- [A Primal-Dual Framework for Transformers and Neural Networks](https://mlsanity.com/paper/2406.13781/) (2024)
- [Transformer Meets Twicing: Harnessing Unattended Residual Information](https://mlsanity.com/paper/2503.00687/) (2025)
- [Transformer with Fourier Integral Attentions](https://mlsanity.com/paper/2206.00206/) (2022)
- [A Semantic and Motion-Aware Spatiotemporal Transformer Network for Action Detection](https://mlsanity.com/paper/2405.08204/) (2024)
- [Robust Filter Attention: Self-Attention as Precision-Weighted State Estimation](https://mlsanity.com/paper/2509.04154/) (2025)
- [Zonkey: A Hierarchical Diffusion Language Model with Differentiable Tokenization and Probabilistic Attention](https://mlsanity.com/paper/2601.21768/) (2026)
- [The Transformer as a Polar State Estimator](https://mlsanity.com/paper/2605.11007/) (2026)

## Abstract (arXiv)

We provide a probabilistic interpretation of attention and show that the standard dot-product attention in transformers is a special case of Maximum A Posteriori (MAP) inference. The proposed approach suggests the use of Expectation Maximization algorithms for online adaptation of key and value model parameters. This approach is useful for cases in which external agents, e.g., annotators, provide inference-time information about the correct values of some tokens, e.g, the semantic category of some pixels, and we need for this new information to propagate to other tokens in a principled manner. We illustrate the approach on an interactive semantic segmentation task in which annotators and models collaborate online to improve annotation efficiency. Using standard benchmarks, we observe that key adaptation boosts model performance (~10% mIoU) in the low feedback regime and value propagation improves model responsiveness in the high feedback regime. A PyTorch layer implementation of our probabilistic attention model will be made publicly available here: https://github.com/apple/ml-probabilistic-attention.