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