Probabilistic Attention for Interactive Segmentation

Prasad Gabbur, Manjot Bilkhu, Javier Movellan

Keywords: Attention, Transformers, Probabilistic model, Gaussian mixture model, Interactive segmentation, Semantic segmentation

TL;DR: A new perspective of attention as a probabilistic generative model with applications to interactive image segmentation.

Abstract: 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 on-line 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 ($\sim10\%$ 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 is available here: https://github.com/apple/ml-probabilistic-attention.

Code Of Conduct: I certify that all co-authors of this work have read and commit to adhering to the NeurIPS Statement on Ethics, Fairness, Inclusivity, and Code of Conduct.

Code: https://github.com/apple/ml-probabilistic-attention

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2106.15338/code)