Probabilistic Attention for Interactive SegmentationDownload PDF

21 May 2021, 20:48 (edited 26 Oct 2021)NeurIPS 2021 SpotlightReaders: Everyone
  • 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.
  • Supplementary Material: pdf
  • 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
11 Replies

Loading