Probabilistic Attention for Interactive Segmentation

Abstract: Prasad Gabbur, Manjot Bilkhu, Javier Movellan

Reviews and author responses:

 Thank you for addressing my questions, and I will keep my rating.

 I would like to thank the authors for the response, which answers my questions. I keep my score and would like to see the paper at NeurIPS. 

 Our sincere thanks for taking time to review our paper and provide feedback. Please find our responses to your concerns below.

* **Motivation for a generative model of attention**

Our main motivation for this work is to improve the effectiveness of annotator feedback on the task of interactive segmentation through the context modeling properties of self-attention. There has been at least one prior work (cited in the paper as [16], Sec. 2.1, line 70) to view the attention operation as a generative model. But it does so only for the queries and with a different motivation. One immediate benefit of the generative model is to enable various inferences, termed as key adaptation and belief propagation in the paper, for improving the efficiency on interactive segmentation. The generative model also has implications in terms of the redundancy of representation afforded by the self-attention mechanism for vision applications. Specifically, we believe that neighboring pixels in the input or some of the hidden representations in the intermediate layers are highly correlated and groups of these pixels can be assumed to arise from an underlying semantic concept. The latent variables could be representative of these concepts potentially leading to more efficient forms of attention by having fewer memory units (latents) than the number of input queries or output values. 

* **Relationship to Variational Attention**

Thank you for pointing us to the interesting work on Variational Attention (VA). The motivation to propose a variational attention mechanism in that work is to alleviate the “bypassing” phenomenon of standard deterministic attention that leads to lower diversity of the generated samples in a Variational Encoder Decoder (VED). Although both our approach and theirs look similar in flavor, there are differences in the parameterizations and the underlying likelihoods. For instance, the per-unit keys and expected values are the same (or shared parameters) in their formulation. The likelihood functions for the values in VA are Gaussians for both the priors ($p(a_j)$) and the posteriors ($q(a_j|x)$). In our case, the posterior value likelihood function ($p(v|q)$) is a mixture of Gaussians, with the mixture components being the per unit value likelihoods ($p(v|u_j)$) and the weights given by the unit posteriors ($p(u_j|q)$). Other than the possibility of inducing a discrete distribution on the attention weights, there are no distributional forms for the queries/keys in VA. We will cite the VA work and the differences to our approach in the final version of the paper.

* **Extension to NLP tasks**

The focus of our work has been to demonstrate the value of probabilistic attention on interactive segmentation as the main motivation. However our approach is applicable to problems in any domain that make use of attention in its various forms. We believe that some domain knowledge is required to design the hyperparameters and distributional forms of the proposed generative model to achieve the best results in the respective domains. We will plan to explore NLP tasks (e.g. GLUE benchmarks) using probabilistic attention in a future work. We will also make our implementation publicly available to make it easy for other practitioners to experiment with our model. 

 Thank you for your review. Please see our comments below specifically addressing the weaknesses above.

* **Why latent memory units?**

One of the main contributions of this work is a new perspective of attention as a probabilistic generative model for queries and values through a set of memory units. From a generative perspective, a memory unit is first sampled from a prior distribution and then a query/value pair is sampled from the likelihoods given the unit. The latents are required to indicate the unknown identity of the unit per query/value pair as in a standard mixture model. These latents also help propagate information from one unit’s output to another during value propagation. While this is just a different perspective on standard attention, this perspective allows us to derive new algorithms for key adaptation and value propagation, which we show to be beneficial on the task that is the main focus of our work--interactive segmentation.

* **Application to tasks other than interactive segmentation**

Interactive segmentation was the motivating task for us to rethink the possibility of self-attention enabling more effective propagation of annotator feedback from one unit to the others. Focusing on this task, we are able to demonstrate in detail, with quantitative results, the benefits provided by different types of inferences (key adaptation and value propagation) from our model. As pointed out above as a strength of our approach, probabilistic attention is a drop-in replacement for any task that uses one or other forms of attention. We are planning to experiment with other problems in computer vision as well as other domains and report results in a future work. At the same time, we will make our implementation public so that practitioners in various fields can tailor the generative model to suit their domains leading to more effective use of the proposed approach. 

 Thank you for your review of our paper.The following are our responses to your comments

* **Comparison to other methods**

The main point of the paper is to propose a novel perspective of attention as inference under a probabilistic generative model. We show that standard attention is equivalent to MAP inference with a constrained version of our model. More importantly, the approach is more general than standard attention and reveals additional inferences using Bayesian calculus to update model parameters online and incorporate/propagate user feedback for some of the outputs. We wanted to highlight these aspects of our model with detailed results as opposed to comparing our approach with the state-of-the-art.

* **Choice of the BoTNet50 architecture**

We chose to experiment with BoTNet50 [46] mainly because of its simplicity in terms of replacing standard attention with our proposed probabilistic attention within the architecture. The BoTNet50 work demonstrates the benefits of using self-attention within computer vision architectures by replacing very few convolutional modules with attention modules in the last bottleneck layer of ResNet50. There has been progress lately on making self-attention feasible throughout the network architecture by working with patches instead of pixels (ViT[17] and more recently Swin [R1], SegFormer[R2], etc.). We are currently experimenting with these architectures by replacing standard attention with ours and are finding that our approach provides benefits to these architectures similar to the ones we report for BoTNet50. These results will be reported in a future work.

**Additional References**
***
* R1: Liu et. al., Swin Transformer: Hierarchical Vision Transformer using Shifted Windows, arXiv:2103.14030
* R2: Xie et. al., SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers, arXiv:2105.15203


 We sincerely thank the Reviewer for taking the time to review our paper. We address the questions and comments below.

1. EM and network training 

    * We use EM to prove that MAP inference of values given queries in our model is equivalent to standard attention under some constraints as described in Sec. 3.3. In this case, only one update is sufficient. However, our model extends to cases beyond standard attention, e.g., using non-zero $\beta_j$’s, where more than one EM iteration may provide improved MAP values.

     * EM is used to derive the online update equations for mixture model parameters (keys and the associated precisions) to adapt them based on the observed queries. In this work, we do not use EM for key adaptation during training but only during inference. This is to demonstrate that a simple online update of the keys allows the models to perform better in the low annotator feedback scenario. We have seen recently that using a single EM iteration during training leads to further performance improvements, most likely due to the way the gradients of queries and keys interact during backpropagation. Although we used zero key prior precisions for experiments in the paper (ML update), we are experimenting with learning the prior precisions as part of backpropagation and plan to report these results in a future work.

    * EM is also used during value propagation to update the per-unit expected values ($\mu_j$’s) in the presence of annotator feedback. In this case, we used one EM iteration during training to learn the fixed value logits (Sec. 4.4, lines 262-263). This was rather a choice of convenience as more EM iterations during training would lead to increased GPU computation and memory. However, during inference we have shown results with 1 and 5 iterations (Fig. 4). All these experiments use a fixed value (per dataset) of prior precision for expected value, chosen using a grid search (Appendix C4, lines 91-92). Learning fixed priors or even input specific priors (Appendix C4, lines 94-95) is a possible future work.

2. It is possible to use our proposed probabilistic attention mechanism in any task that uses one or other forms of attention. It would only require calling our update in place of the current attention update. To this end, we will be making our implementation available as a PyTorch layer. Specifically our implementation is parameter-free allowing users to re-use/define their query, key and value parameterizations and even instantiate any additional ones (e.g. precisions, priors) in any way they choose. 

We are also thankful for pointing out the missing assumption about equal \betas in going from B3 to B4.
Although this assumption is stated in the main paper in Sec. 3.3 (line 137), we will make a note of it in the appendix as well in the final version. 


The paper presents a generative view of attention: N hidden memory units each of which generates observed a query and a value. The memory units densely connected with each other. The paper shows that, given queries, the value of a memory unit, in a certain condition, is computed in a similar way as in attention, using EM. 

Applying this new attention, called probabilistic attention, to interactive image segmentation tasks (where users provide some ground-truth pixels' values on the fly), the paper shows better results than using traditional attention on GrabCut and Berkeley datasets. 

**Contributions** 
1. novel explanation of attention, leading to new attention mechanism relying on probabilistic   **Originality** 

Explaining attention is not new but using probabilistic generative view in this paper is novel and very interesting. It will open a door to a new class of attention layers that have wider probabilistic properties. 

I particularly like the trick of letting beta -> 0 in section 3.3. This is counterintuitive to me as beta -> 0, eq 9 implies that v is more and more independent of all u_j. Turns out, in the end, we have traditional attention. Also, line 142, it is interesting that we just need to use one EM update, exactly like the attention. 


**Quality** 

The paper technically is sound. I found the maths correct (though some steps are missing, such as from B3 to B4 where beta1=beta2=...=beta is crucial). 

I would like to see more in the paper: 
1. What is the effects and use of EM? It's unclear to me whether EM is used together with training the networks for priors. If it is, how? 
2. Interactive image segmentation looks like a good tasks for attention. The authors claim this probabilistic attention is more generic applicable. Is there another successful story? Is it simple just to pickup a task, a model which uses attention, and replace that attention by the probabilistic attention? Is tuning precision prior important?

**Clarity** 

The submission is well written. 

**Significance**

The results shown in the paper is important as now we can have more freedom to extend attention, especially in probabilistic ways.  I don't see any negative societal impact. 

This paper describes an EM-based approach to propagating attention parameters, showing strong performance on an interactive segmentation task.  The paper describes a probabilistic attention model which contains standard transformer attention as a special case, along with an online value propagation method which allows humans to update attention weights in a principled manner. This is evaluated in the domain of interactive segmentation but could seemingly be applied to other domains such as interactive machine translation. This paper is not in my area, but to the best of my knowledge the approach for attention propagation is novel and potentially applicable to a wide number of human-in-the-loop tasks. Although the empirical gains are modest, both the theoretical and empirical sections are clearly written and motivated.

The empirical results show that:
 - Probabilistic attention outperforms standard attention in a low-click interactive segmentation domain and performs similarly to standard attention with a greater number of clicks
 - Unsupervised key adaptation (cf. Section 3.5) leads to significantly improved IoU in a non-interactive domain

The authors also released a Pytorch layer implementing their proposed approach. The paper is seemingly missing a comparison to other methods for interactive segmentation and does not justify its choice of the BoTNet50 architecture as a baseline model. Once again, though, the paper is out of my area and I am not familiar with the state-of-the-art in this domain. I will choose to defer to the other reviewers on my recommendation for acceptance. N/A

This paper offers a new interpretation of standard dot-product
attention as a mixture model. The model assumes memory units are
latent variables that dictate the probabilistic interaction between
query and value vectors. The significance of this probabilistic
interpretation is twofold:

1. MAP inference under a carefully constrained Gaussian mixture model
   case of the model yields exactly the update rule of standard
   dot-product attention. Prior to this the most similar
   interpretation of attention was through modern Hopfield networks
   (Ramsauer et al., 2021).
2. The model admits an EM algorithm which allows for either online or
   offline learning in the precence of partial annotations. 

The paper applies this probabilistic attention approach to interactive
semantic segmentation by incorporating it in various ways within
BotNet. The probabilistic attention either replaces standard attention
or is used within the self-attention at the last part of the
network. Under various settings (e.g. full or axial attention;
factored or full relative positional encoding), the incorporation of
probabilistic attention improves over the corresponding baseline. 

The paper finally introduces a corrective self attention layer in
order to more effectively pass along information from an annotators
corrections of the model output. Results show that such a use of these
corrections is more effective than not.  Strengths:
- The paper presents a novel interpretation of standard dot-product attention as a probabilistic (mixture) model.
- Bayesian inference in the model enables online updating of the
  mixture parameters and corrections on a subset of values to be
  propagated.
- The results on interactive semantic segmentation suggest that the
  approach is a drop in replacement for attention that can improve
  over baselines without.

Weaknesses:
- From the initial description of the mixture model, it was a bit difficult to understand *why* it would make sense to have latent memory units at all. Only later upon thinking what these were being used for within the context of interactive semantic segmentation did this become clear.
- The probabilistic attention mechanism is only applied to one task (interactive semantic segmentation). It would be interesting to see its effectiveness in other settings.
 Yes

Summary and Contributions

The authors study a probabilistic interpretation of attention mechanisms for interactive segmentation.  The goal is to show that self-attention is related to MAP and EM can be used to induce the posterior distribution over attention scores.  The main contributions are: probabilistic interpretation of attention as a generative model, interactive segmentation with the proposed attention model, and classification.  The study shows that the proposed model has competitive results compared to related work.
  Strengths

Clear description of background knowledge, proposed model, and related work needed to understand the proposed model.  
The authors perform a comprehensive comparison across different tasks.

Weaknesses

It is not clear the motivation to  define attention as a generative model. 
It is not discussed the possible difference of the proposed models with variational attention models.

Questions to the Authors
Please address the following questions during the rebuttal:

Which is the relation of the proposed model and variational attention models used in NLP applications? e.g. Variational Attention for Sequence-to-Sequence Models
A possible extra contribution could be the discussion of the proposed model for text classification tasks. e.g. GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding
 Yes, the authors addressed  limitations.


Meta-review: This paper builds and analyzes a probabilistic mixture model for which softmax attention mechanisms are a special case. Reviewers  find the construction enlightening, educational, and with high impact potential. They strongly recommend improving the presentation to make the "story" and motivation much clearer upfront: the reasons for the construction don't reveal themselves until the connections to attention. The constructions seem different, but uthors should discuss the relationship with other probabilstic attention work (e.g. variational attention by Deng et al, but also Heo et al, NeurIPS 19, Bahuleyan et al, Coling 19). The font sizes and line markers/colors in the plot should be changed to improve readability and skimmability of the paper.

Decision: accept