Probabilistic Attention for Interactive Segmentation

1 Introduction . Attention was first introduced as a computational primitive for natural language processing [ 53 ] and has since been widely adopted [ 17 , 63 , 13 , 14 ] as a replacement for recurrent primitives such as LSTMs [ 26 ] . More recently it has been making inroads into computer vision [ 43 , 69 , 57 , 64 , 5 , 19 , 49 ] as a replacement for the long accepted convolution as the main computational primitive . Self-attention based architectures have demonstrated state-of-the-art results in fundamental vision problems including image classification [ 5 , 43 , 69 , 19 , 49 ] , object detection [ 8 , 71 , 57 ] , image and video semantic segmentation [ 55 , 29 , 57 , 41 ] and tracking [ 66 ] to state a few . There are a few different perspectives on the reasons for success of self-attention in computer vision and its superiority over convolution . This includes a view that the self-attention mechanism allows modeling spatially varying dynamic convolution filters [ 32 ] and at the same time enabling parameter independent scaling of receptive fields [ 52 ] . Another includes their ability to capture global context through long range interactions especially when full attention is feasible [ 49 ] at reduced spatial resolution maps or using an approximation of full attention with axial [ 55 ] or criss-cross attention [ 29 ] . A recent work [ 45 ] introduces modern Hopfield networks with continuous states where the update mechanism is shown to be equivalent to the update mechanism of standard dot-product attention [ 53 ] . They show that such a network has the capacity to store exponentially many patterns and retrieve them with high fidelity . In this work , we provide a novel interpretation of attention as a probabilistic generative model for queries and values . Specifically we hypothesize the existence of a bank of probabilistic memory units , each of which maintains a joint probability distribution over queries and values parameterized through keys . A query/value pair is generated by first sampling a unit ( from a prior over units ) followed by sampling the pair from the unit specific joint distribution . This is equivalent to generating the queries and values through a probabilistic mixture model over the units . A particular form for unit joint likelihoods expressed as Gaussians for both the query and value marginals , assuming their independence conditioned on a unit , turns out to be equivalent to traditional 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . dot product attention under a few constraints . As shown in Section 3.3 , maximum likelihood ( ML ) inference for the corresponding value given a query is equivalent to standard dot-product attention . Our probabilistic interpretation provides a systematic framework for online update of mixture model parameters based on a set of observed queries . It also allows propagation of correct values provided by an external agent for some of the units to all other units . Using Bayesian inference in the constrained case , we derive update rules for online unsupervised adaptation ( Section 3.5 ) of query/key likelihood parameters based on a set of observed queries . We also derive update equations for online value propagation ( Section 3.6 ) across units based on fixed externally specified values for a subset of units . The latter is specifically useful for interactive segmentation where a correction provided by an annotator has to be propagated globally to make the process more efficient . We use probabilistic attention in place of standard attention in deep architectures for interactive segmentation both within the backbone and at the network head as a classifier . Specifically we use probabilistic attention updates in the BoTNet50 [ 49 ] architecture and show that adapting keys to incoming queries leads to better model performance in the low annotator feedback regime . Using value propagation within a probabilistic attention layer at the head of the segmentation network leads to a more responsive model through effective feedback propagation in the high feedback regime . We also use both key adaptation and value propagation together and demonstrate the complementary effects of the two in both the low and high annotator feedback regimes . 2 Related Work . 2.1 Attention . Natural language processing has seen the rise [ 3 , 53 ] and widespread adoption of attention in recent years [ 17 , 63 , 13 , 14 , 59 ] . One of the first works on visual attention was on learning to attend to image regions for caption generation [ 60 ] . Since then there has been a steady progress on using attention primitives within vision models for recognition and classification [ 27 , 57 , 5 , 54 , 43 , 69 , 19 , 52 , 49 ] , detection [ 8 , 71 ] , segmentation [ 55 , 29 ] , tracking [ 66 ] and video analysis [ 42 , 41 ] . There have been numerous works interpreting the attention mechanism as a form of computing non-local means [ 7 , 57 ] , approximating dynamic convolution filters [ 32 , 52 ] , and capturing global context through long range interactions [ 27 , 43 , 55 ] . The standard dot-product attention [ 53 ] update was also formulated as emerging from the update rule of modern Hopfield networks [ 45 ] . Our work introduces attention mechanism from a novel perspective as that of inferring from a probabilistic memory bank . To our knowledge , the only work that is closest to our approach is [ 18 ] , which also proposes a similar interpretation but only for the queries in order to study the explaining away effect of attention update . Other notable works modeling attention in a probabilistic framework have done so using neural processes [ 25 ] or variational approximations for modeling alignments as latent variables [ 16 ] or to overcome the bypassing phenomenon of deterministic attention [ 4 ] in sequence tasks . Being different from these approaches , our formulation encapsulates queries and values in a single generative model and provides an interpretation of standard dot-product attention as constrained Bayesian inference . Doubly normalized attention scheme [ 18 ] also emerges as a special case of key adaptation in our framework . 2.2 Interactive Segmentation . Deep neural networks have set state-of-the art in semantic segmentation through the use of fully convolutional architectures [ 36 , 70 , 10 , 11 , 12 , 50 , 56 , 67 ] and more recently using hybrid convolution and self-attention [ 64 ] or stand-alone self-attention architectures [ 55 ] . The input domain of interactive segmentation includes user input in the form of clicks or scribbles in addition to the visual signal ( images or videos ) . The earliest works in interactive segmentation were based on algorithmic approaches for incorporating human inputs into region [ 46 , 51 , 65 ] or boundary [ 40 , 20 ] processing pipelines . [ 44 ] provides a comprehensive survey of interactive segmentation approaches . More recently deep networks have been used to incorporate user feedback to guide their output predictions at the pixel level . Following a similar taxonomy , these can be roughly categorized into region based [ 62 , 61 , 37 , 33 , 6 , 2 , 28 , 68 , 34 , 31 ] or boundary based approaches [ 9 , 1 , 35 , 58 ] . Deep Extreme Cut ( DEXTR ) [ 38 ] demonstrated that user guidance in the form of extreme points could be used in addition to the input channels to accurately localize the object of interest . More recently [ 68 ] argued that three points are sufficient as input guidance to localize the object but additional corrective clicks could be used to further refine the prediction . Other works have used the corrective clicks to adapt the network inputs [ 30 ] , embeddings [ 48 ] or parameters [ 31 ] online . Different from these previous approaches , we use corrective clicks as providing fixed values for a subset of units in the proposed probabilistic attention framework . These values are propagated globally through the attention mechanism to directly and more effectively influence the outputs towards user intended values . 3 Method . We provide a probabilistic interpretation of attention as a generative model for queries and values through a set of memory units . Using this formulation , traditional attention in transformers [ 53 ] reduces to the special case of maximum a posteriori ( MAP ) inference of values given queries , assuming Gaussians for the likelihoods . Using Bayesian inference , we provide a systematic approach to adapt keys online as a locally ML update of the corresponding model parameters . Our formulation also allows to fix the values of certain units and propagate their influence to other units online by conditioning on the fixed unit values . The following sections provide more details on the probabilistic model . 3.1 Probabilistic attention . We assume that there are n memory units , indexed by i , each of which can be queried through a vector qi ∈ Rd to yield an output value vector vi ∈ Rm . The queries and the corresponding values may depend on an input x . For example , each memory unit may represent a pixel xi in an image x . The joint distribution of queries qi and values vi conditioned on the input x is assumed to factorize over memory units p ( q1 : n , v1 : n|x ) = Πni=1pi ( qi , vi | x ) , ( 1 ) where x is the conditioning input , q1 : n = { q1 , · · · , qn } , v1 : n = { v1 , · · · , vn } , and qi ∈ Rd , vi ∈ Rm are the query , value vectors for unit i respectively . The per-unit joint likelihood pi ( qi , vi | x ) is a probabilistic mixture model given by pi ( q , v | x ) = n∑ j=1 pi ( q , v , uj | x ) = n∑ j=1 πi , j ( x ) pi ( q , v | uj , x ) , ( 2 ) where we have dropped the subscript i from q and v for simplicity . In the above , uj indexes unit j , πi , j ( x ) is the probability of activating unit j when unit i is queried , pi ( q , v | uj , x ) is the likelihood of observing the pair ( q , v ) , given the pair is generated through unit j in the mixture , conditioned on the input x . 3.2 Value inference . Using the above model , it is possible to find the most likely value v̂ given a query q to unit i v̂ = argmax v pi ( v|q , x ) . ( 3 ) We use Expectation Maximization ( EM ) [ 15 ] to achieve this : starting with an initial estimate v0 of the most probable value and iterating over the standard EM auxiliary function Qi . Given the latest known estimate vt , the M step produces a new estimate vt+1 that increases Qi by maximizing it w.r.t . vt+1 . This guarantees local maximization of pi ( v|q , x ) . Qi ( v t , vt+1 | x ) = ∑ j wti , j log pi ( uj , q , v t+1 | x ) , ( 4 ) where wti , j = pi ( uj | q , vt , x ) = πi , j ( x ) pi ( q , v t | uj , x ) ∑ j πi , j ( x ) pi ( q , v t | uj , x ) . ( 5 ) The n× n matrix formed by the entries wi , j corresponds to the attention matrix in standard transformers . The optimal value v̂ is obtained by taking the gradient with respect to vt+1 and setting it to zero ∇vt+1Qi ( vt , vt+1 | x ) = ∑ j wi , j∇vt+1 log pi ( q , vt+1 , uj | x ) , ( 6 ) where ∇vt+1 log pi ( q , vt+1 , uj | x ) is the Fisher Score for unit i with respect to vt+1 .

Summary: 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